Skip to content

Commit

Permalink
Add better logging around enumeration process. Refactor to make easier.
Browse files Browse the repository at this point in the history
  • Loading branch information
@pwood
pwood committed Jan 14, 2024
1 parent c232727 commit 4a0b233
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 41 deletions.
111 changes: 74 additions & 37 deletions enumerate_device.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,14 +98,17 @@ func (e enumerateDevice) enumerate(pctx context.Context, n *node) {
return
}

e.logger.LogTrace(ctx, "Running rules against node.")
inv, err = e.runRules(inv)
if err != nil {
e.logger.LogError(ctx, "Failed to run rules against node.", logwrap.Err(err))
return
}

e.logger.LogTrace(ctx, "Grouping endpoints and devices.")
inventoryDevices := e.groupInventoryDevices(inv)
did := e.updateNodeTable(n, inventoryDevices)

did := e.updateNodeTable(ctx, n, inventoryDevices)

for _, id := range inventoryDevices {
d := did[id.deviceId]
Expand Down Expand Up @@ -250,7 +253,10 @@ func (e enumerateDevice) groupInventoryDevices(inv inventory) []inventoryDevice
return outDevices
}

func (e enumerateDevice) updateNodeTable(n *node, inventoryDevices []inventoryDevice) map[uint16]*device {
func (e enumerateDevice) updateNodeTable(ctx context.Context, n *node, inventoryDevices []inventoryDevice) map[uint16]*device {
ctx, end := e.logger.Segment(ctx, "Updating node table.")
defer end()

deviceIdMapping := map[uint16]*device{}
var unsetDevice []*device = nil

Expand Down Expand Up @@ -298,6 +304,7 @@ func (e enumerateDevice) updateNodeTable(n *node, inventoryDevices []inventoryDe
d.deviceIdSet = true
d.m.Unlock()
deviceIdMapping[i.deviceId] = d
e.logger.LogTrace(ctx, "Added new device.", logwrap.Datum("DeviceId", i.deviceId), logwrap.Datum("NewIdentifier", d.Identifier().String()))
}
}

Expand All @@ -308,6 +315,7 @@ func (e enumerateDevice) updateNodeTable(n *node, inventoryDevices []inventoryDe
for _, d := range n.device {
d.m.RLock()
if _, found := deviceIdMapping[d.deviceId]; !found {
e.logger.LogTrace(ctx, "Removing old device.", logwrap.Datum("DeviceId", d.deviceId), logwrap.Datum("OldIdentifier", d.Identifier().String()))
devicesToRemove = append(devicesToRemove, d.address)
}
d.m.RUnlock()
Expand All @@ -322,6 +330,9 @@ func (e enumerateDevice) updateNodeTable(n *node, inventoryDevices []inventoryDe
}

func (e enumerateDevice) updateCapabilitiesOnDevice(ctx context.Context, d *device, id inventoryDevice) map[da.Capability]*capabilities.EnumerationCapability {
ctx, end := e.logger.Segment(ctx, "Enumerating capabilities", logwrap.Datum("Identifier", d.Identifier().String()))
defer end()

errs := map[da.Capability]*capabilities.EnumerationCapability{
capabilities.EnumerateDeviceFlag: {Attached: true},
}
Expand All @@ -334,6 +345,7 @@ func (e enumerateDevice) updateCapabilitiesOnDevice(ctx context.Context, d *devi
for capImplName, settings := range ep.rulesOutput.Capabilities {
cF, found := factory.Mapping[capImplName]
if !found {
e.logger.LogWarn(ctx, "Could not find implementation for capability.", logwrap.Datum("CapabilityImplementation", capImplName))
errs[capabilities.EnumerateDeviceFlag].Errors = append(errs[capabilities.EnumerateDeviceFlag].Errors, fmt.Errorf("could not find capability in rule output: %s", capImplName))
continue
}
Expand All @@ -342,54 +354,28 @@ func (e enumerateDevice) updateCapabilitiesOnDevice(ctx context.Context, d *devi
errs[cF] = &capabilities.EnumerationCapability{Attached: false}
}

if slices.Contains(activeCapabilities, cF) {
errs[cF].Errors = append(errs[cF].Errors, fmt.Errorf("multiple implementations of same category, last attempted: %s", capImplName))
continue
}

c, found := d.capabilities[cF]
if found && c.ImplName() != capImplName {
found = false

if err := c.Detach(ctx, implcaps.NoLongerEnumerated); err != nil {
errs[cF].Errors = append(errs[cF].Errors, fmt.Errorf("failed to detach conflicting capabiltiy: %w", err))
}

delete(d.capabilities, cF)
}

if !found {
if c = e.capabilityFactory(capImplName); c == nil {
errs[cF].Errors = append(errs[cF].Errors, fmt.Errorf("failed to find concrete implementation: %s", capImplName))
continue
}
}

attached, err := c.Attach(ctx, d, implcaps.Enumeration, settings)
ectx, end := e.logger.Segment(ctx, "Enumerating capability.", logwrap.Datum("Endpoint", ep.description.Endpoint), logwrap.Datum("DeviceId", ep.description.DeviceID), logwrap.Datum("CapabilityImplementation", capImplName), logwrap.Datum("Capability", capabilities.StandardNames[cF]))
attached, err := e.enumerateCapabilityOnDevice(ectx, d, capImplName, cF, activeCapabilities, settings)
if err != nil {
errs[cF].Errors = append(errs[cF].Errors, fmt.Errorf("error while attaching: %s: %w", capImplName, err))
errs[cF].Errors = append(errs[cF].Errors, err...)
}
errs[cF].Attached = attached

if !attached {
if err := c.Detach(ctx, implcaps.NoLongerEnumerated); err != nil {
errs[cF].Errors = append(errs[cF].Errors, fmt.Errorf("failed to detach failed attach on capabiltiy: %s: %w", capImplName, err))
}
delete(d.capabilities, cF)
} else {
d.capabilities[cF] = c
if attached {
activeCapabilities = append(activeCapabilities, cF)
}

errs[cF].Attached = attached

activeCapabilities = append(activeCapabilities, cF)
end()
}
}

for k, v := range d.capabilities {
if !slices.Contains(activeCapabilities, k) {
errs[k] = &capabilities.EnumerationCapability{Attached: false}

e.logger.LogInfo(ctx, "Removing redundant capability implementation.", logwrap.Datum("Capability", capabilities.StandardNames[k]))
if err := v.Detach(ctx, implcaps.NoLongerEnumerated); err != nil {
e.logger.LogWarn(ctx, "Failed to detach redundant capability.", logwrap.Datum("RedundantCapabilityImplementationName", v.ImplName()), logwrap.Err(err))
errs[k].Errors = append(errs[k].Errors, fmt.Errorf("failed to detach redundant capabiltiy: %w", err))
}
delete(d.capabilities, k)
Expand All @@ -401,6 +387,57 @@ func (e enumerateDevice) updateCapabilitiesOnDevice(ctx context.Context, d *devi
return errs
}

func (e enumerateDevice) enumerateCapabilityOnDevice(ctx context.Context, d *device, capImplName string, cF da.Capability, activeCapabilities []da.Capability, settings map[string]interface{}) (bool, []error) {
var errs []error

if slices.Contains(activeCapabilities, cF) {
e.logger.LogWarn(ctx, "Multiple capabilities of the same type present on endpoint.")
return false, []error{fmt.Errorf("multiple implementations of same category, last attempted: %s", capImplName)}
}

c, found := d.capabilities[cF]
if found && c.ImplName() != capImplName {
found = false

e.logger.LogInfo(ctx, "Removing redundant capability implementation.", logwrap.Datum("RedundantCapabilityImplementationName", c.ImplName()))

if err := c.Detach(ctx, implcaps.NoLongerEnumerated); err != nil {
e.logger.LogWarn(ctx, "Failed to detach redundant capability.", logwrap.Datum("RedundantCapabilityImplementationName", c.ImplName()), logwrap.Err(err))
errs = append(errs, fmt.Errorf("failed to detach conflicting capabiltiy: %w", err))
}

delete(d.capabilities, cF)
}

if !found {
if c = e.capabilityFactory(capImplName); c == nil {
e.logger.LogError(ctx, "Failed to find implementation of capability.")
return false, []error{fmt.Errorf("failed to find concrete implementation: %s", capImplName)}
}
}

e.logger.LogInfo(ctx, "Attaching capability implementation.")
attached, err := c.Attach(ctx, d, implcaps.Enumeration, settings)
if err != nil {
e.logger.LogWarn(ctx, "Errored while attaching new capability.", logwrap.Err(err), logwrap.Datum("Attached", attached))
errs = append(errs, fmt.Errorf("error while attaching: %s: %w", capImplName, err))
}

if !attached {
e.logger.LogWarn(ctx, "Failed to attach capability implementation.")
if err := c.Detach(ctx, implcaps.NoLongerEnumerated); err != nil {
e.logger.LogWarn(ctx, "Failed to detach failed attaching capability.", logwrap.Err(err))
errs = append(errs, fmt.Errorf("failed to detach failed attach on capabiltiy: %s: %w", capImplName, err))
}
delete(d.capabilities, cF)
} else {
d.capabilities[cF] = c
e.logger.LogInfo(ctx, "Capability attached successfully.")
}

return attached, errs
}

type enumeratedDeviceAttachment struct {
node *node
device *device
Expand Down
8 changes: 4 additions & 4 deletions enumerate_device_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -345,7 +345,7 @@ func Test_enumerateDevice_updateNodeTable(t *testing.T) {
},
}

mapping := ed.updateNodeTable(n, id)
mapping := ed.updateNodeTable(context.Background(), n, id)

assert.Equal(t, d, mapping[expectedDeviceId])
assert.Equal(t, expectedDeviceId, d.deviceId)
Expand All @@ -367,7 +367,7 @@ func Test_enumerateDevice_updateNodeTable(t *testing.T) {
},
}

mapping := ed.updateNodeTable(n, id)
mapping := ed.updateNodeTable(context.Background(), n, id)

assert.Equal(t, d, mapping[existingDeviceId])
assert.Equal(t, existingDeviceId, d.deviceId)
Expand All @@ -389,7 +389,7 @@ func Test_enumerateDevice_updateNodeTable(t *testing.T) {
},
}

mapping := ed.updateNodeTable(n, id)
mapping := ed.updateNodeTable(context.Background(), n, id)

assert.Equal(t, d, mapping[existingDeviceId])
assert.Equal(t, existingDeviceId, d.deviceId)
Expand All @@ -408,7 +408,7 @@ func Test_enumerateDevice_updateNodeTable(t *testing.T) {
d := &device{m: &sync.RWMutex{}, deviceId: unwantedDeviceId, address: address}
n := &node{m: &sync.RWMutex{}, device: map[uint8]*device{0: d}}

mapping := ed.updateNodeTable(n, nil)
mapping := ed.updateNodeTable(context.Background(), n, nil)

assert.Nil(t, mapping[unwantedDeviceId])
})
Expand Down

0 comments on commit 4a0b233

Please sign in to comment.