graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

propagate context to job executors for cancellation on disconnect

Browse source
This commit is contained in:
Graham McIntire 2026-02-12 09:59:45 -06:00
parent 04d8349a6c
commit 94192f402a
No known key found for this signature in database
7 changed files with 62 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
agent.go
View file

@ -222,7 +222,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
slog.Warn("invalid message", "error", err) slog.Warn("invalid message", "error", err)
continue continue
} }
handleMessage(msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh) handleMessage(ctx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
case result := <-snmpResultCh: case result := <-snmpResultCh:
snmpBatch = append(snmpBatch, result) snmpBatch = append(snmpBatch, result)
@ -274,6 +274,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
// handleMessage dispatches incoming channel messages. // handleMessage dispatches incoming channel messages.
func handleMessage( func handleMessage(
ctx context.Context,
msg channelMsg, msg channelMsg,
pools *jobPools, pools *jobPools,
snmpResultCh chan<- *pb.SnmpResult, snmpResultCh chan<- *pb.SnmpResult,
@ -305,7 +306,7 @@ func handleMessage(
} }
slog.Info("received jobs", "count", len(jobList.Jobs)) slog.Info("received jobs", "count", len(jobList.Jobs))
for _, job := range jobList.Jobs { for _, job := range jobList.Jobs {
dispatchJob(job, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh) dispatchJob(ctx, job, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
} }
case "restart": case "restart":
@ -340,6 +341,7 @@ type jobPools struct {
// dispatchJob routes a job to the appropriate worker pool. // dispatchJob routes a job to the appropriate worker pool.
func dispatchJob( func dispatchJob(
ctx context.Context,
job *pb.AgentJob, job *pb.AgentJob,
pools *jobPools, pools *jobPools,
snmpResultCh chan<- *pb.SnmpResult, snmpResultCh chan<- *pb.SnmpResult,
@ -351,13 +353,13 @@ func dispatchJob(
switch job.JobType { switch job.JobType {
case pb.JobType_MIKROTIK: case pb.JobType_MIKROTIK:
pools.mikrotik.submit(func() { executeMikrotikJob(job, mikrotikResultCh) }) pools.mikrotik.submit(func() { executeMikrotikJob(ctx, job, mikrotikResultCh) })
case pb.JobType_TEST_CREDENTIALS: case pb.JobType_TEST_CREDENTIALS:
pools.snmp.submit(func() { executeCredentialTest(job, credTestResultCh) }) pools.snmp.submit(func() { executeCredentialTest(ctx, job, credTestResultCh) })
case pb.JobType_PING: case pb.JobType_PING:
pools.ping.submit(func() { executePingJob(job, monitoringCheckCh) }) pools.ping.submit(func() { executePingJob(ctx, job, monitoringCheckCh) })
default: default:
pools.snmp.submit(func() { executeSnmpJob(job, snmpResultCh) }) pools.snmp.submit(func() { executeSnmpJob(ctx, job, snmpResultCh) })
} }
} }

33
agent_test.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"context"
"encoding/base64" "encoding/base64"
"encoding/json" "encoding/json"
"fmt" "fmt"
@ -100,7 +101,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1) mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Just verify it doesn't panic // Just verify it doesn't panic
}) })
@ -127,7 +128,7 @@ func TestHandleMessage(t *testing.T) {
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
}) })
handleMessage(channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Wait for goroutine to finish // Wait for goroutine to finish
select { select {
case <-snmpCh: case <-snmpCh:
@ -141,7 +142,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1) mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Should log error but not panic // Should log error but not panic
}) })
@ -151,7 +152,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"}) payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"})
handleMessage(channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
}) })
t.Run("invalid protobuf", func(t *testing.T) { t.Run("invalid protobuf", func(t *testing.T) {
@ -160,7 +161,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})}) payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})})
handleMessage(channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
}) })
t.Run("restart", func(t *testing.T) { t.Run("restart", func(t *testing.T) {
@ -174,7 +175,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1) mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
if exitCode != 0 { if exitCode != 0 {
t.Errorf("expected exit code 0, got %d", exitCode) t.Errorf("expected exit code 0, got %d", exitCode)
@ -196,7 +197,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"}) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"})
handleMessage(channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
if calledURL != "https://example.com/agent" { if calledURL != "https://example.com/agent" {
t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL) t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL)
@ -219,7 +220,7 @@ func TestHandleMessage(t *testing.T) {
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
// Missing URL field // Missing URL field
payload, _ := json.Marshal(map[string]string{"checksum": "abc123"}) payload, _ := json.Marshal(map[string]string{"checksum": "abc123"})
handleMessage(channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
if called { if called {
t.Error("selfUpdate should not be called with empty URL") t.Error("selfUpdate should not be called with empty URL")
@ -239,7 +240,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"}) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"})
handleMessage(channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Should log error but not panic // Should log error but not panic
}) })
@ -248,7 +249,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1) mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Should just log and not panic // Should just log and not panic
}) })
@ -274,7 +275,7 @@ func TestHandleMessage(t *testing.T) {
JobType: pb.JobType_DISCOVER, JobType: pb.JobType_DISCOVER,
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}) })
handleMessage(channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
select { select {
case <-snmpCh: case <-snmpCh:
case <-time.After(2 * time.Second): case <-time.After(2 * time.Second):
@ -301,7 +302,7 @@ func TestHandleMessage(t *testing.T) {
JobType: pb.JobType_MIKROTIK, JobType: pb.JobType_MIKROTIK,
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"},
}) })
handleMessage(channelMsg{Event: "backup_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) handleMessage(context.Background(), channelMsg{Event: "backup_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
select { select {
case result := <-mtCh: case result := <-mtCh:
if result.Error != "" { if result.Error != "" {
@ -326,7 +327,7 @@ func TestDispatchJob(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
dispatchJob(&pb.AgentJob{ dispatchJob(context.Background(), &pb.AgentJob{
JobId: "mt1", JobId: "mt1",
JobType: pb.JobType_MIKROTIK, JobType: pb.JobType_MIKROTIK,
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728},
@ -354,7 +355,7 @@ func TestDispatchJob(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
dispatchJob(&pb.AgentJob{ dispatchJob(context.Background(), &pb.AgentJob{
JobId: "tc1", JobId: "tc1",
JobType: pb.JobType_TEST_CREDENTIALS, JobType: pb.JobType_TEST_CREDENTIALS,
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
@ -382,7 +383,7 @@ func TestDispatchJob(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
dispatchJob(&pb.AgentJob{ dispatchJob(context.Background(), &pb.AgentJob{
JobId: "p1", JobId: "p1",
JobType: pb.JobType_PING, JobType: pb.JobType_PING,
SnmpDevice: &pb.SnmpDevice{Ip: "127.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "127.0.0.1"},
@ -414,7 +415,7 @@ func TestDispatchJob(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
dispatchJob(&pb.AgentJob{ dispatchJob(context.Background(), &pb.AgentJob{
JobId: "s1", JobId: "s1",
JobType: pb.JobType_POLL, JobType: pb.JobType_POLL,
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},

2
mikrotik.go
View file

@ -244,7 +244,7 @@ func parseMikrotikAttrs(words []string) map[string]string {
} }
// executeMikrotikJob handles a MikroTik API job including backup-via-SSH. // executeMikrotikJob handles a MikroTik API job including backup-via-SSH.
func executeMikrotikJob(job *pb.AgentJob, resultCh chan<- *pb.MikrotikResult) { func executeMikrotikJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.MikrotikResult) {
dev := job.MikrotikDevice dev := job.MikrotikDevice
if dev == nil { if dev == nil {
slog.Error("job missing mikrotik device", "job_id", job.JobId) slog.Error("job missing mikrotik device", "job_id", job.JobId)

8
snmp.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"context"
"encoding/hex" "encoding/hex"
"fmt" "fmt"
"log/slog" "log/slog"
@ -32,7 +33,7 @@ var snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) {
} }
// executeSnmpJob runs SNMP GET/WALK queries for a job and sends results. // executeSnmpJob runs SNMP GET/WALK queries for a job and sends results.
func executeSnmpJob(job *pb.AgentJob, resultCh chan<- *pb.SnmpResult) { func executeSnmpJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.SnmpResult) {
dev := job.SnmpDevice dev := job.SnmpDevice
if dev == nil { if dev == nil {
slog.Error("job missing snmp device", "job_id", job.JobId) slog.Error("job missing snmp device", "job_id", job.JobId)
@ -53,6 +54,9 @@ func executeSnmpJob(job *pb.AgentJob, resultCh chan<- *pb.SnmpResult) {
oidValues := make(map[string]string, totalOIDs) oidValues := make(map[string]string, totalOIDs)
for _, q := range job.Queries { for _, q := range job.Queries {
if ctx.Err() != nil {
return
}
switch q.QueryType { switch q.QueryType {
case pb.QueryType_GET: case pb.QueryType_GET:
for i := 0; i < len(q.Oids); i += snmpMaxOIDsPerGet { for i := 0; i < len(q.Oids); i += snmpMaxOIDsPerGet {
@ -115,7 +119,7 @@ func executeSnmpJob(job *pb.AgentJob, resultCh chan<- *pb.SnmpResult) {
} }
// executeCredentialTest tests SNMP credentials by reading sysDescr.0. // executeCredentialTest tests SNMP credentials by reading sysDescr.0.
func executeCredentialTest(job *pb.AgentJob, resultCh chan<- *pb.CredentialTestResult) { func executeCredentialTest(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.CredentialTestResult) {
dev := job.SnmpDevice dev := job.SnmpDevice
if dev == nil { if dev == nil {
slog.Error("job missing snmp device", "job_id", job.JobId) slog.Error("job missing snmp device", "job_id", job.JobId)

35
snmp_test.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"context"
"fmt" "fmt"
"testing" "testing"
@ -299,7 +300,7 @@ func (m *mockSnmpQuerier) BulkWalkAll(rootOid string) ([]gosnmp.SnmpPDU, error)
func TestExecuteSnmpJob(t *testing.T) { func TestExecuteSnmpJob(t *testing.T) {
t.Run("nil device", func(t *testing.T) { t.Run("nil device", func(t *testing.T) {
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{JobId: "1"}, ch) executeSnmpJob(context.Background(), &pb.AgentJob{JobId: "1"}, ch)
if len(ch) != 0 { if len(ch) != 0 {
t.Error("expected no result for nil device") t.Error("expected no result for nil device")
} }
@ -313,7 +314,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
}, ch) }, ch)
@ -340,7 +341,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
DeviceId: "dev-1", DeviceId: "dev-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
@ -378,7 +379,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -406,7 +407,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -434,7 +435,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -466,7 +467,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -498,7 +499,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -528,7 +529,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Version: "1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Version: "1"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -564,7 +565,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Version: "2c"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Version: "2c"},
Queries: []*pb.SnmpQuery{ Queries: []*pb.SnmpQuery{
@ -598,7 +599,7 @@ func TestExecuteSnmpJob(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult) // unbuffered, no reader — will be full ch := make(chan *pb.SnmpResult) // unbuffered, no reader — will be full
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "1", JobId: "1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}, ch) }, ch)
@ -632,7 +633,7 @@ func TestExecuteSnmpJobBatchesGets(t *testing.T) {
} }
ch := make(chan *pb.SnmpResult, 1) ch := make(chan *pb.SnmpResult, 1)
executeSnmpJob(&pb.AgentJob{ executeSnmpJob(context.Background(), &pb.AgentJob{
JobId: "batch-test", JobId: "batch-test",
DeviceId: "dev-1", DeviceId: "dev-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
@ -661,7 +662,7 @@ func TestExecuteSnmpJobBatchesGets(t *testing.T) {
func TestExecuteCredentialTest(t *testing.T) { func TestExecuteCredentialTest(t *testing.T) {
t.Run("nil device", func(t *testing.T) { t.Run("nil device", func(t *testing.T) {
ch := make(chan *pb.CredentialTestResult, 1) ch := make(chan *pb.CredentialTestResult, 1)
executeCredentialTest(&pb.AgentJob{JobId: "1"}, ch) executeCredentialTest(context.Background(), &pb.AgentJob{JobId: "1"}, ch)
if len(ch) != 0 { if len(ch) != 0 {
t.Error("expected no result for nil device") t.Error("expected no result for nil device")
} }
@ -675,7 +676,7 @@ func TestExecuteCredentialTest(t *testing.T) {
} }
ch := make(chan *pb.CredentialTestResult, 1) ch := make(chan *pb.CredentialTestResult, 1)
executeCredentialTest(&pb.AgentJob{ executeCredentialTest(context.Background(), &pb.AgentJob{
JobId: "test-1", JobId: "test-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
}, ch) }, ch)
@ -703,7 +704,7 @@ func TestExecuteCredentialTest(t *testing.T) {
} }
ch := make(chan *pb.CredentialTestResult, 1) ch := make(chan *pb.CredentialTestResult, 1)
executeCredentialTest(&pb.AgentJob{ executeCredentialTest(context.Background(), &pb.AgentJob{
JobId: "test-1", JobId: "test-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}, ch) }, ch)
@ -732,7 +733,7 @@ func TestExecuteCredentialTest(t *testing.T) {
} }
ch := make(chan *pb.CredentialTestResult, 1) ch := make(chan *pb.CredentialTestResult, 1)
executeCredentialTest(&pb.AgentJob{ executeCredentialTest(context.Background(), &pb.AgentJob{
JobId: "test-1", JobId: "test-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}, ch) }, ch)
@ -760,7 +761,7 @@ func TestExecuteCredentialTest(t *testing.T) {
} }
ch := make(chan *pb.CredentialTestResult, 1) ch := make(chan *pb.CredentialTestResult, 1)
executeCredentialTest(&pb.AgentJob{ executeCredentialTest(context.Background(), &pb.AgentJob{
JobId: "test-1", JobId: "test-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}, ch) }, ch)

3
ssh.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"context"
"fmt" "fmt"
"log/slog" "log/slog"
"time" "time"
@ -47,7 +48,7 @@ func executeMikrotikBackup(ip string, port uint16, username, password string) (s
} }
// executePingJob pings a device and sends a monitoring check result. // executePingJob pings a device and sends a monitoring check result.
func executePingJob(job *pb.AgentJob, resultCh chan<- *pb.MonitoringCheck) { func executePingJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.MonitoringCheck) {
dev := job.SnmpDevice dev := job.SnmpDevice
if dev == nil { if dev == nil {
slog.Error("job missing device info for ping", "job_id", job.JobId) slog.Error("job missing device info for ping", "job_id", job.JobId)

19
ssh_test.go
View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"context"
"crypto/ecdsa" "crypto/ecdsa"
"crypto/elliptic" "crypto/elliptic"
"crypto/rand" "crypto/rand"
@ -17,7 +18,7 @@ import (
func TestExecutePingJob(t *testing.T) { func TestExecutePingJob(t *testing.T) {
t.Run("nil device", func(t *testing.T) { t.Run("nil device", func(t *testing.T) {
ch := make(chan *pb.MonitoringCheck, 1) ch := make(chan *pb.MonitoringCheck, 1)
executePingJob(&pb.AgentJob{JobId: "p1"}, ch) executePingJob(context.Background(), &pb.AgentJob{JobId: "p1"}, ch)
if len(ch) != 0 { if len(ch) != 0 {
t.Error("expected no result for nil device") t.Error("expected no result for nil device")
} }
@ -31,7 +32,7 @@ func TestExecutePingJob(t *testing.T) {
} }
ch := make(chan *pb.MonitoringCheck, 1) ch := make(chan *pb.MonitoringCheck, 1)
executePingJob(&pb.AgentJob{ executePingJob(context.Background(), &pb.AgentJob{
JobId: "p1", JobId: "p1",
DeviceId: "dev-1", DeviceId: "dev-1",
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
@ -61,7 +62,7 @@ func TestExecutePingJob(t *testing.T) {
} }
ch := make(chan *pb.MonitoringCheck, 1) ch := make(chan *pb.MonitoringCheck, 1)
executePingJob(&pb.AgentJob{ executePingJob(context.Background(), &pb.AgentJob{
JobId: "p2", JobId: "p2",
DeviceId: "dev-2", DeviceId: "dev-2",
SnmpDevice: &pb.SnmpDevice{Ip: "192.168.1.1"}, SnmpDevice: &pb.SnmpDevice{Ip: "192.168.1.1"},
@ -81,7 +82,7 @@ func TestExecutePingJob(t *testing.T) {
func TestExecuteMikrotikJob(t *testing.T) { func TestExecuteMikrotikJob(t *testing.T) {
t.Run("nil device", func(t *testing.T) { t.Run("nil device", func(t *testing.T) {
ch := make(chan *pb.MikrotikResult, 1) ch := make(chan *pb.MikrotikResult, 1)
executeMikrotikJob(&pb.AgentJob{JobId: "m1"}, ch) executeMikrotikJob(context.Background(), &pb.AgentJob{JobId: "m1"}, ch)
if len(ch) != 0 { if len(ch) != 0 {
t.Error("expected no result for nil device") t.Error("expected no result for nil device")
} }
@ -95,7 +96,7 @@ func TestExecuteMikrotikJob(t *testing.T) {
} }
ch := make(chan *pb.MikrotikResult, 1) ch := make(chan *pb.MikrotikResult, 1)
executeMikrotikJob(&pb.AgentJob{ executeMikrotikJob(context.Background(), &pb.AgentJob{
JobId: "m1", JobId: "m1",
DeviceId: "dev-1", DeviceId: "dev-1",
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728},
@ -119,7 +120,7 @@ func TestExecuteMikrotikJob(t *testing.T) {
} }
ch := make(chan *pb.MikrotikResult, 1) ch := make(chan *pb.MikrotikResult, 1)
executeMikrotikJob(&pb.AgentJob{ executeMikrotikJob(context.Background(), &pb.AgentJob{
JobId: "m1", JobId: "m1",
DeviceId: "dev-1", DeviceId: "dev-1",
MikrotikDevice: &pb.MikrotikDevice{ MikrotikDevice: &pb.MikrotikDevice{
@ -154,7 +155,7 @@ func TestExecuteMikrotikJob(t *testing.T) {
} }
ch := make(chan *pb.MikrotikResult, 1) ch := make(chan *pb.MikrotikResult, 1)
executeMikrotikJob(&pb.AgentJob{ executeMikrotikJob(context.Background(), &pb.AgentJob{
JobId: "m1", JobId: "m1",
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728},
MikrotikCommands: []*pb.MikrotikCommand{ MikrotikCommands: []*pb.MikrotikCommand{
@ -180,7 +181,7 @@ func TestExecuteMikrotikJob(t *testing.T) {
} }
ch := make(chan *pb.MikrotikResult, 1) ch := make(chan *pb.MikrotikResult, 1)
executeMikrotikJob(&pb.AgentJob{ executeMikrotikJob(context.Background(), &pb.AgentJob{
JobId: "m1", JobId: "m1",
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728},
MikrotikCommands: []*pb.MikrotikCommand{ MikrotikCommands: []*pb.MikrotikCommand{
@ -203,7 +204,7 @@ func TestExecuteMikrotikJob(t *testing.T) {
} }
ch := make(chan *pb.MikrotikResult, 1) ch := make(chan *pb.MikrotikResult, 1)
executeMikrotikJob(&pb.AgentJob{ executeMikrotikJob(context.Background(), &pb.AgentJob{
JobId: "backup:dev1", JobId: "backup:dev1",
DeviceId: "dev-1", DeviceId: "dev-1",
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"},