1// +build windows
2
3package common
4
5import (
6	"context"
7	"unsafe"
8
9	"github.com/StackExchange/wmi"
10	"golang.org/x/sys/windows"
11)
12
13// for double values
14type PDH_FMT_COUNTERVALUE_DOUBLE struct {
15	CStatus     uint32
16	DoubleValue float64
17}
18
19// for 64 bit integer values
20type PDH_FMT_COUNTERVALUE_LARGE struct {
21	CStatus    uint32
22	LargeValue int64
23}
24
25// for long values
26type PDH_FMT_COUNTERVALUE_LONG struct {
27	CStatus   uint32
28	LongValue int32
29	padding   [4]byte
30}
31
32// windows system const
33const (
34	ERROR_SUCCESS        = 0
35	ERROR_FILE_NOT_FOUND = 2
36	DRIVE_REMOVABLE      = 2
37	DRIVE_FIXED          = 3
38	HKEY_LOCAL_MACHINE   = 0x80000002
39	RRF_RT_REG_SZ        = 0x00000002
40	RRF_RT_REG_DWORD     = 0x00000010
41	PDH_FMT_LONG         = 0x00000100
42	PDH_FMT_DOUBLE       = 0x00000200
43	PDH_FMT_LARGE        = 0x00000400
44	PDH_INVALID_DATA     = 0xc0000bc6
45	PDH_INVALID_HANDLE   = 0xC0000bbc
46	PDH_NO_DATA          = 0x800007d5
47)
48
49var (
50	Modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
51	ModNt       = windows.NewLazySystemDLL("ntdll.dll")
52	ModPdh      = windows.NewLazySystemDLL("pdh.dll")
53	ModPsapi    = windows.NewLazySystemDLL("psapi.dll")
54
55	ProcGetSystemTimes           = Modkernel32.NewProc("GetSystemTimes")
56	ProcNtQuerySystemInformation = ModNt.NewProc("NtQuerySystemInformation")
57	PdhOpenQuery                 = ModPdh.NewProc("PdhOpenQuery")
58	PdhAddCounter                = ModPdh.NewProc("PdhAddCounterW")
59	PdhCollectQueryData          = ModPdh.NewProc("PdhCollectQueryData")
60	PdhGetFormattedCounterValue  = ModPdh.NewProc("PdhGetFormattedCounterValue")
61	PdhCloseQuery                = ModPdh.NewProc("PdhCloseQuery")
62)
63
64type FILETIME struct {
65	DwLowDateTime  uint32
66	DwHighDateTime uint32
67}
68
69// borrowed from net/interface_windows.go
70func BytePtrToString(p *uint8) string {
71	a := (*[10000]uint8)(unsafe.Pointer(p))
72	i := 0
73	for a[i] != 0 {
74		i++
75	}
76	return string(a[:i])
77}
78
79// CounterInfo
80// copied from https://github.com/mackerelio/mackerel-agent/
81type CounterInfo struct {
82	PostName    string
83	CounterName string
84	Counter     windows.Handle
85}
86
87// CreateQuery XXX
88// copied from https://github.com/mackerelio/mackerel-agent/
89func CreateQuery() (windows.Handle, error) {
90	var query windows.Handle
91	r, _, err := PdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&query)))
92	if r != 0 {
93		return 0, err
94	}
95	return query, nil
96}
97
98// CreateCounter XXX
99func CreateCounter(query windows.Handle, pname, cname string) (*CounterInfo, error) {
100	var counter windows.Handle
101	r, _, err := PdhAddCounter.Call(
102		uintptr(query),
103		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(cname))),
104		0,
105		uintptr(unsafe.Pointer(&counter)))
106	if r != 0 {
107		return nil, err
108	}
109	return &CounterInfo{
110		PostName:    pname,
111		CounterName: cname,
112		Counter:     counter,
113	}, nil
114}
115
116// WMIQueryWithContext - wraps wmi.Query with a timed-out context to avoid hanging
117func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, connectServerArgs ...interface{}) error {
118	if _, ok := ctx.Deadline(); !ok {
119		ctxTimeout, cancel := context.WithTimeout(ctx, Timeout)
120		defer cancel()
121		ctx = ctxTimeout
122	}
123
124	errChan := make(chan error, 1)
125	go func() {
126		errChan <- wmi.Query(query, dst, connectServerArgs...)
127	}()
128
129	select {
130	case <-ctx.Done():
131		return ctx.Err()
132	case err := <-errChan:
133		return err
134	}
135}
136