From d41bc015ccc06feedd21c0a87b0a44a02b74d530 Mon Sep 17 00:00:00 2001
From: Simon Rozman <simon@rozman.si>
Date: Fri, 1 Feb 2019 10:58:06 +0100
Subject: [PATCH] Finish support for setupapi.SetupDiGetClassDevsEx()

Signed-off-by: Simon Rozman <simon@rozman.si>
---
 setupapi/setupapi_windows.go       | 50 ++++++++++++++++++++----------
 setupapi/setupapi_windows_test.go  | 39 +++++++++++++++++++++++
 setupapi/zsetupapi_windows.go      | 30 ++++++++++--------
 setupapi/zsetupapi_windows_test.go | 20 ++++++++++++
 4 files changed, 110 insertions(+), 29 deletions(-)
 create mode 100644 setupapi/setupapi_windows_test.go
 create mode 100644 setupapi/zsetupapi_windows_test.go

diff --git a/setupapi/setupapi_windows.go b/setupapi/setupapi_windows.go
index 22cacee..6e81a9e 100644
--- a/setupapi/setupapi_windows.go
+++ b/setupapi/setupapi_windows.go
@@ -6,6 +6,8 @@
 package setupapi
 
 import (
+	"syscall"
+
 	"golang.org/x/sys/windows"
 )
 
