fix: #292 DLL agent

this commit supports creating a DLL with minimal modification to the existing source code
jm33-m0 · Jan 30, 2024 · dddd442 · dddd442
1 parent 4033979
commit dddd442
Show file tree

Hide file tree

Showing 2 changed files with 220 additions and 137 deletions.
diff --git a/core/lib/util/dll_windows.go b/core/lib/util/dll_windows.go
@@ -0,0 +1,80 @@
+//go:build windows
+// +build windows
+
+package util
+
+import (
+	"fmt"
+	"log"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+func IsRunningInDLL() bool {
+	modHandle, _, _ := procGetModuleHandle.Call(0)
+	return modHandle != 0
+}
+
+func ReadDLL(moduleInfo *windows.ModuleInfo, fileName string) (dll_data []byte, err error) {
+	// Allocate a buffer to hold the DLL content
+	dllContent := make([]byte, moduleInfo.SizeOfImage)
+
+	// Read the content of the DLL from memory without specifying a process handle
+	var bytes_read uintptr
+	err = windows.ReadProcessMemory(windows.CurrentProcess(), moduleInfo.BaseOfDll, &dllContent[0],
+		uintptr(moduleInfo.SizeOfImage), &bytes_read)
+	if err != nil {
+		err = fmt.Errorf("failed to get module info of %s: %v", fileName, err)
+		return
+	}
+
+	// Print or process the DLL content as needed
+	dll_data = dllContent
+	return
+}
+
+// Enum all DLLs and get their handles
+func GetAllDLLs() (modules map[string]*windows.ModuleInfo, err error) {
+	modules = make(map[string]*windows.ModuleInfo, 0)
+
+	// Open a handle to the current process
+	processHandle := windows.CurrentProcess()
+
+	// Enumerate the modules (DLLs) loaded in the current process
+	var moduleHandles = make([]windows.Handle, 1024)
+	var neededBytes uint32
+	err = windows.EnumProcessModules(processHandle, &moduleHandles[0], 1024, &neededBytes)
+	if err != nil {
+		err = fmt.Errorf("enum modules: %v", err)
+		return
+	}
+
+	// Calculate the number of modules
+	numModules := int(neededBytes / uint32(unsafe.Sizeof(moduleHandles[0])))
+
+	// Print the file names of the loaded DLLs
+	for i := 0; i < numModules; i++ {
+		// Get the file name of the DLL
+		var fname16 = make([]uint16, windows.MAX_PATH)
+		_, err = windows.GetModuleFileName(moduleHandles[i], &fname16[0], windows.MAX_PATH)
+		if err != nil {
+			log.Printf("get module file name: %v", err)
+			continue
+		}
+		// Convert the UTF-16 encoded file name to a Go string
+		fileName := windows.UTF16ToString(fname16)
+
+		// get module info
+		modinfo := new(windows.ModuleInfo)
+		cb := uint32(unsafe.Sizeof(*modinfo))
+		err = windows.GetModuleInformation(processHandle, moduleHandles[i], modinfo, cb)
+		if err != nil {
+			log.Printf("get modinfo of %s: %v", fileName, err)
+			continue
+		}
+		modules[fileName] = modinfo
+	}
+
+	return
+}
diff --git a/core/lib/util/mem_windows.go b/core/lib/util/mem_windows.go
@@ -1,137 +1,140 @@
-package util
-
-import (
-	"fmt"
-	"log"
-	"os"
-	"syscall"
-	"unsafe"
-)
-
-var (
-	kernel32 = syscall.NewLazyDLL("kernel32.dll")
-	psapi    = syscall.NewLazyDLL("Psapi.dll")
-
-	procOpenProcess        = kernel32.NewProc("OpenProcess")
-	procReadProcessMemory  = kernel32.NewProc("ReadProcessMemory")
-	procWriteProcessMemory = kernel32.NewProc("WriteProcessMemory")
-	procVirtualQuery       = kernel32.NewProc("VirtualQuery")
-	procEnumProcessModules = psapi.NewProc("EnumProcessModulesEx")
-)
-
-const PROCESS_ALL_ACCESS = 0x1F0FFF
-
-func OpenProcess(pid int) uintptr {
-	handle, _, _ := procOpenProcess.Call(uintptr(PROCESS_ALL_ACCESS), uintptr(1), uintptr(pid))
-	return handle
-}
-
-func read_mem(hProcess uintptr, address, size uintptr) []byte {
-	var data = make([]byte, size)
-	var length uint32
-
-	procReadProcessMemory.Call(hProcess, address,
-		uintptr(unsafe.Pointer(&data[0])),
-		size, uintptr(unsafe.Pointer(&length)))
-
-	return data
-}
-
-const (
-	MEM_COMMIT  = 0x1000
-	MEM_RESERVE = 0x2000
-	MEM_FREE    = 0x10000
-)
-
-type MEMORY_BASIC_INFORMATION struct {
-	BaseAddress       uintptr
-	AllocationBase    uintptr
-	AllocationProtect uint32
-	RegionSize        uintptr
-	State             uint32
-	Protect           uint32
-	Type              uint32
-}
-
-func read_self_mem(hProcess uintptr) (mem_data [][]byte, bytes_read int, err error) {
-	// Start with an initial address of 0
-	address := uintptr(0)
-
-	// Loop through the memory regions and print information
-	for {
-		var mbi MEMORY_BASIC_INFORMATION
-		ret, _, _ := procVirtualQuery.Call(address, uintptr(unsafe.Pointer(&mbi)), unsafe.Sizeof(mbi))
-
-		// Check for the end of the memory regions
-		if ret == 0 {
-			break
-		}
-
-		// Move to the next memory region
-		address += mbi.RegionSize
-
-		// Print information about the memory region
-		// log.Printf("BaseAddress: 0x%x, RegionSize: 0x%x, State: %d, Protect: %d, Type: %d\n",
-		// 	mbi.BaseAddress, mbi.RegionSize, mbi.State, mbi.Protect, mbi.Type)
-
-		// if memory is not committed or is read-only, skip it
-		readable := mbi.State == MEM_COMMIT && mbi.Protect&syscall.PAGE_READONLY != 0
-		if !readable {
-			continue
-		}
-
-		// read data from this region
-		data_read := read_mem(hProcess, mbi.BaseAddress, mbi.RegionSize)
-		bytes_read += len(data_read)
-		mem_data = append(mem_data, data_read)
-	}
-
-	return
-}
-
-func write_mem(hProcess uintptr, lpBaseAddress, lpBuffer, nSize uintptr) (int, bool) {
-	var nBytesWritten int
-	ret, _, _ := procWriteProcessMemory.Call(
-		uintptr(hProcess),
-		lpBaseAddress,
-		lpBuffer,
-		nSize,
-		uintptr(unsafe.Pointer(&nBytesWritten)),
-	)
-
-	return nBytesWritten, ret != 0
-}
-
-func getBaseAddress(handle uintptr) uintptr {
-	modules := [1024]uint64{}
-	var needed uintptr
-	procEnumProcessModules.Call(
-		handle,
-		uintptr(unsafe.Pointer(&modules)),
-		uintptr(1024),
-		uintptr(unsafe.Pointer(&needed)),
-		uintptr(0x03),
-	)
-	for i := uintptr(0); i < needed/unsafe.Sizeof(modules[0]); i++ {
-		if i == 0 {
-			return uintptr(modules[i])
-		}
-	}
-	return 0
-}
-
-func crossPlatformDumpSelfMem() (mem_data [][]byte, err error) {
-	// open current process
-	pid := os.Getpid()
-	processHandle := OpenProcess(pid)
-
-	// read memory regions
-	bytes_read := 0
-	mem_data, bytes_read, err = read_self_mem(processHandle)
-	if err != nil {
-		return nil, fmt.Errorf("crossPlatformDumpSelfMem read_self_mem: %v", err)
-	}
-	log.Printf("crossPlatformDumpSelfMem: READ %d bytes from %d memory regions",
-		bytes_read, len(mem_data))
-	return
-}
+//go:build windows
+// +build windows
+
+package util
+
+import (
+	"log"
+	"syscall"
+	"unsafe"
+)
+
+var (
+	kernel32 = syscall.NewLazyDLL("kernel32.dll")
+	psapi    = syscall.NewLazyDLL("Psapi.dll")
+
+	procOpenProcess        = kernel32.NewProc("OpenProcess")
+	procReadProcessMemory  = kernel32.NewProc("ReadProcessMemory")
+	procWriteProcessMemory = kernel32.NewProc("WriteProcessMemory")
+	procVirtualQuery       = kernel32.NewProc("VirtualQuery")
+	procGetModuleFileName  = kernel32.NewProc("GetModuleFileNameW")
+	procGetModuleHandle    = kernel32.NewProc("GetModuleHandleW")
+	procEnumProcessModules = psapi.NewProc("EnumProcessModulesEx")
+)
+
+const PROCESS_ALL_ACCESS = 0x1F0FFF
+
+func OpenProcess(pid int) uintptr {
+	handle, _, _ := procOpenProcess.Call(uintptr(PROCESS_ALL_ACCESS), uintptr(1), uintptr(pid))
+	return handle
+}
+
+func read_mem(hProcess uintptr, address, size uintptr) []byte {
+	var data = make([]byte, size)
+	var length uint32
+
+	procReadProcessMemory.Call(hProcess, address,
+		uintptr(unsafe.Pointer(&data[0])),
+		size, uintptr(unsafe.Pointer(&length)))
+
+	return data
+}
+
+const (
+	MEM_COMMIT  = 0x1000
+	MEM_RESERVE = 0x2000
+	MEM_FREE    = 0x10000
+)
+
+type MEMORY_BASIC_INFORMATION struct {
+	BaseAddress       uintptr
+	AllocationBase    uintptr
+	AllocationProtect uint32
+	RegionSize        uintptr
+	State             uint32
+	Protect           uint32
+	Type              uint32
+}
+
+func read_self_mem(hProcess uintptr) (mem_data [][]byte, bytes_read int, err error) {
+	// Start with an initial address of 0
+	address := uintptr(0)
+
+	// Loop through the memory regions and print information
+	for {
+		var mbi MEMORY_BASIC_INFORMATION
+		ret, _, _ := procVirtualQuery.Call(address, uintptr(unsafe.Pointer(&mbi)), unsafe.Sizeof(mbi))
+
+		// Check for the end of the memory regions
+		if ret == 0 {
+			break
+		}
+
+		// Move to the next memory region
+		address += mbi.RegionSize
+
+		// Print information about the memory region
+		// log.Printf("BaseAddress: 0x%x, RegionSize: 0x%x, State: %d, Protect: %d, Type: %d\n",
+		// 	mbi.BaseAddress, mbi.RegionSize, mbi.State, mbi.Protect, mbi.Type)
+
+		// if memory is not committed or is read-only, skip it
+		readable := mbi.State == MEM_COMMIT && mbi.Protect&syscall.PAGE_READONLY != 0
+		if !readable {
+			continue
+		}
+
+		// read data from this region
+		data_read := read_mem(hProcess, mbi.BaseAddress, mbi.RegionSize)
+		bytes_read += len(data_read)
+		mem_data = append(mem_data, data_read)
+	}
+
+	return
+}
+
+func write_mem(hProcess uintptr, lpBaseAddress, lpBuffer, nSize uintptr) (int, bool) {
+	var nBytesWritten int
+	ret, _, _ := procWriteProcessMemory.Call(
+		uintptr(hProcess),
+		lpBaseAddress,
+		lpBuffer,
+		nSize,
+		uintptr(unsafe.Pointer(&nBytesWritten)),
+	)
+
+	return nBytesWritten, ret != 0
+}
+
+func getBaseAddress(handle uintptr) uintptr {
+	modules := [1024]uint64{}
+	var needed uintptr
+	procEnumProcessModules.Call(
+		handle,
+		uintptr(unsafe.Pointer(&modules)),
+		uintptr(1024),
+		uintptr(unsafe.Pointer(&needed)),
+		uintptr(0x03),
+	)
+	for i := uintptr(0); i < needed/unsafe.Sizeof(modules[0]); i++ {
+		if i == 0 {
+			return uintptr(modules[i])
+		}
+	}
+	return 0
+}
+
+func crossPlatformDumpSelfMem() (mem_data [][]byte, err error) {
+	dlls, err := GetAllDLLs()
+	if err != nil {
+		return
+	}
+	for fileName, dll := range dlls {
+		dll_data, err := ReadDLL(dll, fileName)
+		if err != nil {
+			log.Printf("reading DLL %s: %v", fileName, err)
+			continue
+		}
+		mem_data = append(mem_data, dll_data)
+	}
+	return mem_data, err
+}