Skip to content

Commit ca57b91

Browse files
committed
Use WMI to implement Volume API to reduce PowerShell overhead
1 parent ab1d3ab commit ca57b91

File tree

319 files changed

+53943
-167
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

319 files changed

+53943
-167
lines changed

go.mod

+2
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,12 @@ toolchain go1.22.3
99

1010
require (
1111
github.com/Microsoft/go-winio v0.6.2
12+
github.com/go-ole/go-ole v1.3.0
1213
github.com/google/go-cmp v0.6.0
1314
github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1
1415
github.com/iancoleman/strcase v0.3.0
1516
github.com/kubernetes-csi/csi-proxy/client v1.1.3
17+
github.com/microsoft/wmi v0.25.1
1618
github.com/pkg/errors v0.9.1
1719
github.com/prometheus/client_golang v1.20.5
1820
github.com/sergi/go-diff v1.3.1

go.sum

+5
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
2525
github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
2626
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
2727
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
28+
github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE=
29+
github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78=
2830
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
2931
github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
3032
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -69,6 +71,8 @@ github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0
6971
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
7072
github.com/mauriciopoppe/gengo v0.0.0-20210525224835-9c78f58f3486 h1:+l047vEi0SyAzdVToIaAcfoY5DwwGW+OyqTdH/P3TTg=
7173
github.com/mauriciopoppe/gengo v0.0.0-20210525224835-9c78f58f3486/go.mod h1:xXv3T4UXTLta31wMhVezwVkc26OLei4hMbLeBJbPmxc=
74+
github.com/microsoft/wmi v0.25.1 h1:sQv9hCEHtW5K6yEVL78T6XGRMGxk4aTpcJwCiB5rLN0=
75+
github.com/microsoft/wmi v0.25.1/go.mod h1:1zbdSF0A+5OwTUII5p3hN7/K6KF2m3o27pSG6Y51VU8=
7276
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
7377
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
7478
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
@@ -134,6 +138,7 @@ golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5h
134138
golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
135139
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
136140
golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
141+
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
137142
golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
138143
golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
139144
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

pkg/cim/disk.go

+36
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
//go:build windows
2+
// +build windows
3+
4+
package cim
5+
6+
import (
7+
"fmt"
8+
"strconv"
9+
10+
"github.com/microsoft/wmi/pkg/base/query"
11+
"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
12+
)
13+
14+
// QueryDiskByNumber retrieves disk information for a specific disk identified by its number.
15+
//
16+
// The equivalent WMI query is:
17+
//
18+
// SELECT [selectors] FROM MSFT_Disk
19+
// WHERE DiskNumber = '<diskNumber>'
20+
//
21+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk
22+
// for the WMI class definition.
23+
func QueryDiskByNumber(diskNumber uint32, selectorList []string) (*storage.MSFT_Disk, error) {
24+
diskQuery := query.NewWmiQueryWithSelectList("MSFT_Disk", selectorList, "Number", strconv.Itoa(int(diskNumber)))
25+
instances, err := QueryInstances(WMINamespaceStorage, diskQuery)
26+
if err != nil {
27+
return nil, err
28+
}
29+
30+
disk, err := storage.NewMSFT_DiskEx1(instances[0])
31+
if err != nil {
32+
return nil, fmt.Errorf("failed to query disk %d. error: %v", diskNumber, err)
33+
}
34+
35+
return disk, nil
36+
}

pkg/cim/volume.go

+300
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,300 @@
1+
//go:build windows
2+
// +build windows
3+
4+
package cim
5+
6+
import (
7+
"fmt"
8+
"strconv"
9+
10+
"github.com/microsoft/wmi/pkg/base/query"
11+
"github.com/microsoft/wmi/pkg/errors"
12+
cim "github.com/microsoft/wmi/pkg/wmiinstance"
13+
"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
14+
)
15+
16+
// QueryVolumeByUniqueID retrieves a specific volume by its unique identifier,
17+
// returning the first volume that matches the given volume ID.
18+
//
19+
// The equivalent WMI query is:
20+
//
21+
// SELECT [selectors] FROM MSFT_Volume
22+
//
23+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-volume
24+
// for the WMI class definition.
25+
func QueryVolumeByUniqueID(volumeID string, selectorList []string) (*storage.MSFT_Volume, error) {
26+
var selectors []string
27+
selectors = append(selectors, selectorList...)
28+
selectors = append(selectors, "UniqueId")
29+
volumeQuery := query.NewWmiQueryWithSelectList("MSFT_Volume", selectors)
30+
instances, err := QueryInstances(WMINamespaceStorage, volumeQuery)
31+
if err != nil {
32+
return nil, err
33+
}
34+
35+
for _, instance := range instances {
36+
volume, err := storage.NewMSFT_VolumeEx1(instance)
37+
if err != nil {
38+
return nil, fmt.Errorf("failed to query volume (%s). error: %w", volumeID, err)
39+
}
40+
41+
uniqueID, err := volume.GetPropertyUniqueId()
42+
if err != nil {
43+
return nil, fmt.Errorf("failed to query volume unique ID (%s). error: %w", volumeID, err)
44+
}
45+
46+
if uniqueID == volumeID {
47+
return volume, nil
48+
}
49+
}
50+
51+
return nil, errors.NotFound
52+
}
53+
54+
// ListVolumes retrieves all available volumes on the system.
55+
//
56+
// The equivalent WMI query is:
57+
//
58+
// SELECT [selectors] FROM MSFT_Volume
59+
//
60+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-volume
61+
// for the WMI class definition.
62+
func ListVolumes(selectorList []string) ([]*storage.MSFT_Volume, error) {
63+
diskQuery := query.NewWmiQueryWithSelectList("MSFT_Volume", selectorList)
64+
instances, err := QueryInstances(WMINamespaceStorage, diskQuery)
65+
if IgnoreNotFound(err) != nil {
66+
return nil, err
67+
}
68+
69+
var volumes []*storage.MSFT_Volume
70+
for _, instance := range instances {
71+
volume, err := storage.NewMSFT_VolumeEx1(instance)
72+
if err != nil {
73+
return nil, fmt.Errorf("failed to query volume %v. error: %v", instance, err)
74+
}
75+
76+
volumes = append(volumes, volume)
77+
}
78+
79+
return volumes, nil
80+
}
81+
82+
// ListPartitionsOnDisk retrieves all partitions or a partition with the specified number on a disk.
83+
//
84+
// The equivalent WMI query is:
85+
//
86+
// SELECT [selectors] FROM MSFT_Partition
87+
// WHERE DiskNumber = '<diskNumber>'
88+
// AND PartitionNumber = '<partitionNumber>'
89+
//
90+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition
91+
// for the WMI class definition.
92+
func ListPartitionsOnDisk(diskNumber, partitionNumber uint32, selectorList []string) ([]*storage.MSFT_Partition, error) {
93+
filters := []*query.WmiQueryFilter{
94+
query.NewWmiQueryFilter("DiskNumber", strconv.Itoa(int(diskNumber)), query.Equals),
95+
}
96+
if partitionNumber > 0 {
97+
filters = append(filters, query.NewWmiQueryFilter("PartitionNumber", strconv.Itoa(int(partitionNumber)), query.Equals))
98+
}
99+
return ListPartitionsWithFilters(selectorList, filters...)
100+
}
101+
102+
// ListPartitionsWithFilters retrieves all partitions matching with the conditions specified by query filters.
103+
//
104+
// The equivalent WMI query is:
105+
//
106+
// SELECT [selectors] FROM MSFT_Partition
107+
// WHERE ...
108+
//
109+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition
110+
// for the WMI class definition.
111+
func ListPartitionsWithFilters(selectorList []string, filters ...*query.WmiQueryFilter) ([]*storage.MSFT_Partition, error) {
112+
partitionQuery := query.NewWmiQueryWithSelectList("MSFT_Partition", selectorList)
113+
partitionQuery.Filters = append(partitionQuery.Filters, filters...)
114+
instances, err := QueryInstances(WMINamespaceStorage, partitionQuery)
115+
if IgnoreNotFound(err) != nil {
116+
return nil, err
117+
}
118+
119+
var partitions []*storage.MSFT_Partition
120+
for _, instance := range instances {
121+
part, err := storage.NewMSFT_PartitionEx1(instance)
122+
if err != nil {
123+
return nil, fmt.Errorf("failed to query partition %v. error: %v", instance, err)
124+
}
125+
126+
partitions = append(partitions, part)
127+
}
128+
129+
return partitions, nil
130+
}
131+
132+
// ListPartitionToVolumeMappings builds a mapping between partition and volume with partition Object ID as the key.
133+
//
134+
// The equivalent WMI query is:
135+
//
136+
// SELECT [selectors] FROM MSFT_PartitionToVolume
137+
//
138+
// Partition | Volume
139+
// --------- | ------
140+
// MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
141+
//
142+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume
143+
// for the WMI class definition.
144+
func ListPartitionToVolumeMappings() (map[string]string, error) {
145+
return ListWMIInstanceMappings(WMINamespaceStorage, "MSFT_PartitionToVolume", nil,
146+
mappingObjectRefIndexer("Partition", "MSFT_Partition", "ObjectId"),
147+
mappingObjectRefIndexer("Volume", "MSFT_Volume", "ObjectId"),
148+
)
149+
}
150+
151+
// ListVolumeToPartitionMappings builds a mapping between volume and partition with volume Object ID as the key.
152+
//
153+
// The equivalent WMI query is:
154+
//
155+
// SELECT [selectors] FROM MSFT_PartitionToVolume
156+
//
157+
// Partition | Volume
158+
// --------- | ------
159+
// MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
160+
//
161+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume
162+
// for the WMI class definition.
163+
func ListVolumeToPartitionMappings() (map[string]string, error) {
164+
return ListWMIInstanceMappings(WMINamespaceStorage, "MSFT_PartitionToVolume", nil,
165+
mappingObjectRefIndexer("Volume", "MSFT_Volume", "ObjectId"),
166+
mappingObjectRefIndexer("Partition", "MSFT_Partition", "ObjectId"),
167+
)
168+
}
169+
170+
// FindPartitionsByVolume finds all partitions associated with the given volumes
171+
// using partition-to-volume mapping.
172+
func FindPartitionsByVolume(partitions []*storage.MSFT_Partition, volumes []*storage.MSFT_Volume) ([]*storage.MSFT_Partition, error) {
173+
var partitionInstances []*cim.WmiInstance
174+
for _, part := range partitions {
175+
partitionInstances = append(partitionInstances, part.WmiInstance)
176+
}
177+
178+
var volumeInstances []*cim.WmiInstance
179+
for _, volume := range volumes {
180+
volumeInstances = append(volumeInstances, volume.WmiInstance)
181+
}
182+
183+
partitionToVolumeMappings, err := ListPartitionToVolumeMappings()
184+
if err != nil {
185+
return nil, err
186+
}
187+
188+
filtered, err := FindInstancesByObjectIDMapping(partitionInstances, volumeInstances, partitionToVolumeMappings)
189+
if err != nil {
190+
return nil, err
191+
}
192+
193+
var result []*storage.MSFT_Partition
194+
for _, instance := range filtered {
195+
part, err := storage.NewMSFT_PartitionEx1(instance)
196+
if err != nil {
197+
return nil, fmt.Errorf("failed to query partition %v. error: %v", instance, err)
198+
}
199+
200+
result = append(result, part)
201+
}
202+
203+
return result, nil
204+
}
205+
206+
// FindVolumesByPartition finds all volumes associated with the given partitions
207+
// using volume-to-partition mapping.
208+
func FindVolumesByPartition(volumes []*storage.MSFT_Volume, partitions []*storage.MSFT_Partition) ([]*storage.MSFT_Volume, error) {
209+
var volumeInstances []*cim.WmiInstance
210+
for _, volume := range volumes {
211+
volumeInstances = append(volumeInstances, volume.WmiInstance)
212+
}
213+
214+
var partitionInstances []*cim.WmiInstance
215+
for _, part := range partitions {
216+
partitionInstances = append(partitionInstances, part.WmiInstance)
217+
}
218+
219+
volumeToPartitionMappings, err := ListVolumeToPartitionMappings()
220+
if err != nil {
221+
return nil, err
222+
}
223+
224+
filtered, err := FindInstancesByObjectIDMapping(volumeInstances, partitionInstances, volumeToPartitionMappings)
225+
if err != nil {
226+
return nil, err
227+
}
228+
229+
var result []*storage.MSFT_Volume
230+
for _, instance := range filtered {
231+
volume, err := storage.NewMSFT_VolumeEx1(instance)
232+
if err != nil {
233+
return nil, fmt.Errorf("failed to query volume %v. error: %v", instance, err)
234+
}
235+
236+
result = append(result, volume)
237+
}
238+
239+
return result, nil
240+
}
241+
242+
// GetPartitionByVolumeUniqueID retrieves a specific partition from a volume identified by its unique ID.
243+
func GetPartitionByVolumeUniqueID(volumeID string, partitionSelectorList []string) (*storage.MSFT_Partition, error) {
244+
volume, err := QueryVolumeByUniqueID(volumeID, []string{"ObjectId"})
245+
if err != nil {
246+
return nil, err
247+
}
248+
249+
partitions, err := ListPartitionsWithFilters(partitionSelectorList)
250+
if err != nil {
251+
return nil, err
252+
}
253+
254+
result, err := FindPartitionsByVolume(partitions, []*storage.MSFT_Volume{volume})
255+
if err != nil {
256+
return nil, err
257+
}
258+
259+
return result[0], nil
260+
}
261+
262+
// GetVolumeByDriveLetter retrieves a volume associated with a specific drive letter.
263+
func GetVolumeByDriveLetter(driveLetter string, partitionSelectorList []string) (*storage.MSFT_Volume, error) {
264+
var selectorsForPart []string
265+
selectorsForPart = append(selectorsForPart, partitionSelectorList...)
266+
selectorsForPart = append(selectorsForPart, "ObjectId")
267+
partitions, err := ListPartitionsWithFilters(selectorsForPart, query.NewWmiQueryFilter("DriveLetter", driveLetter, query.Equals))
268+
if err != nil {
269+
return nil, err
270+
}
271+
272+
volumes, err := ListVolumes(partitionSelectorList)
273+
if err != nil {
274+
return nil, err
275+
}
276+
277+
result, err := FindVolumesByPartition(volumes, partitions)
278+
if err != nil {
279+
return nil, err
280+
}
281+
282+
if len(result) == 0 {
283+
return nil, errors.NotFound
284+
}
285+
286+
return result[0], nil
287+
}
288+
289+
// GetPartitionDiskNumber retrieves the disk number associated with a given partition.
290+
//
291+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition
292+
// for the WMI class definitions.
293+
func GetPartitionDiskNumber(part *storage.MSFT_Partition) (uint32, error) {
294+
diskNumber, err := part.GetProperty("DiskNumber")
295+
if err != nil {
296+
return 0, err
297+
}
298+
299+
return uint32(diskNumber.(int32)), nil
300+
}

0 commit comments

Comments
 (0)