Files
overlord/Overlord-Client/cmd/agent/capture/win_bitblt.go
T
2026-04-08 14:14:09 -05:00

578 lines
14 KiB
Go

//go:build windows
package capture
import (
"image"
"log"
"math"
"os"
"strconv"
"sync"
"sync/atomic"
"syscall"
"time"
"unsafe"
)
var (
user32 = syscall.NewLazyDLL("user32.dll")
gdi32 = syscall.NewLazyDLL("gdi32.dll")
shcore = syscall.NewLazyDLL("shcore.dll")
kernel32 = syscall.NewLazyDLL("kernel32.dll")
procGetDC = user32.NewProc("GetDC")
procReleaseDC = user32.NewProc("ReleaseDC")
procGetSystemMetrics = user32.NewProc("GetSystemMetrics")
procCreateCompatibleDC = gdi32.NewProc("CreateCompatibleDC")
procCreateDCW = gdi32.NewProc("CreateDCW")
procDeleteDC = gdi32.NewProc("DeleteDC")
procCreateCompatibleBitmap = gdi32.NewProc("CreateCompatibleBitmap")
procCreateDIBSection = gdi32.NewProc("CreateDIBSection")
procDeleteObject = gdi32.NewProc("DeleteObject")
procSelectObject = gdi32.NewProc("SelectObject")
procBitBlt = gdi32.NewProc("BitBlt")
procStretchBlt = gdi32.NewProc("StretchBlt")
procSetStretchBltMode = gdi32.NewProc("SetStretchBltMode")
procSetProcessDpiAwareness = shcore.NewProc("SetProcessDpiAwareness")
)
const (
SM_CXSCREEN = 0
SM_CYSCREEN = 1
SM_XVIRTUALSCREEN = 76
SM_YVIRTUALSCREEN = 77
SM_CXVIRTUALSCREEN = 78
SM_CYVIRTUALSCREEN = 79
SRCCOPY = 0x00CC0020
CAPTUREBLT = 0x40000000
BI_RGB = 0
DIB_RGB_COLORS = 0
PROCESS_PER_MONITOR_DPI_AWARE = 2
COLORONCOLOR = 3
)
type bitmapInfoHeader struct {
biSize uint32
biWidth int32
biHeight int32
biPlanes uint16
biBitCount uint16
biCompression uint32
biSizeImage uint32
biXPelsPerMeter int32
biYPelsPerMeter int32
biClrUsed uint32
biClrImportant uint32
}
type bitmapInfo struct {
bmiHeader bitmapInfoHeader
redMask uint32
greenMask uint32
blueMask uint32
alphaMask uint32
}
func getSystemMetric(i int32) int32 {
r, _, _ := procGetSystemMetrics.Call(uintptr(i))
return int32(r)
}
func getDC(hwnd uintptr) uintptr {
dc, _, _ := procGetDC.Call(hwnd)
return dc
}
func releaseDC(hwnd, hdc uintptr) {
_, _, _ = procReleaseDC.Call(hwnd, hdc)
}
func createCompatibleDC(hdc uintptr) uintptr {
r, _, _ := procCreateCompatibleDC.Call(hdc)
return r
}
func createDisplayDC() uintptr {
displayName := syscall.StringToUTF16Ptr("DISPLAY")
r, _, _ := procCreateDCW.Call(uintptr(unsafe.Pointer(displayName)), 0, 0, 0)
return r
}
func deleteDC(hdc uintptr) {
_, _, _ = procDeleteDC.Call(hdc)
}
func createCompatibleBitmap(hdc uintptr, w, h int32) uintptr {
r, _, _ := procCreateCompatibleBitmap.Call(hdc, uintptr(w), uintptr(h))
return r
}
func createDIBSection(hdc uintptr, bmi *bitmapInfo, usage uint32, bits *unsafe.Pointer) uintptr {
r, _, _ := procCreateDIBSection.Call(hdc, uintptr(unsafe.Pointer(bmi)), uintptr(usage), uintptr(unsafe.Pointer(bits)), 0, 0)
return r
}
func deleteObject(obj uintptr) {
_, _, _ = procDeleteObject.Call(obj)
}
func selectObject(hdc, hgdiobj uintptr) uintptr {
r, _, _ := procSelectObject.Call(hdc, hgdiobj)
return r
}
func bitBlt(hdcDest uintptr, x, y, cx, cy int32, hdcSrc uintptr, x1, y1 int32, rop uint32) bool {
r, _, _ := procBitBlt.Call(hdcDest, uintptr(x), uintptr(y), uintptr(cx), uintptr(cy), hdcSrc, uintptr(x1), uintptr(y1), uintptr(rop))
return r != 0
}
func stretchBlt(hdcDest uintptr, xDest, yDest, wDest, hDest int32, hdcSrc uintptr, xSrc, ySrc, wSrc, hSrc int32, rop uint32) bool {
r, _, _ := procStretchBlt.Call(hdcDest, uintptr(xDest), uintptr(yDest), uintptr(wDest), uintptr(hDest), hdcSrc, uintptr(xSrc), uintptr(ySrc), uintptr(wSrc), uintptr(hSrc), uintptr(rop))
return r != 0
}
func setStretchBltMode(hdc uintptr, mode int32) {
_, _, _ = procSetStretchBltMode.Call(hdc, uintptr(mode))
}
func bitBltWithFallback(hdcMem, hdcScreen uintptr, bounds image.Rectangle, w, h int) bool {
return bitBlt(hdcMem, 0, 0, int32(w), int32(h), hdcScreen, int32(bounds.Min.X), int32(bounds.Min.Y), SRCCOPY)
}
func stretchBltCapture(hdcMem, hdcScreen uintptr, bounds image.Rectangle, srcW, srcH, dstW, dstH int) bool {
setStretchBltMode(hdcMem, COLORONCOLOR)
return stretchBlt(hdcMem, 0, 0, int32(dstW), int32(dstH), hdcScreen, int32(bounds.Min.X), int32(bounds.Min.Y), int32(srcW), int32(srcH), SRCCOPY)
}
var dpiAwareOnce sync.Once
func setDPIAware() {
dpiAwareOnce.Do(func() {
if procSetProcessDpiAwareness.Find() == nil {
procSetProcessDpiAwareness.Call(uintptr(PROCESS_PER_MONITOR_DPI_AWARE))
}
})
}
var (
captureCount atomic.Int64
bitbltNs atomic.Int64
convertNs atomic.Int64
lastCaptureLogNs atomic.Int64
scaleOnce sync.Once
cachedScale float64
maxResHeight atomic.Int64 // 0 = default (1080), -1 = native, >0 = specific max height
lastCapScale atomic.Uint64
state capState
captureMu sync.Mutex
stateCreatedAt time.Time
stateRefreshInterval = 30 * time.Second
)
var captureDisplayFn = func(display int) (*image.RGBA, error) {
if useDesktopDuplication() {
img, err := captureDisplayDXGI(display)
if err == nil && img != nil {
return img, nil
}
if err != nil {
log.Printf("capture: dxgi duplication failed: %v (falling back to bitblt)", err)
}
}
return captureDisplayBitBlt(display)
}
func captureDisplayBitBlt(display int) (*image.RGBA, error) {
captureMu.Lock()
defer captureMu.Unlock()
setDPIAware()
maxDisplays := displayCount()
if display < 0 || display >= maxDisplays {
display = 0
}
mons := monitorList()
if display >= len(mons) {
display = 0
}
mon := mons[display]
bounds := resolveBounds(mon)
srcW := bounds.Dx()
srcH := bounds.Dy()
if srcW <= 0 || srcH <= 0 {
ResetMonitorCache()
mons = monitorList()
if display >= len(mons) {
display = 0
}
if display < len(mons) {
mon = mons[display]
bounds = resolveBounds(mon)
srcW = bounds.Dx()
srcH = bounds.Dy()
}
if srcW <= 0 || srcH <= 0 {
return nil, syscall.EINVAL
}
}
userScale := effectiveScale(srcW, srcH)
dstW := int(float64(srcW) * userScale)
dstH := int(float64(srcH) * userScale)
if dstW <= 0 || dstH <= 0 {
dstW = srcW
dstH = srcH
}
useStretch := dstW > 0 && dstH > 0 && (dstW < srcW || dstH < srcH)
capW := srcW
capH := srcH
if useStretch {
capW = dstW
capH = dstH
}
hdcScreen := createDisplayDC()
fromCreateDC := hdcScreen != 0
if hdcScreen == 0 {
hdcScreen = getDC(0)
}
if hdcScreen == 0 {
return nil, syscall.EINVAL
}
closeDC := func() {
if hdcScreen == 0 {
return
}
if fromCreateDC {
deleteDC(hdcScreen)
} else {
releaseDC(0, hdcScreen)
}
hdcScreen = 0
}
defer closeDC()
if capW <= 0 || capH <= 0 {
return nil, syscall.EINVAL
}
captureWithDC := func() (*image.RGBA, time.Duration, time.Duration, time.Duration, error) {
hdcMem, _, buf, stride, err := state.ensure(hdcScreen, capW, capH)
if err != nil {
return nil, 0, 0, 0, err
}
bitStart := time.Now()
var captured bool
if useStretch {
captured = stretchBltCapture(hdcMem, hdcScreen, bounds, srcW, srcH, capW, capH)
} else {
captured = bitBltWithFallback(hdcMem, hdcScreen, bounds, capW, capH)
}
if !captured {
state.reset()
hdcMem, _, buf, stride, err = state.ensure(hdcScreen, capW, capH)
if err != nil {
return nil, 0, 0, 0, err
}
if useStretch {
captured = stretchBltCapture(hdcMem, hdcScreen, bounds, srcW, srcH, capW, capH)
} else {
captured = bitBltWithFallback(hdcMem, hdcScreen, bounds, capW, capH)
}
if !captured {
return nil, 0, 0, 0, syscall.EINVAL
}
}
bitDur := time.Since(bitStart)
cursorDrawn := false
if useStretch {
sx := float64(capW) / float64(srcW)
sy := float64(capH) / float64(srcH)
cursorDrawn = DrawCursorOnDCScaled(hdcMem, bounds, sx, sy)
} else {
cursorDrawn = DrawCursorOnDC(hdcMem, bounds)
}
if len(buf) == 0 {
return nil, 0, 0, 0, syscall.EINVAL
}
dibDur := time.Duration(0)
convStart := time.Now()
swapRB(buf)
img := &image.RGBA{Pix: buf, Stride: stride, Rect: image.Rect(0, 0, capW, capH)}
convDur := time.Since(convStart)
if !useStretch && !cursorDrawn {
DrawCursorOnImage(img, bounds)
}
if dstW != capW || dstH != capH {
img = resizeNearest(img, dstW, dstH)
} else {
// Detach from the reusable DIB buffer to avoid concurrent mutation during encode.
stable := GetRGBA(img.Rect.Dx(), img.Rect.Dy())
copy(stable.Pix, img.Pix)
img = stable
}
return img, bitDur, dibDur, convDur, nil
}
if !stateCreatedAt.IsZero() && time.Since(stateCreatedAt) > stateRefreshInterval {
state.reset()
}
img, bitDur, dibDur, convDur, err := captureWithDC()
if err != nil && fromCreateDC {
state.reset()
closeDC()
hdcScreen = getDC(0)
fromCreateDC = false
if hdcScreen == 0 {
return nil, err
}
img, bitDur, dibDur, convDur, err = captureWithDC()
}
if err != nil {
state.reset()
return nil, err
}
logCaptureTimings(bitDur, dibDur, convDur)
return img, nil
}
func clampToVirtual(bounds image.Rectangle) image.Rectangle {
vx := int(getSystemMetric(SM_XVIRTUALSCREEN))
vy := int(getSystemMetric(SM_YVIRTUALSCREEN))
vw := int(getSystemMetric(SM_CXVIRTUALSCREEN))
vh := int(getSystemMetric(SM_CYVIRTUALSCREEN))
if vw <= 0 || vh <= 0 {
return bounds
}
virtual := image.Rect(vx, vy, vx+vw, vy+vh)
inter := bounds.Intersect(virtual)
if inter.Empty() {
return bounds
}
return inter
}
func resolveBounds(mon monitorDesc) image.Rectangle {
virtualW := int(getSystemMetric(SM_CXVIRTUALSCREEN))
virtualH := int(getSystemMetric(SM_CYVIRTUALSCREEN))
if virtualW > 0 && virtualH > 0 {
return clampToVirtual(mon.rect)
}
return captureBounds(mon)
}
func captureBounds(mon monitorDesc) image.Rectangle {
bounds := mon.rect
if mon.physW > 0 && mon.physH > 0 {
phys := image.Rect(mon.posX, mon.posY, mon.posX+mon.physW, mon.posY+mon.physH)
if phys.Dx() > 0 && phys.Dy() > 0 {
if mon.scale != 1.0 || phys.Dx() != bounds.Dx() || phys.Dy() != bounds.Dy() {
bounds = phys
}
}
}
return bounds
}
func swapRB(pix []byte) {
for i := 0; i+3 < len(pix); i += 4 {
pix[i], pix[i+2] = pix[i+2], pix[i]
}
}
func logCaptureTimings(bitDur, dibDur, convDur time.Duration) {
captureCount.Add(1)
bitbltNs.Add(bitDur.Nanoseconds())
convertNs.Add(convDur.Nanoseconds())
nowNs := time.Now().UnixNano()
last := lastCaptureLogNs.Load()
if last != 0 && time.Duration(nowNs-last) < 5*time.Second {
return
}
if !lastCaptureLogNs.CompareAndSwap(last, nowNs) {
return
}
frames := captureCount.Swap(0)
if frames == 0 {
return
}
avg := func(totalNs atomic.Int64) float64 {
return float64(totalNs.Swap(0)) / 1e6 / float64(frames)
}
bitAvg := avg(bitbltNs)
convAvg := avg(convertNs)
log.Printf("capture: win bitblt avg bitblt=%.2fms convert=%.2fms frames=%d", bitAvg, convAvg, frames)
}
func resizeNearest(src *image.RGBA, w, h int) *image.RGBA {
dst := GetRGBA(w, h)
srcW := src.Bounds().Dx()
srcH := src.Bounds().Dy()
if srcW <= 0 || srcH <= 0 || w <= 0 || h <= 0 {
return dst
}
xOff := make([]int, w)
for x := 0; x < w; x++ {
xOff[x] = (x * srcW / w) * 4
}
srcPix := src.Pix
dstPix := dst.Pix
srcStride := src.Stride
dstStride := dst.Stride
for y := 0; y < h; y++ {
sp := (y * srcH / h) * srcStride
dp := y * dstStride
for x := 0; x < w; x++ {
*(*uint32)(unsafe.Pointer(&dstPix[dp+x*4])) = *(*uint32)(unsafe.Pointer(&srcPix[sp+xOff[x]]))
}
}
return dst
}
type capState struct {
mu sync.Mutex
hdcMem uintptr
hbmp uintptr
buf []byte
stride int
w int
h int
bits unsafe.Pointer
}
func (s *capState) ensure(hdcScreen uintptr, w, h int) (uintptr, uintptr, []byte, int, error) {
s.mu.Lock()
defer s.mu.Unlock()
if s.hdcMem == 0 {
s.hdcMem = createCompatibleDC(hdcScreen)
if s.hdcMem == 0 {
return 0, 0, nil, 0, syscall.EINVAL
}
stateCreatedAt = time.Now()
}
if s.w != w || s.h != h || s.hbmp == 0 {
bmi := bitmapInfo{
bmiHeader: bitmapInfoHeader{
biSize: uint32(unsafe.Sizeof(bitmapInfoHeader{})),
biWidth: int32(w),
biHeight: -int32(h),
biPlanes: 1,
biBitCount: 32,
biCompression: BI_RGB,
},
}
var bits unsafe.Pointer
newBmp := createDIBSection(s.hdcMem, &bmi, DIB_RGB_COLORS, &bits)
if newBmp == 0 || bits == nil {
return 0, 0, nil, 0, syscall.EINVAL
}
if selectObject(s.hdcMem, newBmp) == 0 {
deleteObject(newBmp)
return 0, 0, nil, 0, syscall.EINVAL
}
if s.hbmp != 0 {
deleteObject(s.hbmp)
}
s.hbmp = newBmp
s.w = w
s.h = h
s.stride = w * 4
s.bits = bits
s.buf = unsafe.Slice((*byte)(bits), s.stride*h)
}
return s.hdcMem, s.hbmp, s.buf, s.stride, nil
}
func (s *capState) reset() {
s.mu.Lock()
defer s.mu.Unlock()
if s.hbmp != 0 {
deleteObject(s.hbmp)
s.hbmp = 0
}
if s.hdcMem != 0 {
deleteDC(s.hdcMem)
s.hdcMem = 0
}
s.buf = nil
s.stride = 0
s.w = 0
s.h = 0
s.bits = nil
stateCreatedAt = time.Time{}
}
func captureScale() float64 {
scaleOnce.Do(func() {
s := 1.0
if env := os.Getenv("OVERLORD_RD_SCALE"); env != "" {
if v, err := strconv.ParseFloat(env, 64); err == nil && v > 0.2 && v <= 1.5 {
s = v
}
}
cachedScale = s
})
return cachedScale
}
func effectiveScale(srcW, srcH int) float64 {
s := captureScale()
maxH := int(maxResHeight.Load())
if maxH == 0 {
maxH = 1080
} else if maxH < 0 {
storeLastScale(s)
return s
}
if srcH > maxH {
resCap := float64(maxH) / float64(srcH)
if resCap < s {
s = resCap
}
}
storeLastScale(s)
return s
}
func storeLastScale(s float64) {
lastCapScale.Store(math.Float64bits(s))
}
func SetMaxResolution(maxH int) {
maxResHeight.Store(int64(maxH))
log.Printf("capture: max resolution set to %d", maxH)
}
func BypassResolutionCap() (restore func()) {
old := maxResHeight.Swap(-1)
return func() { maxResHeight.Store(old) }
}
func EffectiveScaleForInput() float64 {
bits := lastCapScale.Load()
if bits == 0 {
return 1.0
}
return math.Float64frombits(bits)
}