Skip to content

Commit

Permalink
Merge pull request #1367 from FrankSpitulski/feat/implement-label-win…
Browse files Browse the repository at this point in the history 
…dows

[disk][windows] Support LabelWithContext & SerialNumberWithContext - rebased
  • Loading branch information
@shirou
shirou committed Nov 30, 2022
2 parents 9fc4d70 + d7c7912 commit af44f9d
Showing 1 changed file with 114 additions and 56 deletions.
170 changes: 114 additions & 56 deletions disk/disk_windows.go
View file
Expand Up @@ -6,12 +6,15 @@ package disk
import (
"bytes"
"context"
"errors"
"fmt"
"strconv"
"syscall"
"unsafe"

"github.com/shirou/gopsutil/v3/internal/common"
"golang.org/x/sys/windows"

"github.com/shirou/gopsutil/v3/internal/common"
"golang.org/x/sys/windows/registry"
)

Expand Down Expand Up @@ -95,55 +98,27 @@ func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, erro
}
for _, v := range lpBuffer {
if v >= 65 && v <= 90 {
path := string(v) + ":"
typepath, _ := windows.UTF16PtrFromString(path)
typeret, _, _ := procGetDriveType.Call(uintptr(unsafe.Pointer(typepath)))
if typeret == 0 {
err := windows.GetLastError()
warnings.Add(err)
i, err := getVolumeInformation(string(v), &warnings)
if err != nil && !errors.Is(err, errDeviceNotReady) && !errors.Is(err, errInvalidDriveType) {
continue
}
// 2: DRIVE_REMOVABLE 3: DRIVE_FIXED 4: DRIVE_REMOTE 5: DRIVE_CDROM

if typeret == 2 || typeret == 3 || typeret == 4 || typeret == 5 {
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 {
if typeret == 5 || typeret == 2 {
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")
}

d := PartitionStat{
Mountpoint: path,
Device: path,
Fstype: string(bytes.Replace(lpFileSystemNameBuffer, []byte("\x00"), []byte(""), -1)),
Opts: opts,
}
ret = append(ret, d)
opts := []string{"rw"}
if i.FileSystemFlags&fileReadOnlyVolume != 0 {
opts = []string{"ro"}
}
if i.FileSystemFlags&fileFileCompression != 0 {
opts = append(opts, "compress")
}

path := string(v) + ":"
d := PartitionStat{
Mountpoint: path,
Device: path,
Fstype: i.FileSystemName,
Opts: opts,
}
ret = append(ret, d)
}
}
return ret, warnings.Reference()
Expand All @@ -161,7 +136,8 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC
}
for _, v := range lpBuffer[:lpBufferLen] {
if 'A' <= v && v <= 'Z' {
path := string(rune(v)) + ":"
vStr := string(rune(v))
path := vStr + ":"
typepath, _ := windows.UTF16PtrFromString(path)
typeret := windows.GetDriveType(typepath)
if typeret == 0 {
Expand All @@ -186,24 +162,106 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC
if err != nil {
return drivemap, err
}

i, err := getVolumeInformation(vStr, nil)
if err != nil {
return nil, err
}

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,
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,
SerialNumber: strconv.FormatInt(i.SerialNumber, 10),
Label: i.Label,
}
}
}
return drivemap, nil
}

func SerialNumberWithContext(ctx context.Context, name string) (string, error) {
return "", common.ErrNotImplementedError
i, err := getVolumeInformation(name, nil)
if err != nil {
return "", err
}

return strconv.FormatInt(i.SerialNumber, 10), nil
}

func LabelWithContext(ctx context.Context, name string) (string, error) {
return "", common.ErrNotImplementedError
i, err := getVolumeInformation(name, nil)
if err != nil {
return "", err
}

return i.Label, nil
}

type volumeInformation struct {
Label string
SerialNumber int64
MaximumComponentLength int64
FileSystemFlags int64
FileSystemName string
}

var (
errDeviceNotReady = errors.New("device not ready")
errInvalidDriveType = errors.New("invalid drive type specified")
)

// getVolumeInformation returns all the information gathered from GetVolumeInformationW
func getVolumeInformation(name string, warnings *common.Warnings) (*volumeInformation, error) {
path := name + ":"
typepath, _ := windows.UTF16PtrFromString(path)
typeret, _, _ := procGetDriveType.Call(uintptr(unsafe.Pointer(typepath)))
if typeret == 0 {
err := windows.GetLastError()
if warnings != nil {
warnings.Add(err)
}
return nil, err
}

if typeret == windows.DRIVE_REMOVABLE || typeret == windows.DRIVE_FIXED || typeret == windows.DRIVE_REMOTE || typeret == windows.DRIVE_CDROM {
lpVolumeNameBuffer := make([]byte, 256)
lpVolumeSerialNumber := int64(0)
lpMaximumComponentLength := int64(0)
lpFileSystemFlags := int64(0)
lpFileSystemNameBuffer := make([]byte, 256)
volpath, _ := windows.UTF16PtrFromString(name + ":/")
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 {
if typeret == 5 || typeret == 2 {
return nil, errDeviceNotReady // device is not ready will happen if there is no disk in the drive
}
if warnings != nil {
warnings.Add(err)
}
return nil, err
}

return &volumeInformation{
Label: string(bytes.Replace(lpVolumeNameBuffer, []byte("\x00"), []byte(""), -1)),
SerialNumber: lpVolumeSerialNumber,
MaximumComponentLength: lpMaximumComponentLength,
FileSystemFlags: lpFileSystemFlags,
FileSystemName: string(bytes.Replace(lpFileSystemNameBuffer, []byte("\x00"), []byte(""), -1)),
}, nil
}

return nil, errInvalidDriveType
}

0 comments on commit af44f9d

Please sign in to comment.