setupapi: Add support for SetupDi(Get|Set)DeviceRegistryProperty()

Signed-off-by: Simon Rozman <simon@rozman.si>

4 files changed, 242 insertions, 24 deletions

diff --git a/setupapi/setupapi_windows.go b/setupapi/setupapi_windows.go
index d87520b..0a833ed 100644
--- a/setupapi/setupapi_windows.go
+++ b/setupapi/setupapi_windows.go
@@ -6,6 +6,8 @@
 package setupapi
 
 import (
+	"encoding/binary"
+	"fmt"
 	"syscall"
 	"unsafe"
 
@@ -144,6 +146,95 @@ func (DeviceInfoSet DevInfo) OpenDevRegKey(DeviceInfoData *SP_DEVINFO_DATA, Scop
 	return SetupDiOpenDevRegKey(DeviceInfoSet, DeviceInfoData, Scope, HwProfile, KeyType, samDesired)
 }
 
+//sys	setupDiGetDeviceRegistryProperty(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP, PropertyRegDataType *uint32, PropertyBuffer *byte, PropertyBufferSize uint32, RequiredSize *uint32) (err error) = setupapi.SetupDiGetDeviceRegistryPropertyW
+
+// SetupDiGetDeviceRegistryProperty function retrieves a specified Plug and Play device property.
+func SetupDiGetDeviceRegistryProperty(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP) (value interface{}, err error) {
+	buf := make([]byte, 0x100)
+	var dataType, bufLen uint32
+	err = setupDiGetDeviceRegistryProperty(DeviceInfoSet, DeviceInfoData, Property, &dataType, &buf[0], uint32(cap(buf)), &bufLen)
+	if err == nil {
+		// The buffer was sufficiently big.
+		return getRegistryValue(buf[:bufLen], dataType)
+	}
+
+	if errWin, ok := err.(syscall.Errno); ok && errWin == windows.ERROR_INSUFFICIENT_BUFFER {
+		// The buffer was too small. Now that we got the required size, create another one big enough and retry.
+		buf = make([]byte, bufLen)
+		err = setupDiGetDeviceRegistryProperty(DeviceInfoSet, DeviceInfoData, Property, &dataType, &buf[0], uint32(cap(buf)), &bufLen)
+		if err == nil {
+			return getRegistryValue(buf[:bufLen], dataType)
+		}
+	}
+
+	return
+}
+
+func getRegistryValue(buf []byte, dataType uint32) (interface{}, error) {
+	switch dataType {
+	case windows.REG_SZ:
+		return windows.UTF16ToString(toUTF16(buf)), nil
+	case windows.REG_EXPAND_SZ:
+		return registry.ExpandString(windows.UTF16ToString(toUTF16(buf)))
+	case windows.REG_BINARY:
+		return buf, nil
+	case windows.REG_DWORD_LITTLE_ENDIAN:
+		return binary.LittleEndian.Uint32(buf), nil
+	case windows.REG_DWORD_BIG_ENDIAN:
+		return binary.BigEndian.Uint32(buf), nil
+	case windows.REG_MULTI_SZ:
+		bufW := toUTF16(buf)
+		a := []string{}
+		for i := 0; i < len(bufW); {
+			j := i + wcslen(bufW[i:])
+			if i < j {
+				a = append(a, windows.UTF16ToString(bufW[i:j]))
+			}
+			i = j + 1
+		}
+		return a, nil
+	case windows.REG_QWORD_LITTLE_ENDIAN:
+		return binary.LittleEndian.Uint64(buf), nil
+	default:
+		return nil, fmt.Errorf("Unsupported registry value type: %v", dataType)
+	}
+}
+
+func toUTF16(buf []byte) []uint16 {
+	sl := struct {
+		addr *uint16
+		len  int
+		cap  int
+	}{(*uint16)(unsafe.Pointer(&buf[0])), len(buf) / 2, cap(buf) / 2}
+	return *(*[]uint16)(unsafe.Pointer(&sl))
+}
+
+func wcslen(str []uint16) int {
+	for i := 0; i < len(str); i++ {
+		if str[i] == 0 {
+			return i
+		}
+	}
+	return len(str)
+}
+
+// GetDeviceRegistryProperty method retrieves a specified Plug and Play device property.
+func (DeviceInfoSet DevInfo) GetDeviceRegistryProperty(DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP) (value interface{}, err error) {
+	return SetupDiGetDeviceRegistryProperty(DeviceInfoSet, DeviceInfoData, Property)
+}
+
+//sys	setupDiSetDeviceRegistryProperty(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP, PropertyBuffer *byte, PropertyBufferSize uint32) (err error) = setupapi.SetupDiSetDeviceRegistryPropertyW
+
+// SetupDiSetDeviceRegistryProperty function sets a Plug and Play device property for a device.
+func SetupDiSetDeviceRegistryProperty(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP, PropertyBuffer []byte) (err error) {
+	return setupDiSetDeviceRegistryProperty(DeviceInfoSet, DeviceInfoData, Property, &PropertyBuffer[0], uint32(len(PropertyBuffer)))
+}
+
+// SetDeviceRegistryProperty function sets a Plug and Play device property for a device.
+func (DeviceInfoSet DevInfo) SetDeviceRegistryProperty(DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP, PropertyBuffer []byte) (err error) {
+	return SetupDiSetDeviceRegistryProperty(DeviceInfoSet, DeviceInfoData, Property, PropertyBuffer)
+}
+
 //sys	setupDiGetDeviceInstallParams(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, DeviceInstallParams *_SP_DEVINSTALL_PARAMS) (err error) = setupapi.SetupDiGetDeviceInstallParamsW
 
 // SetupDiGetDeviceInstallParams function retrieves device installation parameters for a device information set or a particular device information element.
@@ -248,9 +339,9 @@ func SetupDiClassGuidsFromNameEx(ClassName string, MachineName string) (ClassGui
 		return
 	}
 
-	const bufLen = 4
-	var buf [bufLen]windows.GUID
-	var bufCount uint32
+	const bufCapacity = 4
+	var buf [bufCapacity]windows.GUID
+	var bufLen uint32
 
 	var machineNameUTF16 *uint16
 	if MachineName != "" {
@@ -260,18 +351,18 @@ func SetupDiClassGuidsFromNameEx(ClassName string, MachineName string) (ClassGui
 		}
 	}
 
-	err = setupDiClassGuidsFromNameEx(classNameUTF16, &buf[0], bufLen, &bufCount, machineNameUTF16, 0)
+	err = setupDiClassGuidsFromNameEx(classNameUTF16, &buf[0], bufCapacity, &bufLen, machineNameUTF16, 0)
 	if err == nil {
 		// The GUID array was sufficiently big. Return its slice.
-		return buf[:bufCount], nil
+		return buf[:bufLen], nil
 	}
 
 	if errWin, ok := err.(syscall.Errno); ok && errWin == windows.ERROR_INSUFFICIENT_BUFFER {
 		// The GUID array was too small. Now that we got the required size, create another one big enough and retry.
-		buf := make([]windows.GUID, bufCount)
-		err = setupDiClassGuidsFromNameEx(classNameUTF16, &buf[0], bufCount, &bufCount, machineNameUTF16, 0)
+		buf := make([]windows.GUID, bufLen)
+		err = setupDiClassGuidsFromNameEx(classNameUTF16, &buf[0], bufLen, &bufLen, machineNameUTF16, 0)
 		if err == nil {
-			return buf[:bufCount], nil
+			return buf[:bufLen], nil
 		}
 	}
 
diff --git a/setupapi/setupapi_windows_test.go b/setupapi/setupapi_windows_test.go
index 81a0a09..93ae42a 100644
--- a/setupapi/setupapi_windows_test.go
+++ b/setupapi/setupapi_windows_test.go
@@ -179,6 +179,56 @@ func TestSetupDiOpenDevRegKey(t *testing.T) {
 	}
 }
 