@@ -16,26 +18,42 @@ const (
 type DIGCF uint32
 
 const (
-	Default         DIGCF = 0x00000001
-	Present         DIGCF = 0x00000002
-	AllClasses      DIGCF = 0x00000004
-	Profile         DIGCF = 0x00000008
-	DeviceInterface DIGCF = 0x00000010
-	InterfaceDevice DIGCF = 0x00000010
+	DIGCF_DEFAULT         DIGCF = 0x00000001 // only valid with DIGCF_DEVICEINTERFACE
+	DIGCF_PRESENT         DIGCF = 0x00000002
+	DIGCF_ALLCLASSES      DIGCF = 0x00000004
+	DIGCF_PROFILE         DIGCF = 0x00000008
+	DIGCF_DEVICEINTERFACE DIGCF = 0x00000010
 )
 
 type DevInfo windows.Handle
 
-//sys	setupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator *string, hwndParent uintptr, Flags uint32, DeviceInfoSet uintptr, MachineName string, reserved uint32) (handle windows.Handle, err error) = setupapi.SetupDiGetClassDevsExW
-
-// The SetupDiGetClassDevsEx function returns a handle to a device information set that contains requested device information elements for a local or a remote computer.
-func SetupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator string, hwndParent uintptr, Flags DIGCF, DeviceInfoSet DevInfo, MachineName string) (DevInfo, error) {
-	enumerator := &Enumerator
-
-	if Enumerator == "" {
-		enumerator = nil
+// The SetupDiDestroyDeviceInfoList function deletes a device information set and frees all associated memory.
+func (h DevInfo) Close() error {
+	if h != DevInfo(windows.InvalidHandle) {
+		return SetupDiDestroyDeviceInfoList(h)
 	}
 
-	h, err := setupDiGetClassDevsEx(ClassGuid, enumerator, hwndParent, uint32(Flags), uintptr(DeviceInfoSet), MachineName, 0)
-	return DevInfo(h), err
+	return nil
+}
+
+//sys	setupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator *uint16, hwndParent uintptr, Flags DIGCF, DeviceInfoSet DevInfo, MachineName *uint16, reserved uintptr) (handle DevInfo, err error) [failretval==DevInfo(windows.InvalidHandle)] = setupapi.SetupDiGetClassDevsExW
+//sys	SetupDiDestroyDeviceInfoList(DeviceInfoSet DevInfo) (err error) = setupapi.SetupDiDestroyDeviceInfoList
+
+// The SetupDiGetClassDevsEx function returns a handle to a device information set that contains requested device information elements for a local or a remote computer.
+func SetupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator string, hwndParent uintptr, Flags DIGCF, DeviceInfoSet DevInfo, MachineName string) (handle DevInfo, err error) {
+	var _p0 *uint16
+	if Enumerator != "" {
+		_p0, err = syscall.UTF16PtrFromString(Enumerator)
+		if err != nil {
+			return
+		}
+	}
+	var _p1 *uint16
+	if MachineName != "" {
+		_p1, err = syscall.UTF16PtrFromString(MachineName)
+		if err != nil {
+			return
+		}
+	}
+	return setupDiGetClassDevsEx(ClassGuid, _p0, hwndParent, Flags, DeviceInfoSet, _p1, 0)
 }
diff --git a/setupapi/setupapi_windows_test.go b/setupapi/setupapi_windows_test.go
new file mode 100644
index 0000000..709dd18
--- /dev/null
+++ b/setupapi/setupapi_windows_test.go
@@ -0,0 +1,39 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2019 WireGuard LLC. All Rights Reserved.
+ */
+
+package setupapi
+
+import (
+	"syscall"
+	"testing"
+
+	"golang.org/x/sys/windows"
+)
+
+func TestSetupDiGetClassDevsEx(t *testing.T) {
+	guidDeviceClassNet := windows.GUID{0x4d36e972, 0xe325, 0x11ce, [8]byte{0xbf, 0xc1, 0x08, 0x00, 0x2b, 0xe1, 0x03, 0x18}}
+
+	compName, err := windows.ComputerName()
+	if err != nil {
+		t.Errorf("Error getting computer name: %s", err.Error())
+	}
+
+	dev_info_list, err := SetupDiGetClassDevsEx(&guidDeviceClassNet, "PCI", 0, DIGCF_PRESENT, DevInfo(0), compName)
+	if err == nil {
+		dev_info_list.Close()
+	} else {
+		t.Errorf("Error calling SetupDiGetClassDevsEx: %s", err.Error())
+	}
+
+	dev_info_list, err = SetupDiGetClassDevsEx(nil, "", 0, DIGCF_PRESENT, DevInfo(0), "")
+	if err == nil {
+		dev_info_list.Close()
+		t.Errorf("SetupDiGetClassDevsEx(nil, ...) should fail")
+	} else {
+		if errWin, ok := err.(syscall.Errno); !ok || errWin != 87 /*ERROR_INVALID_PARAMETER*/ {
+			t.Errorf("SetupDiGetClassDevsEx(nil, ...) should fail with ERROR_INVALID_PARAMETER")
+		}
+	}
+}
diff --git a/setupapi/zsetupapi_windows.go b/setupapi/zsetupapi_windows.go
index 84ead2d..bce8c08 100644
--- a/setupapi/zsetupapi_windows.go
+++ b/setupapi/zsetupapi_windows.go
@@ -39,22 +39,26 @@ func errnoErr(e syscall.Errno) error {
 var (
 	modsetupapi = windows.NewLazySystemDLL("setupapi.dll")
 
-	procSetupDiGetClassDevsExW = modsetupapi.NewProc("SetupDiGetClassDevsExW")
+	procSetupDiGetClassDevsExW       = modsetupapi.NewProc("SetupDiGetClassDevsExW")
+	procSetupDiDestroyDeviceInfoList = modsetupapi.NewProc("SetupDiDestroyDeviceInfoList")
 )
 
-func setupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator *string, hwndParent uintptr, Flags uint32, DeviceInfoSet uintptr, MachineName string, reserved uint32) (handle windows.Handle, err error) {
-	var _p0 *uint16
-	_p0, err = syscall.UTF16PtrFromString(MachineName)
-	if err != nil {
-		return
-	}
-	return _setupDiGetClassDevsEx(ClassGuid, Enumerator, hwndParent, Flags, DeviceInfoSet, _p0, reserved)
-}
-
-func _setupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator *string, hwndParent uintptr, Flags uint32, DeviceInfoSet uintptr, MachineName *uint16, reserved uint32) (handle windows.Handle, err error) {
+func setupDiGetClassDevsEx(ClassGuid *windows.GUID, Enumerator *uint16, hwndParent uintptr, Flags DIGCF, DeviceInfoSet DevInfo, MachineName *uint16, reserved uintptr) (handle DevInfo, err error) {
 	r0, _, e1 := syscall.Syscall9(procSetupDiGetClassDevsExW.Addr(), 7, uintptr(unsafe.Pointer(ClassGuid)), uintptr(unsafe.Pointer(Enumerator)), uintptr(hwndParent), uintptr(Flags), uintptr(DeviceInfoSet), uintptr(unsafe.Pointer(MachineName)), uintptr(reserved), 0, 0)
-	handle = windows.Handle(r0)
-	if handle == 0 {
+	handle = DevInfo(r0)
+	if handle == DevInfo(windows.InvalidHandle) {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func SetupDiDestroyDeviceInfoList(DeviceInfoSet DevInfo) (err error) {
+	r1, _, e1 := syscall.Syscall(procSetupDiDestroyDeviceInfoList.Addr(), 1, uintptr(DeviceInfoSet), 0, 0)
+	if r1 == 0 {
 		if e1 != 0 {
 			err = errnoErr(e1)
 		} else {
diff --git a/setupapi/zsetupapi_windows_test.go b/setupapi/zsetupapi_windows_test.go
new file mode 100644
index 0000000..09b9195
--- /dev/null
+++ b/setupapi/zsetupapi_windows_test.go
@@ -0,0 +1,20 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2019 WireGuard LLC. All Rights Reserved.
+ */
+
+package setupapi
+
+import (
+	"syscall"
+	"testing"
+
+	"golang.org/x/sys/windows"
+)
+
+func TestSetupDiDestroyDeviceInfoList(t *testing.T) {
+	err := SetupDiDestroyDeviceInfoList(DevInfo(windows.InvalidHandle))
+	if errWin, ok := err.(syscall.Errno); !ok || errWin != 6 /*ERROR_INVALID_HANDLE*/ {
+		t.Errorf("SetupDiDestroyDeviceInfoList(nil, ...) should fail with ERROR_INVALID_HANDLE")
+	}
+}