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