+func TestSetupDiGetDeviceRegistryProperty(t *testing.T) {
+	devInfoList, err := SetupDiGetClassDevsEx(&deviceClassNetGUID, "", 0, DIGCF_PRESENT, DevInfo(0), "")
+	if err != nil {
+		t.Errorf("Error calling SetupDiGetClassDevsEx: %s", err.Error())
+	}
+	defer devInfoList.Close()
+
+	for i := 0; true; i++ {
+		data, err := devInfoList.EnumDeviceInfo(i)
+		if err != nil {
+			if errWin, ok := err.(syscall.Errno); ok && errWin == 259 /*ERROR_NO_MORE_ITEMS*/ {
+				break
+			}
+			continue
+		}
+
+		val, err := devInfoList.GetDeviceRegistryProperty(data, SPDRP_CLASS)
+		if err != nil {
+			t.Errorf("Error calling SetupDiGetDeviceRegistryProperty(SPDRP_CLASS): %s", err.Error())
+		} else if class, ok := val.(string); !ok || strings.ToLower(class) != "net" {
+			t.Errorf("SetupDiGetDeviceRegistryProperty(SPDRP_CLASS) should return \"Net\"")
+		}
+
+		val, err = devInfoList.GetDeviceRegistryProperty(data, SPDRP_CLASSGUID)
+		if err != nil {
+			t.Errorf("Error calling SetupDiGetDeviceRegistryProperty(SPDRP_CLASSGUID): %s", err.Error())
+		} /* TODO: Parse GUID string: else if classGUID, ok := val.(string); !ok || parseGUID(classGUID) != deviceClassNetGUID {
+			t.Errorf("SetupDiGetDeviceRegistryProperty(SPDRP_CLASSGUID) should return %x", deviceClassNetGUID)
+		}*/
+
+		val, err = devInfoList.GetDeviceRegistryProperty(data, SPDRP_COMPATIBLEIDS)
+		if err != nil {
+			// Some devices have no SPDRP_COMPATIBLEIDS.
+			if errWin, ok := err.(syscall.Errno); !ok || errWin != 13 /*windows.ERROR_INVALID_DATA*/ {
+				t.Errorf("Error calling SetupDiGetDeviceRegistryProperty(SPDRP_COMPATIBLEIDS): %s", err.Error())
+			}
+		}
+
+		val, err = devInfoList.GetDeviceRegistryProperty(data, SPDRP_CONFIGFLAGS)
+		if err != nil {
+			t.Errorf("Error calling SetupDiGetDeviceRegistryProperty(SPDRP_CONFIGFLAGS): %s", err.Error())
+		}
+
+		val, err = devInfoList.GetDeviceRegistryProperty(data, SPDRP_DEVICE_POWER_DATA)
+		if err != nil {
+			t.Errorf("Error calling SetupDiGetDeviceRegistryProperty(SPDRP_DEVICE_POWER_DATA): %s", err.Error())
+		}
+	}
+}
+
 func TestSetupDiGetDeviceInstallParams(t *testing.T) {
 	devInfoList, err := SetupDiGetClassDevsEx(&deviceClassNetGUID, "", 0, DIGCF_PRESENT, DevInfo(0), "")
 	if err != nil {
diff --git a/setupapi/types_windows.go b/setupapi/types_windows.go
index 83a4753..c93d3aa 100644
--- a/setupapi/types_windows.go
+++ b/setupapi/types_windows.go
@@ -269,3 +269,54 @@ const (
 	DIREG_DRV  DIREG = 0x00000002 // Open/Create/Delete driver key
 	DIREG_BOTH DIREG = 0x00000004 // Delete both driver and Device key
 )
+
+//
+// SPDRP specifies device registry property codes
+// (Codes marked as read-only (R) may only be used for
+// SetupDiGetDeviceRegistryProperty)
+//
+// These values should cover the same set of registry properties
+// as defined by the CM_DRP codes in cfgmgr32.h.
+//
+// Note that SPDRP codes are zero based while CM_DRP codes are one based!
+//
+type SPDRP uint32
+
+const (
+	SPDRP_DEVICEDESC                  SPDRP = 0x00000000 // DeviceDesc (R/W)
+	SPDRP_HARDWAREID                  SPDRP = 0x00000001 // HardwareID (R/W)
+	SPDRP_COMPATIBLEIDS               SPDRP = 0x00000002 // CompatibleIDs (R/W)
+	SPDRP_SERVICE                     SPDRP = 0x00000004 // Service (R/W)
+	SPDRP_CLASS                       SPDRP = 0x00000007 // Class (R--tied to ClassGUID)
+	SPDRP_CLASSGUID                   SPDRP = 0x00000008 // ClassGUID (R/W)
+	SPDRP_DRIVER                      SPDRP = 0x00000009 // Driver (R/W)
+	SPDRP_CONFIGFLAGS                 SPDRP = 0x0000000A // ConfigFlags (R/W)
+	SPDRP_MFG                         SPDRP = 0x0000000B // Mfg (R/W)
+	SPDRP_FRIENDLYNAME                SPDRP = 0x0000000C // FriendlyName (R/W)
+	SPDRP_LOCATION_INFORMATION        SPDRP = 0x0000000D // LocationInformation (R/W)
+	SPDRP_PHYSICAL_DEVICE_OBJECT_NAME SPDRP = 0x0000000E // PhysicalDeviceObjectName (R)
+	SPDRP_CAPABILITIES                SPDRP = 0x0000000F // Capabilities (R)
+	SPDRP_UI_NUMBER                   SPDRP = 0x00000010 // UiNumber (R)
+	SPDRP_UPPERFILTERS                SPDRP = 0x00000011 // UpperFilters (R/W)
+	SPDRP_LOWERFILTERS                SPDRP = 0x00000012 // LowerFilters (R/W)
+	SPDRP_BUSTYPEGUID                 SPDRP = 0x00000013 // BusTypeGUID (R)
+	SPDRP_LEGACYBUSTYPE               SPDRP = 0x00000014 // LegacyBusType (R)
+	SPDRP_BUSNUMBER                   SPDRP = 0x00000015 // BusNumber (R)
+	SPDRP_ENUMERATOR_NAME             SPDRP = 0x00000016 // Enumerator Name (R)
+	SPDRP_SECURITY                    SPDRP = 0x00000017 // Security (R/W, binary form)
+	SPDRP_SECURITY_SDS                SPDRP = 0x00000018 // Security (W, SDS form)
+	SPDRP_DEVTYPE                     SPDRP = 0x00000019 // Device Type (R/W)
+	SPDRP_EXCLUSIVE                   SPDRP = 0x0000001A // Device is exclusive-access (R/W)
+	SPDRP_CHARACTERISTICS             SPDRP = 0x0000001B // Device Characteristics (R/W)
+	SPDRP_ADDRESS                     SPDRP = 0x0000001C // Device Address (R)
+	SPDRP_UI_NUMBER_DESC_FORMAT       SPDRP = 0x0000001D // UiNumberDescFormat (R/W)
+	SPDRP_DEVICE_POWER_DATA           SPDRP = 0x0000001E // Device Power Data (R)
+	SPDRP_REMOVAL_POLICY              SPDRP = 0x0000001F // Removal Policy (R)
+	SPDRP_REMOVAL_POLICY_HW_DEFAULT   SPDRP = 0x00000020 // Hardware Removal Policy (R)
+	SPDRP_REMOVAL_POLICY_OVERRIDE     SPDRP = 0x00000021 // Removal Policy Override (RW)
+	SPDRP_INSTALL_STATE               SPDRP = 0x00000022 // Device Install State (R)
+	SPDRP_LOCATION_PATHS              SPDRP = 0x00000023 // Device Location Paths (R)
+	SPDRP_BASE_CONTAINERID            SPDRP = 0x00000024 // Base ContainerID (R)
+
+	SPDRP_MAXIMUM_PROPERTY SPDRP = 0x00000025 // Upper bound on ordinals
+)
diff --git a/setupapi/zsetupapi_windows.go b/setupapi/zsetupapi_windows.go
index 3a0b441..c9aa72a 100644
--- a/setupapi/zsetupapi_windows.go
+++ b/setupapi/zsetupapi_windows.go
@@ -39,22 +39,24 @@ func errnoErr(e syscall.Errno) error {
 var (
 	modsetupapi = windows.NewLazySystemDLL("setupapi.dll")
 
-	procSetupDiCreateDeviceInfoListExW  = modsetupapi.NewProc("SetupDiCreateDeviceInfoListExW")
-	procSetupDiGetDeviceInfoListDetailW = modsetupapi.NewProc("SetupDiGetDeviceInfoListDetailW")
-	procSetupDiCreateDeviceInfoW        = modsetupapi.NewProc("SetupDiCreateDeviceInfoW")
-	procSetupDiEnumDeviceInfo           = modsetupapi.NewProc("SetupDiEnumDeviceInfo")
-	procSetupDiDestroyDeviceInfoList    = modsetupapi.NewProc("SetupDiDestroyDeviceInfoList")
-	procSetupDiGetClassDevsExW          = modsetupapi.NewProc("SetupDiGetClassDevsExW")
-	procSetupDiCallClassInstaller       = modsetupapi.NewProc("SetupDiCallClassInstaller")
-	procSetupDiOpenDevRegKey            = modsetupapi.NewProc("SetupDiOpenDevRegKey")
-	procSetupDiGetDeviceInstallParamsW  = modsetupapi.NewProc("SetupDiGetDeviceInstallParamsW")
-	procSetupDiGetClassInstallParamsW   = modsetupapi.NewProc("SetupDiGetClassInstallParamsW")
-	procSetupDiSetDeviceInstallParamsW  = modsetupapi.NewProc("SetupDiSetDeviceInstallParamsW")
-	procSetupDiSetClassInstallParamsW   = modsetupapi.NewProc("SetupDiSetClassInstallParamsW")
-	procSetupDiClassNameFromGuidExW     = modsetupapi.NewProc("SetupDiClassNameFromGuidExW")
-	procSetupDiClassGuidsFromNameExW    = modsetupapi.NewProc("SetupDiClassGuidsFromNameExW")
-	procSetupDiGetSelectedDevice        = modsetupapi.NewProc("SetupDiGetSelectedDevice")
-	procSetupDiSetSelectedDevice        = modsetupapi.NewProc("SetupDiSetSelectedDevice")
+	procSetupDiCreateDeviceInfoListExW    = modsetupapi.NewProc("SetupDiCreateDeviceInfoListExW")
+	procSetupDiGetDeviceInfoListDetailW   = modsetupapi.NewProc("SetupDiGetDeviceInfoListDetailW")
+	procSetupDiCreateDeviceInfoW          = modsetupapi.NewProc("SetupDiCreateDeviceInfoW")
+	procSetupDiEnumDeviceInfo             = modsetupapi.NewProc("SetupDiEnumDeviceInfo")
+	procSetupDiDestroyDeviceInfoList      = modsetupapi.NewProc("SetupDiDestroyDeviceInfoList")
+	procSetupDiGetClassDevsExW            = modsetupapi.NewProc("SetupDiGetClassDevsExW")
+	procSetupDiCallClassInstaller         = modsetupapi.NewProc("SetupDiCallClassInstaller")
+	procSetupDiOpenDevRegKey              = modsetupapi.NewProc("SetupDiOpenDevRegKey")
+	procSetupDiGetDeviceRegistryPropertyW = modsetupapi.NewProc("SetupDiGetDeviceRegistryPropertyW")
+	procSetupDiSetDeviceRegistryPropertyW = modsetupapi.NewProc("SetupDiSetDeviceRegistryPropertyW")
+	procSetupDiGetDeviceInstallParamsW    = modsetupapi.NewProc("SetupDiGetDeviceInstallParamsW")
+	procSetupDiGetClassInstallParamsW     = modsetupapi.NewProc("SetupDiGetClassInstallParamsW")
+	procSetupDiSetDeviceInstallParamsW    = modsetupapi.NewProc("SetupDiSetDeviceInstallParamsW")
+	procSetupDiSetClassInstallParamsW     = modsetupapi.NewProc("SetupDiSetClassInstallParamsW")
+	procSetupDiClassNameFromGuidExW       = modsetupapi.NewProc("SetupDiClassNameFromGuidExW")
+	procSetupDiClassGuidsFromNameExW      = modsetupapi.NewProc("SetupDiClassGuidsFromNameExW")
+	procSetupDiGetSelectedDevice          = modsetupapi.NewProc("SetupDiGetSelectedDevice")
+	procSetupDiSetSelectedDevice          = modsetupapi.NewProc("SetupDiSetSelectedDevice")
 )
 
 func setupDiCreateDeviceInfoListEx(ClassGUID *windows.GUID, hwndParent uintptr, MachineName *uint16, Reserved uintptr) (handle DevInfo, err error) {
@@ -156,6 +158,30 @@ func setupDiOpenDevRegKey(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA
 	return
 }
 
+func setupDiGetDeviceRegistryProperty(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP, PropertyRegDataType *uint32, PropertyBuffer *byte, PropertyBufferSize uint32, RequiredSize *uint32) (err error) {
+	r1, _, e1 := syscall.Syscall9(procSetupDiGetDeviceRegistryPropertyW.Addr(), 7, uintptr(DeviceInfoSet), uintptr(unsafe.Pointer(DeviceInfoData)), uintptr(Property), uintptr(unsafe.Pointer(PropertyRegDataType)), uintptr(unsafe.Pointer(PropertyBuffer)), uintptr(PropertyBufferSize), uintptr(unsafe.Pointer(RequiredSize)), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func setupDiSetDeviceRegistryProperty(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, Property SPDRP, PropertyBuffer *byte, PropertyBufferSize uint32) (err error) {
+	r1, _, e1 := syscall.Syscall6(procSetupDiSetDeviceRegistryPropertyW.Addr(), 5, uintptr(DeviceInfoSet), uintptr(unsafe.Pointer(DeviceInfoData)), uintptr(Property), uintptr(unsafe.Pointer(PropertyBuffer)), uintptr(PropertyBufferSize), 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
 func setupDiGetDeviceInstallParams(DeviceInfoSet DevInfo, DeviceInfoData *SP_DEVINFO_DATA, DeviceInstallParams *_SP_DEVINSTALL_PARAMS) (err error) {
 	r1, _, e1 := syscall.Syscall(procSetupDiGetDeviceInstallParamsW.Addr(), 3, uintptr(DeviceInfoSet), uintptr(unsafe.Pointer(DeviceInfoData)), uintptr(unsafe.Pointer(DeviceInstallParams)))
 	if r1 == 0 {