mi: validate unmarshal target before querying and unit-test setField (#2520)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/mi/char16_test.go

@@ -65,9 +65,3 @@ func TestGetValue_CHAR16(t *testing.T) {
 		})
 	}
 }
-
-// Note: The CHAR16 unmarshal branch in unmarshalInstance (session.go) handles
-// string and uint fields. Integration testing of this path requires a WMI class
-// with a CHAR16 property, which is extremely rare. The setUintField and
-// setIntField helpers used by the branch are covered by the signed/unsigned
-// integer tests (Test_MI_QueryUnmarshal_SignedInt, Test_MI_Unmarshal_TypeMismatch).

internal/mi/operation.go

@@ -201,31 +201,44 @@ func (o *Operation) GetInstance() (*Instance, bool, error) {
 }
 
 func (o *Operation) Unmarshal(dst any) error {
-	return o.unmarshal(dst, false)
-}
-
-// unmarshal iterates over the operation's instances and populates dst.
-// skipMissing controls how missing elements are handled (see unmarshalInstance).
-func (o *Operation) unmarshal(dst any, skipMissing bool) error {
 	if o == nil || o.ft == nil {
 		return ErrNotInitialized
 	}
 
+	dv, err := structSlice(dst)
+	if err != nil {
+		return err
+	}
+
+	return o.unmarshal(dv, false)
+}
+
+// structSlice checks that dst is a non-nil pointer to a slice of structs and resets
+// that slice to empty. Callers run it before starting a query, so invalid input is
+// rejected without a WMI round trip.
+func structSlice(dst any) (reflect.Value, error) {
 	dv := reflect.ValueOf(dst)
 	if dv.Kind() != reflect.Pointer || dv.IsNil() {
-		return ErrInvalidEntityType
+		return reflect.Value{}, ErrInvalidEntityType
 	}
 
 	dv = dv.Elem()
 
-	elemType := dv.Type().Elem()
-
-	if dv.Kind() != reflect.Slice || elemType.Kind() != reflect.Struct {
-		return ErrInvalidEntityType
+	if dv.Kind() != reflect.Slice || dv.Type().Elem().Kind() != reflect.Struct {
+		return reflect.Value{}, ErrInvalidEntityType
 	}
 
 	dv.Set(reflect.MakeSlice(dv.Type(), 0, 0))
 
+	return dv, nil
+}
+
+// unmarshal iterates over the operation's instances and appends them to dv,
+// which must come from structSlice.
+// skipMissing controls how missing elements are handled (see unmarshalInstance).
+func (o *Operation) unmarshal(dv reflect.Value, skipMissing bool) error {
+	elemType := dv.Type().Elem()
+
 	for {
 		instance, moreResults, err := o.GetInstance()
 		if err != nil {

internal/mi/session.go

@@ -209,84 +209,80 @@ func unmarshalInstance(instance *Instance, structType reflect.Type, structValue
 			return fmt.Errorf("failed to get element %s: %w", miTag, err)
 		}
 
-		switch element.valueType {
-		case ValueTypeBOOLEAN:
-			if field.Kind() != reflect.Bool {
-				return fieldTypeError(miTag, field, "boolean")
-			}
-
-			field.SetBool(element.value == 1)
-		case ValueTypeUINT8, ValueTypeUINT16, ValueTypeUINT32, ValueTypeUINT64:
-			if err := setUintField(miTag, field, uint64(element.value)); err != nil {
-				return err
-			}
-		case ValueTypeSINT8:
-			if err := setIntField(miTag, field, int64(int8(element.value))); err != nil {
-				return err
-			}
-		case ValueTypeSINT16:
-			if err := setIntField(miTag, field, int64(int16(element.value))); err != nil {
-				return err
-			}
-		case ValueTypeSINT32:
-			if err := setIntField(miTag, field, int64(int32(element.value))); err != nil {
-				return err
-			}
-		case ValueTypeSINT64:
-			if err := setIntField(miTag, field, int64(element.value)); err != nil {
-				return err
-			}
-		case ValueTypeSTRING:
-			if field.Kind() != reflect.String {
-				return fieldTypeError(miTag, field, "string")
-			}
-
-			if element.value == 0 {
-				continue
-			}
-
-			stringValue := windows.UTF16PtrToString((*uint16)(unsafe.Pointer(element.value)))
-
-			field.SetString(stringValue)
-		case ValueTypeCHAR16:
-			// A CHAR16 is a single UTF-16 code unit, not a pointer to a string.
-			// Note: lone UTF-16 surrogates (0xD800–0xDFFF) are replaced with U+FFFD
-			// by string(rune(...)). Use a uint16 Go field when the raw code unit is needed.
-			if field.Kind() == reflect.String {
-				field.SetString(string(rune(element.value)))
-
-				continue
-			}
-
-			if err := setUintField(miTag, field, uint64(element.value)); err != nil {
-				return err
-			}
-		case ValueTypeREAL32:
-			if field.Kind() != reflect.Float32 && field.Kind() != reflect.Float64 {
-				return fieldTypeError(miTag, field, "float")
-			}
-
-			field.SetFloat(float64(math.Float32frombits(uint32(element.value))))
-		case ValueTypeREAL64:
-			if field.Kind() != reflect.Float64 {
-				return fieldTypeError(miTag, field, "float")
-			}
-
-			field.SetFloat(math.Float64frombits(uint64(element.value)))
-		case ValueTypeUINT16A:
-			if field.Type() != reflect.TypeFor[[]uint16]() {
-				return fmt.Errorf("cannot unmarshal UINT16A into field of type %s, expected []uint16", field.Type())
-			}
-
-			field.Set(reflect.ValueOf(element.getUint16Array()))
-		default:
-			return fmt.Errorf("unsupported value type: %d", element.valueType)
+		if err := setField(miTag, field, element); err != nil {
+			return err
 		}
 	}
 
 	return nil
 }
 
+// setField assigns the value of element to field, rejecting Go field types that
+// cannot hold the MI value instead of letting the reflect package panic.
+func setField(miTag string, field reflect.Value, element *Element) error {
+	switch element.valueType {
+	case ValueTypeBOOLEAN:
+		if field.Kind() != reflect.Bool {
+			return fieldTypeError(miTag, field, "boolean")
+		}
+
+		field.SetBool(element.value == 1)
+	case ValueTypeUINT8, ValueTypeUINT16, ValueTypeUINT32, ValueTypeUINT64:
+		return setUintField(miTag, field, uint64(element.value))
+	case ValueTypeSINT8:
+		return setIntField(miTag, field, int64(int8(element.value)))
+	case ValueTypeSINT16:
+		return setIntField(miTag, field, int64(int16(element.value)))
+	case ValueTypeSINT32:
+		return setIntField(miTag, field, int64(int32(element.value)))
+	case ValueTypeSINT64:
+		return setIntField(miTag, field, int64(element.value))
+	case ValueTypeSTRING:
+		if field.Kind() != reflect.String {
+			return fieldTypeError(miTag, field, "string")
+		}
+
+		if element.value == 0 {
+			return nil
+		}
+
+		field.SetString(windows.UTF16PtrToString((*uint16)(unsafe.Pointer(element.value))))
+	case ValueTypeCHAR16:
+		// A CHAR16 is a single UTF-16 code unit, not a pointer to a string.
+		// Note: lone UTF-16 surrogates (0xD800–0xDFFF) are replaced with U+FFFD
+		// by string(rune(...)). Use a uint16 Go field when the raw code unit is needed.
+		if field.Kind() == reflect.String {
+			field.SetString(string(rune(uint16(element.value))))
+
+			return nil
+		}
+
+		return setUintField(miTag, field, uint64(uint16(element.value)))
+	case ValueTypeREAL32:
+		if field.Kind() != reflect.Float32 && field.Kind() != reflect.Float64 {
+			return fieldTypeError(miTag, field, "float")
+		}
+
+		field.SetFloat(float64(math.Float32frombits(uint32(element.value))))
+	case ValueTypeREAL64:
+		if field.Kind() != reflect.Float64 {
+			return fieldTypeError(miTag, field, "float")
+		}
+
+		field.SetFloat(math.Float64frombits(uint64(element.value)))
+	case ValueTypeUINT16A:
+		if field.Type() != reflect.TypeFor[[]uint16]() {
+			return fmt.Errorf("cannot unmarshal UINT16A into field of type %s, expected []uint16", field.Type())
+		}
+
+		field.Set(reflect.ValueOf(element.getUint16Array()))
+	default:
+		return fmt.Errorf("unsupported value type: %d", element.valueType)
+	}
+
+	return nil
+}
+
 // setUintField assigns value to a numeric field, rejecting fields that cannot hold it
 // instead of letting the reflect package panic or silently wrap.
 func setUintField(miTag string, field reflect.Value, value uint64) error {
@@ -366,6 +362,11 @@ func (s *Session) QueryUnmarshal(dst any,
 		return ErrNotInitialized
 	}
 
+	dv, err := structSlice(dst)
+	if err != nil {
+		return err
+	}
+
 	operation := &Operation{}
 
 	if operationOptions == nil {
@@ -392,7 +393,7 @@ func (s *Session) QueryUnmarshal(dst any,
 		_ = operation.Close()
 	}()
 
-	return operation.unmarshal(dst, true)
+	return operation.unmarshal(dv, true)
 }
 
 // Query queries for a set of instances based on a query expression.

internal/mi/unmarshal_test.go

@@ -0,0 +1,131 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build windows
+
+package mi
+
+import (
+	"math"
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// newField returns a settable reflect.Value holding the zero value of T.
+func newField[T any]() reflect.Value {
+	return reflect.New(reflect.TypeFor[T]()).Elem()
+}
+
+func TestSetField(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name    string
+		element Element
+		field   reflect.Value
+		want    any
+		wantErr bool
+	}{
+		{"boolean", Element{value: 1, valueType: ValueTypeBOOLEAN}, newField[bool](), true, false},
+		{"boolean_into_int", Element{value: 1, valueType: ValueTypeBOOLEAN}, newField[int](), nil, true},
+
+		{"uint32_into_uint64", Element{value: math.MaxUint32, valueType: ValueTypeUINT32}, newField[uint64](), uint64(math.MaxUint32), false},
+		{"uint64_overflows_uint32", Element{value: math.MaxUint32 + 1, valueType: ValueTypeUINT64}, newField[uint32](), nil, true},
+		{"uint32_into_string", Element{value: 1, valueType: ValueTypeUINT32}, newField[string](), nil, true},
+
+		{"sint8_sign_extends", Element{value: 0xFF, valueType: ValueTypeSINT8}, newField[int64](), int64(-1), false},
+		{"sint16_sign_extends", Element{value: 0xFFFF, valueType: ValueTypeSINT16}, newField[int64](), int64(-1), false},
+		{"sint32_sign_extends", Element{value: 0xFFFFFFFF, valueType: ValueTypeSINT32}, newField[int64](), int64(-1), false},
+		{"sint16_positive_into_uint16", Element{value: 42, valueType: ValueTypeSINT16}, newField[uint16](), uint16(42), false},
+		{"sint16_negative_into_uint16", Element{value: 0xFFFF, valueType: ValueTypeSINT16}, newField[uint16](), nil, true},
+		{"sint32_overflows_int8", Element{value: 1000, valueType: ValueTypeSINT32}, newField[int8](), nil, true},
+
+		{"real32_into_float32", Element{value: uintptr(math.Float32bits(1.5)), valueType: ValueTypeREAL32}, newField[float32](), float32(1.5), false},
+		{"real32_into_float64", Element{value: uintptr(math.Float32bits(-2.25)), valueType: ValueTypeREAL32}, newField[float64](), float64(-2.25), false},
+		{"real64_into_float64", Element{value: uintptr(math.Float64bits(9.875)), valueType: ValueTypeREAL64}, newField[float64](), float64(9.875), false},
+		{"real64_into_float32", Element{value: uintptr(math.Float64bits(1)), valueType: ValueTypeREAL64}, newField[float32](), nil, true},
+
+		{"string_null_keeps_zero_value", Element{value: 0, valueType: ValueTypeSTRING}, newField[string](), "", false},
+		{"string_into_int", Element{value: 0, valueType: ValueTypeSTRING}, newField[int](), nil, true},
+
+		{"char16_into_string", Element{value: 'A', valueType: ValueTypeCHAR16}, newField[string](), "A", false},
+		{"char16_euro_into_string", Element{value: 0x20AC, valueType: ValueTypeCHAR16}, newField[string](), "€", false},
+		{"char16_surrogate_into_string", Element{value: 0xD83D, valueType: ValueTypeCHAR16}, newField[string](), "�", false},
+		{"char16_into_uint16", Element{value: 0xD83D, valueType: ValueTypeCHAR16}, newField[uint16](), uint16(0xD83D), false},
+		{"char16_overflows_uint8", Element{value: 0x20AC, valueType: ValueTypeCHAR16}, newField[uint8](), nil, true},
+		{"char16_into_bool", Element{value: 'A', valueType: ValueTypeCHAR16}, newField[bool](), nil, true},
+
+		{"uint16a_into_wrong_slice", Element{value: 0, valueType: ValueTypeUINT16A}, newField[[]uint32](), nil, true},
+		{"unsupported_type", Element{value: 0, valueType: ValueTypeDATETIME}, newField[string](), nil, true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := setField("Val", tt.field, &tt.element)
+			if tt.wantErr {
+				require.Error(t, err)
+
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tt.want, tt.field.Interface())
+		})
+	}
+}
+
+func TestStructSlice(t *testing.T) {
+	t.Parallel()
+
+	type row struct {
+		Name string `mi:"Name"`
+	}
+
+	t.Run("resets_valid_slice", func(t *testing.T) {
+		t.Parallel()
+
+		dst := []row{{Name: "stale"}}
+
+		dv, err := structSlice(&dst)
+		require.NoError(t, err)
+		require.Equal(t, 0, dv.Len())
+		require.Empty(t, dst)
+	})
+
+	invalid := []struct {
+		name string
+		dst  any
+	}{
+		{"nil", nil},
+		{"slice_not_pointer", []row{}},
+		{"nil_pointer", (*[]row)(nil)},
+		{"pointer_to_int", new(int)},
+		{"pointer_to_struct", new(row)},
+		{"pointer_to_slice_of_int", new([]int)},
+	}
+
+	for _, tt := range invalid {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			_, err := structSlice(tt.dst)
+			require.ErrorIs(t, err, ErrInvalidEntityType)
+		})
+	}
+}

internal/utils/testutils/testutils.go

@@ -116,16 +116,7 @@ func TestCollector[C collector.Collector, V any](t *testing.T, fn func(*V) C, co
 // is only accepted here (Build path) so that a mistyped class in production
 // code still fails Collect.
 func unsupportedBuild(err error) bool {
-	return errors.Is(err, mi.MI_RESULT_INVALID_CLASS) ||
-		errors.Is(err, mi.MI_RESULT_INVALID_NAMESPACE) ||
-		errors.Is(err, mi.MI_RESULT_INVALID_QUERY) ||
-		errors.Is(err, pdh.NewPdhError(pdh.CstatusNoCounter)) ||
-		errors.Is(err, pdh.NewPdhError(pdh.CstatusNoObject)) ||
-		errors.Is(err, pdh.ErrPerformanceCounterNotInitialized) ||
-		errors.Is(err, pdh.ErrNoData) ||
-		errors.Is(err, update.ErrUpdateServiceDisabled) ||
-		errors.Is(err, os.ErrNotExist) ||
-		errors.Is(err, windows.Errno(2151088411))
+	return errors.Is(err, mi.MI_RESULT_INVALID_CLASS) || unsupportedCollect(err)
 }
 
 // unsupportedCollect returns true for Collect-path errors that indicate a