Skip to content

Commit 384f0db

Browse files
andyzhangxandyzhangx
committed
feat: Use WMI to implement Volume API to reduce PowerShell overhead
1 parent a1dfb14 commit 384f0db

File tree

3 files changed

+695
-0
lines changed

3 files changed

+695
-0
lines changed

‎pkg/os/cim/disk.go

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

‎pkg/os/cim/volume.go

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

0 commit comments

Comments
 (0)