headers: fail Build on missing native exports and assert MI/COM layouts (#2739)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/ole/com_windows.go

@@ -27,6 +27,7 @@ import (
 	"unicode/utf8"
 	"unsafe"
 
+	"github.com/prometheus-community/windows_exporter/internal/headers/lazyproc"
 	"golang.org/x/sys/windows"
 )
 
@@ -40,12 +41,21 @@ var (
 	sysAllocStringLen = oleaut32.NewProc("SysAllocStringLen")
 	sysFreeString     = oleaut32.NewProc("SysFreeString")
 	sysStringLen      = oleaut32.NewProc("SysStringLen")
+
+	loadOle32    = lazyproc.Once(ole32, coInitializeEx, coUninitialize, coCreateInstance)
+	loadOleaut32 = lazyproc.Once(oleaut32, sysAllocStringLen, sysFreeString, sysStringLen)
 )
 
 // Initialize initializes the MTA and disables legacy OLE 1.0 DDE. The caller
 // must lock its OS thread before calling and release all objects before
 // Uninitialize on that same thread. S_FALSE is successful initialization.
+// A missing ole32 or oleaut32 export fails Initialize instead of panicking
+// in a later call.
 func Initialize() error {
+	if err := errors.Join(loadOle32(), loadOleaut32()); err != nil {
+		return fmt.Errorf("initialize COM: %w", err)
+	}
+
 	hr, _, _ := coInitializeEx.Call(0, windows.COINIT_MULTITHREADED|windows.COINIT_DISABLE_OLE1DDE)
 	if err := ResultError(hr); err != nil {
 		return fmt.Errorf("initialize COM: %w", err)
@@ -88,15 +98,37 @@ type nativeScalar interface {
 // nativeObject requires the common COM interface layout without extra Go data.
 type nativeObject interface{ ~struct{ Object } }
 
-// Method returns the address of the zero-based native vtable slot.
-func (o *Object) Method(slot uintptr) uintptr {
-	return *(*uintptr)(unsafe.Add(unsafe.Pointer(o.VTable), slot*unsafe.Sizeof(uintptr(0))))
+// ErrNilObject reports a call on a nil interface pointer, an interface without
+// a vtable, or an empty vtable slot. Dereferencing them would crash the process
+// instead of failing the call.
+var ErrNilObject = errors.New("COM interface or vtable slot is nil")
+
+// Method returns the address of the zero-based native vtable slot. It returns
+// ErrNilObject instead of dereferencing a nil interface, a nil vtable or
+// calling an empty slot.
+func (o *Object) Method(slot uintptr) (uintptr, error) {
+	if o == nil || o.VTable == nil {
+		return 0, ErrNilObject
+	}
+
+	fn := *(*uintptr)(unsafe.Add(unsafe.Pointer(o.VTable), slot*unsafe.Sizeof(uintptr(0))))
+	if fn == 0 {
+		return 0, fmt.Errorf("COM vtable slot %d: %w", slot, ErrNilObject)
+	}
+
+	return fn, nil
 }
 
 // Release relinquishes an owned interface reference. Do not call it on borrowed
-// iterator items. Like IUnknown::Release, it does not return an HRESULT.
+// iterator items. Like IUnknown::Release, it does not return an HRESULT, and
+// releasing a nil interface does nothing.
 func (o *Object) Release() {
-	_ = Call(o.Method(2), uintptr(unsafe.Pointer(o)))
+	fn, err := o.Method(2)
+	if err != nil {
+		return
+	}
+
+	_ = Call(fn, uintptr(unsafe.Pointer(o)))
 	runtime.KeepAlive(o)
 }
 
@@ -108,7 +140,12 @@ func (o *Object) Release() {
 func (o *Object) Get[T nativeScalar](slot uintptr) (T, error) {
 	var value T
 
-	hr := Call(o.Method(slot), uintptr(unsafe.Pointer(o)), uintptr(unsafe.Pointer(&value)))
+	fn, err := o.Method(slot)
+	if err != nil {
+		return value, err
+	}
+
+	hr := Call(fn, uintptr(unsafe.Pointer(o)), uintptr(unsafe.Pointer(&value)))
 	runtime.KeepAlive(o)
 
 	return value, ResultError(hr)
@@ -122,8 +159,13 @@ func (o *Object) Get[T nativeScalar](slot uintptr) (T, error) {
 func (o *Object) GetArg[T nativeScalar](slot, arg uintptr) (T, error) {
 	var value T
 
+	fn, err := o.Method(slot)
+	if err != nil {
+		return value, err
+	}
+
 	hr := Call(
-		o.Method(slot),
+		fn,
 		uintptr(unsafe.Pointer(o)),
 		arg,
 		uintptr(unsafe.Pointer(&value)),
@@ -138,7 +180,12 @@ func (o *Object) GetArg[T nativeScalar](slot, arg uintptr) (T, error) {
 func (o *Object) GetObject[T nativeObject](slot uintptr) (*T, error) {
 	var value *T
 
-	hr := Call(o.Method(slot), uintptr(unsafe.Pointer(o)), uintptr(unsafe.Pointer(&value)))
+	fn, err := o.Method(slot)
+	if err != nil {
+		return nil, err
+	}
+
+	hr := Call(fn, uintptr(unsafe.Pointer(o)), uintptr(unsafe.Pointer(&value)))
 	runtime.KeepAlive(o)
 
 	if err := ResultError(hr); err != nil {
@@ -160,8 +207,13 @@ func (o *Object) GetObject[T nativeObject](slot uintptr) (*T, error) {
 func (o *Object) GetObjectArg[T nativeObject](slot, arg uintptr) (*T, error) {
 	var value *T
 
+	fn, err := o.Method(slot)
+	if err != nil {
+		return nil, err
+	}
+
 	hr := Call(
-		o.Method(slot),
+		fn,
 		uintptr(unsafe.Pointer(o)),
 		arg,
 		uintptr(unsafe.Pointer(&value)),
@@ -181,7 +233,12 @@ func (o *Object) GetObjectArg[T nativeObject](slot, arg uintptr) (*T, error) {
 
 // Put calls a native setter with a scalar argument.
 func (o *Object) Put(slot, arg uintptr) error {
-	hr := Call(o.Method(slot), uintptr(unsafe.Pointer(o)), arg)
+	fn, err := o.Method(slot)
+	if err != nil {
+		return err
+	}
+
+	hr := Call(fn, uintptr(unsafe.Pointer(o)), arg)
 	runtime.KeepAlive(o)
 
 	return ResultError(hr)

internal/ole/com_windows_test.go

@@ -264,3 +264,39 @@ func TestVariantLayout(t *testing.T) {
 	require.Equal(t, uintptr(24), unsafe.Sizeof(Variant{}))
 	require.Equal(t, uintptr(8), unsafe.Offsetof(Variant{}.Value))
 }
+
+// TestNilObject checks that calls on a nil interface, a nil vtable or an empty
+// vtable slot fail instead of crashing the process.
+func TestNilObject(t *testing.T) {
+	var methods [8]uintptr
+
+	for name, object := range map[string]*Object{
+		"nil interface": nil,
+		"nil vtable":    {},
+		"empty slot":    {VTable: &methods[0]},
+	} {
+		t.Run(name, func(t *testing.T) {
+			_, err := object.Method(7)
+			require.ErrorIs(t, err, ErrNilObject)
+
+			_, err = object.Get[int32](7)
+			require.ErrorIs(t, err, ErrNilObject)
+
+			_, err = object.GetArg[int32](7, 0)
+			require.ErrorIs(t, err, ErrNilObject)
+
+			_, err = object.GetObject[nativeTestObject](7)
+			require.ErrorIs(t, err, ErrNilObject)
+
+			_, err = object.GetObjectArg[nativeTestObject](7, 0)
+			require.ErrorIs(t, err, ErrNilObject)
+
+			require.ErrorIs(t, object.Put(7, 0), ErrNilObject)
+
+			_, err = object.String(7)
+			require.ErrorIs(t, err, ErrNilObject)
+
+			require.NotPanics(t, object.Release)
+		})
+	}
+}

internal/ole/taskschd/taskschd_windows.go

@@ -67,8 +67,13 @@ func (s *TaskService) Connect() error {
 	domain := ole.NewEmptyVariant()
 	password := ole.NewEmptyVariant()
 
+	fn, err := s.Method(taskServiceConnect)
+	if err != nil {
+		return err
+	}
+
 	hr := ole.Call(
-		s.Method(taskServiceConnect),
+		fn,
 		uintptr(unsafe.Pointer(s)),
 		uintptr(unsafe.Pointer(server)),
 		uintptr(unsafe.Pointer(user)),