Skip to content

Commit

Permalink
invoke method experiment
Browse files Browse the repository at this point in the history
Signed-off-by: Michael Engl <[email protected]>
  • Loading branch information
englm committed Jan 16, 2025
1 parent f73a74b commit 99f1ed9
Show file tree
Hide file tree
Showing 4 changed files with 127 additions and 2 deletions.
54 changes: 52 additions & 2 deletions internal/collector/dhcp/dhcp.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package dhcp

import (
"errors"
"fmt"
"log/slog"

Expand All @@ -35,7 +36,8 @@ var ConfigDefaults = Config{}

// A Collector is a Prometheus Collector perflib DHCP metrics.
type Collector struct {
config Config
config Config
miSession *mi.Session

perfDataCollector *pdh.Collector
perfDataObject []perfDataCounterValues
Expand Down Expand Up @@ -67,6 +69,25 @@ type Collector struct {
requestsTotal *prometheus.Desc
}

type ScopeStatistics struct {
AddressesFree uint32 `mi:"AddressesFree"`
AddressesInUse uint32 `mi:"AddressesInUse"`
PendingOffers uint32 `mi:"PendingOffers"`
ScopeId string `mi:"ScopeId"`
SuperscopeName string `mi:"SuperscopeName"`
ReservedAddress uint32 `mi:"ReservedAddress"`
AddressesFreeOnThisServer uint32 `mi:"AddressesFreeOnThisServer"`
AddressesFreeOnPartnerServer uint32 `mi:"AddressesFreeOnPartnerServer"`
AddressesInUseOnThisServer uint32 `mi:"AddressesInUseOnThisServer"`
AddressesInUseOnPartnerServer uint32 `mi:"AddressesInUseOnPartnerServer"`
// PercentageInUse real32 `mi:PercentageInUse`
}

type ScopeStatisticsResponse struct {
CmdletOutput []ScopeStatistics `mi:"cmdletOutput"`
ReturnValue uint32 `mi:"ReturnValue"`
}

func New(config *Config) *Collector {
if config == nil {
config = &ConfigDefaults
Expand All @@ -93,9 +114,15 @@ func (c *Collector) Close() error {
return nil
}

func (c *Collector) Build(_ *slog.Logger, _ *mi.Session) error {
func (c *Collector) Build(_ *slog.Logger, miSession *mi.Session) error {
var err error

if miSession == nil {
return errors.New("miSession is nil")
}

c.miSession = miSession

c.perfDataCollector, err = pdh.NewCollector[perfDataCounterValues](pdh.CounterTypeRaw, "DHCP Server", nil)
if err != nil {
return fmt.Errorf("failed to create DHCP Server collector: %w", err)
Expand Down Expand Up @@ -261,6 +288,29 @@ func (c *Collector) Collect(ch chan<- prometheus.Metric) error {
return fmt.Errorf("failed to collect DHCP Server metrics: %w", err)
}

class, err := mi.NewClass("PS_DhcpServerv4ScopeStatistics")
if err != nil {
return err
}
method, err := mi.NewMethod("Get")
if err != nil {
return err
}

var res []ScopeStatisticsResponse

err = c.miSession.InvokeUnmarshal(
&res,
mi.OperationFlagsDefaultRTTI,
&mi.OperationOptions{},
mi.NamespaceRootWindowsDHCP,
class,
method,
)
if err != nil {
return err
}

ch <- prometheus.MustNewConstMetric(
c.packetsReceivedTotal,
prometheus.CounterValue,
Expand Down
2 changes: 2 additions & 0 deletions internal/mi/callbacks.go
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,8 @@ func (o *OperationUnmarshalCallbacks) InstanceResult(
field.SetString(stringValue)
case ValueTypeREAL32, ValueTypeREAL64:
field.SetFloat(float64(element.value))
case ValueTypeINSTANCEA:
// todo
default:
o.errCh <- fmt.Errorf("unsupported value type: %d", element.valueType)

Expand Down
60 changes: 60 additions & 0 deletions internal/mi/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -253,3 +253,63 @@ func (s *Session) Query(dst any, namespaceName Namespace, queryExpression Query)

return nil
}

// Invokes a method in the provider
//
// https://learn.microsoft.com/en-us/windows/win32/api/mi/nf-mi-mi_session_invoke
func (s *Session) InvokeUnmarshal(dst any, flags OperationFlags, operationOptions *OperationOptions, namespaceName Namespace, className Class, methodName Method) error {
if s == nil || s.ft == nil {
return ErrNotInitialized
}

operation := &Operation{}

if operationOptions == nil {
operationOptions = s.defaultOperationOptions
}

errCh := make(chan error, 1)

operationCallbacks, err := NewUnmarshalOperationsCallbacks(dst, errCh)
if err != nil {
return err
}

r0, _, _ := syscall.SyscallN(
s.ft.Invoke,
uintptr(unsafe.Pointer(s)),
uintptr(flags),
uintptr(unsafe.Pointer(operationOptions)),
uintptr(unsafe.Pointer(namespaceName)),
uintptr(unsafe.Pointer(className)),
uintptr(unsafe.Pointer(methodName)),
0,
0,
uintptr(unsafe.Pointer(operationCallbacks)),
uintptr(unsafe.Pointer(operation)),
)

if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
return result
}

errs := make([]error, 0)

// We need an active go routine to prevent a
// fatal error: all goroutines are asleep - deadlock!
// ref: https://github.com/golang/go/issues/55015
// go time.Sleep(5 * time.Second)

for {
if err, ok := <-errCh; err != nil {
errs = append(errs, err)
} else if !ok {
break
}
}

// KeepAlive is used to ensure that the callbacks are not garbage collected before the operation is closed.
runtime.KeepAlive(operationCallbacks.CallbackContext)

return errors.Join(errs...)
}
13 changes: 13 additions & 0 deletions internal/mi/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,23 @@ func NewNamespace(namespace string) (Namespace, error) {
var (
NamespaceRootCIMv2 = utils.Must(NewNamespace("root/CIMv2"))
NamespaceRootWindowsFSRM = utils.Must(NewNamespace("root/microsoft/windows/fsrm"))
NamespaceRootWindowsDHCP = utils.Must(NewNamespace("root/microsoft/windows/dhcp"))
NamespaceRootWebAdministration = utils.Must(NewNamespace("root/WebAdministration"))
NamespaceRootMSCluster = utils.Must(NewNamespace("root/MSCluster"))
)

type Method *uint16

func NewMethod(methodName string) (Method, error) {
return windows.UTF16PtrFromString(methodName)
}

type Class *uint16

func NewClass(ClassName string) (Class, error) {
return windows.UTF16PtrFromString(ClassName)
}

type Query *uint16

func NewQuery(query string) (Query, error) {
Expand Down

0 comments on commit 99f1ed9

Please sign in to comment.