Use WMI to implement Volume API to reduce PowerShell overhead

go.mod

@@ -9,10 +9,12 @@ toolchain go1.22.3
 
 require (
 	github.com/Microsoft/go-winio v0.6.2
+	github.com/go-ole/go-ole v1.3.0
 	github.com/google/go-cmp v0.6.0
 	github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1
 	github.com/iancoleman/strcase v0.3.0
 	github.com/kubernetes-csi/csi-proxy/client v1.1.3
+	github.com/microsoft/wmi v0.25.1
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.20.5
 	github.com/sergi/go-diff v1.3.1

go.sum

@@ -25,6 +25,8 @@ github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
 github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
 github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE=
+github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -69,6 +71,8 @@ github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0
 github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/mauriciopoppe/gengo v0.0.0-20210525224835-9c78f58f3486 h1:+l047vEi0SyAzdVToIaAcfoY5DwwGW+OyqTdH/P3TTg=
 github.com/mauriciopoppe/gengo v0.0.0-20210525224835-9c78f58f3486/go.mod h1:xXv3T4UXTLta31wMhVezwVkc26OLei4hMbLeBJbPmxc=
+github.com/microsoft/wmi v0.25.1 h1:sQv9hCEHtW5K6yEVL78T6XGRMGxk4aTpcJwCiB5rLN0=
+github.com/microsoft/wmi v0.25.1/go.mod h1:1zbdSF0A+5OwTUII5p3hN7/K6KF2m3o27pSG6Y51VU8=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
@@ -134,6 +138,7 @@ golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

pkg/cim/disk.go

@@ -0,0 +1,33 @@
+package cim
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/microsoft/wmi/pkg/base/query"
+	"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
+)
+
+// QueryDiskByNumber retrieves disk information for a specific disk identified by its number.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_Disk
+//	  WHERE DiskNumber = '<diskNumber>'
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk
+// for the WMI class definition.
+func QueryDiskByNumber(diskNumber uint32, selectorList []string) (*storage.MSFT_Disk, error) {
+	diskQuery := query.NewWmiQueryWithSelectList("MSFT_Disk", selectorList, "Number", strconv.Itoa(int(diskNumber)))
+	instances, err := QueryInstances(WMINamespaceStorage, diskQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	disk, err := storage.NewMSFT_DiskEx1(instances[0])
+	if err != nil {
+		return nil, fmt.Errorf("failed to query disk %d. error: %v", diskNumber, err)
+	}
+
+	return disk, nil
+}

pkg/cim/volume.go

@@ -0,0 +1,297 @@
+package cim
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/microsoft/wmi/pkg/base/query"
+	"github.com/microsoft/wmi/pkg/errors"
+	cim "github.com/microsoft/wmi/pkg/wmiinstance"
+	"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
+)
+
+// QueryVolumeByUniqueID retrieves a specific volume by its unique identifier,
+// returning the first volume that matches the given volume ID.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_Volume
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-volume
+// for the WMI class definition.
+func QueryVolumeByUniqueID(volumeID string, selectorList []string) (*storage.MSFT_Volume, error) {
+	var selectors []string
+	selectors = append(selectors, selectorList...)
+	selectors = append(selectors, "UniqueId")
+	volumeQuery := query.NewWmiQueryWithSelectList("MSFT_Volume", selectors)
+	instances, err := QueryInstances(WMINamespaceStorage, volumeQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, instance := range instances {
+		volume, err := storage.NewMSFT_VolumeEx1(instance)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query volume (%s). error: %w", volumeID, err)
+		}
+
+		uniqueID, err := volume.GetPropertyUniqueId()
+		if err != nil {
+			return nil, fmt.Errorf("failed to query volume unique ID (%s). error: %w", volumeID, err)
+		}
+
+		if uniqueID == volumeID {
+			return volume, nil
+		}
+	}
+
+	return nil, errors.NotFound
+}
+
+// ListVolumes retrieves all available volumes on the system.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_Volume
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-volume
+// for the WMI class definition.
+func ListVolumes(selectorList []string) ([]*storage.MSFT_Volume, error) {
+	diskQuery := query.NewWmiQueryWithSelectList("MSFT_Volume", selectorList)
+	instances, err := QueryInstances(WMINamespaceStorage, diskQuery)
+	if IgnoreNotFound(err) != nil {
+		return nil, err
+	}
+
+	var volumes []*storage.MSFT_Volume
+	for _, instance := range instances {
+		volume, err := storage.NewMSFT_VolumeEx1(instance)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query volume %v. error: %v", instance, err)
+		}
+
+		volumes = append(volumes, volume)
+	}
+
+	return volumes, nil
+}
+
+// ListPartitionsOnDisk retrieves all partitions or a partition with the specified number on a disk.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_Partition
+//	  WHERE DiskNumber = '<diskNumber>'
+//	    AND PartitionNumber = '<partitionNumber>'
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition
+// for the WMI class definition.
+func ListPartitionsOnDisk(diskNumber, partitionNumber uint32, selectorList []string) ([]*storage.MSFT_Partition, error) {
+	filters := []*query.WmiQueryFilter{
+		query.NewWmiQueryFilter("DiskNumber", strconv.Itoa(int(diskNumber)), query.Equals),
+	}
+	if partitionNumber > 0 {
+		filters = append(filters, query.NewWmiQueryFilter("PartitionNumber", strconv.Itoa(int(partitionNumber)), query.Equals))
+	}
+	return ListPartitionsWithFilters(selectorList, filters...)
+}
+
+// ListPartitionsWithFilters retrieves all partitions matching with the conditions specified by query filters.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_Partition
+//	  WHERE ...
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition
+// for the WMI class definition.
+func ListPartitionsWithFilters(selectorList []string, filters ...*query.WmiQueryFilter) ([]*storage.MSFT_Partition, error) {
+	partitionQuery := query.NewWmiQueryWithSelectList("MSFT_Partition", selectorList)
+	partitionQuery.Filters = append(partitionQuery.Filters, filters...)
+	instances, err := QueryInstances(WMINamespaceStorage, partitionQuery)
+	if IgnoreNotFound(err) != nil {
+		return nil, err
+	}
+
+	var partitions []*storage.MSFT_Partition
+	for _, instance := range instances {
+		part, err := storage.NewMSFT_PartitionEx1(instance)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query partition %v. error: %v", instance, err)
+		}
+
+		partitions = append(partitions, part)
+	}
+
+	return partitions, nil
+}
+
+// ListPartitionToVolumeMappings builds a mapping between partition and volume with partition Object ID as the key.
+//
+// The equivalent WMI query is:
+//
+//		SELECT [selectors] FROM MSFT_PartitionToVolume
+//
+//	 Partition                                                               | Volume
+//	 ---------                                                               | ------
+//	 MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume
+// for the WMI class definition.
+func ListPartitionToVolumeMappings() (map[string]string, error) {
+	return ListWMIInstanceMappings(WMINamespaceStorage, "MSFT_PartitionToVolume", nil,
+		mappingObjectRefIndexer("Partition", "MSFT_Partition", "ObjectId"),
+		mappingObjectRefIndexer("Volume", "MSFT_Volume", "ObjectId"),
+	)
+}
+
+// ListVolumeToPartitionMappings builds a mapping between volume and partition with volume Object ID as the key.
+//
+// The equivalent WMI query is:
+//
+//		SELECT [selectors] FROM MSFT_PartitionToVolume
+//
+//	 Partition                                                               | Volume
+//	 ---------                                                               | ------
+//	 MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume
+// for the WMI class definition.
+func ListVolumeToPartitionMappings() (map[string]string, error) {
+	return ListWMIInstanceMappings(WMINamespaceStorage, "MSFT_PartitionToVolume", nil,
+		mappingObjectRefIndexer("Volume", "MSFT_Volume", "ObjectId"),
+		mappingObjectRefIndexer("Partition", "MSFT_Partition", "ObjectId"),
+	)
+}
+
+// FindPartitionsByVolume finds all partitions associated with the given volumes
+// using partition-to-volume mapping.
+func FindPartitionsByVolume(partitions []*storage.MSFT_Partition, volumes []*storage.MSFT_Volume) ([]*storage.MSFT_Partition, error) {
+	var partitionInstances []*cim.WmiInstance
+	for _, part := range partitions {
+		partitionInstances = append(partitionInstances, part.WmiInstance)
+	}
+
+	var volumeInstances []*cim.WmiInstance
+	for _, volume := range volumes {
+		volumeInstances = append(volumeInstances, volume.WmiInstance)
+	}
+
+	partitionToVolumeMappings, err := ListPartitionToVolumeMappings()
+	if err != nil {
+		return nil, err
+	}
+
+	filtered, err := FindInstancesByObjectIDMapping(partitionInstances, volumeInstances, partitionToVolumeMappings)
+	if err != nil {
+		return nil, err
+	}
+
+	var result []*storage.MSFT_Partition
+	for _, instance := range filtered {
+		part, err := storage.NewMSFT_PartitionEx1(instance)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query partition %v. error: %v", instance, err)
+		}
+
+		result = append(result, part)
+	}
+
+	return result, nil
+}
+
+// FindVolumesByPartition finds all volumes associated with the given partitions
+// using volume-to-partition mapping.
+func FindVolumesByPartition(volumes []*storage.MSFT_Volume, partitions []*storage.MSFT_Partition) ([]*storage.MSFT_Volume, error) {
+	var volumeInstances []*cim.WmiInstance
+	for _, volume := range volumes {
+		volumeInstances = append(volumeInstances, volume.WmiInstance)
+	}
+
+	var partitionInstances []*cim.WmiInstance
+	for _, part := range partitions {
+		partitionInstances = append(partitionInstances, part.WmiInstance)
+	}
+
+	volumeToPartitionMappings, err := ListVolumeToPartitionMappings()
+	if err != nil {
+		return nil, err
+	}
+
+	filtered, err := FindInstancesByObjectIDMapping(volumeInstances, partitionInstances, volumeToPartitionMappings)
+	if err != nil {
+		return nil, err
+	}
+
+	var result []*storage.MSFT_Volume
+	for _, instance := range filtered {
+		volume, err := storage.NewMSFT_VolumeEx1(instance)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query volume %v. error: %v", instance, err)
+		}
+
+		result = append(result, volume)
+	}
+
+	return result, nil
+}
+
+// GetPartitionByVolumeUniqueID retrieves a specific partition from a volume identified by its unique ID.
+func GetPartitionByVolumeUniqueID(volumeID string, partitionSelectorList []string) (*storage.MSFT_Partition, error) {
+	volume, err := QueryVolumeByUniqueID(volumeID, []string{"ObjectId"})
+	if err != nil {
+		return nil, err
+	}
+
+	partitions, err := ListPartitionsWithFilters(partitionSelectorList)
+	if err != nil {
+		return nil, err
+	}
+
+	result, err := FindPartitionsByVolume(partitions, []*storage.MSFT_Volume{volume})
+	if err != nil {
+		return nil, err
+	}
+
+	return result[0], nil
+}
+
+// GetVolumeByDriveLetter retrieves a volume associated with a specific drive letter.
+func GetVolumeByDriveLetter(driveLetter string, partitionSelectorList []string) (*storage.MSFT_Volume, error) {
+	var selectorsForPart []string
+	selectorsForPart = append(selectorsForPart, partitionSelectorList...)
+	selectorsForPart = append(selectorsForPart, "ObjectId")
+	partitions, err := ListPartitionsWithFilters(selectorsForPart, query.NewWmiQueryFilter("DriveLetter", driveLetter, query.Equals))
+	if err != nil {
+		return nil, err
+	}
+
+	volumes, err := ListVolumes(partitionSelectorList)
+	if err != nil {
+		return nil, err
+	}
+
+	result, err := FindVolumesByPartition(volumes, partitions)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(result) == 0 {
+		return nil, errors.NotFound
+	}
+
+	return result[0], nil
+}
+
+// GetPartitionDiskNumber retrieves the disk number associated with a given partition.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition
+// for the WMI class definitions.
+func GetPartitionDiskNumber(part *storage.MSFT_Partition) (uint32, error) {
+	diskNumber, err := part.GetProperty("DiskNumber")
+	if err != nil {
+		return 0, err
+	}
+
+	return uint32(diskNumber.(int32)), nil
+}