Skip to content

Commit 8162276

Browse files
Monster0506Monster0506
committed
Add unit tests for Agent command processing and execution
1 parent 576dd90 commit 8162276

2 files changed

Lines changed: 295 additions & 184 deletions

File tree

‎internal/agent/agent.go‎

Lines changed: 182 additions & 184 deletions
Original file line numberDiff line numberDiff line change
@@ -1,184 +1,182 @@
1-
package agent
2-
3-
import (
4-
"context"
5-
"sync"
6-
"time"
7-
8-
core "github.com/monster0506/meshexec/internal"
9-
"github.com/monster0506/meshexec/internal/logging"
10-
"github.com/monster0506/meshexec/internal/messages"
11-
)
12-
13-
// Agent implements core.Agent
14-
type Agent struct {
15-
mesh core.MeshNode
16-
security SignVerify
17-
exec core.CommandExecutor
18-
targetEval core.TargetEvaluator
19-
device core.DeviceInfo
20-
logger *logging.Logger
21-
22-
cancel context.CancelFunc
23-
mu sync.Mutex
24-
run bool
25-
}
26-
27-
type SignVerify interface {
28-
SignMeshMessage(msg *core.MeshMessage) error
29-
VerifyMeshMessage(msg *core.MeshMessage) error
30-
}
31-
32-
func New(mesh core.MeshNode, security SignVerify, exec core.CommandExecutor, target core.TargetEvaluator, device core.DeviceInfo, logger *logging.Logger) *Agent {
33-
if logger == nil {
34-
logger = logging.NewLogger("info")
35-
}
36-
return &Agent{mesh: mesh, security: security, exec: exec, targetEval: target, device: device, logger: logger}
37-
}
38-
39-
func (a *Agent) Start(ctx context.Context) error {
40-
a.mu.Lock()
41-
if a.run {
42-
a.mu.Unlock()
43-
return nil
44-
}
45-
runCtx, cancel := context.WithCancel(ctx)
46-
a.cancel = cancel
47-
a.run = true
48-
a.mu.Unlock()
49-
50-
cmdCh := a.mesh.Subscribe(core.MessageTypeCommand)
51-
go func() {
52-
for {
53-
select {
54-
case <-runCtx.Done():
55-
return
56-
case msg, ok := <-cmdCh:
57-
if !ok {
58-
return
59-
}
60-
_ = a.ProcessCommand(msg)
61-
}
62-
}
63-
}()
64-
return nil
65-
}
66-
67-
func (a *Agent) Stop() error {
68-
a.mu.Lock()
69-
if !a.run {
70-
a.mu.Unlock()
71-
return nil
72-
}
73-
if a.cancel != nil {
74-
a.cancel()
75-
}
76-
a.run = false
77-
a.mu.Unlock()
78-
return nil
79-
}
80-
81-
func (a *Agent) ProcessCommand(msg *core.MeshMessage) error {
82-
if msg == nil {
83-
return core.NewExecutionError("nil_message", "received nil command message", nil)
84-
}
85-
86-
// verify signature if available
87-
if a.security != nil {
88-
if err := a.security.VerifyMeshMessage(msg); err != nil {
89-
if a.logger != nil {
90-
a.logger.Warn("Rejected command due to invalid signature", map[string]interface{}{"id": msg.ID, "error": err.Error()})
91-
}
92-
return core.NewSecurityError("signature_invalid", "invalid message signature", map[string]interface{}{"id": msg.ID})
93-
}
94-
}
95-
96-
// target matching
97-
if a.targetEval != nil {
98-
// Treat empty target as broadcast
99-
matched := true
100-
if len(msg.Target) > 0 {
101-
// Combine simple target list as OR of tokens (device names or tags)
102-
// For now, evaluate against device name equality or delegate to evaluator on first token
103-
expr := msg.Target[0]
104-
ok, err := a.targetEval.Evaluate(expr, &a.device)
105-
if err != nil {
106-
return core.NewTargetingError("evaluation_failed", "failed to evaluate target expression", map[string]interface{}{"expr": expr, "error": err.Error()})
107-
}
108-
matched = ok
109-
}
110-
if !matched {
111-
if a.logger != nil {
112-
a.logger.Debug("Command not applicable to this device", map[string]interface{}{"id": msg.ID})
113-
}
114-
return nil
115-
}
116-
}
117-
118-
// Safety validation and execution
119-
var execRes *core.ExecutionResult
120-
var execErr error
121-
if a.exec != nil {
122-
cmdLine := buildCommandLine(msg.Command, msg.Payload)
123-
if err := a.exec.ValidateCommand(cmdLine); err != nil {
124-
return err
125-
}
126-
timeout := 30 * time.Second
127-
ectx, cancel := context.WithTimeout(context.Background(), timeout)
128-
defer cancel()
129-
execRes, execErr = a.exec.Execute(ectx, cmdLine)
130-
}
131-
if execRes == nil {
132-
execRes = &core.ExecutionResult{Status: "failed", ExitCode: -1, Stderr: "no executor"}
133-
}
134-
if execErr != nil {
135-
execRes.Status = "failed"
136-
if execRes.Stderr == "" {
137-
execRes.Stderr = execErr.Error()
138-
}
139-
} else {
140-
if execRes.Status == "" {
141-
if execRes.ExitCode == 0 {
142-
execRes.Status = "success"
143-
} else {
144-
execRes.Status = "failed"
145-
}
146-
}
147-
}
148-
execRes.Device = a.device.Name
149-
150-
// Create and sign result
151-
mh := messages.NewMessageHandlerWithLevel("none")
152-
result := mh.CreateResultMessage(msg.ID, *execRes, a.device.Name)
153-
if a.security != nil {
154-
_ = a.security.SignMeshMessage(&result.MeshMessage)
155-
}
156-
// Publish result
157-
return a.mesh.SendMessage(&result.MeshMessage)
158-
}
159-
160-
func (a *Agent) ExecuteCommand(cmd string) (*core.ExecutionResult, error) {
161-
if a.exec == nil {
162-
return &core.ExecutionResult{Status: "failed", ExitCode: -1, Stderr: "no executor"}, nil
163-
}
164-
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
165-
defer cancel()
166-
return a.exec.Execute(ctx, cmd)
167-
}
168-
169-
func (a *Agent) ValidateCommand(msg *core.MeshMessage) error {
170-
if a.exec == nil {
171-
return nil
172-
}
173-
cmdLine := buildCommandLine(msg.Command, msg.Payload)
174-
return a.exec.ValidateCommand(cmdLine)
175-
}
176-
177-
func buildCommandLine(command string, payload []byte) string {
178-
if command == "" {
179-
return ""
180-
}
181-
return command
182-
}
183-
184-
1+
package agent
2+
3+
import (
4+
"context"
5+
"sync"
6+
"time"
7+
8+
core "github.com/monster0506/meshexec/internal"
9+
"github.com/monster0506/meshexec/internal/logging"
10+
"github.com/monster0506/meshexec/internal/messages"
11+
)
12+
13+
// Agent implements core.Agent
14+
type Agent struct {
15+
mesh core.MeshNode
16+
security SignVerify
17+
exec core.CommandExecutor
18+
targetEval core.TargetEvaluator
19+
device core.DeviceInfo
20+
logger *logging.Logger
21+
22+
cancel context.CancelFunc
23+
mu sync.Mutex
24+
run bool
25+
}
26+
27+
type SignVerify interface {
28+
SignMeshMessage(msg *core.MeshMessage) error
29+
VerifyMeshMessage(msg *core.MeshMessage) error
30+
}
31+
32+
func New(mesh core.MeshNode, security SignVerify, exec core.CommandExecutor, target core.TargetEvaluator, device core.DeviceInfo, logger *logging.Logger) *Agent {
33+
if logger == nil {
34+
logger = logging.NewLogger("info")
35+
}
36+
return &Agent{mesh: mesh, security: security, exec: exec, targetEval: target, device: device, logger: logger}
37+
}
38+
39+
func (a *Agent) Start(ctx context.Context) error {
40+
a.mu.Lock()
41+
if a.run {
42+
a.mu.Unlock()
43+
return nil
44+
}
45+
runCtx, cancel := context.WithCancel(ctx)
46+
a.cancel = cancel
47+
a.run = true
48+
a.mu.Unlock()
49+
50+
cmdCh := a.mesh.Subscribe(core.MessageTypeCommand)
51+
go func() {
52+
for {
53+
select {
54+
case <-runCtx.Done():
55+
return
56+
case msg, ok := <-cmdCh:
57+
if !ok {
58+
return
59+
}
60+
_ = a.ProcessCommand(msg)
61+
}
62+
}
63+
}()
64+
return nil
65+
}
66+
67+
func (a *Agent) Stop() error {
68+
a.mu.Lock()
69+
if !a.run {
70+
a.mu.Unlock()
71+
return nil
72+
}
73+
if a.cancel != nil {
74+
a.cancel()
75+
}
76+
a.run = false
77+
a.mu.Unlock()
78+
return nil
79+
}
80+
81+
func (a *Agent) ProcessCommand(msg *core.MeshMessage) error {
82+
if msg == nil {
83+
return core.NewExecutionError("nil_message", "received nil command message", nil)
84+
}
85+
86+
// verify signature if available
87+
if a.security != nil {
88+
if err := a.security.VerifyMeshMessage(msg); err != nil {
89+
if a.logger != nil {
90+
a.logger.Warn("Rejected command due to invalid signature", map[string]interface{}{"id": msg.ID, "error": err.Error()})
91+
}
92+
return core.NewSecurityError("signature_invalid", "invalid message signature", map[string]interface{}{"id": msg.ID})
93+
}
94+
}
95+
96+
// target matching
97+
if a.targetEval != nil {
98+
// Treat empty target as broadcast
99+
matched := true
100+
if len(msg.Target) > 0 {
101+
// Combine simple target list as OR of tokens (device names or tags)
102+
// For now, evaluate against device name equality or delegate to evaluator on first token
103+
expr := msg.Target[0]
104+
ok, err := a.targetEval.Evaluate(expr, &a.device)
105+
if err != nil {
106+
return core.NewTargetingError("evaluation_failed", "failed to evaluate target expression", map[string]interface{}{"expr": expr, "error": err.Error()})
107+
}
108+
matched = ok
109+
}
110+
if !matched {
111+
if a.logger != nil {
112+
a.logger.Debug("Command not applicable to this device", map[string]interface{}{"id": msg.ID})
113+
}
114+
return nil
115+
}
116+
}
117+
118+
// Safety validation and execution
119+
var execRes *core.ExecutionResult
120+
var execErr error
121+
if a.exec != nil {
122+
cmdLine := buildCommandLine(msg.Command, msg.Payload)
123+
if err := a.exec.ValidateCommand(cmdLine); err != nil {
124+
return err
125+
}
126+
timeout := 30 * time.Second
127+
ectx, cancel := context.WithTimeout(context.Background(), timeout)
128+
defer cancel()
129+
execRes, execErr = a.exec.Execute(ectx, cmdLine)
130+
}
131+
if execRes == nil {
132+
execRes = &core.ExecutionResult{Status: "failed", ExitCode: -1, Stderr: "no executor"}
133+
}
134+
if execErr != nil {
135+
execRes.Status = "failed"
136+
if execRes.Stderr == "" {
137+
execRes.Stderr = execErr.Error()
138+
}
139+
} else {
140+
if execRes.Status == "" {
141+
if execRes.ExitCode == 0 {
142+
execRes.Status = "success"
143+
} else {
144+
execRes.Status = "failed"
145+
}
146+
}
147+
}
148+
execRes.Device = a.device.Name
149+
150+
// Create and sign result
151+
mh := messages.NewMessageHandlerWithLevel("none")
152+
result := mh.CreateResultMessage(msg.ID, *execRes, a.device.Name)
153+
if a.security != nil {
154+
_ = a.security.SignMeshMessage(&result.MeshMessage)
155+
}
156+
// Publish result
157+
return a.mesh.SendMessage(&result.MeshMessage)
158+
}
159+
160+
func (a *Agent) ExecuteCommand(cmd string) (*core.ExecutionResult, error) {
161+
if a.exec == nil {
162+
return &core.ExecutionResult{Status: "failed", ExitCode: -1, Stderr: "no executor"}, nil
163+
}
164+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
165+
defer cancel()
166+
return a.exec.Execute(ctx, cmd)
167+
}
168+
169+
func (a *Agent) ValidateCommand(msg *core.MeshMessage) error {
170+
if a.exec == nil {
171+
return nil
172+
}
173+
cmdLine := buildCommandLine(msg.Command, msg.Payload)
174+
return a.exec.ValidateCommand(cmdLine)
175+
}
176+
177+
func buildCommandLine(command string, payload []byte) string {
178+
if command == "" {
179+
return ""
180+
}
181+
return command
182+
}

0 commit comments

Comments
 (0)