mi: make timeout configurable for build functions (#2377)

internal/mi/application.go

@@ -205,6 +205,16 @@ func (application *Application) NewOperationOptions() (*OperationOptions, error)
 	return operationOptions, nil
 }
 
+// MustNewOperationOptions is the panicking version of NewOperationOptions.
+func (application *Application) MustNewOperationOptions() *OperationOptions {
+	operationOptions, err := application.NewOperationOptions()
+	if err != nil {
+		panic(fmt.Sprintf("failed to create operation options: %v", err))
+	}
+
+	return operationOptions
+}
+
 // NewDestinationOptions creates an DestinationOptions object that can be used with the Application.NewSession function.
 //
 // https://learn.microsoft.com/en-us/windows/win32/api/mi/nf-mi-mi_application_newdestinationoptions

internal/mi/mi_test.go

@@ -257,7 +257,7 @@ func Test_MI_FD_Leak(t *testing.T) {
 	for range 300 {
 		var processes []win32Process
 
-		err := session.Query(&processes, mi.NamespaceRootCIMv2, queryPrinter)
+		err := session.Query(&processes, mi.NamespaceRootCIMv2, queryPrinter, -1)
 		require.NoError(t, err)
 
 		currentFileHandle, err = testutils.GetProcessHandleCount(windows.CurrentProcess())
@@ -287,3 +287,48 @@ func Test_MI_FD_Leak(t *testing.T) {
 
 	t.Log("Current File Handle Count: ", currentFileHandle)
 }
+
+func Test_MI_QueryTimeout(t *testing.T) {
+	application, err := mi.ApplicationInitialize()
+	require.NoError(t, err)
+	require.NotEmpty(t, application)
+
+	destinationOptions, err := application.NewDestinationOptions()
+	require.NoError(t, err)
+	require.NotEmpty(t, destinationOptions)
+
+	err = destinationOptions.SetTimeout(1 * time.Second)
+	require.NoError(t, err)
+
+	err = destinationOptions.SetLocale(mi.LocaleEnglish)
+	require.NoError(t, err)
+
+	session, err := application.NewSession(destinationOptions)
+	require.NoError(t, err)
+	require.NotEmpty(t, session)
+
+	operationOptions, err := application.NewOperationOptions()
+	require.NoError(t, err)
+	require.NotEmpty(t, operationOptions)
+
+	err = operationOptions.SetTimeout(1 * time.Millisecond)
+	require.NoError(t, err)
+
+	operation, err := session.QueryInstances(mi.OperationFlagsStandardRTTI, operationOptions, mi.NamespaceRootCIMv2, mi.QueryDialectWQL, "select Name from win32_process where handle = 0")
+	require.NoError(t, err)
+	require.NotEmpty(t, operation)
+
+	instance, moreResults, err := operation.GetInstance()
+	require.ErrorIs(t, err, mi.MI_RESULT_INVALID_OPERATION_TIMEOUT)
+	require.False(t, moreResults)
+	require.Empty(t, instance)
+
+	err = operation.Close()
+	require.NoError(t, err)
+
+	err = session.Close()
+	require.NoError(t, err)
+
+	err = application.Close()
+	require.NoError(t, err)
+}

internal/mi/operation.go

@@ -21,6 +21,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"strings"
 	"syscall"
 	"time"
 	"unsafe"
@@ -150,8 +151,35 @@ func (o *Operation) GetInstance() (*Instance, bool, error) {
 		uintptr(unsafe.Pointer(&errorMessageUTF16)),
 		uintptr(unsafe.Pointer(&errorDetails)),
 	)
+
+	//nolint:nestif
 	if !errors.Is(instanceResult, MI_RESULT_OK) {
-		return nil, false, fmt.Errorf("instance result: %w (%s)", instanceResult, windows.UTF16PtrToString(errorMessageUTF16))
+		errorMessage := strings.TrimSpace(windows.UTF16PtrToString(errorMessageUTF16))
+
+		// We need a language neutral way to detect an operation timeout, because MI_RESULT_OPERATION_TIMED_OUT
+		// is not returned by the API, but instead we get MI_RESULT_INVALID_OPERATION_TIMEOUT with a specific error code
+		// in the error details.
+		if errorDetails != nil {
+			count, _ := errorDetails.GetElementCount()
+			if count != 0 {
+				errorCodeRaw, err := errorDetails.GetElement("error_Code")
+				if err == nil {
+					errorCodeValue, _ := errorCodeRaw.GetValue()
+
+					errorCode, ok := errorCodeValue.(uint32)
+					if ok && errorCode == 262148 {
+						instanceResult = MI_RESULT_INVALID_OPERATION_TIMEOUT
+						errorMessage = ""
+					}
+				}
+			}
+		}
+
+		if errorMessage != "" {
+			errorMessage = fmt.Sprintf(" (%s)", errorMessage)
+		}
+
+		return nil, false, fmt.Errorf("instance result: %w%s", instanceResult, errorMessage)
 	}
 
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
@@ -273,6 +301,20 @@ func (o *OperationOptions) SetTimeout(timeout time.Duration) error {
 	return nil
 }
 
+func (o *OperationOptions) Close() error {
+	if o == nil || o.ft == nil {
+		return ErrNotInitialized
+	}
+
+	r0, _, _ := syscall.SyscallN(o.ft.Clone, uintptr(unsafe.Pointer(o)))
+
+	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
+		return result
+	}
+
+	return nil
+}
+
 func (o *OperationOptions) Delete() error {
 	if o == nil || o.ft == nil {
 		return ErrNotInitialized

internal/mi/session.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"reflect"
 	"syscall"
+	"time"
 	"unsafe"
 
 	"golang.org/x/sys/windows"
@@ -229,7 +230,7 @@ func (s *Session) QueryUnmarshal(dst any,
 	)
 
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
-		return result
+		return fmt.Errorf("failed to query instances: %w", result)
 	}
 
 	defer func() {
@@ -308,11 +309,28 @@ func (s *Session) QueryUnmarshal(dst any,
 }
 
 // Query queries for a set of instances based on a query expression.
-func (s *Session) Query(dst any, namespaceName Namespace, queryExpression Query) error {
-	err := s.QueryUnmarshal(dst, OperationFlagsStandardRTTI, nil, namespaceName, QueryDialectWQL, queryExpression)
-	if err != nil {
-		return fmt.Errorf("WMI query failed: %w", err)
+//
+//nolint:nestif
+func (s *Session) Query(dst any, namespaceName Namespace, queryExpression Query, queryTimeout time.Duration) error {
+	var operationOptions *OperationOptions
+
+	if queryTimeout >= 0 {
+		app, err := s.GetApplication()
+		if err != nil {
+			return fmt.Errorf("failed to get application: %w", err)
+		}
+
+		operationOptions, err = app.NewOperationOptions()
+		if err != nil {
+			return fmt.Errorf("failed to create operation options: %w", err)
+		}
+
+		if queryTimeout > 0 {
+			if err = operationOptions.SetTimeout(queryTimeout); err != nil {
+				return fmt.Errorf("failed to set timeout: %w", err)
+			}
+		}
 	}
 
-	return nil
+	return s.QueryUnmarshal(dst, OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
 }