From 31830002c204af593829a72943f778bd9f0dd369 Mon Sep 17 00:00:00 2001
From: James O'Doherty <james@odoherty.family>
Date: Sun, 19 Feb 2023 15:08:12 -0500
Subject: [PATCH] Allow missing and null pointer fields to be nil.

This commit ensures missing fields for pointer field types remain nil:

  * This allows one to distinguish the field wasn't present from the
    field being NULL when the PtrNil option is false.
  * This also makes it more convenient to reuse structs with pointer
    field types with various encoders like the standard library JSON
    marshaler, where you may want to omit a field if it didn't exist in
    the queried results.

The expectations around this change are captured in the
TestMissingFields test.

In addition, it also fixes an edge case where even with PtrNil set to
true, pointer field types would always return with a pointer to a zero
value instead of being nil even when the query property's value is
VT_NULL.

This fix is checked by the TestNullPointerField test.
---
 swbemservices.go      |  1 +
 swbemservices_test.go |  7 ++--
 wmi.go                | 13 +++----
 wmi_test.go           | 82 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 94 insertions(+), 9 deletions(-)

diff --git a/swbemservices.go b/swbemservices.go
index 3ff8756..a250c84 100644
--- a/swbemservices.go
+++ b/swbemservices.go
@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package wmi
diff --git a/swbemservices_test.go b/swbemservices_test.go
index 1d619af..30a1f8d 100644
--- a/swbemservices_test.go
+++ b/swbemservices_test.go
@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package wmi
@@ -103,8 +104,8 @@ func WbemGetMemoryUsageMB(s *SWbemServices) (float64, float64, float64) {
 //Run all benchmarks (should run for at least 60s to get a stable number):
 //go test -run=NONE -bench=Version -benchtime=120s
 
-//Individual benchmarks:
-//go test -run=NONE -bench=NewVersion -benchtime=120s
+// Individual benchmarks:
+// go test -run=NONE -bench=NewVersion -benchtime=120s
 func BenchmarkNewVersion(b *testing.B) {
 	s, err := InitializeSWbemServices(DefaultClient)
 	if err != nil {
@@ -128,7 +129,7 @@ func BenchmarkNewVersion(b *testing.B) {
 	}
 }
 
-//go test -run=NONE -bench=OldVersion -benchtime=120s
+// go test -run=NONE -bench=OldVersion -benchtime=120s
 func BenchmarkOldVersion(b *testing.B) {
 	var dst []Win32_OperatingSystem
 	q := CreateQuery(&dst, "")
diff --git a/wmi.go b/wmi.go
index b4bb4f0..26c3581 100644
--- a/wmi.go
+++ b/wmi.go
@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 /*
@@ -20,7 +21,6 @@ Example code to print names of running processes:
 			println(i, v.Name)
 		}
 	}
-
 */
 package wmi
 
@@ -338,11 +338,6 @@ func (c *Client) loadEntity(dst interface{}, src *ole.IDispatch) (errFieldMismat
 		f := v.Field(i)
 		of := f
 		isPtr := f.Kind() == reflect.Ptr
-		if isPtr {
-			ptr := reflect.New(f.Type().Elem())
-			f.Set(ptr)
-			f = f.Elem()
-		}
 		n := v.Type().Field(i).Name
 		if n[0] < 'A' || n[0] > 'Z' {
 			continue
@@ -367,6 +362,12 @@ func (c *Client) loadEntity(dst interface{}, src *ole.IDispatch) (errFieldMismat
 		}
 		defer prop.Clear()
 
+		if isPtr && !(c.PtrNil && prop.VT == 0x1) {
+			ptr := reflect.New(f.Type().Elem())
+			f.Set(ptr)
+			f = f.Elem()
+		}
+
 		if prop.VT == 0x1 { //VT_NULL
 			continue
 		}
diff --git a/wmi_test.go b/wmi_test.go
index 32c4dcd..6587f2c 100644
--- a/wmi_test.go
+++ b/wmi_test.go
@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package wmi
@@ -39,6 +40,87 @@ func TestFieldMismatch(t *testing.T) {
 	}
 }
 
+func TestMissingFields(t *testing.T) {
+	type s struct {
+		Name           string
+		Missing        uint32
+		MissingPointer *uint32
+	}
+
+	var dst []s
+
+	client := &Client{
+		AllowMissingFields: true,
+	}
+	err := client.Query("SELECT Name FROM Win32_Process", &dst)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := range dst {
+		if dst[i].Missing != 0 {
+			t.Fatal("Expected Missing field to be 0")
+		}
+		if dst[i].MissingPointer != nil {
+			t.Fatal("Expected MissingPointer field to be nil")
+		}
+	}
+
+	// NonePtrZero and PtrNil should only affect the behavior of fields that
+	// exist as result properties, not missing fields.
+	client = &Client{
+		NonePtrZero:        true,
+		PtrNil:             true,
+		AllowMissingFields: true,
+	}
+	dst = []s{}
+	err = client.Query("SELECT Name FROM Win32_Process", &dst)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := range dst {
+		if dst[i].Missing != 0 {
+			t.Fatal("Expected Missing field to be 0")
+		}
+		if dst[i].MissingPointer != nil {
+			t.Fatal("Expected MissingPointer field to be nil")
+		}
+	}
+}
+
+func TestNullPointerField(t *testing.T) {
+	type s struct {
+		Name   string
+		Status *string
+	}
+
+	var dst []s
+
+	client := &Client{}
+	err := client.Query("SELECT Name, Status FROM Win32_Process WHERE Status IS NULL", &dst)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := range dst {
+		if dst[i].Status == nil {
+			t.Fatal("Expected Status field to not be nil")
+		}
+	}
+
+	client = &Client{
+		PtrNil: true,
+	}
+	dst = []s{}
+	err = client.Query("SELECT Name, Status FROM Win32_Process WHERE Status IS NULL", &dst)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := range dst {
+		if dst[i].Status != nil {
+			t.Fatal("Expected Status field to be nil")
+		}
+	}
+}
+
 func TestStrings(t *testing.T) {
 	printed := false
 	f := func() {