Use WMI to implement Disk API to reduce PowerShell overhead

Author: laozc

pkg/cim/disk.go

@@ -28,6 +28,14 @@ const (
 	ErrorCodeCreatePartitionAccessPathAlreadyInUse = 42002
 )
 
+var (
+	DiskSelectorListForDiskNumberAndLocation = []string{"Number", "Location"}
+	DiskSelectorListForPartitionStyle        = []string{"PartitionStyle"}
+	DiskSelectorListForPathAndSerialNumber   = []string{"Path", "SerialNumber"}
+	DiskSelectorListForIsOffline             = []string{"IsOffline"}
+	DiskSelectorListForSize                  = []string{"Size"}
+)
+
 // QueryDiskByNumber retrieves disk information for a specific disk identified by its number.
 //
 // The equivalent WMI query is:
@@ -52,6 +60,43 @@ func QueryDiskByNumber(diskNumber uint32, selectorList []string) (*storage.MSFT_
 	return disk, nil
 }
 
+// ListDisks retrieves information about all available disks.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_Disk
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk
+// for the WMI class definition.
+func ListDisks(selectorList []string) ([]*storage.MSFT_Disk, error) {
+	diskQuery := query.NewWmiQueryWithSelectList("MSFT_Disk", selectorList)
+	instances, err := QueryInstances(WMINamespaceStorage, diskQuery)
+	if IgnoreNotFound(err) != nil {
+		return nil, err
+	}
+
+	var disks []*storage.MSFT_Disk
+	for _, instance := range instances {
+		disk, err := storage.NewMSFT_DiskEx1(instance)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query disk %v. error: %v", instance, err)
+		}
+
+		disks = append(disks, disk)
+	}
+
+	return disks, nil
+}
+
+// InitializeDisk initializes a RAW disk with a particular partition style.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/initialize-msft-disk
+// for the WMI method definition.
+func InitializeDisk(disk *storage.MSFT_Disk, partitionStyle int) (int, error) {
+	result, err := disk.InvokeMethodWithReturn("Initialize", int32(partitionStyle))
+	return int(result), err
+}
+
 // RefreshDisk Refreshes the cached disk layout information.
 //
 // Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-refresh
@@ -61,3 +106,85 @@ func RefreshDisk(disk *storage.MSFT_Disk) (int, string, error) {
 	result, err := disk.InvokeMethodWithReturn("Refresh", &status)
 	return int(result), status, err
 }
+
+// CreatePartition creates a partition on a disk.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/createpartition-msft-disk
+// for the WMI method definition.
+func CreatePartition(disk *storage.MSFT_Disk, params ...interface{}) (int, error) {
+	result, err := disk.InvokeMethodWithReturn("CreatePartition", params...)
+	return int(result), err
+}
+
+// SetDiskState takes a disk online or offline.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-online and
+// https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-offline
+// for the WMI method definition.
+func SetDiskState(disk *storage.MSFT_Disk, online bool) (int, string, error) {
+	method := "Offline"
+	if online {
+		method = "Online"
+	}
+
+	var status string
+	result, err := disk.InvokeMethodWithReturn(method, &status)
+	return int(result), status, err
+}
+
+// RescanDisks rescans all changes by updating the internal cache of software objects (that is, Disks, Partitions, Volumes)
+// for the storage setting.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-storagesetting-updatehoststoragecache
+// for the WMI method definition.
+func RescanDisks() (int, error) {
+	result, _, err := InvokeCimMethod(WMINamespaceStorage, "MSFT_StorageSetting", "UpdateHostStorageCache", nil)
+	return result, err
+}
+
+// GetDiskNumber returns the number of a disk.
+func GetDiskNumber(disk *storage.MSFT_Disk) (uint32, error) {
+	number, err := disk.GetProperty("Number")
+	if err != nil {
+		return 0, err
+	}
+	return uint32(number.(int32)), err
+}
+
+// GetDiskLocation returns the location of a disk.
+func GetDiskLocation(disk *storage.MSFT_Disk) (string, error) {
+	return disk.GetPropertyLocation()
+}
+
+// GetDiskPartitionStyle returns the partition style of a disk.
+func GetDiskPartitionStyle(disk *storage.MSFT_Disk) (int32, error) {
+	retValue, err := disk.GetProperty("PartitionStyle")
+	if err != nil {
+		return 0, err
+	}
+	return retValue.(int32), err
+}
+
+// IsDiskOffline returns whether a disk is offline.
+func IsDiskOffline(disk *storage.MSFT_Disk) (bool, error) {
+	return disk.GetPropertyIsOffline()
+}
+
+// GetDiskSize returns the size of a disk.
+func GetDiskSize(disk *storage.MSFT_Disk) (int64, error) {
+	sz, err := disk.GetProperty("Size")
+	if err != nil {
+		return -1, err
+	}
+	return strconv.ParseInt(sz.(string), 10, 64)
+}
+
+// GetDiskPath returns the path of a disk.
+func GetDiskPath(disk *storage.MSFT_Disk) (string, error) {
+	return disk.GetPropertyPath()
+}
+
+// GetDiskSerialNumber returns the serial number of a disk.
+func GetDiskSerialNumber(disk *storage.MSFT_Disk) (string, error) {
+	return disk.GetPropertySerialNumber()
+}