diff --git a/internal/common/common_windows.go b/internal/common/common_windows.go index 2471bae..23417c2 100644 --- a/internal/common/common_windows.go +++ b/internal/common/common_windows.go @@ -6,6 +6,7 @@ import ( "context" "fmt" "path/filepath" + "reflect" "strings" "syscall" "unsafe" @@ -48,11 +49,18 @@ const ( PDH_INVALID_DATA = 0xc0000bc6 PDH_INVALID_HANDLE = 0xC0000bbc PDH_NO_DATA = 0x800007d5 + + STATUS_BUFFER_OVERFLOW = 0x80000005 + STATUS_BUFFER_TOO_SMALL = 0xC0000023 + STATUS_INFO_LENGTH_MISMATCH = 0xC0000004 ) const ( ProcessBasicInformation = 0 ProcessWow64Information = 26 + ProcessQueryInformation = windows.PROCESS_DUP_HANDLE | windows.PROCESS_QUERY_INFORMATION + + SystemExtendedHandleInformationClass = 64 ) var ( @@ -227,3 +235,66 @@ func ConvertDOSPath(p string) string { } return p } + +type NtStatus uint32 + +func (s NtStatus) Error() error { + if s == 0 { + return nil + } + return fmt.Errorf("NtStatus 0x%08x", uint32(s)) +} + +func (s NtStatus) IsError() bool { + return s>>30 == 3 +} + +type SystemExtendedHandleTableEntryInformation struct { + Object uintptr + UniqueProcessId uintptr + HandleValue uintptr + GrantedAccess uint32 + CreatorBackTraceIndex uint16 + ObjectTypeIndex uint16 + HandleAttributes uint32 + Reserved uint32 +} + +type SystemExtendedHandleInformation struct { + NumberOfHandles uintptr + Reserved uintptr + Handles [1]SystemExtendedHandleTableEntryInformation +} + +// CallWithExpandingBuffer https://github.com/hillu/go-ntdll +func CallWithExpandingBuffer(fn func() NtStatus, buf *[]byte, resultLength *uint32) NtStatus { + for { + if st := fn(); st == STATUS_BUFFER_OVERFLOW || st == STATUS_BUFFER_TOO_SMALL || st == STATUS_INFO_LENGTH_MISMATCH { + if int(*resultLength) <= cap(*buf) { + (*reflect.SliceHeader)(unsafe.Pointer(buf)).Len = int(*resultLength) + } else { + *buf = make([]byte, int(*resultLength)) + } + continue + } else { + if !st.IsError() { + *buf = (*buf)[:int(*resultLength)] + } + return st + } + } +} + +func NtQuerySystemInformation( + SystemInformationClass uint32, + SystemInformation *byte, + SystemInformationLength uint32, + ReturnLength *uint32, +) NtStatus { + r0, _, _ := ProcNtQuerySystemInformation.Call( + uintptr(SystemInformationClass), + uintptr(unsafe.Pointer(SystemInformation)), + uintptr(SystemInformationLength), + uintptr(unsafe.Pointer(ReturnLength))) + return NtStatus(r0) +} diff --git a/process/process_windows.go b/process/process_windows.go index c88f949..99d90b8 100644 --- a/process/process_windows.go +++ b/process/process_windows.go @@ -9,8 +9,10 @@ import ( "fmt" "io" "os" + "reflect" "strings" "syscall" + "unicode/utf16" "unsafe" "github.com/shirou/gopsutil/cpu" @@ -603,7 +605,79 @@ func (p *Process) ChildrenWithContext(ctx context.Context) ([]*Process, error) { } func (p *Process) OpenFilesWithContext(ctx context.Context) ([]OpenFilesStat, error) { - return nil, common.ErrNotImplementedError + files := make([]OpenFilesStat, 0) + fileExists := make(map[string]bool) + + process, err := windows.OpenProcess(common.ProcessQueryInformation, false, uint32(p.Pid)) + if err != nil { + return nil, err + } + + buffer := make([]byte, 1024) + var size uint32 + + st := common.CallWithExpandingBuffer( + func() common.NtStatus { + return common.NtQuerySystemInformation( + common.SystemExtendedHandleInformationClass, + &buffer[0], + uint32(len(buffer)), + &size, + ) + }, + &buffer, + &size, + ) + if st.IsError() { + return nil, st.Error() + } + + handlesList := (*common.SystemExtendedHandleInformation)(unsafe.Pointer(&buffer[0])) + handles := make([]common.SystemExtendedHandleTableEntryInformation, int(handlesList.NumberOfHandles)) + hdr := (*reflect.SliceHeader)(unsafe.Pointer(&handles)) + hdr.Data = uintptr(unsafe.Pointer(&handlesList.Handles[0])) + + currentProcess, err := windows.GetCurrentProcess() + if err != nil { + return nil, err + } + + for _, handle := range handles { + var file uintptr + if int32(handle.UniqueProcessId) != p.Pid { + continue + } + if windows.DuplicateHandle(process, windows.Handle(handle.HandleValue), currentProcess, (*windows.Handle)(&file), + 0, true, windows.DUPLICATE_SAME_ACCESS) != nil { + continue + } + fileType, _ := windows.GetFileType(windows.Handle(file)) + if fileType != windows.FILE_TYPE_DISK { + continue + } + + var buf [syscall.MAX_LONG_PATH]uint16 + n, err := windows.GetFinalPathNameByHandle(windows.Handle(file), &buf[0], syscall.MAX_LONG_PATH, 0) + if err != nil { + continue + } + + fileName := string(utf16.Decode(buf[:n])) + fileInfo, _ := os.Stat(fileName) + if fileInfo.IsDir() { + continue + } + + if _, exists := fileExists[fileName]; !exists { + files = append(files, OpenFilesStat{ + Path: fileName, + Fd: uint64(file), + }) + fileExists[fileName] = true + } + } + + return files, nil } func (p *Process) ConnectionsWithContext(ctx context.Context) ([]net.ConnectionStat, error) {