Merge pull request #1166 from easyops-cn/kestrel/process

shirou · web-flow · commit 0ce33dbe6490 · 2021-11-13T20:51:25.000+09:00
feat(process): implement the 'OpenFilesWithContext' function of the windows system
diff --git 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_test.go b/process/process_test.go
@@ -657,6 +657,9 @@ func Test_CPUTimes(t *testing.T) {
 }
 
 func Test_OpenFiles(t *testing.T) {
+	fp, err := os.Open("process_test.go")
+	defer fp.Close()
+
 	pid := os.Getpid()
 	p, err := NewProcess(int32(pid))
 	skipIfNotImplementedErr(t, err)
diff --git 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"
@@ -667,7 +669,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) {
diff --git a/v3/internal/common/common_windows.go b/v3/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/v3/process/process_test.go b/v3/process/process_test.go
@@ -659,6 +659,9 @@ func Test_CPUTimes(t *testing.T) {
 }
 
 func Test_OpenFiles(t *testing.T) {
+	fp, err := os.Open("process_test.go")
+	defer fp.Close()
+
 	pid := os.Getpid()
 	p, err := NewProcess(int32(pid))
 	skipIfNotImplementedErr(t, err)
diff --git a/v3/process/process_windows.go b/v3/process/process_windows.go
@@ -9,8 +9,10 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"reflect"
 	"strings"
 	"syscall"
+	"unicode/utf16"
 	"unsafe"
 
 	"github.com/shirou/gopsutil/v3/cpu"
@@ -654,7 +656,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) {

Original file line number	Diff line number	Diff line change
`@@ -657,6 +657,9 @@ func Test_CPUTimes(t *testing.T) {`
`657`	`657`	`}`
`658`	`658`
`659`	`659`	`func Test_OpenFiles(t *testing.T) {`
	`660`	`+ fp, err := os.Open("process_test.go")`
	`661`	`+ defer fp.Close()`
	`662`	`+`
`660`	`663`	`pid := os.Getpid()`
`661`	`664`	`p, err := NewProcess(int32(pid))`
`662`	`665`	`skipIfNotImplementedErr(t, err)`
Original file line number	Diff line number	Diff line change
`@@ -659,6 +659,9 @@ func Test_CPUTimes(t *testing.T) {`
`659`	`659`	`}`
`660`	`660`
`661`	`661`	`func Test_OpenFiles(t *testing.T) {`
	`662`	`+ fp, err := os.Open("process_test.go")`
	`663`	`+ defer fp.Close()`
	`664`	`+`
`662`	`665`	`pid := os.Getpid()`
`663`	`666`	`p, err := NewProcess(int32(pid))`
`664`	`667`	`skipIfNotImplementedErr(t, err)`