/
pdhmocks_windows.go
214 lines (189 loc) · 6.11 KB
/
pdhmocks_windows.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.
//go:build windows
package pdhutil
import (
"fmt"
"regexp"
"strings"
)
var activeAvailableCounters AvailableCounters
type mockCounter struct {
path string
machine string
class string
instance string
counter string
}
type mockQuery struct {
counters map[int]mockCounter
}
var openQueries = make(map[int]mockQuery)
var openQueriesIndex = 0
var counterIndex = 0 // index of counter into the query must be global, because
// RemoveCounter can be called on just a counter index
// class -> counter -> instance -> values
var counterValues = make(map[string]map[string]map[string][]float64)
func mockCounterFromString(path string) mockCounter {
// Example: \\.\LogicalDisk(HarddiskVolume2)\Current Disk Queue Length
// Example: \\.\Memory\Available Bytes
r := regexp.MustCompile(`\\\\([^\\]+)\\([^\\\(]+)(?:\(([^\\\)]+)\))?\\(.+)`)
res := r.FindStringSubmatch(path)
return mockCounter{
path: path,
machine: res[1],
class: res[2],
instance: res[3],
counter: res[4],
}
}
func mockPdhOpenQuery(szDataSource uintptr, dwUserData uintptr, phQuery *PDH_HQUERY) uint32 { //nolint:revive // TODO fix revive unused-parameter
var mq mockQuery
mq.counters = make(map[int]mockCounter)
openQueriesIndex++
openQueries[openQueriesIndex] = mq
*phQuery = PDH_HQUERY(uintptr(openQueriesIndex))
return 0
}
func mockPdhAddEnglishCounter(hQuery PDH_HQUERY, szFullCounterPath string, dwUserData uintptr, phCounter *PDH_HCOUNTER) uint32 { //nolint:revive // TODO fix revive unused-parameter
ndx := int(hQuery)
var thisQuery mockQuery
var ok bool
if thisQuery, ok = openQueries[ndx]; !ok {
return uint32(PDH_INVALID_PATH)
}
counterIndex++
mc := mockCounterFromString(szFullCounterPath)
thisQuery.counters[counterIndex] = mc
*phCounter = PDH_HCOUNTER(uintptr(counterIndex))
return 0
}
func mockPdhCollectQueryData(hQuery PDH_HQUERY) uint32 { //nolint:revive // TODO fix revive unused-parameter
return 0
}
func mockPdhCloseQuery(hQuery PDH_HQUERY) uint32 {
iQuery := int(hQuery)
if _, ok := openQueries[iQuery]; !ok {
return PDH_INVALID_HANDLE
}
delete(openQueries, iQuery)
return 0
}
func mockCounterFromHandle(hCounter PDH_HCOUNTER) (mockCounter, error) {
// check to see that it's a valid counter
ndx := int(hCounter)
var ctr mockCounter
var ok bool
for _, query := range openQueries {
if ctr, ok = query.counters[ndx]; ok {
break
}
}
if !ok {
return ctr, fmt.Errorf("Invalid handle")
}
return ctr, nil
}
func mockPdhGetFormattedCounterArray(hCounter PDH_HCOUNTER, format uint32) (outItems []PdhCounterValueItem, err error) { //nolint:revive // TODO fix revive unused-parameter
ctr, err := mockCounterFromHandle(hCounter)
if err != nil {
return nil, err
}
if classMap, ok := counterValues[ctr.class]; ok {
if instMap, ok := classMap[ctr.counter]; ok {
for inst, vals := range instMap {
if len(vals) > 0 {
outItems = append(outItems,
PdhCounterValueItem{
instance: inst,
value: PdhCounterValue{
CStatus: PDH_CSTATUS_NEW_DATA,
Double: vals[0],
},
},
)
instMap[inst] = vals[1:]
}
}
return outItems, nil
}
}
return nil, NewErrPdhInvalidInstance("Invalid counter instance")
}
func mockpdhGetFormattedCounterValueFloat(hCounter PDH_HCOUNTER) (val float64, err error) {
ctr, err := mockCounterFromHandle(hCounter)
if err != nil {
return 0, err
}
if classMap, ok := counterValues[ctr.class]; ok {
if instMap, ok := classMap[ctr.counter]; ok {
if vals, ok := instMap[ctr.instance]; ok {
if len(vals) > 0 {
val, instMap[ctr.instance] = vals[0], vals[1:]
return val, nil
}
}
}
}
return 0, NewErrPdhInvalidInstance("Invalid counter instance")
}
func mockpdhMakeCounterPath(machine string, object string, instance string, counter string) (path string, err error) {
var inst string
if len(instance) != 0 {
inst = fmt.Sprintf("(%s)", instance)
}
if len(machine) == 0 {
machine = "."
}
path = fmt.Sprintf("\\\\%s\\%s%s\\%s", machine, object, inst, counter)
return
}
// SetupTesting initializes the PDH libarary with the mock functions rather than the real thing
func SetupTesting(counterstringsfile, countersfile string) { //nolint:revive // TODO fix revive unused-parameter
activeAvailableCounters, _ = ReadCounters(countersfile)
// For testing
pfnPdhOpenQuery = mockPdhOpenQuery
pfnPdhAddEnglishCounter = mockPdhAddEnglishCounter
pfnPdhCollectQueryData = mockPdhCollectQueryData
pfnPdhGetFormattedCounterValueFloat = mockpdhGetFormattedCounterValueFloat
pfnPdhGetFormattedCounterArray = mockPdhGetFormattedCounterArray
pfnPdhCloseQuery = mockPdhCloseQuery
pfnPdhMakeCounterPath = mockpdhMakeCounterPath
}
// SetQueryReturnValue provides an entry point for tests to set expected values for a
// given counter
func SetQueryReturnValue(counter string, val float64) {
mc := mockCounterFromString(counter)
// class -> counter
counterMap, ok := counterValues[mc.class]
if !ok {
counterMap = make(map[string]map[string][]float64)
counterValues[mc.class] = counterMap
}
// counter -> instance
instMap, ok := counterMap[mc.counter]
if !ok {
instMap = make(map[string][]float64)
counterMap[mc.counter] = instMap
}
// instance -> value list
instMap[mc.instance] = append(instMap[mc.instance], val)
}
// RemoveCounterInstance removes a specific instance from the table of available instances
func RemoveCounterInstance(clss, inst string) {
for idx, val := range activeAvailableCounters.instancesByClass[clss] {
if strings.EqualFold(inst, val) {
activeAvailableCounters.instancesByClass[clss] =
append(activeAvailableCounters.instancesByClass[clss][:idx],
activeAvailableCounters.instancesByClass[clss][idx+1:]...)
return
}
}
}
// AddCounterInstance adds a specific instance to the table of available instances
func AddCounterInstance(clss, inst string) {
activeAvailableCounters.instancesByClass[clss] =
append(activeAvailableCounters.instancesByClass[clss], inst)
}