Ensure COM threading apartment for API calls

Author: laozc

pkg/cim/wmi.go

@@ -4,13 +4,16 @@
 package cim
 
 import (
+	"errors"
 	"fmt"
+	"runtime"
 
 	"github.com/go-ole/go-ole"
 	"github.com/go-ole/go-ole/oleutil"
 	"github.com/microsoft/wmi/pkg/base/query"
 	wmierrors "github.com/microsoft/wmi/pkg/errors"
 	cim "github.com/microsoft/wmi/pkg/wmiinstance"
+	"golang.org/x/sys/windows"
 	"k8s.io/klog/v2"
 )
 
@@ -248,3 +251,39 @@ func IgnoreNotFound(err error) error {
 	}
 	return err
 }
+
+// WithCOMThread runs the given function `fn` on a locked OS thread
+// with COM initialized using COINIT_MULTITHREADED.
+//
+// This is necessary for using COM/OLE APIs directly (e.g., via go-ole),
+// because COM requires that initialization and usage occur on the same thread.
+//
+// It performs the following steps:
+//   - Locks the current goroutine to its OS thread
+//   - Calls ole.CoInitializeEx with COINIT_MULTITHREADED
+//   - Executes the user-provided function
+//   - Uninitializes COM
+//   - Unlocks the thread
+//
+// If COM initialization fails, or if the user's function returns an error,
+// that error is returned by WithCOMThread.
+func WithCOMThread(fn func() error) error {
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+
+	if err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED); err != nil {
+		var oleError *ole.OleError
+		if errors.As(err, &oleError) && oleError != nil && oleError.Code() == uintptr(windows.S_FALSE) {
+			klog.V(10).Infof("COM library has been already initialized for the calling thread, proceeding to the function with no error")
+			err = nil
+		}
+		if err != nil {
+			return err
+		}
+	} else {
+		klog.V(10).Infof("COM library is initialized for the calling thread")
+	}
+	defer ole.CoUninitialize()
+
+	return fn()
+}