Skip to content

Commit 13f6143

Browse files
laozclaozc
committed
Use associator to find WMI instances for volume APIs
1 parent 24f2ecd commit 13f6143

File tree

4 files changed

+72
-232
lines changed

4 files changed

+72
-232
lines changed

‎docs/IMPLEMENTATION.md‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,20 @@ func CallMethod(disp *ole.IDispatch, name string, params ...interface{}) (result
121121
}
122122
```
123123

124+
### Association
125+
126+
Association can be used to retrieve all instances that are associated with
127+
a particular source instance.
128+
129+
There are a few Association classes in WMI.
130+
131+
For example, association class [MSFT_PartitionToVolume](https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume)
132+
can be used to retrieve a volume (`MSFT_Volume`) from a partition (`MSFT_Partition`), and vice versa.
133+
134+
```go
135+
collection, err := part.GetAssociated("MSFT_PartitionToVolume", "MSFT_Volume", "Volume", "Partition")
136+
```
137+
124138
<a name="debug-powershell"></a>
125139
## Debug with PowerShell
126140

@@ -181,6 +195,13 @@ PS C:\Users\Administrator> $vol.FileSystem
181195
NTFS
182196
```
183197

198+
### Association
199+
200+
```powershell
201+
PS C:\Users\Administrator> $partition = (Get-CimInstance -Namespace root\Microsoft\Windows\Storage -ClassName MSFT_Partition -Filter "DiskNumber = 0")[0]
202+
PS C:\Users\Administrator> Get-CimAssociatedInstance -InputObject $partition -Association MSFT_PartitionToVolume
203+
```
204+
184205
### Call Class Method
185206

186207
You may get Class Methods for a single CIM class using `$class.CimClassMethods`.

‎pkg/cim/volume.go‎

Lines changed: 45 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import (
99

1010
"github.com/microsoft/wmi/pkg/base/query"
1111
"github.com/microsoft/wmi/pkg/errors"
12-
cim "github.com/microsoft/wmi/pkg/wmiinstance"
1312
"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
1413
)
1514

