From aef770ea83ef451477b644a2d2038cebfc4c554e Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Wed, 4 Mar 2026 16:06:29 -0600 Subject: [PATCH] test: achieve 92%+ coverage with comprehensive integration and edge case tests Cover WebSocket session lifecycle (result channels, SNMP batch flushing, heartbeats, restart), MikroTik TLS TOFU verification mismatch, worker pool exhaustion, plaintext WebSocket rejection, HTTP/TCP/DNS check execution, IPv4 fallback, and customer network edge cases. Make heartbeat intervals and toWebSocketURL exit injectable for testing. --- agent.go | 6 +- agent_test.go | 817 +++++++++++++++++++++++++++ checks_test.go | 1372 +++++++++++++++++++++++++++++++++++++++++++++ hostkeys_test.go | 21 +- main.go | 2 +- main_test.go | 38 ++ mikrotik_test.go | 131 +++++ ping_test.go | 8 + websocket_test.go | 54 ++ 9 files changed, 2442 insertions(+), 7 deletions(-) create mode 100644 checks_test.go diff --git a/agent.go b/agent.go index 251c7ab..af1af63 100644 --- a/agent.go +++ b/agent.go @@ -25,6 +25,8 @@ var doSelfUpdate = selfUpdate var errRestartRequested = fmt.Errorf("restart requested") var joinTimeout = 10 * time.Second +var heartbeatInterval = 60 * time.Second +var channelHeartbeatInterval = 25 * time.Second const maxJobPayloadBytes = 4 << 20 // 4 MB — well above any legitimate job list @@ -224,9 +226,9 @@ func runSession(ctx context.Context, baseURL, token string) error { } }() - heartbeatTicker := time.NewTicker(60 * time.Second) + heartbeatTicker := time.NewTicker(heartbeatInterval) defer heartbeatTicker.Stop() - channelHeartbeatTicker := time.NewTicker(25 * time.Second) + channelHeartbeatTicker := time.NewTicker(channelHeartbeatInterval) defer channelHeartbeatTicker.Stop() flushTicker := time.NewTicker(100 * time.Millisecond) defer flushTicker.Stop() diff --git a/agent_test.go b/agent_test.go index de9e1f3..843bc35 100644 --- a/agent_test.go +++ b/agent_test.go @@ -3,8 +3,10 @@ package main import ( "context" "encoding/base64" + "encoding/binary" "encoding/json" "fmt" + "io" "net" "strings" "sync" @@ -284,6 +286,59 @@ func TestHandleMessage(t *testing.T) { // Should just log and not panic }) + t.Run("check_jobs valid", func(t *testing.T) { + checkList := &pb.CheckList{Checks: []*pb.Check{ + {Id: "c1", CheckType: "tcp", TimeoutMs: 1000, + Config: &pb.Check_Tcp{Tcp: &pb.TcpCheckConfig{Host: "127.0.0.1", Port: 1}}}, + }} + bin, _ := proto.Marshal(checkList) + payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString(bin)}) + + 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: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) + select { + case <-checkCh: + case <-time.After(5 * time.Second): + t.Error("timed out waiting for check result") + } + }) + + t.Run("check_jobs invalid 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: "check_jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) + // Should log error but not panic + }) + + t.Run("check_jobs 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: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) + // Should log error but not panic + }) + + t.Run("check_jobs 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: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh) + // Should log error but not panic + }) + t.Run("discovery_job event", func(t *testing.T) { origDial := snmpDial defer func() { snmpDial = origDial }() @@ -698,3 +753,765 @@ func makeTextFrame(payload []byte) []byte { frame = append(frame, payload...) return frame } + +func TestDispatchJobCancelledContext(t *testing.T) { + // submit returns false when context is cancelled, triggering "pool full" log + p := testPools(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) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately + + // All dispatch types should handle ctx cancellation gracefully + dispatchJob(ctx, &pb.AgentJob{ + JobId: "s1", JobType: pb.JobType_POLL, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, p, snmpCh, mtCh, credCh, monCh, checkCh) + + dispatchJob(ctx, &pb.AgentJob{ + JobId: "m1", JobType: pb.JobType_MIKROTIK, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1"}, + }, p, snmpCh, mtCh, credCh, monCh, checkCh) + + dispatchJob(ctx, &pb.AgentJob{ + JobId: "tc1", JobType: pb.JobType_TEST_CREDENTIALS, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, p, snmpCh, mtCh, credCh, monCh, checkCh) + + dispatchJob(ctx, &pb.AgentJob{ + JobId: "p1", JobType: pb.JobType_PING, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, p, snmpCh, mtCh, credCh, monCh, checkCh) +} + +// fakeWSServer is a test helper that sets up a WebSocket server for runSession tests. +type fakeWSServer struct { + ln net.Listener + conn net.Conn +} + +func newFakeWSServer(t *testing.T) *fakeWSServer { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ln.Close() }) + return &fakeWSServer{ln: ln} +} + +func (s *fakeWSServer) addr() string { return s.ln.Addr().String() } + +// acceptAndJoin accepts one WS connection and responds with a successful join. +func (s *fakeWSServer) acceptAndJoin(t *testing.T) { + t.Helper() + conn, err := s.ln.Accept() + if err != nil { + t.Logf("accept: %v", err) + return + } + s.conn = conn + + // Read HTTP upgrade + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + key := extractWSKey(string(buf[:n])) + 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 + frameBuf := make([]byte, 4096) + _, _ = conn.Read(frameBuf) + + // Send join OK + reply, _ := json.Marshal(channelMsg{ + Topic: "agent:agent-0", + Event: "phx_reply", + Payload: json.RawMessage(`{"status":"ok"}`), + Ref: strPtr("1"), + }) + _, _ = conn.Write(makeTextFrame(reply)) +} + +// sendEvent sends a channel message to the connected client. +func (s *fakeWSServer) sendEvent(event string, payload json.RawMessage) { + msg, _ := json.Marshal(channelMsg{ + Topic: "agent:agent-0", + Event: event, + Payload: payload, + }) + _, _ = s.conn.Write(makeTextFrame(msg)) +} + +// close shuts down the server connection. +func (s *fakeWSServer) close() { + if s.conn != nil { + _ = s.conn.Close() + } +} + +func TestRunSessionCtxCancel(t *testing.T) { + srv := newFakeWSServer(t) + + go srv.acceptAndJoin(t) + + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan error, 1) + go func() { + done <- runSession(ctx, "ws://"+srv.addr(), "token") + }() + + // Give the session time to enter the main loop + time.Sleep(200 * time.Millisecond) + cancel() + + select { + case err := <-done: + if err != nil { + t.Errorf("expected nil error on ctx cancel, got: %v", err) + } + case <-time.After(5 * time.Second): + t.Error("runSession did not exit after ctx cancel") + } + srv.close() +} + +func TestRunSessionReadError(t *testing.T) { + srv := newFakeWSServer(t) + + go func() { + srv.acceptAndJoin(t) + // Small delay, then close to trigger read error + time.Sleep(200 * time.Millisecond) + srv.close() + }() + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + if err == nil { + t.Error("expected read error") + } + if !strings.Contains(err.Error(), "read:") { + t.Errorf("expected 'read:' in error, got: %v", err) + } +} + +func TestRunSessionInvalidMessage(t *testing.T) { + srv := newFakeWSServer(t) + + go func() { + srv.acceptAndJoin(t) + // Send invalid JSON — should be logged but not crash + time.Sleep(100 * time.Millisecond) + _, _ = srv.conn.Write(makeTextFrame([]byte("not json"))) + // Then close to end session + time.Sleep(100 * time.Millisecond) + srv.close() + }() + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + // Should end with read error from close, not crash + if err == nil { + t.Error("expected error") + } +} + +func TestRunSessionConnectError(t *testing.T) { + err := runSession(context.Background(), "ws://127.0.0.1:1", "token") + if err == nil { + t.Error("expected connect error") + } + if !strings.Contains(err.Error(), "connect:") { + t.Errorf("expected 'connect:' in error, got: %v", err) + } +} + +func TestRunSessionJoinWriteError(t *testing.T) { + // Server accepts WS upgrade but closes immediately after + 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 + } + // Read HTTP upgrade + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + key := extractWSKey(string(buf[:n])) + 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)) + // Close immediately so join write may fail + _ = conn.Close() + }() + + err = runSession(context.Background(), "ws://"+ln.Addr().String(), "token") + if err == nil { + t.Error("expected error") + } +} + +func TestRunSessionJoinUnmarshalError(t *testing.T) { + // Server sends binary garbage as 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) + key := extractWSKey(string(buf[:n])) + 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)) + frameBuf := make([]byte, 4096) + _, _ = conn.Read(frameBuf) + // Send invalid JSON as reply + _, _ = conn.Write(makeTextFrame([]byte("{invalid json"))) + time.Sleep(time.Second) + }() + + err = runSession(context.Background(), "ws://"+ln.Addr().String(), "token") + if err == nil { + t.Error("expected unmarshal error") + } + if !strings.Contains(err.Error(), "join reply unmarshal") { + t.Errorf("expected 'join reply unmarshal' in error, got: %v", err) + } +} + +func TestRunSessionReadErrorDuringJoin(t *testing.T) { + // Server sends upgrade then closes before sending 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 + } + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + key := extractWSKey(string(buf[:n])) + 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)) + frameBuf := make([]byte, 4096) + _, _ = conn.Read(frameBuf) + // Close without sending reply + _ = conn.Close() + }() + + err = runSession(context.Background(), "ws://"+ln.Addr().String(), "token") + if err == nil { + t.Error("expected read during join error") + } + if !strings.Contains(err.Error(), "read during join") { + t.Errorf("expected 'read during join' in error, got: %v", err) + } +} + +func TestRunAgentReconnectOnError(t *testing.T) { + // Server that fails first connection then succeeds, then sends restart + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer func() { _ = ln.Close() }() + + connCount := 0 + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + connCount++ + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + key := extractWSKey(string(buf[:n])) + accept := computeAcceptKey(key) + + if connCount == 1 { + // First connection: upgrade then close immediately + 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)) + frameBuf := make([]byte, 4096) + _, _ = conn.Read(frameBuf) + _ = conn.Close() + } else { + // Second connection: proper session with restart + 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)) + frameBuf := make([]byte, 4096) + _, _ = conn.Read(frameBuf) + reply, _ := json.Marshal(channelMsg{ + Topic: "agent:agent-0", + Event: "phx_reply", + Payload: json.RawMessage(`{"status":"ok"}`), + Ref: strPtr("1"), + }) + _, _ = conn.Write(makeTextFrame(reply)) + time.Sleep(50 * time.Millisecond) + restart, _ := json.Marshal(channelMsg{ + Topic: "agent:agent-0", + Event: "restart", + Payload: json.RawMessage(`{}`), + }) + _, _ = conn.Write(makeTextFrame(restart)) + time.Sleep(time.Second) + _ = conn.Close() + } + } + }() + + origExit := osExit + defer func() { osExit = origExit }() + exitCalled := make(chan int, 1) + osExit = func(code int) { exitCalled <- code } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + done := make(chan struct{}) + go func() { + runAgent(ctx, "ws://"+ln.Addr().String(), "token") + close(done) + }() + + select { + case code := <-exitCalled: + if code != 0 { + t.Errorf("expected exit code 0, got %d", code) + } + case <-time.After(10 * time.Second): + t.Error("runAgent did not reconnect and restart") + } +} + +func TestRunAgentContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately + + done := make(chan struct{}) + go func() { + runAgent(ctx, "ws://127.0.0.1:1", "token") + close(done) + }() + + select { + case <-done: + // runAgent returned as expected + case <-time.After(5 * time.Second): + t.Error("runAgent did not return after context cancellation") + } +} + +func TestRunAgentRestart(t *testing.T) { + origExit := osExit + defer func() { osExit = origExit }() + + exitCalled := make(chan int, 1) + osExit = func(code int) { + exitCalled <- code + } + + // Start a fake WebSocket server that accepts the join and sends restart + 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 + frameBuf := make([]byte, 4096) + _, _ = conn.Read(frameBuf) + + // Send join OK reply + reply, _ := json.Marshal(channelMsg{ + Topic: "agent:agent-0", + Event: "phx_reply", + Payload: json.RawMessage(`{"status":"ok"}`), + Ref: strPtr("1"), + }) + _, _ = conn.Write(makeTextFrame(reply)) + + // Small delay then send restart event + time.Sleep(50 * time.Millisecond) + restart, _ := json.Marshal(channelMsg{ + Topic: "agent:agent-0", + Event: "restart", + Payload: json.RawMessage(`{}`), + }) + _, _ = conn.Write(makeTextFrame(restart)) + + // Keep connection open briefly + time.Sleep(2 * time.Second) + }() + + addr := ln.Addr().String() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + done := make(chan struct{}) + go func() { + runAgent(ctx, "ws://"+addr, "test-token") + close(done) + }() + + select { + case code := <-exitCalled: + if code != 0 { + t.Errorf("expected exit code 0, got %d", code) + } + case <-time.After(5 * time.Second): + t.Error("osExit was not called after restart") + } +} + +// readMaskedFrame reads a single masked WebSocket frame from the server side. +func readMaskedFrame(conn net.Conn) ([]byte, error) { + var header [2]byte + if _, err := io.ReadFull(conn, header[:]); err != nil { + return nil, err + } + masked := header[1]&0x80 != 0 + length := uint64(header[1] & 0x7F) + switch length { + case 126: + var ext [2]byte + if _, err := io.ReadFull(conn, ext[:]); err != nil { + return nil, err + } + length = uint64(binary.BigEndian.Uint16(ext[:])) + case 127: + var ext [8]byte + if _, err := io.ReadFull(conn, ext[:]); err != nil { + return nil, err + } + length = binary.BigEndian.Uint64(ext[:]) + } + var maskKey [4]byte + if masked { + if _, err := io.ReadFull(conn, maskKey[:]); err != nil { + return nil, err + } + } + payload := make([]byte, length) + if length > 0 { + if _, err := io.ReadFull(conn, payload); err != nil { + return nil, err + } + } + if masked { + for i := range payload { + payload[i] ^= maskKey[i%4] + } + } + return payload, nil +} + +// drainFrames reads and discards client frames until stop channel is closed or error. +func drainFrames(conn net.Conn, stop <-chan struct{}) { + for { + select { + case <-stop: + return + default: + } + _ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond)) + if _, err := readMaskedFrame(conn); err != nil { + select { + case <-stop: + return + default: + continue // timeout, retry + } + } + } +} + +func TestRunSessionProcessesJobResults(t *testing.T) { + // Mock external dependencies for fast execution + origSnmpDial := snmpDial + origMtDial := mikrotikDial + origPing := doPing + defer func() { + snmpDial = origSnmpDial + mikrotikDial = origMtDial + doPing = origPing + }() + + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + }, func() {}, nil + } + mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { + return nil, fmt.Errorf("unreachable") + } + doPing = func(ip string, timeoutMs int) (float64, error) { + return 1.5, nil + } + + srv := newFakeWSServer(t) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + srv.acceptAndJoin(t) + + // Drain client frames so writes don't block + stopDrain := make(chan struct{}) + go drainFrames(srv.conn, stopDrain) + + time.Sleep(100 * time.Millisecond) + + // SNMP job → snmpResultCh → snmpBatch → flushSnmpBatch → sendBinaryResult + srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{ + JobId: "s1", JobType: pb.JobType_POLL, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + })) + + // Mikrotik job → mikrotikResultCh → sendBinaryResult("mikrotik_result") + srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{ + JobId: "m1", JobType: pb.JobType_MIKROTIK, + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, + })) + + // Ping job → monitoringCheckCh → sendBinaryResult("monitoring_check") + srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{ + JobId: "p1", JobType: pb.JobType_PING, DeviceId: "dev1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + })) + + // Credential test → credTestResultCh → sendBinaryResult("credential_test_result") + srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{ + JobId: "ct1", JobType: pb.JobType_TEST_CREDENTIALS, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + })) + + // Check job → checkResultCh → sendBinaryResult("check_result") + checkList := &pb.CheckList{Checks: []*pb.Check{ + {Id: "c1", CheckType: "tcp", TimeoutMs: 500, + Config: &pb.Check_Tcp{Tcp: &pb.TcpCheckConfig{Host: "127.0.0.1", Port: 1}}}, + }} + bin, _ := proto.Marshal(checkList) + checkPayload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString(bin)}) + srv.sendEvent("check_jobs", checkPayload) + + // Wait for results to flow through all channels + time.Sleep(2 * time.Second) + close(stopDrain) + srv.close() + }() + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + <-serverDone + if err == nil { + t.Error("expected error after server close") + } +} + +func TestRunSessionSnmpBatchThreshold(t *testing.T) { + origSnmpDial := snmpDial + defer func() { snmpDial = origSnmpDial }() + + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + }, func() {}, nil + } + + srv := newFakeWSServer(t) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + srv.acceptAndJoin(t) + + stopDrain := make(chan struct{}) + go drainFrames(srv.conn, stopDrain) + + time.Sleep(100 * time.Millisecond) + + // Send 55 SNMP jobs to trigger batch threshold (>=50) + jobs := make([]*pb.AgentJob, 55) + for i := range jobs { + jobs[i] = &pb.AgentJob{ + JobId: fmt.Sprintf("s%d", i), + JobType: pb.JobType_POLL, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + } + } + srv.sendEvent("jobs", makeJobPayload(jobs...)) + + time.Sleep(3 * time.Second) + close(stopDrain) + srv.close() + }() + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + <-serverDone + if err == nil { + t.Error("expected error after server close") + } +} + +func TestExecuteCheckPoolFull(t *testing.T) { + // Use a pool with 1 worker and 1 queue slot, then fill both + p := &jobPools{ + snmp: newWorkerPool(4), + mikrotik: newWorkerPool(4), + ping: newWorkerPool(4), + checks: newWorkerPool(1), + } + t.Cleanup(func() { p.snmp.stop(); p.mikrotik.stop(); p.ping.stop(); p.checks.stop() }) + + done := make(chan struct{}) + started := make(chan struct{}) + + // Block the single worker + p.checks.submit(context.Background(), func() { + close(started) + <-done + }) + <-started + + // Fill the queue slot + p.checks.submit(context.Background(), func() { <-done }) + + // Pool is now full — submit with cancelled ctx will fail + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + checkCh := make(chan *pb.CheckResult, 1) + check := &pb.Check{Id: "c1", CheckType: "tcp", TimeoutMs: 1000} + executeCheck(ctx, check, p, checkCh) + // Should log "check rejected (pool full)" but not panic + close(done) +} + +func TestExecuteCheckCtxDoneInClosure(t *testing.T) { + // Tests the ctx.Done path inside executeCheck's closure: + // Submit succeeds, check runs, but result channel is blocked so ctx.Done fires. + p := testPools(t) + ctx, cancel := context.WithCancel(context.Background()) + checkCh := make(chan *pb.CheckResult) // unbuffered, no reader + + check := &pb.Check{Id: "c1", CheckType: "tcp", TimeoutMs: 100, + Config: &pb.Check_Tcp{Tcp: &pb.TcpCheckConfig{Host: "127.0.0.1", Port: 1}}} + + executeCheck(ctx, check, p, checkCh) + + // Wait for the check to complete (TCP to port 1 fails fast) + time.Sleep(500 * time.Millisecond) + // Cancel ctx so the closure's select picks ctx.Done instead of blocked channel send + cancel() + time.Sleep(100 * time.Millisecond) +} + +func TestRunSessionRestartInMainLoop(t *testing.T) { + srv := newFakeWSServer(t) + + go func() { + srv.acceptAndJoin(t) + stopDrain := make(chan struct{}) + go drainFrames(srv.conn, stopDrain) + + time.Sleep(100 * time.Millisecond) + // Send restart event — exercised in the main loop select + srv.sendEvent("restart", json.RawMessage(`{}`)) + time.Sleep(time.Second) + close(stopDrain) + srv.close() + }() + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + if err != errRestartRequested { + t.Errorf("expected errRestartRequested, got: %v", err) + } +} + +func TestRunSessionHeartbeats(t *testing.T) { + // Use short heartbeat intervals to exercise heartbeat paths + origHB := heartbeatInterval + origCHB := channelHeartbeatInterval + defer func() { + heartbeatInterval = origHB + channelHeartbeatInterval = origCHB + }() + heartbeatInterval = 100 * time.Millisecond + channelHeartbeatInterval = 100 * time.Millisecond + + srv := newFakeWSServer(t) + go func() { + srv.acceptAndJoin(t) + stopDrain := make(chan struct{}) + go drainFrames(srv.conn, stopDrain) + // Let heartbeats fire a few times + time.Sleep(500 * time.Millisecond) + close(stopDrain) + srv.close() + }() + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + if err == nil { + t.Error("expected error after server close") + } +} + +func TestRunAgentCancelDuringRetry(t *testing.T) { + // Connects to a port nothing listens on, then cancel during retry delay + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan struct{}) + go func() { + runAgent(ctx, "ws://127.0.0.1:1", "token") + close(done) + }() + + // Wait for first connection attempt to fail and retry delay to start + time.Sleep(1500 * time.Millisecond) + cancel() + + select { + case <-done: + // runAgent returned after cancel during retry + case <-time.After(5 * time.Second): + t.Error("runAgent did not return after cancel during retry") + } +} diff --git a/checks_test.go b/checks_test.go new file mode 100644 index 0000000..f2cd47d --- /dev/null +++ b/checks_test.go @@ -0,0 +1,1372 @@ +package main + +import ( + "bufio" + "context" + "crypto/tls" + "fmt" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/towerops-app/towerops-agent/pb" +) + +// --------------------------------------------------------------------------- +// ExecuteCheck routing tests +// --------------------------------------------------------------------------- + +func TestExecuteCheck_UnknownCheckType(t *testing.T) { + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-1", + CheckType: "foobar", + TimeoutMs: 5000, + }) + if result.Status != 3 { + t.Fatalf("expected status 3, got %d", result.Status) + } + if result.CheckId != "chk-1" { + t.Fatalf("expected CheckId chk-1, got %s", result.CheckId) + } + if !strings.Contains(result.Output, "Unknown check type") { + t.Fatalf("expected unknown check type message, got %s", result.Output) + } +} + +func TestExecuteCheck_MissingHTTPConfig(t *testing.T) { + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-2", + CheckType: "http", + TimeoutMs: 5000, + // No Config set + }) + if result.Status != 3 { + t.Fatalf("expected status 3, got %d", result.Status) + } + if !strings.Contains(result.Output, "Missing HTTP config") { + t.Fatalf("expected missing HTTP config message, got %s", result.Output) + } +} + +func TestExecuteCheck_MissingTCPConfig(t *testing.T) { + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-3", + CheckType: "tcp", + TimeoutMs: 5000, + }) + if result.Status != 3 { + t.Fatalf("expected status 3, got %d", result.Status) + } + if !strings.Contains(result.Output, "Missing TCP config") { + t.Fatalf("expected missing TCP config message, got %s", result.Output) + } +} + +func TestExecuteCheck_MissingDNSConfig(t *testing.T) { + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-4", + CheckType: "dns", + TimeoutMs: 5000, + }) + if result.Status != 3 { + t.Fatalf("expected status 3, got %d", result.Status) + } + if !strings.Contains(result.Output, "Missing DNS config") { + t.Fatalf("expected missing DNS config message, got %s", result.Output) + } +} + +func TestExecuteCheck_SetsCheckIdAndTimestamp(t *testing.T) { + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-ts", + CheckType: "unknown", + TimeoutMs: 5000, + }) + if result.CheckId != "chk-ts" { + t.Fatalf("expected CheckId chk-ts, got %s", result.CheckId) + } + if result.Timestamp == 0 { + t.Fatal("expected non-zero timestamp") + } + if result.ResponseTimeMs < 0 { + t.Fatalf("expected non-negative response time, got %f", result.ResponseTimeMs) + } +} + +func TestExecuteCheck_HTTPRouting(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + })) + defer srv.Close() + + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-route-http", + CheckType: "http", + TimeoutMs: 5000, + Config: &pb.Check_Http{ + Http: &pb.HttpCheckConfig{ + Url: srv.URL, + }, + }, + }) + if result.Status != 0 { + t.Fatalf("expected status 0, got %d: %s", result.Status, result.Output) + } +} + +func TestExecuteCheck_TCPRouting(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + _, portStr, _ := net.SplitHostPort(ln.Addr().String()) + + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-route-tcp", + CheckType: "tcp", + TimeoutMs: 5000, + Config: &pb.Check_Tcp{ + Tcp: &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: parsePort(portStr), + }, + }, + }) + if result.Status != 0 { + t.Fatalf("expected status 0, got %d: %s", result.Status, result.Output) + } +} + +func TestExecuteCheck_DNSRouting(t *testing.T) { + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-route-dns", + CheckType: "dns", + TimeoutMs: 5000, + Config: &pb.Check_Dns{ + Dns: &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "A", + }, + }, + }) + // DNS for localhost may or may not resolve depending on system config, + // but the routing should work regardless. + if result.Status == 3 { + t.Fatalf("expected routing to DNS handler, got status 3 (UNKNOWN): %s", result.Output) + } +} + +// --------------------------------------------------------------------------- +// HTTP check tests +// --------------------------------------------------------------------------- + +func TestHTTPCheck_SuccessfulGET(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + fmt.Fprint(w, "OK") + })) + defer srv.Close() + + status, output, rt := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } + if !strings.Contains(output, "HTTP 200 OK") { + t.Fatalf("expected HTTP 200 OK in output, got %s", output) + } + if rt < 0 { + t.Fatalf("expected non-negative response time, got %f", rt) + } +} + +func TestHTTPCheck_CustomMethod_POST(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("expected POST, got %s", r.Method) + w.WriteHeader(405) + return + } + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Method: "post", + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_CustomExpectedStatus(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(201) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + ExpectedStatus: 201, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_WrongStatusCode(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(500) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } + if !strings.Contains(output, "500") && !strings.Contains(output, "expected 200") { + t.Fatalf("expected status code mismatch message, got %s", output) + } +} + +func TestHTTPCheck_DefaultMethodIsGET(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" { + t.Errorf("expected default GET, got %s", r.Method) + w.WriteHeader(405) + return + } + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Method: "", // should default to GET + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_DefaultExpectedStatusIs200(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + ExpectedStatus: 0, // should default to 200 + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_VerifySslFalseWithSelfSigned(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + })) + defer srv.Close() + + // With VerifySsl=false (InsecureSkipVerify=true), the self-signed cert should be accepted + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + VerifySsl: false, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 with VerifySsl=false, got %d: %s", status, output) + } +} + +func TestHTTPCheck_VerifySslTrueRejectsSelfSigned(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + })) + defer srv.Close() + + // With VerifySsl=true (InsecureSkipVerify=false), self-signed cert should fail + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + VerifySsl: true, + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2 with VerifySsl=true on self-signed, got %d: %s", status, output) + } + if !strings.Contains(output, "Request failed") { + t.Fatalf("expected request failed message, got %s", output) + } +} + +func TestHTTPCheck_FollowRedirectsTrue(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/redirect" { + http.Redirect(w, r, "/final", http.StatusMovedPermanently) + return + } + w.WriteHeader(200) + fmt.Fprint(w, "final page") + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL + "/redirect", + FollowRedirects: true, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 after following redirect, got %d: %s", status, output) + } + if !strings.Contains(output, "200") { + t.Fatalf("expected final 200 status in output, got %s", output) + } +} + +func TestHTTPCheck_FollowRedirectsFalse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/redirect" { + http.Redirect(w, r, "/final", http.StatusMovedPermanently) + return + } + w.WriteHeader(200) + })) + defer srv.Close() + + // When FollowRedirects is false, we should see the 301 directly + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL + "/redirect", + FollowRedirects: false, + ExpectedStatus: 301, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 (matching 301), got %d: %s", status, output) + } +} + +func TestHTTPCheck_FollowRedirectsFalse_DefaultExpects200(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "/final", http.StatusMovedPermanently) + })) + defer srv.Close() + + // FollowRedirects=false and default expected=200, but we get 301 → CRITICAL + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL + "/redirect", + FollowRedirects: false, + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2 (301 != 200), got %d: %s", status, output) + } +} + +func TestHTTPCheck_RegexMatchSucceeds(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + fmt.Fprint(w, "Hello World! Version 1.2.3") + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Regex: `Version \d+\.\d+\.\d+`, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 for regex match, got %d: %s", status, output) + } +} + +func TestHTTPCheck_RegexMatchFails(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + fmt.Fprint(w, "Hello World!") + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Regex: `Version \d+`, + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2 for regex mismatch, got %d: %s", status, output) + } + if !strings.Contains(output, "does not match") { + t.Fatalf("expected 'does not match' in output, got %s", output) + } +} + +func TestHTTPCheck_InvalidRegex(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + fmt.Fprint(w, "some body") + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Regex: `[invalid`, + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2 for invalid regex, got %d: %s", status, output) + } + if !strings.Contains(output, "Invalid regex") { + t.Fatalf("expected 'Invalid regex' in output, got %s", output) + } +} + +func TestHTTPCheck_CustomHeaders(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Custom") != "test-value" { + t.Errorf("expected X-Custom header to be test-value, got %s", r.Header.Get("X-Custom")) + w.WriteHeader(400) + return + } + if r.Header.Get("Authorization") != "Bearer abc123" { + t.Errorf("expected Authorization header, got %s", r.Header.Get("Authorization")) + w.WriteHeader(401) + return + } + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Headers: map[string]string{ + "X-Custom": "test-value", + "Authorization": "Bearer abc123", + }, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_UnreachableServer(t *testing.T) { + // Use a non-routable address to guarantee failure + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: "http://192.0.2.1:1/unreachable", + }, 1000) + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } + if !strings.Contains(output, "Request failed") { + t.Fatalf("expected 'Request failed' in output, got %s", output) + } +} + +func TestHTTPCheck_InvalidURL(t *testing.T) { + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: "://not-a-url", + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } + if !strings.Contains(output, "Failed to create request") { + t.Fatalf("expected 'Failed to create request' in output, got %s", output) + } +} + +func TestHTTPCheck_RequestWithBody(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + buf := make([]byte, 1024) + n, _ := r.Body.Read(buf) + body := string(buf[:n]) + if body != `{"key":"value"}` { + t.Errorf("expected JSON body, got %s", body) + w.WriteHeader(400) + return + } + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Method: "POST", + Body: `{"key":"value"}`, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_ContextCancellation(t *testing.T) { + // Server that hangs forever + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + defer srv.Close() + + ctx, cancel := context.WithCancel(context.Background()) + // Cancel after a short delay + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + status, output, _ := executeHTTPCheck(ctx, &pb.HttpCheckConfig{ + Url: srv.URL, + }, 30000) // long timeout so the context cancel hits first + + if status != 2 { + t.Fatalf("expected status 2 on context cancel, got %d: %s", status, output) + } + if !strings.Contains(output, "Request failed") { + t.Fatalf("expected 'Request failed' in output, got %s", output) + } +} + +func TestHTTPCheck_LargeResponseBodyWithRegex(t *testing.T) { + // Generate a large body (500KB) with a marker at the end + bigBody := strings.Repeat("a", 500*1024) + "MARKER_FOUND_HERE" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + fmt.Fprint(w, bigBody) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Regex: `MARKER_FOUND_HERE`, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 for regex match in large body, got %d: %s", status, output) + } +} + +func TestHTTPCheck_SlowServerTimeout(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(500 * time.Millisecond) + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + }, 50) // 50ms timeout, server takes 500ms + + if status != 2 { + t.Fatalf("expected status 2 for timeout, got %d: %s", status, output) + } +} + +func TestHTTPCheck_EmptyBody_NoRegex(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +// --------------------------------------------------------------------------- +// TCP check tests +// --------------------------------------------------------------------------- + +func TestTCPCheck_PortOpen(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + // Accept connections in background so dial doesn't hang + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + conn.Close() + } + }() + + port := parsePort(portFromListener(ln)) + + status, output, rt := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } + if !strings.Contains(output, fmt.Sprintf("TCP port %d open", port)) { + t.Fatalf("expected port open message, got %s", output) + } + if rt < 0 { + t.Fatalf("expected non-negative response time, got %f", rt) + } +} + +func TestTCPCheck_PortClosed(t *testing.T) { + // Find a port that's definitely not listening by binding and immediately closing + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + port := parsePort(portFromListener(ln)) + ln.Close() // close immediately so port is refused + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + }, 2000) + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } + if !strings.Contains(output, "Connection failed") { + t.Fatalf("expected 'Connection failed' in output, got %s", output) + } +} + +func TestTCPCheck_SendExpectSuccess(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + // Echo server + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + scanner := bufio.NewScanner(c) + if scanner.Scan() { + line := scanner.Text() + fmt.Fprintf(c, "ECHO:%s\n", line) + } + }(conn) + } + }() + + port := parsePort(portFromListener(ln)) + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + Send: "hello\n", + Expect: "ECHO:hello", + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestTCPCheck_SendExpectMismatch(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + buf := make([]byte, 1024) + c.Read(buf) + fmt.Fprint(c, "WRONG_RESPONSE") + }(conn) + } + }() + + port := parsePort(portFromListener(ln)) + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + Send: "hello", + Expect: "EXPECTED_VALUE", + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } + if !strings.Contains(output, "Unexpected response") { + t.Fatalf("expected 'Unexpected response' in output, got %s", output) + } +} + +func TestTCPCheck_SendWithEmptyExpect(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + buf := make([]byte, 1024) + c.Read(buf) + // Don't send anything back + }(conn) + } + }() + + port := parsePort(portFromListener(ln)) + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + Send: "data\n", + Expect: "", // no expect check + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 (send only, no expect), got %d: %s", status, output) + } +} + +func TestTCPCheck_IPv6Localhost(t *testing.T) { + ln, err := net.Listen("tcp6", "[::1]:0") + if err != nil { + t.Skip("IPv6 not available on this system") + } + defer ln.Close() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + conn.Close() + } + }() + + port := parsePort(portFromListener(ln)) + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "::1", + Port: port, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 for IPv6, got %d: %s", status, output) + } +} + +func TestTCPCheck_ReadTimeoutOnExpect(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + // Server that accepts and reads but never writes back + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + buf := make([]byte, 1024) + c.Read(buf) + // Intentionally never respond - hold connection open + time.Sleep(10 * time.Second) + }(conn) + } + }() + + port := parsePort(portFromListener(ln)) + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + Send: "hello", + Expect: "response", + }, 200) // short timeout + + if status != 2 { + t.Fatalf("expected status 2 for read timeout, got %d: %s", status, output) + } + if !strings.Contains(output, "Receive failed") { + t.Fatalf("expected 'Receive failed' in output, got %s", output) + } +} + +func TestTCPCheck_VeryShortTimeout(t *testing.T) { + // Use TEST-NET address that won't respond + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "192.0.2.1", + Port: 80, + }, 1) // 1ms timeout - should fail + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } +} + +func TestTCPCheck_BinaryDataSendExpect(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + // Server that echoes binary data back + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + buf := make([]byte, 4096) + n, err := c.Read(buf) + if err != nil { + return + } + c.Write(buf[:n]) + }(conn) + } + }() + + port := parsePort(portFromListener(ln)) + + // Send some binary-ish data + sendData := "BIN\x00\x01\x02DATA" + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + Send: sendData, + Expect: sendData, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 for binary echo, got %d: %s", status, output) + } +} + +// --------------------------------------------------------------------------- +// DNS check tests +// --------------------------------------------------------------------------- + +func TestDNSCheck_ARecord(t *testing.T) { + status, output, rt := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "A", + }, 5000) + + // localhost should resolve on most systems, but skip if it doesn't + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available for localhost") + } + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } + if !strings.Contains(output, "Resolved to:") { + t.Fatalf("expected 'Resolved to:' in output, got %s", output) + } + if rt < 0 { + t.Fatalf("expected non-negative response time, got %f", rt) + } +} + +func TestDNSCheck_AAAARecord(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "AAAA", + }, 5000) + + // AAAA for localhost may or may not exist - just verify it doesn't crash + // and returns a valid status + if status != 0 && status != 2 { + t.Fatalf("expected status 0 or 2, got %d: %s", status, output) + } +} + +func TestDNSCheck_CNAMERecord(t *testing.T) { + // Use a well-known domain that has a CNAME + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "www.google.com", + RecordType: "CNAME", + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available") + } + // CNAME lookup may return the hostname itself if no CNAME exists + if status != 0 && status != 2 { + t.Fatalf("expected status 0 or 2, got %d: %s", status, output) + } +} + +func TestDNSCheck_MXRecord(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "google.com", + RecordType: "MX", + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available") + } + if status != 0 { + t.Fatalf("expected status 0 for MX lookup, got %d: %s", status, output) + } + if !strings.Contains(output, "Resolved to:") { + t.Fatalf("expected 'Resolved to:' in output, got %s", output) + } +} + +func TestDNSCheck_TXTRecord(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "google.com", + RecordType: "TXT", + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available") + } + if status != 0 { + t.Fatalf("expected status 0 for TXT lookup, got %d: %s", status, output) + } +} + +func TestDNSCheck_UnsupportedRecordType(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "example.com", + RecordType: "SRV", + }, 5000) + + if status != 3 { + t.Fatalf("expected status 3, got %d: %s", status, output) + } + if !strings.Contains(output, "Unsupported record type") { + t.Fatalf("expected 'Unsupported record type' in output, got %s", output) + } +} + +func TestDNSCheck_ExpectedMatchesResult(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "A", + Expected: "127.0.0.1", + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available for localhost") + } + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestDNSCheck_ExpectedDoesNotMatch(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "A", + Expected: "10.99.99.99", + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available for localhost") + } + // If localhost resolved, the expected won't match + if status != 2 { + t.Fatalf("expected status 2 for mismatched expected, got %d: %s", status, output) + } + if !strings.Contains(output, "Expected '10.99.99.99'") { + t.Fatalf("expected mismatch message, got %s", output) + } +} + +func TestDNSCheck_NonexistentDomain(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "this-domain-does-not-exist-towerops-test.invalid", + RecordType: "A", + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2 for NXDOMAIN, got %d: %s", status, output) + } +} + +func TestDNSCheck_DefaultRecordTypeIsA(t *testing.T) { + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "", // should default to A + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available for localhost") + } + // Should behave the same as explicitly specifying "A" + if status != 0 { + t.Fatalf("expected status 0 with default record type, got %d: %s", status, output) + } +} + +func TestDNSCheck_CustomDNSServer(t *testing.T) { + // Test with Google's public DNS + conn, err := net.DialTimeout("udp", "8.8.8.8:53", 2*time.Second) + if err != nil { + t.Skip("Cannot reach 8.8.8.8:53, skipping custom DNS server test") + } + conn.Close() + + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "example.com", + RecordType: "A", + Server: "8.8.8.8", + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 with custom DNS server, got %d: %s", status, output) + } +} + +func TestDNSCheck_RecordTypeCaseInsensitive(t *testing.T) { + // The code does strings.ToUpper, so lowercase should work + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "localhost", + RecordType: "a", // lowercase + }, 5000) + + if status == 2 && strings.Contains(output, "DNS query failed") { + t.Skip("DNS resolution not available for localhost") + } + if status != 0 { + t.Fatalf("expected status 0 with lowercase record type, got %d: %s", status, output) + } +} + +func TestDNSCheck_NoRecordsFound(t *testing.T) { + // AAAA for a domain that likely only has A records + // Use a known domain that almost certainly won't have AAAA + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "this-domain-does-not-exist-towerops-test.invalid", + RecordType: "AAAA", + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2, got %d: %s", status, output) + } +} + +func TestDNSCheck_VeryShortTimeout(t *testing.T) { + // Use a custom DNS server with 1ms timeout - should timeout + status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{ + Hostname: "example.com", + RecordType: "A", + Server: "8.8.8.8", + }, 1) // 1ms timeout + + // Should fail due to timeout + if status != 2 { + // On very fast networks this could actually succeed, so just verify it doesn't crash + t.Logf("DNS with 1ms timeout returned status %d: %s (may succeed on fast networks)", status, output) + } +} + +func TestTCPCheck_SendFailsOnClosedConnection(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + // Server accepts and immediately closes the connection + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + conn.Close() + } + }() + + port := parsePort(portFromListener(ln)) + + // Send a large payload to trigger a write error on a closed connection. + // The first small write might succeed (kernel buffer), but a large write + // after the peer has closed should fail with a broken pipe or similar. + largePayload := strings.Repeat("x", 1024*1024) // 1MB + // Give the server time to close the connection + time.Sleep(50 * time.Millisecond) + + status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + Send: largePayload, + Expect: "something", + }, 5000) + + // This may hit either "Send failed" or "Receive failed" depending on timing + if status != 2 { + t.Fatalf("expected status 2 for write to closed conn, got %d: %s", status, output) + } +} + +// --------------------------------------------------------------------------- +// Edge case tests +// --------------------------------------------------------------------------- + +func TestHTTPCheck_MethodUppercased(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "PUT" { + t.Errorf("expected PUT, got %s", r.Method) + w.WriteHeader(405) + return + } + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Method: "put", // lowercase + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_ResponseTimeIsPositive(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(10 * time.Millisecond) + w.WriteHeader(200) + })) + defer srv.Close() + + _, _, rt := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + }, 5000) + + if rt <= 0 { + t.Fatalf("expected positive response time, got %f", rt) + } +} + +func TestTCPCheck_ResponseTimeReported(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + conn.Close() + } + }() + + port := parsePort(portFromListener(ln)) + _, _, rt := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{ + Host: "127.0.0.1", + Port: port, + }, 5000) + + if rt < 0 { + t.Fatalf("expected non-negative response time, got %f", rt) + } +} + +func TestExecuteCheck_ResponseTimeFallback(t *testing.T) { + // When status/output are set directly (unknown type), responseTimeMs should be calculated + result := ExecuteCheck(context.Background(), &pb.Check{ + Id: "chk-fallback", + CheckType: "invalid", + TimeoutMs: 5000, + }) + + if result.ResponseTimeMs < 0 { + t.Fatalf("expected non-negative response time, got %f", result.ResponseTimeMs) + } +} + +func TestHTTPCheck_TLSServerNoVerify(t *testing.T) { + // Create a TLS server with custom cert + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + fmt.Fprint(w, "secure") + })) + srv.TLS = &tls.Config{} + srv.StartTLS() + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + VerifySsl: false, + Regex: "secure", + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_HeadMethod(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "HEAD" { + t.Errorf("expected HEAD, got %s", r.Method) + } + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Method: "HEAD", + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_DeleteMethod(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "DELETE" { + t.Errorf("expected DELETE, got %s", r.Method) + } + w.WriteHeader(204) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Method: "DELETE", + ExpectedStatus: 204, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0, got %d: %s", status, output) + } +} + +func TestHTTPCheck_RegexOnEmptyBody(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + // No body written + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Regex: "something", + }, 5000) + + if status != 2 { + t.Fatalf("expected status 2 (regex no match on empty body), got %d: %s", status, output) + } +} + +func TestHTTPCheck_MultipleRedirects(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/a": + http.Redirect(w, r, "/b", http.StatusFound) + case "/b": + http.Redirect(w, r, "/c", http.StatusFound) + case "/c": + w.WriteHeader(200) + fmt.Fprint(w, "final") + } + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL + "/a", + FollowRedirects: true, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 after multiple redirects, got %d: %s", status, output) + } +} + +func TestHTTPCheck_EmptyHeadersMap(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + })) + defer srv.Close() + + status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{ + Url: srv.URL, + Headers: map[string]string{}, + }, 5000) + + if status != 0 { + t.Fatalf("expected status 0 with empty headers, got %d: %s", status, output) + } +} + +// --------------------------------------------------------------------------- +// Helper functions +// --------------------------------------------------------------------------- + +// portFromListener extracts the port string from a net.Listener's address. +func portFromListener(ln net.Listener) string { + _, port, _ := net.SplitHostPort(ln.Addr().String()) + return port +} + +// parsePort converts a port string to uint32 for use in proto configs. +func parsePort(s string) uint32 { + var port uint32 + fmt.Sscanf(s, "%d", &port) + return port +} diff --git a/hostkeys_test.go b/hostkeys_test.go index b36586c..7ecc781 100644 --- a/hostkeys_test.go +++ b/hostkeys_test.go @@ -1,7 +1,7 @@ package main import ( - "os" + "crypto/x509" "path/filepath" "sync" "testing" @@ -116,7 +116,20 @@ func TestHostKeyStoreSaveError(t *testing.T) { } func TestTlsCertFingerprint(t *testing.T) { - // Just verify it doesn't panic with nil - we test with real certs in integration - // The function is simple enough: sha256 of cert.Raw - _ = os.Getenv("dummy") // placeholder + cert := &x509.Certificate{ + Raw: []byte("test certificate data"), + } + fp := tlsCertFingerprint(cert) + if fp == "" { + t.Error("expected non-empty fingerprint") + } + // Verify it's a hex-encoded SHA256 (64 chars) + if len(fp) != 64 { + t.Errorf("expected 64-char hex fingerprint, got %d chars: %s", len(fp), fp) + } + // Same input should produce same output + fp2 := tlsCertFingerprint(cert) + if fp != fp2 { + t.Error("fingerprint not deterministic") + } } diff --git a/main.go b/main.go index af71359..bd9a258 100644 --- a/main.go +++ b/main.go @@ -93,7 +93,7 @@ func toWebSocketURL(rawURL string) string { if strings.HasPrefix(result, "ws://") && !insecureFlag { slog.Error("plaintext ws:// connection rejected — use wss:// or pass --insecure to allow") - os.Exit(1) + osExit(1) } return result } diff --git a/main_test.go b/main_test.go index bb7b7ab..3b57722 100644 --- a/main_test.go +++ b/main_test.go @@ -39,6 +39,14 @@ func TestSanitizeURL(t *testing.T) { } } +func TestIsFlagSet(t *testing.T) { + // In tests, flags are not set via flag.Parse on the default flag set + // so isFlagSet should return false for any arbitrary name. + if isFlagSet("nonexistent-flag-xyz") { + t.Error("expected isFlagSet to return false for unset flag") + } +} + func TestToWebSocketURL(t *testing.T) { // Enable insecure for testing plaintext conversions origInsecure := insecureFlag @@ -62,3 +70,33 @@ func TestToWebSocketURL(t *testing.T) { } } } + +func TestToWebSocketURLRejectsPlaintext(t *testing.T) { + origInsecure := insecureFlag + origExit := osExit + defer func() { + insecureFlag = origInsecure + osExit = origExit + }() + + insecureFlag = false + exitCode := -1 + osExit = func(code int) { exitCode = code } + + tests := []struct { + name string + input string + }{ + {"ws:// scheme", "ws://localhost:4000"}, + {"http:// converts to ws://", "http://localhost:4000"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + exitCode = -1 + toWebSocketURL(tt.input) + if exitCode != 1 { + t.Errorf("expected osExit(1), got %d", exitCode) + } + }) + } +} diff --git a/mikrotik_test.go b/mikrotik_test.go index bc8105e..1872287 100644 --- a/mikrotik_test.go +++ b/mikrotik_test.go @@ -2,10 +2,21 @@ package main import ( "bytes" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "encoding/pem" "fmt" "io" + "math/big" "net" + "path/filepath" + "strings" + "sync" "testing" + "time" ) func TestEncodeLength(t *testing.T) { @@ -467,6 +478,78 @@ func TestMikrotikConnectSSL(t *testing.T) { } } +func TestMikrotikConnectSSLWithServer(t *testing.T) { + // Generate a self-signed cert at runtime + cert := generateTestCert(t) + tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}} + ln, err := tls.Listen("tcp", "127.0.0.1:0", tlsConfig) + if err != nil { + t.Fatal(err) + } + defer func() { _ = ln.Close() }() + + // Reset host key store for this test + origStore := globalHostKeys + defer func() { + hostKeysOnce = sync.Once{} + globalHostKeys = origStore + }() + hostKeysOnce = sync.Once{} + t.Setenv("TOWEROPS_HOST_KEYS_FILE", filepath.Join(t.TempDir(), "hosts.json")) + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer func() { _ = conn.Close() }() + sc := &mikrotikClient{conn: conn} + _, _ = sc.readSentence() + _ = sc.writeSentence([]string{"!done"}) + _, _ = sc.readSentence() + _ = sc.writeSentence([]string{"!fatal"}) + }() + + _, port, _ := net.SplitHostPort(ln.Addr().String()) + var portNum uint32 + _, _ = fmt.Sscanf(port, "%d", &portNum) + + client, err := mikrotikConnect("127.0.0.1", portNum, "admin", "pass", true) + if err != nil { + t.Fatalf("expected TLS connection to succeed: %v", err) + } + _ = client.close() +} + +func generateTestCert(t *testing.T) tls.Certificate { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatal(err) + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + NotBefore: time.Now(), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + t.Fatal(err) + } + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER}) + keyDER, err := x509.MarshalECPrivateKey(key) + if err != nil { + t.Fatal(err) + } + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}) + tlsCert, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + t.Fatal(err) + } + return tlsCert +} + func TestReadResponseEmptySentence(t *testing.T) { // An empty sentence (just the terminator byte) should be skipped var buf bytes.Buffer @@ -636,3 +719,51 @@ func (rwc *readWriteCloser) Close() error { } return nil } + +func TestMikrotikConnectSSLTOFUMismatch(t *testing.T) { + cert := generateTestCert(t) + tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}} + ln, err := tls.Listen("tcp", "127.0.0.1:0", tlsConfig) + if err != nil { + t.Fatal(err) + } + defer func() { _ = ln.Close() }() + + // Reset host key store and pre-populate with wrong fingerprint + origStore := globalHostKeys + defer func() { + hostKeysOnce = sync.Once{} + globalHostKeys = origStore + }() + hostKeysOnce = sync.Once{} + t.Setenv("TOWEROPS_HOST_KEYS_FILE", filepath.Join(t.TempDir(), "hosts.json")) + + _, port, _ := net.SplitHostPort(ln.Addr().String()) + var portNum uint32 + _, _ = fmt.Sscanf(port, "%d", &portNum) + + // Pre-register wrong fingerprint so TOFU verification fails + store := getHostKeyStore() + store.keys["tls:"+net.JoinHostPort("127.0.0.1", port)] = "wrong_fingerprint" + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer func() { _ = conn.Close() }() + // Complete TLS handshake so client gets peer certificates + if tlsConn, ok := conn.(*tls.Conn); ok { + _ = tlsConn.Handshake() + } + time.Sleep(time.Second) + }() + + _, err = mikrotikConnect("127.0.0.1", portNum, "admin", "pass", true) + if err == nil { + t.Error("expected TOFU verification failure") + } + if !strings.Contains(err.Error(), "TOFU verification failed") { + t.Errorf("expected 'TOFU verification failed' in error, got: %v", err) + } +} diff --git a/ping_test.go b/ping_test.go index e5b2fb5..5a91d4d 100644 --- a/ping_test.go +++ b/ping_test.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "runtime" "testing" ) @@ -65,6 +66,13 @@ func TestIcmpPingInvalidIP(t *testing.T) { } } +func TestErrICMPUnavailableError(t *testing.T) { + err := &errICMPUnavailable{err: fmt.Errorf("permission denied")} + if err.Error() != "permission denied" { + t.Errorf("got %q, want %q", err.Error(), "permission denied") + } +} + func TestParsePingTime(t *testing.T) { tests := []struct { name string diff --git a/websocket_test.go b/websocket_test.go index 659cea3..7bf9410 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -657,6 +657,60 @@ func TestWSDialIPv4Fallback(t *testing.T) { } } +func TestWSDialTLSPath(t *testing.T) { + // Verify the TLS dial path is exercised (connection will fail, but we hit the code path) + origTLS := tlsDial + defer func() { tlsDial = origTLS }() + + called := false + tlsDial = func(network, addr string) (net.Conn, error) { + called = true + return nil, fmt.Errorf("tls dial: connection refused") + } + + _, err := WSDial("wss://127.0.0.1:9999/path") + if err == nil { + t.Error("expected TLS dial error") + } + if !called { + t.Error("expected tlsDial to be called for wss:// URL") + } +} + +func TestWSDialMissingAcceptHeader(t *testing.T) { + // Server sends 101 but without Sec-WebSocket-Accept header + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer func() { _ = ln.Close() }() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer func() { _ = c.Close() }() + buf := make([]byte, 4096) + _, _ = c.Read(buf) + resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n" + _, _ = c.Write([]byte(resp)) + }(conn) + } + }() + + addr := ln.Addr().String() + _, err = WSDial("ws://" + addr + "/socket") + if err == nil { + t.Error("expected error for missing accept header") + } + if !strings.Contains(err.Error(), "missing Sec-WebSocket-Accept") { + t.Errorf("expected 'missing Sec-WebSocket-Accept' in error, got: %v", err) + } +} + func TestWSDialDefaultPorts(t *testing.T) { // Test that ws:// defaults to port 80 — will fail to connect but verifies URL parsing _, err := WSDial("ws://127.0.0.1/path")