package main import ( "context" "encoding/base64" "encoding/json" "fmt" "net" "strings" "sync" "testing" "time" "github.com/gosnmp/gosnmp" "github.com/towerops-app/towerops-agent/pb" "google.golang.org/protobuf/proto" ) func TestChannelMsgSerialization(t *testing.T) { msg := channelMsg{ Topic: "agent:123", Event: "phx_join", Payload: json.RawMessage(`{"token":"test"}`), Ref: strPtr("1"), } data, err := json.Marshal(msg) if err != nil { t.Fatal(err) } s := string(data) checks := []string{"agent:123", "phx_join", "token", "test"} for _, c := range checks { if !contains(s, c) { t.Errorf("expected %q in JSON output %q", c, s) } } } func TestChannelMsgDeserialization(t *testing.T) { raw := `{"topic":"agent:123","event":"phx_reply","payload":{"status":"ok"},"ref":"1"}` var msg channelMsg if err := json.Unmarshal([]byte(raw), &msg); err != nil { t.Fatal(err) } if msg.Topic != "agent:123" { t.Errorf("topic: got %q, want %q", msg.Topic, "agent:123") } if msg.Event != "phx_reply" { t.Errorf("event: got %q, want %q", msg.Event, "phx_reply") } if msg.Ref == nil || *msg.Ref != "1" { t.Errorf("ref: got %v, want %q", msg.Ref, "1") } } func TestChannelMsgNullRef(t *testing.T) { raw := `{"topic":"agent:123","event":"job","payload":{},"ref":null}` var msg channelMsg if err := json.Unmarshal([]byte(raw), &msg); err != nil { t.Fatal(err) } if msg.Ref != nil { t.Errorf("expected nil ref, got %q", *msg.Ref) } } func contains(s, substr string) bool { return len(s) >= len(substr) && searchString(s, substr) } func searchString(s, substr string) bool { for i := 0; i <= len(s)-len(substr); i++ { if s[i:i+len(substr)] == substr { return true } } return false } func testPools(t *testing.T) *jobPools { t.Helper() p := &jobPools{ snmp: newWorkerPool(4), mikrotik: newWorkerPool(4), ping: newWorkerPool(4), checks: newWorkerPool(4), } t.Cleanup(func() { p.snmp.stop(); p.mikrotik.stop(); p.ping.stop(); p.checks.stop() }) return p } // makeJobPayload creates a base64-encoded protobuf job list payload. func makeJobPayload(jobs ...*pb.AgentJob) json.RawMessage { list := &pb.AgentJobList{Jobs: jobs} bin, _ := proto.Marshal(list) payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString(bin)}) return payload } func TestHandleMessage(t *testing.T) { t.Run("phx_reply", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) handleMessage(context.Background(), channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) // Just verify it doesn't panic }) t.Run("jobs valid protobuf", func(t *testing.T) { origDial := snmpDial defer func() { snmpDial = origDial }() snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { return &mockSnmpQuerier{ getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { return &gosnmp.SnmpPacket{}, nil }, }, func() {}, nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload := makeJobPayload(&pb.AgentJob{ JobId: "j1", JobType: pb.JobType_POLL, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, }) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) // Wait for goroutine to finish select { case <-snmpCh: case <-time.After(2 * time.Second): t.Error("timed out waiting for snmp result") } }) t.Run("invalid payload json", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) // Should log error but not panic }) t.Run("invalid base64", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"}) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) }) t.Run("invalid protobuf", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})}) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) }) t.Run("restart", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) shouldEnd := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) if !shouldEnd { t.Error("expected handleMessage to return true for restart") } }) t.Run("update success", func(t *testing.T) { origUpdate := doSelfUpdate defer func() { doSelfUpdate = origUpdate }() var calledURL string doSelfUpdate = func(url, checksum string) error { calledURL = url return nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"}) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) if calledURL != "https://example.com/agent" { t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL) } }) t.Run("update invalid payload", func(t *testing.T) { origUpdate := doSelfUpdate defer func() { doSelfUpdate = origUpdate }() called := false doSelfUpdate = func(url, checksum string) error { called = true return nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) // Missing URL field payload, _ := json.Marshal(map[string]string{"checksum": "abc123"}) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) if called { t.Error("selfUpdate should not be called with empty URL") } }) t.Run("update error", func(t *testing.T) { origUpdate := doSelfUpdate defer func() { doSelfUpdate = origUpdate }() doSelfUpdate = func(url, checksum string) error { return fmt.Errorf("download failed") } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"}) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) // Should log error but not panic }) t.Run("update missing checksum", func(t *testing.T) { origUpdate := doSelfUpdate defer func() { doSelfUpdate = origUpdate }() called := false doSelfUpdate = func(url, checksum string) error { called = true return nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"}) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) if called { t.Error("selfUpdate should not be called with empty checksum") } }) t.Run("unknown event", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) handleMessage(context.Background(), channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) // Should just log and not panic }) t.Run("discovery_job event", func(t *testing.T) { origDial := snmpDial defer func() { snmpDial = origDial }() snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { return &mockSnmpQuerier{ getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { return &gosnmp.SnmpPacket{}, nil }, }, func() {}, nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload := makeJobPayload(&pb.AgentJob{ JobId: "d1", JobType: pb.JobType_DISCOVER, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, }) handleMessage(context.Background(), channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) select { case <-snmpCh: case <-time.After(2 * time.Second): t.Error("timed out waiting for discovery result") } }) t.Run("backup_job event", func(t *testing.T) { origDial := mikrotikDial origSSH := sshBackup defer func() { mikrotikDial = origDial; sshBackup = origSSH }() sshBackup = func(ip string, port uint16, username, password string) (string, error) { return "/ip address\nadd address=10.0.0.1/24", nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload := makeJobPayload(&pb.AgentJob{ JobId: "backup:dev1", JobType: pb.JobType_MIKROTIK, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, }) handleMessage(context.Background(), channelMsg{Event: "backup_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) select { case result := <-mtCh: if result.Error != "" { t.Errorf("unexpected error: %s", result.Error) } case <-time.After(2 * time.Second): t.Error("timed out waiting for backup result") } }) } func TestHandleMessageRejectsOversizedPayload(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) // Create a binary payload larger than maxJobPayloadBytes oversized := make([]byte, maxJobPayloadBytes+1) encoded := base64.StdEncoding.EncodeToString(oversized) payload, _ := json.Marshal(map[string]string{"binary": encoded}) handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) // Verify no jobs were dispatched select { case <-snmpCh: t.Error("expected no SNMP result for oversized payload") case <-time.After(100 * time.Millisecond): // Good — nothing dispatched } } func TestBufferPoolZeroesOnReturn(t *testing.T) { pool := &sync.Pool{ New: func() any { b := make([]byte, 0, 64) return &b }, } // Get a buffer, write some data, return it bp := pool.Get().(*[]byte) *bp = append((*bp)[:0], []byte("sensitive credentials data here")...) full := (*bp)[:cap(*bp)] // Simulate the zeroing that sendBinaryResult should do zeroBytes(full) *bp = full[:0] pool.Put(bp) // Get the buffer back and verify it's zeroed bp2 := pool.Get().(*[]byte) full2 := (*bp2)[:cap(*bp2)] for i, b := range full2 { if b != 0 { t.Errorf("byte[%d] = %d, expected 0 — pool buffer not zeroed", i, b) break } } } func TestZeroBytes(t *testing.T) { b := []byte{1, 2, 3, 4, 5} zeroBytes(b) for i, v := range b { if v != 0 { t.Errorf("byte[%d] = %d, want 0", i, v) } } // Nil/empty should not panic zeroBytes(nil) zeroBytes([]byte{}) } func TestNextBackoff(t *testing.T) { maxDelay := 60 * time.Second // Test doubling with jitter for i := 0; i < 100; i++ { current := 2 * time.Second next := nextBackoff(current, maxDelay) doubled := current * 2 maxWithJitter := doubled + doubled/4 if next < doubled || next > maxWithJitter { t.Errorf("nextBackoff(%v) = %v, want in [%v, %v]", current, next, doubled, maxWithJitter) } } // Test cap at max for i := 0; i < 100; i++ { next := nextBackoff(30*time.Second, maxDelay) if next > maxDelay+maxDelay/4 { t.Errorf("nextBackoff(30s) = %v, exceeded max+jitter", next) } } // Test that already-at-max stays at max (with jitter) for i := 0; i < 100; i++ { next := nextBackoff(maxDelay, maxDelay) maxWithJitter := maxDelay + maxDelay/4 if next < maxDelay || next > maxWithJitter { t.Errorf("nextBackoff(max) = %v, want in [%v, %v]", next, maxDelay, maxWithJitter) } } } func TestDispatchJob(t *testing.T) { t.Run("MIKROTIK", func(t *testing.T) { origDial := mikrotikDial defer func() { mikrotikDial = origDial }() mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { return nil, fmt.Errorf("not reachable") } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) dispatchJob(context.Background(), &pb.AgentJob{ JobId: "mt1", JobType: pb.JobType_MIKROTIK, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, }, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) select { case result := <-mtCh: if result.Error == "" { t.Error("expected error from unreachable device") } case <-time.After(2 * time.Second): t.Error("timed out") } }) t.Run("TEST_CREDENTIALS", func(t *testing.T) { origDial := snmpDial defer func() { snmpDial = origDial }() snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { return nil, nil, fmt.Errorf("refused") } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) dispatchJob(context.Background(), &pb.AgentJob{ JobId: "tc1", JobType: pb.JobType_TEST_CREDENTIALS, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, }, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) select { case result := <-credCh: if result.Success { t.Error("expected failure") } case <-time.After(2 * time.Second): t.Error("timed out") } }) t.Run("PING", func(t *testing.T) { origPing := doPing defer func() { doPing = origPing }() doPing = func(ip string, timeoutMs int) (float64, error) { return 5.5, nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) dispatchJob(context.Background(), &pb.AgentJob{ JobId: "p1", JobType: pb.JobType_PING, SnmpDevice: &pb.SnmpDevice{Ip: "127.0.0.1"}, }, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) select { case result := <-monCh: if result.Status != "success" { t.Errorf("expected success, got %q", result.Status) } case <-time.After(2 * time.Second): t.Error("timed out") } }) t.Run("default SNMP", func(t *testing.T) { origDial := snmpDial defer func() { snmpDial = origDial }() snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { return &mockSnmpQuerier{ getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { return &gosnmp.SnmpPacket{}, nil }, }, func() {}, nil } snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) dispatchJob(context.Background(), &pb.AgentJob{ JobId: "s1", JobType: pb.JobType_POLL, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, }, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) select { case <-snmpCh: case <-time.After(2 * time.Second): t.Error("timed out") } }) } func TestRunSessionRejectsFailedJoin(t *testing.T) { origTimeout := joinTimeout defer func() { joinTimeout = origTimeout }() joinTimeout = 2 * time.Second // Start a fake WebSocket server that accepts the upgrade then sends a phx_error join reply ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } defer func() { _ = ln.Close() }() go func() { conn, err := ln.Accept() if err != nil { return } defer func() { _ = conn.Close() }() // Read HTTP upgrade request buf := make([]byte, 4096) n, _ := conn.Read(buf) reqStr := string(buf[:n]) // Extract key and compute accept key := extractWSKey(reqStr) accept := computeAcceptKey(key) // Send valid 101 upgrade resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n" _, _ = conn.Write([]byte(resp)) // Read the join message (masked WebSocket frame) — just consume it frameBuf := make([]byte, 4096) _, _ = conn.Read(frameBuf) // Send phx_error reply as an unmasked text frame reply, _ := json.Marshal(channelMsg{ Topic: "agent:agent-0", Event: "phx_reply", Payload: json.RawMessage(`{"status":"error","response":{"reason":"invalid token"}}`), Ref: strPtr("1"), }) frame := makeTextFrame(reply) _, _ = conn.Write(frame) // Keep connection open for a bit time.Sleep(time.Second) }() addr := ln.Addr().String() err = runSession(context.Background(), "ws://"+addr, "bad-token") if err == nil { t.Fatal("expected error from rejected join") } if !strings.Contains(err.Error(), "join rejected") { t.Errorf("expected 'join rejected' in error, got: %v", err) } } func TestRunSessionJoinTimeout(t *testing.T) { origTimeout := joinTimeout defer func() { joinTimeout = origTimeout }() joinTimeout = 500 * time.Millisecond // Server that upgrades but never sends a join reply ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } defer func() { _ = ln.Close() }() go func() { conn, err := ln.Accept() if err != nil { return } defer func() { _ = conn.Close() }() buf := make([]byte, 4096) n, _ := conn.Read(buf) reqStr := string(buf[:n]) key := extractWSKey(reqStr) accept := computeAcceptKey(key) resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n" _, _ = conn.Write([]byte(resp)) // Read join frame but never reply frameBuf := make([]byte, 4096) _, _ = conn.Read(frameBuf) time.Sleep(5 * time.Second) }() addr := ln.Addr().String() start := time.Now() err = runSession(context.Background(), "ws://"+addr, "token") elapsed := time.Since(start) if err == nil { t.Fatal("expected timeout error") } if !strings.Contains(err.Error(), "join timeout") { t.Errorf("expected 'join timeout' in error, got: %v", err) } if elapsed > 3*time.Second { t.Errorf("took too long (%v), timeout didn't trigger", elapsed) } } // extractWSKey extracts the Sec-WebSocket-Key from a raw HTTP request. func extractWSKey(req string) string { for _, line := range strings.Split(req, "\r\n") { lower := strings.ToLower(line) if strings.HasPrefix(lower, "sec-websocket-key: ") { return strings.TrimSpace(line[len("Sec-WebSocket-Key: "):]) } } return "" } // makeTextFrame creates an unmasked WebSocket text frame (server→client). func makeTextFrame(payload []byte) []byte { length := len(payload) var frame []byte frame = append(frame, 0x81) // FIN + text if length <= 125 { frame = append(frame, byte(length)) } else if length <= 65535 { frame = append(frame, 126, byte(length>>8), byte(length)) } frame = append(frame, payload...) return frame }