Merge pull request #52 from ksubrmnn/disknumber

GetDiskNumberWithID API

Author: k8s-ci-robot

internal/os/disk/api.go

@@ -1,15 +1,29 @@
+// +build windows
+
 package disk
 
 import (
 	"encoding/json"
 	"fmt"
 	"os/exec"
+	"strconv"
 	"strings"
+	"syscall"
+	"unsafe"
 
 	shared "github.com/kubernetes-csi/csi-proxy/internal/shared/disk"
 	"k8s.io/klog"
 )
 
+var (
+	kernel32DLL = syscall.NewLazyDLL("kernel32.dll")
+)
+
+const (
+	IOCTL_STORAGE_GET_DEVICE_NUMBER = 0x2D1080
+	IOCTL_STORAGE_QUERY_PROPERTY    = 0x002d1400
+)
+
 // Implements the OS API calls related to Disk Devices. All code here should be very simple
 // pass-through to the OS APIs or cmdlets. Any logic around the APIs/cmdlet invocation
 // should go in internal/server/filesystem/disk.go so that logic can be easily unit-tested
@@ -123,11 +137,83 @@ func (APIImplementor) CreatePartition(diskID string) error {
 	return nil
 }
 
-func (APIImplementor) GetDiskNumberByName(diskName string) (string, error) {
-	out, err := exec.Command("DiskUtil.exe", "-GetDiskNumberWithId", diskName).CombinedOutput()
-	outString := string(out)
+func (imp APIImplementor) GetDiskNumberByName(diskName string) (string, error) {
+	diskNumber, err := imp.GetDiskNumberWithID(diskName)
+	return strconv.FormatUint(uint64(diskNumber), 10), err
+}
+
+func (APIImplementor) DiskNumber(disk syscall.Handle) (uint32, error) {
+	var bytes uint32
+	devNum := StorageDeviceNumber{}
+	buflen := uint32(unsafe.Sizeof(devNum.DeviceType)) + uint32(unsafe.Sizeof(devNum.DeviceNumber)) + uint32(unsafe.Sizeof(devNum.PartitionNumber))
+
+	err := syscall.DeviceIoControl(disk, IOCTL_STORAGE_GET_DEVICE_NUMBER, nil, 0, (*byte)(unsafe.Pointer(&devNum)), buflen, &bytes, nil)
+
+	return devNum.DeviceNumber, err
+}
+
+func (APIImplementor) DiskHasPage83ID(disk syscall.Handle, matchID string) (bool, error) {
+	query := StoragePropertyQuery{}
+
+	bufferSize := uint32(4 * 1024)
+	buffer := make([]byte, 4*1024)
+	var size uint32
+	var n uint32
+	var m uint16
+
+	query.QueryType = PropertyStandardQuery
+	query.PropertyID = StorageDeviceIDProperty
+
+	querySize := uint32(unsafe.Sizeof(query.PropertyID)) + uint32(unsafe.Sizeof(query.QueryType)) + uint32(unsafe.Sizeof(query.Byte))
+	querySize = uint32(unsafe.Sizeof(query))
+	err := syscall.DeviceIoControl(disk, IOCTL_STORAGE_QUERY_PROPERTY, (*byte)(unsafe.Pointer(&query)), querySize, (*byte)(unsafe.Pointer(&buffer[0])), bufferSize, &size, nil)
+	if err != nil {
+		return false, fmt.Errorf("IOCTL_STORAGE_QUERY_PROPERTY failed: %v", err)
+	}
+
+	devIDDesc := (*StorageDeviceIDDescriptor)(unsafe.Pointer(&buffer[0]))
+
+	pID := (*StorageIdentifier)(unsafe.Pointer(&devIDDesc.Identifiers[0]))
+
+	page83ID := []byte{}
+	byteSize := unsafe.Sizeof(byte(0))
+	for n = 0; n < devIDDesc.NumberOfIdentifiers; n++ {
+		if pID.CodeSet == StorageIDCodeSetASCII && pID.Association == StorageIDAssocDevice {
+			for m = 0; m < pID.IdentifierSize; m++ {
+				page83ID = append(page83ID, *(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(&pID.Identifier[0])) + byteSize*uintptr(m))))
+			}
+
+			page83IDString := string(page83ID)
+			if strings.Contains(page83IDString, matchID) {
+				return true, nil
+			}
+		}
+		pID = (*StorageIdentifier)(unsafe.Pointer(uintptr(unsafe.Pointer(pID)) + byteSize*uintptr(pID.NextOffset)))
+	}
+	return false, nil
+}
+
+func (imp APIImplementor) GetDiskNumberWithID(page83ID string) (uint32, error) {
+	out, err := exec.Command("powershell.exe", "(get-disk | select Path) | ConvertTo-Json").CombinedOutput()
 	if err != nil {
-		return "", fmt.Errorf("error getting disk number by name %s: %v, output %v", diskName, err, out)
+		return 0, fmt.Errorf("Could not query disk paths")
 	}
-	return strings.TrimSpace(outString), nil
+
+	outString := string(out)
+	diskPaths := []DiskPath{}
+	json.Unmarshal([]byte(outString), &diskPaths)
+
+	for i := range diskPaths {
+		h, err := syscall.Open(diskPaths[i].Path, syscall.O_RDONLY, 0)
+		if err != nil {
+			return 0, err
+		}
+
+		found, err := imp.DiskHasPage83ID(h, page83ID)
+		if found {
+			return imp.DiskNumber(h)
+		}
+	}
+
+	return 0, fmt.Errorf("Could not find disk with Page83 ID %s", page83ID)
 }

