[disk][windows] Support LabelWithContext & SerialNumberWithContext

Add support for LabelWithContext & SerialNumberWithContext, additionally returning this data back for the IOCountersWithContext function.
pull/1367/head
Matthew Palmer 3 years ago committed by Frank Spitulski
parent 400a453e67
commit b515acb162

@ -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"
)
@ -93,57 +96,29 @@ 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))
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 {
if i.FileSystemFlags&fileReadOnlyVolume != 0 {
opts = []string{"ro"}
}
if lpFileSystemFlags&fileFileCompression != 0 {
if i.FileSystemFlags&fileFileCompression != 0 {
opts = append(opts, "compress")
}
path := string(v) + ":"
d := PartitionStat{
Mountpoint: path,
Device: path,
Fstype: string(bytes.Replace(lpFileSystemNameBuffer, []byte("\x00"), []byte(""), -1)),
Fstype: i.FileSystemName,
Opts: opts,
}
ret = append(ret, d)
}
}
}
return ret, warnings.Reference()
}
@ -159,7 +134,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 {
@ -184,6 +160,12 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC
if err != nil {
return drivemap, err
}
i, err := getVolumeInformation(vStr)
if err != nil {
return nil, err
}
drivemap[path] = IOCountersStat{
ReadBytes: uint64(diskPerformance.BytesRead),
WriteBytes: uint64(diskPerformance.BytesWritten),
@ -192,6 +174,8 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC
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,
}
}
}
@ -199,9 +183,76 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC
}
func SerialNumberWithContext(ctx context.Context, name string) (string, error) {
return "", common.ErrNotImplementedError
i, err := getVolumeInformation(name)
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)
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) (*volumeInformation, error) {
path := name + ":"
typepath, _ := windows.UTF16PtrFromString(path)
typeret, _, _ := procGetDriveType.Call(uintptr(unsafe.Pointer(typepath)))
if typeret == 0 {
return nil, windows.GetLastError()
}
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
}
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
}

Loading…
Cancel
Save