Use WMI to implement Volume API to reduce PowerShell overhead (library version)

go.mod

@@ -1,9 +1,21 @@
 module github.com/kubernetes-csi/csi-proxy/v2
 
-go 1.16
+go 1.22.0
+
+toolchain go1.22.3
 
 require (
-	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/stretchr/testify v1.5.1
+	github.com/go-ole/go-ole v1.3.0
+	github.com/microsoft/wmi v0.25.1
+	github.com/stretchr/testify v1.7.0
+	golang.org/x/sys v0.25.0
 	k8s.io/klog/v2 v2.9.0
 )
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/go-logr/logr v0.4.0 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0 // indirect
+)

go.sum

@@ -3,14 +3,24 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/go-logr/logr v0.4.0 h1:K7/B1jt6fIBQVd4Owv2MqGQClcgf0R266+7C/QjRcLc=
 github.com/go-logr/logr v0.4.0/go.mod h1:z6/tIYblkpsD+a4lm/fGIIU9mZ+XfAiaFtq7xTgseGU=
+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/microsoft/wmi v0.25.1 h1:sQv9hCEHtW5K6yEVL78T6XGRMGxk4aTpcJwCiB5rLN0=
+github.com/microsoft/wmi v0.25.1/go.mod h1:1zbdSF0A+5OwTUII5p3hN7/K6KF2m3o27pSG6Y51VU8=
+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/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.5.1 h1:nOGnQDM7FYENwehXlg/kFVnos3rEvtKTjRvOWSzb6H4=
-github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34=
+golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
-gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.0 h1:hjy8E9ON/egN1tAYqKb61G10WtihqetD4sz2H+8nIeA=
+gopkg.in/yaml.v3 v3.0.0/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 k8s.io/klog/v2 v2.9.0 h1:D7HV+n1V57XeZ0m6tdRkfknthUaM06VFbWldOFh8kzM=
 k8s.io/klog/v2 v2.9.0/go.mod h1:hy9LJ/NvuK+iVyP4Ehqva4HxZG/oXyIS3n3Jmire4Ec=

pkg/cim/disk.go

@@ -0,0 +1,36 @@
+//go:build windows
+// +build windows
+
+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,300 @@
+//go:build windows
+// +build windows
+
+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
+}