mi: use generic methods for Query, QueryUnmarshal and Unmarshal (#2518)

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

internal/mi/operation.go

@@ -215,45 +215,66 @@ func (o *Operation) GetInstance() (*Instance, bool, error) {
 	return instance, moreResults == True, nil
 }
 
-func (o *Operation) Unmarshal(dst any) error {
+func (o *Operation) Unmarshal[T any](dst *[]T) error {
 	if o == nil || o.ft == nil {
 		return ErrNotInitialized
 	}
 
-	dv, err := structSlice(dst)
+	fields, err := prepareUnmarshal(dst)
 	if err != nil {
 		return err
 	}
 
-	return o.unmarshal(dv, false)
+	return o.unmarshal(dst, fields, 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 reflect.Value{}, ErrInvalidEntityType
-	}
-
-	dv = dv.Elem()
-
-	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
+// miField is a struct field that is populated from the MI element named by its `mi` tag.
+type miField struct {
+	index int
+	tag   string
 }
 
-// unmarshal iterates over the operation's instances and appends them to dv,
-// which must come from structSlice.
+// prepareUnmarshal checks that dst is non-nil and T is a struct, resets *dst to empty
+// and returns the `mi`-tagged fields of T. Callers run it before starting a query,
+// so invalid input is rejected without a WMI round trip.
+func prepareUnmarshal[T any](dst *[]T) ([]miField, error) {
+	if dst == nil {
+		return nil, ErrInvalidEntityType
+	}
+
+	fields, err := miFieldsOf[T]()
+	if err != nil {
+		return nil, err
+	}
+
+	*dst = (*dst)[:0]
+
+	return fields, nil
+}
+
+// miFieldsOf returns the fields of T that carry an `mi` tag.
+// It returns ErrInvalidEntityType if T is not a struct.
+func miFieldsOf[T any]() ([]miField, error) {
+	structType := reflect.TypeFor[T]()
+	if structType.Kind() != reflect.Struct {
+		return nil, ErrInvalidEntityType
+	}
+
+	fields := make([]miField, 0, structType.NumField())
+
+	for i := range structType.NumField() {
+		if miTag := structType.Field(i).Tag.Get("mi"); miTag != "" {
+			fields = append(fields, miField{index: i, tag: miTag})
+		}
+	}
+
+	return fields, nil
+}
+
+// unmarshal iterates over the operation's instances and appends them to dst.
+// fields must come from prepareUnmarshal.
 // skipMissing controls how missing elements are handled (see unmarshalInstance).
-func (o *Operation) unmarshal(dv reflect.Value, skipMissing bool) error {
-	elemType := dv.Type().Elem()
-
+func (o *Operation) unmarshal[T any](dst *[]T, fields []miField, skipMissing bool) error {
 	for {
 		instance, moreResults, err := o.GetInstance()
 		if err != nil {
@@ -274,13 +295,13 @@ func (o *Operation) unmarshal(dv reflect.Value, skipMissing bool) error {
 			break
 		}
 
-		elemValue := reflect.New(elemType).Elem()
+		var elem T
 
-		if err := unmarshalInstance(instance, elemType, elemValue, skipMissing); err != nil {
+		if err := unmarshalInstance(instance, fields, reflect.ValueOf(&elem).Elem(), skipMissing); err != nil {
 			return err
 		}
 
-		dv.Set(reflect.Append(dv, elemValue))
+		*dst = append(*dst, elem)
 
 		if !moreResults {
 			break

internal/mi/session.go

@@ -160,6 +160,14 @@ func (s *Session) QueryInstances(flags OperationFlags, operationOptions *Operati
 		return nil, err
 	}
 
+	return s.queryInstances(flags, operationOptions, namespaceName, queryDialect, queryExpressionUTF16)
+}
+
+// queryInstances starts an MI_Session_QueryInstances operation. Nil operationOptions
+// fall back to the session defaults. The caller must close the returned operation.
+func (s *Session) queryInstances(flags OperationFlags, operationOptions *OperationOptions, namespaceName Namespace,
+	queryDialect QueryDialect, queryExpression Query,
+) (*Operation, error) {
 	operation := &Operation{}
 
 	if operationOptions == nil {
@@ -173,7 +181,7 @@ func (s *Session) QueryInstances(flags OperationFlags, operationOptions *Operati
 		uintptr(unsafe.Pointer(operationOptions)),
 		uintptr(unsafe.Pointer(namespaceName)),
 		uintptr(unsafe.Pointer(queryDialect)),
-		uintptr(unsafe.Pointer(queryExpressionUTF16)),
+		uintptr(unsafe.Pointer(queryExpression)),
 		0,
 		uintptr(unsafe.Pointer(operation)),
 	)
@@ -191,25 +199,18 @@ func (s *Session) QueryInstances(flags OperationFlags, operationOptions *Operati
 // QueryUnmarshal passes true, tolerating classes whose MOF gained fields the query did not
 // select. Operation.Unmarshal passes false, so a missing element is reported as an error
 // instead of being silently left at its zero value.
-func unmarshalInstance(instance *Instance, structType reflect.Type, structValue reflect.Value, skipMissing bool) error {
-	for i := range structType.NumField() {
-		field := structValue.Field(i)
-
-		miTag := structType.Field(i).Tag.Get("mi")
-		if miTag == "" {
-			continue
-		}
-
-		element, err := instance.GetElement(miTag)
+func unmarshalInstance(instance *Instance, fields []miField, structValue reflect.Value, skipMissing bool) error {
+	for _, f := range fields {
+		element, err := instance.GetElement(f.tag)
 		if err != nil {
 			if skipMissing && errors.Is(err, MI_RESULT_NO_SUCH_PROPERTY) {
 				continue
 			}
 
-			return fmt.Errorf("failed to get element %s: %w", miTag, err)
+			return fmt.Errorf("failed to get element %s: %w", f.tag, err)
 		}
 
-		if err := setField(miTag, field, element); err != nil {
+		if err := setField(f.tag, structValue.Field(f.index), element); err != nil {
 			return err
 		}
 	}
@@ -354,7 +355,7 @@ func fieldTypeError(miTag string, field reflect.Value, miType string) error {
 // QueryUnmarshal queries for a set of instances based on a query expression.
 //
 // https://learn.microsoft.com/en-us/windows/win32/api/mi/nf-mi-mi_session_queryinstances
-func (s *Session) QueryUnmarshal(dst any,
+func (s *Session) QueryUnmarshal[T any](dst *[]T,
 	flags OperationFlags, operationOptions *OperationOptions,
 	namespaceName Namespace, queryDialect QueryDialect, queryExpression Query,
 ) error {
@@ -362,42 +363,25 @@ func (s *Session) QueryUnmarshal(dst any,
 		return ErrNotInitialized
 	}
 
-	dv, err := structSlice(dst)
+	fields, err := prepareUnmarshal(dst)
 	if err != nil {
 		return err
 	}
 
-	operation := &Operation{}
-
-	if operationOptions == nil {
-		operationOptions = s.defaultOperationOptions
-	}
-
-	r0, _, _ := syscall.SyscallN(
-		s.ft.QueryInstances,
-		uintptr(unsafe.Pointer(s)),
-		uintptr(flags),
-		uintptr(unsafe.Pointer(operationOptions)),
-		uintptr(unsafe.Pointer(namespaceName)),
-		uintptr(unsafe.Pointer(queryDialect)),
-		uintptr(unsafe.Pointer(queryExpression)),
-		0,
-		uintptr(unsafe.Pointer(operation)),
-	)
-
-	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
-		return fmt.Errorf("failed to query instances: %w", result)
+	operation, err := s.queryInstances(flags, operationOptions, namespaceName, queryDialect, queryExpression)
+	if err != nil {
+		return fmt.Errorf("failed to query instances: %w", err)
 	}
 
 	defer func() {
 		_ = operation.Close()
 	}()
 
-	return operation.unmarshal(dv, true)
+	return operation.unmarshal(dst, fields, true)
 }
 
 // Query queries for a set of instances based on a query expression.
-func (s *Session) Query(dst any, namespaceName Namespace, queryExpression Query, queryTimeout time.Duration) error {
+func (s *Session) Query[T any](dst *[]T, namespaceName Namespace, queryExpression Query, queryTimeout time.Duration) error {
 	operationOptions, err := s.newOperationOptions(queryTimeout)
 	if err != nil {
 		return err
@@ -434,26 +418,11 @@ func (s *Session) QueryFunc(namespaceName Namespace, queryExpression Query, quer
 		defer func() {
 			_ = operationOptions.Delete()
 		}()
-	} else {
-		operationOptions = s.defaultOperationOptions
 	}
 
-	operation := &Operation{}
-
-	r0, _, _ := syscall.SyscallN(
-		s.ft.QueryInstances,
-		uintptr(unsafe.Pointer(s)),
-		uintptr(OperationFlagsStandardRTTI),
-		uintptr(unsafe.Pointer(operationOptions)),
-		uintptr(unsafe.Pointer(namespaceName)),
-		uintptr(unsafe.Pointer(QueryDialectWQL)),
-		uintptr(unsafe.Pointer(queryExpression)),
-		0,
-		uintptr(unsafe.Pointer(operation)),
-	)
-
-	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
-		return fmt.Errorf("failed to query instances: %w", result)
+	operation, err := s.queryInstances(OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
+	if err != nil {
+		return fmt.Errorf("failed to query instances: %w", err)
 	}
 
 	// Close cancels the operation if fn aborted it early.

internal/mi/unmarshal_test.go

@@ -90,11 +90,13 @@ func TestSetField(t *testing.T) {
 	}
 }
 
-func TestStructSlice(t *testing.T) {
+func TestPrepareUnmarshal(t *testing.T) {
 	t.Parallel()
 
 	type row struct {
-		Name string `mi:"Name"`
+		Name     string `mi:"Name"`
+		Internal int
+		ID       uint32 `mi:"ProcessId"`
 	}
 
 	t.Run("resets_valid_slice", func(t *testing.T) {
@@ -102,30 +104,30 @@ func TestStructSlice(t *testing.T) {
 
 		dst := []row{{Name: "stale"}}
 
-		dv, err := structSlice(&dst)
+		fields, err := prepareUnmarshal(&dst)
 		require.NoError(t, err)
-		require.Equal(t, 0, dv.Len())
 		require.Empty(t, dst)
+		require.Equal(t, []miField{{index: 0, tag: "Name"}, {index: 2, tag: "ProcessId"}}, fields)
 	})
 
-	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)},
-	}
+	t.Run("nil_pointer", func(t *testing.T) {
+		t.Parallel()
 
-	for _, tt := range invalid {
-		t.Run(tt.name, func(t *testing.T) {
-			t.Parallel()
+		_, err := prepareUnmarshal((*[]row)(nil))
+		require.ErrorIs(t, err, ErrInvalidEntityType)
+	})
 
-			_, err := structSlice(tt.dst)
-			require.ErrorIs(t, err, ErrInvalidEntityType)
-		})
-	}
+	t.Run("slice_of_int", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := prepareUnmarshal(new([]int))
+		require.ErrorIs(t, err, ErrInvalidEntityType)
+	})
+
+	t.Run("slice_of_struct_pointer", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := prepareUnmarshal(new([]*row))
+		require.ErrorIs(t, err, ErrInvalidEntityType)
+	})
 }