From 311c16dbaf567f322bdec539effcc98515e444d5 Mon Sep 17 00:00:00 2001 From: Dani Louca Date: Tue, 3 Oct 2023 17:10:11 -0400 Subject: [PATCH] Fetch stats from folder mounted volumes Signed-off-by: Dani Louca --- disk/disk_windows.go | 283 ++++++++++++++++++++++++++++++++++----------------- go.mod | 1 + go.sum | 2 + 3 files changed, 194 insertions(+), 92 deletions(-) diff --git a/disk/disk_windows.go b/disk/disk_windows.go index 0abc844..3d22a3d 100644 --- a/disk/disk_windows.go +++ b/disk/disk_windows.go @@ -7,7 +7,10 @@ import ( "bytes" "context" "errors" + "fmt" + "reflect" "syscall" + "unicode/utf16" "unsafe" "golang.org/x/sys/windows" @@ -16,10 +19,17 @@ import ( "github.com/shirou/gopsutil/v4/internal/common" ) +const volumeNameBufferLength = uint32(windows.MAX_PATH + 1) +const volumePathBufferLength = volumeNameBufferLength + var ( - procGetDiskFreeSpaceExW = common.Modkernel32.NewProc("GetDiskFreeSpaceExW") - procGetLogicalDriveStringsW = common.Modkernel32.NewProc("GetLogicalDriveStringsW") - procGetVolumeInformation = common.Modkernel32.NewProc("GetVolumeInformationW") + procGetDiskFreeSpaceExW = common.Modkernel32.NewProc("GetDiskFreeSpaceExW") + procGetLogicalDriveStringsW = common.Modkernel32.NewProc("GetLogicalDriveStringsW") + procGetVolumeInformation = common.Modkernel32.NewProc("GetVolumeInformationW") + procFindFirstVolumeW = common.Modkernel32.NewProc("FindFirstVolumeW") + procFindNextVolumeW = common.Modkernel32.NewProc("FindNextVolumeW") + procFindVolumeClose = common.Modkernel32.NewProc("FindVolumeClose") + procGetVolumePathNamesForVolumeNameW = common.Modkernel32.NewProc("GetVolumePathNamesForVolumeNameW") ) var ( @@ -81,84 +91,85 @@ func UsageWithContext(_ context.Context, path string) (*UsageStat, error) { } // PartitionsWithContext returns disk partitions. -// Since GetVolumeInformation doesn't have a timeout, this method uses context to set deadline by users. +// It uses procGetLogicalDriveStringsW to get drives with drive letters and procFindFirstVolumeW to get volumes without drive letters. +// Since the api calls don't have a timeout, this method uses context to set deadline by users. func PartitionsWithContext(ctx context.Context, _ bool) ([]PartitionStat, error) { - warnings := Warnings{ - Verbose: true, - } - - var errLogicalDrives error + warnings := Warnings{Verbose: true} + var errInitialCall error retChan := make(chan PartitionStat) quitChan := make(chan struct{}) defer close(quitChan) + processedPaths := make(map[string]struct{}) getPartitions := func() { defer close(retChan) + // Get drives with drive letters (including remote drives, ex: SMB shares) lpBuffer := make([]byte, 254) - - diskret, _, err := procGetLogicalDriveStringsW.Call( + if diskret, _, err := procGetLogicalDriveStringsW.Call( uintptr(len(lpBuffer)), - uintptr(unsafe.Pointer(&lpBuffer[0]))) - if diskret == 0 { - errLogicalDrives = err + uintptr(unsafe.Pointer(&lpBuffer[0]))); diskret == 0 { + errInitialCall = err return } for _, v := range lpBuffer { if v >= 65 && v <= 90 { path := string(v) + ":" - typepath, _ := windows.UTF16PtrFromString(path) - typeret := windows.GetDriveType(typepath) - switch typeret { - case windows.DRIVE_UNKNOWN: - continue - case windows.DRIVE_REMOVABLE, - windows.DRIVE_FIXED, - windows.DRIVE_REMOTE, - windows.DRIVE_CDROM: - lpVolumeNameBuffer := make([]byte, 256) - lpVolumeSerialNumber := int64(0) - lpMaximumComponentLength := int64(0) - lpFileSystemFlags := int64(0) - lpFileSystemNameBuffer := make([]byte, 256) - volpath, _ := windows.UTF16PtrFromString(string(v) + ":/") - driveret, _, err := procGetVolumeInformation.Call( - uintptr(unsafe.Pointer(volpath)), - uintptr(unsafe.Pointer(&lpVolumeNameBuffer[0])), - uintptr(len(lpVolumeNameBuffer)), - uintptr(unsafe.Pointer(&lpVolumeSerialNumber)), - uintptr(unsafe.Pointer(&lpMaximumComponentLength)), - uintptr(unsafe.Pointer(&lpFileSystemFlags)), - uintptr(unsafe.Pointer(&lpFileSystemNameBuffer[0])), - uintptr(len(lpFileSystemNameBuffer))) - if driveret == 0 { - switch typeret { - case windows.DRIVE_REMOVABLE, windows.DRIVE_REMOTE, windows.DRIVE_CDROM: - continue // device is not ready will happen if there is no disk in the drive - } - warnings.Add(err) - continue - } - opts := []string{"rw"} - if lpFileSystemFlags&fileReadOnlyVolume != 0 { - opts = []string{"ro"} - } - if lpFileSystemFlags&fileFileCompression != 0 { - opts = append(opts, "compress") + if partitionStat, warning := buildPartitionStat(path); warning == nil { + processedPaths[partitionStat.Mountpoint+"\\"] = struct{}{} + select { + case retChan <- partitionStat: + case <-quitChan: + return } + } else { + warnings.Add(warning) + } + } + } + // Get volumes without drive letters (ex: mounted folders with no drive letter) + volNameBuf := make([]uint16, volumeNameBufferLength) + nextVolHandle, _, err := procFindFirstVolumeW.Call( + uintptr(unsafe.Pointer(&volNameBuf[0])), + uintptr(volumeNameBufferLength)) + if windows.Handle(nextVolHandle) == windows.InvalidHandle { + errInitialCall = fmt.Errorf("failed to get first-volume: %w", err) + return + } + defer procFindVolumeClose.Call(nextVolHandle) + for { + mounts, err := getVolumePaths(volNameBuf) + if err != nil { + warnings.Add(fmt.Errorf("failed to find paths for volume %s", windows.UTF16ToString(volNameBuf))) + continue + } + + for _, mount := range mounts { + if _, ok := processedPaths[mount]; ok { + continue + } + if partitionStat, warning := buildPartitionStat(mount); warning == nil { select { - case retChan <- PartitionStat{ - Mountpoint: path, - Device: path, - Fstype: string(bytes.ReplaceAll(lpFileSystemNameBuffer, []byte("\x00"), []byte(""))), - Opts: opts, - }: + case retChan <- partitionStat: case <-quitChan: return } + } else { + warnings.Add(warning) } } + + volNameBuf = make([]uint16, volumeNameBufferLength) + if volRet, _, err := procFindNextVolumeW.Call( + nextVolHandle, + uintptr(unsafe.Pointer(&volNameBuf[0])), + uintptr(volumeNameBufferLength)); err != nil && volRet == 0 { + if errno, ok := err.(syscall.Errno); ok && errno == windows.ERROR_NO_MORE_FILES { + break + } + warnings.Add(fmt.Errorf("failed to find next volume: %w", err)) + } } } @@ -169,18 +180,72 @@ func PartitionsWithContext(ctx context.Context, _ bool) ([]PartitionStat, error) select { case p, ok := <-retChan: if !ok { - if errLogicalDrives != nil { - return ret, errLogicalDrives + if errInitialCall != nil { + return ret, errInitialCall } return ret, warnings.Reference() } - ret = append(ret, p) + if !reflect.DeepEqual(p, PartitionStat{}) { + ret = append(ret, p) + } case <-ctx.Done(): return ret, ctx.Err() } } } +func buildPartitionStat(path string) (PartitionStat, error) { + typePath, _ := windows.UTF16PtrFromString(path) + driveType := windows.GetDriveType(typePath) + + if driveType == windows.DRIVE_UNKNOWN { + return PartitionStat{}, windows.GetLastError() + } + + if driveType == windows.DRIVE_REMOVABLE || driveType == windows.DRIVE_FIXED || + driveType == windows.DRIVE_REMOTE || driveType == windows.DRIVE_CDROM { + volPath, _ := windows.UTF16PtrFromString(path + "/") + volumeName := make([]byte, 256) + fsName := make([]byte, 256) + var serialNumber, maxComponentLength, fsFlags int64 + + ret, _, err := procGetVolumeInformation.Call( + uintptr(unsafe.Pointer(volPath)), + uintptr(unsafe.Pointer(&volumeName[0])), + uintptr(len(volumeName)), + uintptr(unsafe.Pointer(&serialNumber)), + uintptr(unsafe.Pointer(&maxComponentLength)), + uintptr(unsafe.Pointer(&fsFlags)), + uintptr(unsafe.Pointer(&fsName[0])), + uintptr(len(fsName)), + ) + + if ret == 0 { + if driveType == windows.DRIVE_REMOVABLE || driveType == windows.DRIVE_REMOTE || driveType == windows.DRIVE_CDROM { + return PartitionStat{}, nil // Device not ready + } + return PartitionStat{}, err + } + + opts := []string{"rw"} + if fsFlags&fileReadOnlyVolume != 0 { + opts = []string{"ro"} + } + if fsFlags&fileFileCompression != 0 { + opts = append(opts, "compress") + } + + return PartitionStat{ + Mountpoint: path, + Device: path, + Fstype: string(bytes.ReplaceAll(fsName, []byte("\x00"), []byte(""))), + Opts: opts, + }, nil + } + + return PartitionStat{}, nil +} + func IOCountersWithContext(_ context.Context, names ...string) (map[string]IOCountersStat, error) { // https://github.com/giampaolo/psutil/blob/544e9daa4f66a9f80d7bf6c7886d693ee42f0a13/psutil/arch/windows/disk.c#L83 drivemap := make(map[string]IOCountersStat, 0) @@ -192,41 +257,40 @@ func IOCountersWithContext(_ context.Context, names ...string) (map[string]IOCou return drivemap, err } for _, v := range lpBuffer[:lpBufferLen] { - if 'A' > v || v > 'Z' { - continue - } - path := string(rune(v)) + ":" - typepath, _ := windows.UTF16PtrFromString(path) - typeret := windows.GetDriveType(typepath) - if typeret != windows.DRIVE_FIXED { - continue - } - szDevice := `\\.\` + path - const IOCTL_DISK_PERFORMANCE = 0x70020 - h, err := windows.CreateFile(syscall.StringToUTF16Ptr(szDevice), 0, windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE, nil, windows.OPEN_EXISTING, 0, 0) - if err != nil { - if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) { + if 'A' <= v && v <= 'Z' { + path := string(rune(v)) + ":" + typepath, _ := windows.UTF16PtrFromString(path) + typeret := windows.GetDriveType(typepath) + if typeret != windows.DRIVE_FIXED { continue } - return drivemap, err - } - defer windows.CloseHandle(h) + szDevice := `\\.\` + path + const IOCTL_DISK_PERFORMANCE = 0x70020 + h, err := windows.CreateFile(syscall.StringToUTF16Ptr(szDevice), 0, windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE, nil, windows.OPEN_EXISTING, 0, 0) + if err != nil { + if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) { + continue + } + return drivemap, err + } + defer windows.CloseHandle(h) - var diskPerformanceSize uint32 - err = windows.DeviceIoControl(h, IOCTL_DISK_PERFORMANCE, nil, 0, (*byte)(unsafe.Pointer(&diskPerformance)), uint32(unsafe.Sizeof(diskPerformance)), &diskPerformanceSize, nil) - if err != nil { - return drivemap, err - } + var diskPerformanceSize uint32 + err = windows.DeviceIoControl(h, IOCTL_DISK_PERFORMANCE, nil, 0, (*byte)(unsafe.Pointer(&diskPerformance)), uint32(unsafe.Sizeof(diskPerformance)), &diskPerformanceSize, nil) + if err != nil { + return drivemap, err + } - if len(names) == 0 || common.StringsHas(names, path) { - drivemap[path] = IOCountersStat{ - ReadBytes: uint64(diskPerformance.BytesRead), - WriteBytes: uint64(diskPerformance.BytesWritten), - ReadCount: uint64(diskPerformance.ReadCount), - WriteCount: uint64(diskPerformance.WriteCount), - ReadTime: uint64(diskPerformance.ReadTime / 10000 / 1000), // convert to ms: https://github.com/giampaolo/psutil/issues/1012 - WriteTime: uint64(diskPerformance.WriteTime / 10000 / 1000), - Name: path, + if len(names) == 0 || common.StringsHas(names, path) { + drivemap[path] = IOCountersStat{ + ReadBytes: uint64(diskPerformance.BytesRead), + WriteBytes: uint64(diskPerformance.BytesWritten), + ReadCount: uint64(diskPerformance.ReadCount), + WriteCount: uint64(diskPerformance.WriteCount), + ReadTime: uint64(diskPerformance.ReadTime / 10000 / 1000), // convert to ms: https://github.com/giampaolo/psutil/issues/1012 + WriteTime: uint64(diskPerformance.WriteTime / 10000 / 1000), + Name: path, + } } } } @@ -240,3 +304,38 @@ func SerialNumberWithContext(_ context.Context, _ string) (string, error) { func LabelWithContext(_ context.Context, _ string) (string, error) { return "", common.ErrNotImplementedError } + +// getVolumePaths returns the path for the given volume name. +func getVolumePaths(volNameBuf []uint16) ([]string, error) { + volPathsBuf := make([]uint16, volumePathBufferLength) + returnLen := uint32(0) + if result, _, err := procGetVolumePathNamesForVolumeNameW.Call( + uintptr(unsafe.Pointer(&volNameBuf[0])), + uintptr(unsafe.Pointer(&volPathsBuf[0])), + uintptr(volumePathBufferLength), + uintptr(unsafe.Pointer(&returnLen))); err != nil && result == 0 { + return nil, err + } + return split0(volPathsBuf, int(returnLen)), nil +} + +// split0 iterates through s16 upto `end` and slices `s16` into sub-slices separated by the null character (uint16(0)). +// split0 converts the sub-slices between the null characters into strings then returns them in a slice. +func split0(s16 []uint16, end int) []string { + if end > len(s16) { + end = len(s16) + } + + from, ss := 0, make([]string, 0) + + for to := 0; to < end; to++ { + if s16[to] == 0 { + if from < to && s16[from] != 0 { + ss = append(ss, string(utf16.Decode(s16[from:to]))) + } + from = to + 1 + } + } + + return ss +} diff --git a/go.mod b/go.mod index 8662ed2..9f6989a 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/ebitengine/purego v0.8.4 github.com/google/go-cmp v0.7.0 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 + github.com/pkg/errors v0.9.1 github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c github.com/stretchr/testify v1.10.0 github.com/tklauser/go-sysconf v0.3.12 diff --git a/go.sum b/go.sum index b14b648..a38601e 100644 --- a/go.sum +++ b/go.sum @@ -9,6 +9,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=