internal/os/disk/types.go

@@ -0,0 +1,114 @@
+package disk
+
+type StorageDeviceNumber struct {
+	DeviceType      DeviceType
+	DeviceNumber    uint32
+	PartitionNumber uint32
+}
+type DeviceType uint32
+
+type StoragePropertyID uint32
+
+const (
+	StorageDeviceProperty                  StoragePropertyID = 0
+	StorageAdapterProperty                                   = 1
+	StorageDeviceIDProperty                                  = 2
+	StorageDeviceUniqueIDProperty                            = 3
+	StorageDeviceWriteCacheProperty                          = 4
+	StorageMiniportProperty                                  = 5
+	StorageAccessAlignmentProperty                           = 6
+	StorageDeviceSeekPenaltyProperty                         = 7
+	StorageDeviceTrimProperty                                = 8
+	StorageDeviceWriteAggregationProperty                    = 9
+	StorageDeviceDeviceTelemetryProperty                     = 10
+	StorageDeviceLBProvisioningProperty                      = 11
+	StorageDevicePowerProperty                               = 12
+	StorageDeviceCopyOffloadProperty                         = 13
+	StorageDeviceResiliencyProperty                          = 14
+	StorageDeviceMediumProductType                           = 15
+	StorageAdapterRpmbProperty                               = 16
+	StorageAdapterCryptoProperty                             = 17
+	StorageDeviceIoCapabilityProperty                        = 18
+	StorageAdapterProtocolSpecificProperty                   = 19
+	StorageDeviceProtocolSpecificProperty                    = 20
+	StorageAdapterTemperatureProperty                        = 21
+	StorageDeviceTemperatureProperty                         = 22
+	StorageAdapterPhysicalTopologyProperty                   = 23
+	StorageDevicePhysicalTopologyProperty                    = 24
+	StorageDeviceAttributesProperty                          = 25
+	StorageDeviceManagementStatus                            = 26
+	StorageAdapterSerialNumberProperty                       = 27
+	StorageDeviceLocationProperty                            = 28
+	StorageDeviceNumaProperty                                = 29
+	StorageDeviceZonedDeviceProperty                         = 30
+	StorageDeviceUnsafeShutdownCount                         = 31
+	StorageDeviceEnduranceProperty                           = 32
+)
+
+type StorageQueryType uint32
+
+const (
+	PropertyStandardQuery StorageQueryType = iota
+	PropertyExistsQuery
+	PropertyMaskQuery
+	PropertyQueryMaxDefined
+)
+
+type StoragePropertyQuery struct {
+	PropertyID StoragePropertyID
+	QueryType  StorageQueryType
+	Byte       []AdditionalParameters
+}
+
+type AdditionalParameters byte
+
+type StorageDeviceIDDescriptor struct {
+	Version             uint32
+	Size                uint32
+	NumberOfIdentifiers uint32
+	Identifiers         [1]byte
+}
+
+type StorageIdentifierCodeSet uint32
+
+const (
+	StorageIDCodeSetReserved StorageIdentifierCodeSet = 0
+	StorageIDCodeSetBinary                            = 1
+	StorageIDCodeSetASCII                             = 2
+	StorageIDCodeSetUtf8                              = 3
+)
+
+type StorageIdentifierType uint32
+
+const (
+	StorageIdTypeVendorSpecific           StorageIdentifierType = 0
+	StorageIDTypeVendorID                                       = 1
+	StorageIDTypeEUI64                                          = 2
+	StorageIDTypeFCPHName                                       = 3
+	StorageIDTypePortRelative                                   = 4
+	StorageIDTypeTargetPortGroup                                = 5
+	StorageIDTypeLogicalUnitGroup                               = 6
+	StorageIDTypeMD5LogicalUnitIdentifier                       = 7
+	StorageIDTypeScsiNameString                                 = 8
+)
+
+type StorageAssociationType uint32
+
+const (
+	StorageIDAssocDevice StorageAssociationType = 0
+	StorageIDAssocPort                          = 1
+	StorageIDAssocTarget                        = 2
+)
+
+type StorageIdentifier struct {
+	CodeSet        StorageIdentifierCodeSet
+	Type           StorageIdentifierType
+	IdentifierSize uint16
+	NextOffset     uint16
+	Association    StorageAssociationType
+	Identifier     [1]byte
+}
+
+type DiskPath struct {
+	Path string `json:"Path"`
+}