Move WMI Volume and Partition functions to cim package

Author: laozc

pkg/cim/volume.go

@@ -7,9 +7,23 @@ import (
 	"fmt"
 	"strconv"
 
+	"github.com/go-ole/go-ole"
 	"github.com/microsoft/wmi/pkg/base/query"
 	"github.com/microsoft/wmi/pkg/errors"
 	"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
+	"k8s.io/klog/v2"
+)
+
+const (
+	FileSystemUnknown = 0
+)
+
+var (
+	VolumeSelectorListForFileSystemType = []string{"FileSystemType"}
+	VolumeSelectorListForStats          = []string{"UniqueId", "SizeRemaining", "Size"}
+	VolumeSelectorListUniqueID          = []string{"UniqueId"}
+
+	PartitionSelectorListObjectID = []string{"ObjectId"}
 )
 
 // QueryVolumeByUniqueID retrieves a specific volume by its unique identifier,
@@ -78,6 +92,68 @@ func ListVolumes(selectorList []string) ([]*storage.MSFT_Volume, error) {
 	return volumes, nil
 }
 
+// FormatVolume formats the specified volume.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/format-msft-volume
+// for the WMI method definition.
+func FormatVolume(volume *storage.MSFT_Volume, params ...interface{}) (int, error) {
+	result, err := volume.InvokeMethodWithReturn("Format", params...)
+	return int(result), err
+}
+
+// FlushVolume flushes the cached data in the volume's file system to disk.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-volume-flush
+// for the WMI method definition.
+func FlushVolume(volume *storage.MSFT_Volume) (int, error) {
+	result, err := volume.Flush()
+	return int(result), err
+}
+
+// GetVolumeUniqueID returns the unique ID (object ID) of a volume.
+func GetVolumeUniqueID(volume *storage.MSFT_Volume) (string, error) {
+	return volume.GetPropertyUniqueId()
+}
+
+// GetVolumeFileSystemType returns the file system type of a volume.
+func GetVolumeFileSystemType(volume *storage.MSFT_Volume) (int32, error) {
+	fsType, err := volume.GetProperty("FileSystemType")
+	if err != nil {
+		return 0, err
+	}
+	return fsType.(int32), nil
+}
+
+// GetVolumeSize returns the size of a volume.
+func GetVolumeSize(volume *storage.MSFT_Volume) (int64, error) {
+	volumeSizeVal, err := volume.GetProperty("Size")
+	if err != nil {
+		return -1, err
+	}
+
+	volumeSize, err := strconv.ParseInt(volumeSizeVal.(string), 10, 64)
+	if err != nil {
+		return -1, err
+	}
+
+	return volumeSize, err
+}
+
+// GetVolumeSizeRemaining returns the remaining size of a volume.
+func GetVolumeSizeRemaining(volume *storage.MSFT_Volume) (int64, error) {
+	volumeSizeRemainingVal, err := volume.GetProperty("SizeRemaining")
+	if err != nil {
+		return -1, err
+	}
+
+	volumeSizeRemaining, err := strconv.ParseInt(volumeSizeRemainingVal.(string), 10, 64)
+	if err != nil {
+		return -1, err
+	}
+
+	return volumeSizeRemaining, err
+}
+
 // ListPartitionsOnDisk retrieves all partitions or a partition with the specified number on a disk.
 //
 // The equivalent WMI query is:
@@ -262,6 +338,55 @@ func SetPartitionState(part *storage.MSFT_Partition, online bool) (int, string,
 	return int(result), status, err
 }
 
+// GetPartitionSupportedSize retrieves the minimum and maximum sizes that the partition can be resized to using the ResizePartition method.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-getsupportedsizes
+// for the WMI method definition.
+func GetPartitionSupportedSize(part *storage.MSFT_Partition) (result int, sizeMin, sizeMax int64, status string, err error) {
+	sizeMin = -1
+	sizeMax = -1
+
+	var sizeMinVar, sizeMaxVar ole.VARIANT
+	invokeResult, err := part.InvokeMethodWithReturn("GetSupportedSize", &sizeMinVar, &sizeMaxVar, &status)
+	if invokeResult != 0 || err != nil {
+		result = int(invokeResult)
+	}
+	klog.V(5).Infof("got sizeMin (%v) sizeMax (%v) from partition (%v), status: %s", sizeMinVar, sizeMaxVar, part, status)
+
+	sizeMin, err = strconv.ParseInt(sizeMinVar.ToString(), 10, 64)
+	if err != nil {
+		return
+	}
+
+	sizeMax, err = strconv.ParseInt(sizeMaxVar.ToString(), 10, 64)
+	return
+}
+
+// ResizePartition resizes a partition.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-resize
+// for the WMI method definition.
+func ResizePartition(part *storage.MSFT_Partition, size int64) (int, string, error) {
+	var status string
+	result, err := part.InvokeMethodWithReturn("Resize", strconv.Itoa(int(size)), &status)
+	return int(result), status, err
+}
+
+// GetPartitionSize returns the size of a partition.
+func GetPartitionSize(part *storage.MSFT_Partition) (int64, error) {
+	sizeProp, err := part.GetProperty("Size")
+	if err != nil {
+		return -1, err
+	}
+
+	size, err := strconv.ParseInt(sizeProp.(string), 10, 64)
+	if err != nil {
+		return -1, err
+	}
+
+	return size, err
+}
+
 // FilterForPartitionOnDisk creates a WMI query filter to query a disk by its number.
 func FilterForPartitionOnDisk(diskNumber uint32) *query.WmiQueryFilter {
 	return query.NewWmiQueryFilter("DiskNumber", strconv.Itoa(int(diskNumber)), query.Equals)

pkg/cim/wmi.go

@@ -9,7 +9,7 @@ import (
 	"github.com/go-ole/go-ole"
 	"github.com/go-ole/go-ole/oleutil"
 	"github.com/microsoft/wmi/pkg/base/query"
-	"github.com/microsoft/wmi/pkg/errors"
+	wmierrors "github.com/microsoft/wmi/pkg/errors"
 	cim "github.com/microsoft/wmi/pkg/wmiinstance"
 	"k8s.io/klog/v2"
 )
@@ -61,7 +61,7 @@ func QueryFromWMI(namespace string, query *query.WmiQuery, handler InstanceHandl
 	}
 
 	if len(instances) == 0 {
-		return errors.NotFound
+		return wmierrors.NotFound
 	}
 
 	var cont bool
@@ -95,7 +95,7 @@ func executeClassMethodParam(classInst *cim.WmiInstance, method *cim.WmiMethod,
 
 	iDispatchInstance := classInst.GetIDispatch()
 	if iDispatchInstance == nil {
-		return nil, errors.Wrapf(errors.InvalidInput, "InvalidInstance")
+		return nil, wmierrors.Wrapf(wmierrors.InvalidInput, "InvalidInstance")
 	}
 	rawResult, err := iDispatchInstance.GetProperty("Methods_")
 	if err != nil {
@@ -235,10 +235,15 @@ func InvokeCimMethod(namespace, class, methodName string, inputParameters map[st
 	return int(result.ReturnValue), outputParameters, nil
 }
 
+// IsNotFound returns true if it's a "not found" error.
+func IsNotFound(err error) bool {
+	return wmierrors.IsNotFound(err)
+}
+
 // IgnoreNotFound returns nil if the error is nil or a "not found" error,
 // otherwise returns the original error.
 func IgnoreNotFound(err error) error {
-	if err == nil || errors.IsNotFound(err) {
+	if err == nil || IsNotFound(err) {
 		return nil
 	}
 	return err