@@ -129,131 +128,84 @@ func ListPartitionsWithFilters(selectorList []string, filters ...*query.WmiQuery
129128
return partitions, nil
130129
}
131130

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:
131+
// FindPartitionsByVolume finds all partitions associated with the given volumes
132+
// using MSFT_PartitionToVolume association.
154133
//
155-
// SELECT [selectors] FROM MSFT_PartitionToVolume
134+
// WMI association MSFT_PartitionToVolume:
156135
//
157-
// Partition | Volume
158-
// --------- | ------
159-
// MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
136+
// Partition | Volume
137+
// --------- | ------
138+
// MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
160139
//
161140
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume
162141
// 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-
142+
func FindPartitionsByVolume(volumes []*storage.MSFT_Volume) ([]*storage.MSFT_Partition, error) {
193143
var result []*storage.MSFT_Partition
194-
for _, instance := range filtered {
195-
part, err := storage.NewMSFT_PartitionEx1(instance)
144+
for _, vol := range volumes {
145+
collection, err := vol.GetAssociated("MSFT_PartitionToVolume", "MSFT_Partition", "Partition", "Volume")
196146
if err != nil {
197-
return nil, fmt.Errorf("failed to query partition %v. error: %v", instance, err)
147+
return nil, fmt.Errorf("failed to query associated partition for %v. error: %v", vol, err)
198148
}
199149

200-
result = append(result, part)
150+
for _, instance := range collection {
151+
part, err := storage.NewMSFT_PartitionEx1(instance)
152+
if err != nil {
153+
return nil, fmt.Errorf("failed to query partition %v. error: %v", instance, err)
154+
}
155+
156+
result = append(result, part)
157+
}
201158
}
202159

203160
return result, nil
204161
}
205162

206163
// 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-
164+
// using MSFT_PartitionToVolume association.
165+
//
166+
// WMI association MSFT_PartitionToVolume:
167+
//
168+
// Partition | Volume
169+
// --------- | ------
170+
// MSFT_Partition (ObjectId = "{1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Win...) | MSFT_Volume (ObjectId = "{1}\\WIN-8E2EVAQ9QS...
171+
//
172+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume
173+
// for the WMI class definition.
174+
func FindVolumesByPartition(partitions []*storage.MSFT_Partition) ([]*storage.MSFT_Volume, error) {
229175
var result []*storage.MSFT_Volume
230-
for _, instance := range filtered {
231-
volume, err := storage.NewMSFT_VolumeEx1(instance)
176+
for _, part := range partitions {
177+
collection, err := part.GetAssociated("MSFT_PartitionToVolume", "MSFT_Volume", "Volume", "Partition")
232178
if err != nil {
233-
return nil, fmt.Errorf("failed to query volume %v. error: %v", instance, err)
179+
return nil, fmt.Errorf("failed to query associated volumes for %v. error: %v", part, err)
234180
}
235181

236-
result = append(result, volume)
182+
for _, instance := range collection {
183+
volume, err := storage.NewMSFT_VolumeEx1(instance)
184+
if err != nil {
185+
return nil, fmt.Errorf("failed to query volume %v. error: %v", instance, err)
186+
}
187+
188+
result = append(result, volume)
189+
}
237190
}
238191

239192
return result, nil
240193
}
241194

242195
// GetPartitionByVolumeUniqueID retrieves a specific partition from a volume identified by its unique ID.
243-
func GetPartitionByVolumeUniqueID(volumeID string, partitionSelectorList []string) (*storage.MSFT_Partition, error) {
196+
func GetPartitionByVolumeUniqueID(volumeID string) (*storage.MSFT_Partition, error) {
244197
volume, err := QueryVolumeByUniqueID(volumeID, []string{"ObjectId"})
245198
if err != nil {
246199
return nil, err
247200
}
248201

249-
partitions, err := ListPartitionsWithFilters(partitionSelectorList)
202+
result, err := FindPartitionsByVolume([]*storage.MSFT_Volume{volume})
250203
if err != nil {
251204
return nil, err
252205
}
253206

254-
result, err := FindPartitionsByVolume(partitions, []*storage.MSFT_Volume{volume})
255-
if err != nil {
256-
return nil, err
207+
if len(result) == 0 {
208+
return nil, errors.NotFound
257209
}
258210

259211
return result[0], nil
@@ -269,12 +221,7 @@ func GetVolumeByDriveLetter(driveLetter string, partitionSelectorList []string)
269221
return nil, err
270222
}
271223

272-
volumes, err := ListVolumes(partitionSelectorList)
273-
if err != nil {
274-
return nil, err
275-
}
276-
277-
result, err := FindVolumesByPartition(volumes, partitions)
224+
result, err := FindVolumesByPartition(partitions)
278225
if err != nil {
279226
return nil, err
280227
}

‎pkg/cim/wmi.go‎

Lines changed: 2 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ package cim
55

66
import (
77
"fmt"
8-
"strings"
98

109
"github.com/go-ole/go-ole"
1110
"github.com/go-ole/go-ole/oleutil"
@@ -16,21 +15,18 @@ import (
1615
)
1716

1817
const (
19-
WMINamespaceRoot = "Root\\CimV2"
18+
WMINamespaceCimV2 = "Root\\CimV2"
2019
WMINamespaceStorage = "Root\\Microsoft\\Windows\\Storage"
2120
WMINamespaceSmb = "Root\\Microsoft\\Windows\\Smb"
2221
)
2322

2423
type InstanceHandler func(instance *cim.WmiInstance) (bool, error)
2524

26-
// An InstanceIndexer provides index key to a WMI Instance in a map
27-
type InstanceIndexer func(instance *cim.WmiInstance) (string, error)
28-
2925
// NewWMISession creates a new local WMI session for the given namespace, defaulting
3026
// to root namespace if none specified.
3127
func NewWMISession(namespace string) (*cim.WmiSession, error) {
3228
if namespace == "" {
33-
namespace = WMINamespaceRoot
29+
namespace = WMINamespaceCimV2
3430
}
3531

3632
sessionManager := cim.NewWmiSessionManager()
@@ -247,122 +243,3 @@ func IgnoreNotFound(err error) error {
247243
}
248244
return err
249245
}
250-
251-
// parseObjectRef extracts the object ID from a WMI object reference string.
252-
// The result string is in this format
253-
// {1}\\WIN-8E2EVAQ9QSB\ROOT/Microsoft/Windows/Storage/Providers_v2\WSP_Partition.ObjectId="{b65bb3cd-da86-11ee-854b-806e6f6e6963}:PR:{00000000-0000-0000-0000-100000000000}\\?\scsi#disk&ven_vmware&prod_virtual_disk#4&2c28f6c4&0&000000#{53f56307-b6bf-11d0-94f2-00a0c91efb8b}"
254-
// from an escape string
255-
func parseObjectRef(input, objectClass, refName string) (string, error) {
256-
tokens := strings.Split(input, fmt.Sprintf("%s.%s=", objectClass, refName))
257-
if len(tokens) < 2 {
258-
return "", fmt.Errorf("invalid object ID value: %s", input)
259-
}
260-
261-
objectID := tokens[1]
262-
objectID = strings.ReplaceAll(objectID, "\\\"", "\"")
263-
objectID = strings.ReplaceAll(objectID, "\\\\", "\\")
264-
objectID = objectID[1 : len(objectID)-1]
265-
return objectID, nil
266-
}
267-
268-
// ListWMIInstanceMappings queries WMI instances and creates a map using custom indexing functions
269-
// to extract keys and values from each instance.
270-
func ListWMIInstanceMappings(namespace, mappingClassName string, selectorList []string, keyIndexer InstanceIndexer, valueIndexer InstanceIndexer) (map[string]string, error) {
271-
q := query.NewWmiQueryWithSelectList(mappingClassName, selectorList)
272-
mappingInstances, err := QueryInstances(namespace, q)
273-
if err != nil {
274-
return nil, err
275-
}
276-
277-
result := make(map[string]string)
278-
for _, mapping := range mappingInstances {
279-
key, err := keyIndexer(mapping)
280-
if err != nil {
281-
return nil, err
282-
}
283-
284-
value, err := valueIndexer(mapping)
285-
if err != nil {
286-
return nil, err
287-
}
288-
289-
result[key] = value
290-
}
291-
292-
return result, nil
293-
}
294-
295-
// FindInstancesByMapping filters instances based on a mapping relationship,
296-
// matching instances through custom indexing and mapping functions.
297-
func FindInstancesByMapping(instanceToFind []*cim.WmiInstance, instanceToFindIndex InstanceIndexer, associatedInstances []*cim.WmiInstance, associatedInstanceIndexer InstanceIndexer, instanceMappings map[string]string) ([]*cim.WmiInstance, error) {
298-
associatedInstanceObjectIDMapping := map[string]*cim.WmiInstance{}
299-
for _, inst := range associatedInstances {
300-
key, err := associatedInstanceIndexer(inst)
301-
if err != nil {
302-
return nil, err
303-
}
304-
305-
associatedInstanceObjectIDMapping[key] = inst
306-
}
307-
308-
var filtered []*cim.WmiInstance
309-
for _, inst := range instanceToFind {
310-
key, err := instanceToFindIndex(inst)
311-
if err != nil {
312-
return nil, err
313-
}
314-
315-
valueObjectID, ok := instanceMappings[key]
316-
if !ok {
317-
continue
318-
}
319-
320-
_, ok = associatedInstanceObjectIDMapping[strings.ToUpper(valueObjectID)]
321-
if !ok {
322-
continue
323-
}
324-
filtered = append(filtered, inst)
325-
}
326-
327-
if len(filtered) == 0 {
328-
return nil, errors.NotFound
329-
}
330-
331-
return filtered, nil
332-
}
333-
334-
// mappingObjectRefIndexer indexes an WMI object by the Object ID reference from a specified property.
335-
func mappingObjectRefIndexer(propertyName, className, refName string) InstanceIndexer {
336-
return func(instance *cim.WmiInstance) (string, error) {
337-
valueVal, err := instance.GetProperty(propertyName)
338-
if err != nil {
339-
return "", err
340-
}
341-
342-
refValue, err := parseObjectRef(valueVal.(string), className, refName)
343-
return strings.ToUpper(refValue), err
344-
}
345-
}
346-
347-
// stringPropertyIndexer indexes a WMI object from a string property.
348-
func stringPropertyIndexer(propertyName string) InstanceIndexer {
349-
return func(instance *cim.WmiInstance) (string, error) {
350-
valueVal, err := instance.GetProperty(propertyName)
351-
if err != nil {
352-
return "", err
353-
}
354-
355-
return strings.ToUpper(valueVal.(string)), err
356-
}
357-
}
358-
359-
var (
360-
// objectIDPropertyIndexer indexes a WMI object from its ObjectId property.
361-
objectIDPropertyIndexer = stringPropertyIndexer("ObjectId")
362-
)
363-
364-
// FindInstancesByObjectIDMapping filters instances based on ObjectId mapping
365-
// between two sets of WMI instances.
366-
func FindInstancesByObjectIDMapping(instanceToFind []*cim.WmiInstance, associatedInstances []*cim.WmiInstance, instanceMappings map[string]string) ([]*cim.WmiInstance, error) {
367-
return FindInstancesByMapping(instanceToFind, objectIDPropertyIndexer, associatedInstances, objectIDPropertyIndexer, instanceMappings)
368-
}

0 commit comments

Comments
 (0)