diff --git a/internal/mi/operation.go b/internal/mi/operation.go index 5c67dd0d..f6d5639f 100644 --- a/internal/mi/operation.go +++ b/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 diff --git a/internal/mi/session.go b/internal/mi/session.go index be6afa07..50c0f34c 100644 --- a/internal/mi/session.go +++ b/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. diff --git a/internal/mi/unmarshal_test.go b/internal/mi/unmarshal_test.go index d79e742e..03f17dcd 100644 --- a/internal/mi/unmarshal_test.go +++ b/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) + }) }