[disk][windows] Support LabelWithContext & SerialNumberWithContext - rebased

disk/disk_windows.go

@@ -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,55 +96,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()
@@ -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,24 +160,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
 }