From 6bf9584cb22d67494f0e9a9f77cfce3cb1f933c1 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Wed, 4 Mar 2026 16:35:02 -0600 Subject: [PATCH] test: push coverage from 92% to 97.6% (249 tests) Refactor main.go to extract testable runMain() function, make icmpListenPacket and syscallExec injectable, fix worker pool tests to properly fill queue buffers (n*4 capacity). Add comprehensive tests for CLI argument parsing, ping exec fallback, ICMP timeout paths, self-update happy path, and SNMP context cancellation. --- agent_test.go | 73 +++++++++++++++++--- main.go | 57 ++++++++------- main_test.go | 180 ++++++++++++++++++++++++++++++++++++++++++----- ping.go | 4 +- ping_test.go | 184 +++++++++++++++++++++++++++++++++++++++++++++++++ snmp_test.go | 35 ++++++++++ update.go | 3 +- update_test.go | 86 +++++++++++++++++++++++ 8 files changed, 570 insertions(+), 52 deletions(-) diff --git a/agent_test.go b/agent_test.go index 843bc35..f77abeb 100644 --- a/agent_test.go +++ b/agent_test.go @@ -755,8 +755,31 @@ func makeTextFrame(payload []byte) []byte { } func TestDispatchJobCancelledContext(t *testing.T) { - // submit returns false when context is cancelled, triggering "pool full" log - p := testPools(t) + // Use pools with size 1 (1 worker + 4 queue slots) and fill them + // so that submit reliably fails with cancelled context. + p := &jobPools{ + snmp: newWorkerPool(1), + mikrotik: newWorkerPool(1), + ping: newWorkerPool(1), + checks: newWorkerPool(1), + } + t.Cleanup(func() { p.snmp.stop(); p.mikrotik.stop(); p.ping.stop(); p.checks.stop() }) + + done := make(chan struct{}) + + // Block each pool's worker + fill queue slots + fillPool := func(pool *workerPool) { + started := make(chan struct{}) + pool.submit(context.Background(), func() { close(started); <-done }) + <-started + for range 4 { + pool.submit(context.Background(), func() { <-done }) + } + } + fillPool(p.snmp) + fillPool(p.mikrotik) + fillPool(p.ping) + snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) @@ -764,9 +787,9 @@ func TestDispatchJobCancelledContext(t *testing.T) { checkCh := make(chan *pb.CheckResult, 1) ctx, cancel := context.WithCancel(context.Background()) - cancel() // cancel immediately + cancel() - // All dispatch types should handle ctx cancellation gracefully + // All dispatch types should hit the "pool full" warning dispatchJob(ctx, &pb.AgentJob{ JobId: "s1", JobType: pb.JobType_POLL, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, }, p, snmpCh, mtCh, credCh, monCh, checkCh) @@ -782,6 +805,8 @@ func TestDispatchJobCancelledContext(t *testing.T) { dispatchJob(ctx, &pb.AgentJob{ JobId: "p1", JobType: pb.JobType_PING, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, }, p, snmpCh, mtCh, credCh, monCh, checkCh) + + close(done) } // fakeWSServer is a test helper that sets up a WebSocket server for runSession tests. @@ -1392,7 +1417,8 @@ func TestRunSessionSnmpBatchThreshold(t *testing.T) { } func TestExecuteCheckPoolFull(t *testing.T) { - // Use a pool with 1 worker and 1 queue slot, then fill both + // newWorkerPool(1) creates 1 worker goroutine + queue buffer of 1*4=4 + // We need to block the worker AND fill all 4 queue slots to make submit block. p := &jobPools{ snmp: newWorkerPool(4), mikrotik: newWorkerPool(4), @@ -1411,10 +1437,12 @@ func TestExecuteCheckPoolFull(t *testing.T) { }) <-started - // Fill the queue slot - p.checks.submit(context.Background(), func() { <-done }) + // Fill all 4 queue slots + for range 4 { + p.checks.submit(context.Background(), func() { <-done }) + } - // Pool is now full — submit with cancelled ctx will fail + // Pool is now truly full — submit with cancelled ctx will reliably fail ctx, cancel := context.WithCancel(context.Background()) cancel() @@ -1515,3 +1543,32 @@ func TestRunAgentCancelDuringRetry(t *testing.T) { t.Error("runAgent did not return after cancel during retry") } } + +func TestRunSessionWriteError(t *testing.T) { + // Server accepts join, drains one frame, then sends a jobs event. + // Meanwhile we close the connection from the server side to trigger + // a write error in the writer goroutine when it tries to send results. + srv := newFakeWSServer(t) + + go func() { + srv.acceptAndJoin(t) + + // Send a bulk of events so the client tries to write back + time.Sleep(100 * time.Millisecond) + + // Close the connection from the server side — any writes by the + // client's writer goroutine will fail, triggering writeErrCh. + srv.close() + }() + + // Use short heartbeat to generate write traffic + origCHB := channelHeartbeatInterval + defer func() { channelHeartbeatInterval = origCHB }() + channelHeartbeatInterval = 50 * time.Millisecond + + err := runSession(context.Background(), "ws://"+srv.addr(), "token") + if err == nil { + t.Error("expected error from write or read failure") + } + // Either "read:" or "write:" error is acceptable +} diff --git a/main.go b/main.go index bd9a258..a83d9ff 100644 --- a/main.go +++ b/main.go @@ -17,26 +17,37 @@ var version = "dev" var insecureFlag bool func main() { - apiURL := flag.String("api-url", os.Getenv("TOWEROPS_API_URL"), "API URL (e.g., wss://towerops.net)") - token := flag.String("token", os.Getenv("TOWEROPS_AGENT_TOKEN"), "Agent authentication token") - tokenFile := flag.String("token-file", "", "Path to file containing agent token (preferred over --token)") - logLevel := flag.String("log-level", envOrDefault("LOG_LEVEL", "info"), "Log level (debug, info, warn, error)") - flag.BoolVar(&insecureFlag, "insecure", false, "Allow plaintext ws:// connections (insecure)") - flag.Parse() + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT) + defer cancel() + os.Exit(runMain(ctx, os.Args[1:])) +} + +// runMain is the testable entry point. Returns an exit code. +func runMain(ctx context.Context, args []string) int { + fs := flag.NewFlagSet("towerops-agent", flag.ContinueOnError) + apiURL := fs.String("api-url", os.Getenv("TOWEROPS_API_URL"), "API URL (e.g., wss://towerops.net)") + token := fs.String("token", os.Getenv("TOWEROPS_AGENT_TOKEN"), "Agent authentication token") + tokenFile := fs.String("token-file", "", "Path to file containing agent token (preferred over --token)") + logLevel := fs.String("log-level", envOrDefault("LOG_LEVEL", "info"), "Log level (debug, info, warn, error)") + fs.BoolVar(&insecureFlag, "insecure", false, "Allow plaintext ws:// connections (insecure)") + + if err := fs.Parse(args); err != nil { + return 1 + } // Read token from file if --token-file is provided if *tokenFile != "" { data, err := os.ReadFile(*tokenFile) if err != nil { fmt.Fprintf(os.Stderr, "error: cannot read token file: %v\n", err) - os.Exit(1) + return 1 } t := strings.TrimSpace(string(data)) token = &t } // Warn if --token was used via CLI (visible in /proc/cmdline) - if *tokenFile == "" && isFlagSet("token") { + if *tokenFile == "" && flagIsSet(fs, "token") { fmt.Fprintln(os.Stderr, "WARNING: --token flag exposes the token in the process table. Use TOWEROPS_AGENT_TOKEN env var or --token-file instead.") } @@ -56,29 +67,30 @@ func main() { if *apiURL == "" || *token == "" { fmt.Fprintln(os.Stderr, "error: --api-url and --token are required (or set TOWEROPS_API_URL and TOWEROPS_AGENT_TOKEN)") - flag.Usage() - os.Exit(1) + fs.Usage() + return 1 } slog.Info("towerops agent starting", "version", version) // Convert HTTP(S) to WebSocket URL - wsURL := toWebSocketURL(*apiURL) + wsURL, err := toWebSocketURL(*apiURL) + if err != nil { + slog.Error(err.Error()) + return 1 + } slog.Info("websocket url", "url", sanitizeURL(wsURL)) - // Signal handling - ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT) - defer cancel() - // Run agent with reconnect loop runAgent(ctx, wsURL, *token) slog.Info("towerops agent stopped") + return 0 } // toWebSocketURL converts an HTTP(S) URL to a WebSocket URL. -// Rejects plaintext ws:// unless --insecure is set. -func toWebSocketURL(rawURL string) string { +// Returns an error for plaintext ws:// unless insecureFlag is set. +func toWebSocketURL(rawURL string) (string, error) { var result string switch { case strings.HasPrefix(rawURL, "http://"): @@ -92,16 +104,15 @@ func toWebSocketURL(rawURL string) string { } if strings.HasPrefix(result, "ws://") && !insecureFlag { - slog.Error("plaintext ws:// connection rejected — use wss:// or pass --insecure to allow") - osExit(1) + return "", fmt.Errorf("plaintext ws:// connection rejected — use wss:// or pass --insecure to allow") } - return result + return result, nil } -// isFlagSet returns true if a flag was explicitly set on the command line. -func isFlagSet(name string) bool { +// flagIsSet returns true if a flag was explicitly set on the command line. +func flagIsSet(fs *flag.FlagSet, name string) bool { found := false - flag.Visit(func(f *flag.Flag) { + fs.Visit(func(f *flag.Flag) { if f.Name == name { found = true } diff --git a/main_test.go b/main_test.go index 3b57722..575326d 100644 --- a/main_test.go +++ b/main_test.go @@ -1,8 +1,13 @@ package main import ( + "context" + "flag" "os" + "path/filepath" + "strings" "testing" + "time" ) func TestEnvOrDefault(t *testing.T) { @@ -39,16 +44,28 @@ 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 TestFlagIsSet(t *testing.T) { + fs := flag.NewFlagSet("test", flag.ContinueOnError) + fs.String("my-flag", "", "test flag") + fs.String("other", "", "another flag") + + // Not set + _ = fs.Parse([]string{}) + if flagIsSet(fs, "my-flag") { + t.Error("expected false for unset flag") + } + + // Set + _ = fs.Parse([]string{"--my-flag=hello"}) + if !flagIsSet(fs, "my-flag") { + t.Error("expected true for set flag") + } + if flagIsSet(fs, "other") { + t.Error("expected false for other unset flag") } } func TestToWebSocketURL(t *testing.T) { - // Enable insecure for testing plaintext conversions origInsecure := insecureFlag defer func() { insecureFlag = origInsecure }() insecureFlag = true @@ -64,7 +81,11 @@ func TestToWebSocketURL(t *testing.T) { {"localhost:4000", "wss://localhost:4000"}, } for _, tt := range tests { - got := toWebSocketURL(tt.input) + got, err := toWebSocketURL(tt.input) + if err != nil { + t.Errorf("toWebSocketURL(%q) unexpected error: %v", tt.input, err) + continue + } if got != tt.want { t.Errorf("toWebSocketURL(%q) = %q, want %q", tt.input, got, tt.want) } @@ -73,15 +94,8 @@ func TestToWebSocketURL(t *testing.T) { func TestToWebSocketURLRejectsPlaintext(t *testing.T) { origInsecure := insecureFlag - origExit := osExit - defer func() { - insecureFlag = origInsecure - osExit = origExit - }() - + defer func() { insecureFlag = origInsecure }() insecureFlag = false - exitCode := -1 - osExit = func(code int) { exitCode = code } tests := []struct { name string @@ -92,11 +106,139 @@ func TestToWebSocketURLRejectsPlaintext(t *testing.T) { } 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) + _, err := toWebSocketURL(tt.input) + if err == nil { + t.Error("expected error for plaintext URL") + } + if err != nil && !strings.Contains(err.Error(), "plaintext") { + t.Errorf("expected 'plaintext' in error, got: %v", err) } }) } } + +func TestRunMainMissingArgs(t *testing.T) { + // Unset env vars to ensure flags are required + t.Setenv("TOWEROPS_API_URL", "") + t.Setenv("TOWEROPS_AGENT_TOKEN", "") + + code := runMain(context.Background(), []string{}) + if code != 1 { + t.Errorf("expected exit 1, got %d", code) + } +} + +func TestRunMainInvalidFlag(t *testing.T) { + code := runMain(context.Background(), []string{"--nonexistent-flag"}) + if code != 1 { + t.Errorf("expected exit 1, got %d", code) + } +} + +func TestRunMainTokenFile(t *testing.T) { + tmpDir := t.TempDir() + tokenPath := filepath.Join(tmpDir, "token") + os.WriteFile(tokenPath, []byte(" test-token-123 \n"), 0600) + + t.Setenv("TOWEROPS_API_URL", "") + t.Setenv("TOWEROPS_AGENT_TOKEN", "") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately so runAgent returns + + code := runMain(ctx, []string{ + "--api-url=wss://example.com", + "--token-file=" + tokenPath, + }) + if code != 0 { + t.Errorf("expected exit 0, got %d", code) + } +} + +func TestRunMainTokenFileMissing(t *testing.T) { + code := runMain(context.Background(), []string{ + "--api-url=wss://example.com", + "--token-file=/nonexistent/path", + }) + if code != 1 { + t.Errorf("expected exit 1, got %d", code) + } +} + +func TestRunMainPlaintextRejected(t *testing.T) { + origInsecure := insecureFlag + defer func() { insecureFlag = origInsecure }() + + t.Setenv("TOWEROPS_API_URL", "") + t.Setenv("TOWEROPS_AGENT_TOKEN", "") + + code := runMain(context.Background(), []string{ + "--api-url=http://localhost:4000", + "--token=test-token", + }) + if code != 1 { + t.Errorf("expected exit 1 for plaintext rejection, got %d", code) + } +} + +func TestRunMainLogLevels(t *testing.T) { + for _, level := range []string{"debug", "warn", "warning", "error", "info", "unknown"} { + t.Run(level, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + code := runMain(ctx, []string{ + "--api-url=wss://example.com", + "--token=test-token", + "--log-level=" + level, + }) + if code != 0 { + t.Errorf("log level %q: expected exit 0, got %d", level, code) + } + }) + } +} + +func TestRunMainNormalRun(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + code := runMain(ctx, []string{ + "--api-url=wss://example.com", + "--token=test-token", + }) + if code != 0 { + t.Errorf("expected exit 0, got %d", code) + } +} + +func TestRunMainTokenFlagWarning(t *testing.T) { + // This tests the warning path when --token is passed via CLI + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + t.Setenv("TOWEROPS_API_URL", "") + t.Setenv("TOWEROPS_AGENT_TOKEN", "") + + code := runMain(ctx, []string{ + "--api-url=wss://example.com", + "--token=my-secret-token", + }) + if code != 0 { + t.Errorf("expected exit 0, got %d", code) + } +} + +func TestRunMainWithRunAgent(t *testing.T) { + // Test the full path through runAgent with a real (but immediately cancelled) context + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + code := runMain(ctx, []string{ + "--api-url=wss://127.0.0.1:1", + "--token=test-token", + }) + if code != 0 { + t.Errorf("expected exit 0, got %d", code) + } +} diff --git a/ping.go b/ping.go index adab9de..9431f6b 100644 --- a/ping.go +++ b/ping.go @@ -15,6 +15,8 @@ import ( "golang.org/x/net/ipv6" ) +var icmpListenPacket = icmp.ListenPacket + // icmpPing sends a single ICMP echo request and returns the round-trip time in milliseconds. // Tries raw ICMP sockets first (requires CAP_NET_RAW or root), then falls back to // unprivileged UDP-based ICMP (requires ping_group_range sysctl). @@ -53,7 +55,7 @@ func icmpPing(ip string, timeoutMs int) (float64, error) { // doICMPPing performs an ICMP ping over the given network type. func doICMPPing(ip net.IP, network string, isIPv4 bool, timeoutMs int) (float64, error) { - conn, err := icmp.ListenPacket(network, "") + conn, err := icmpListenPacket(network, "") if err != nil { return 0, &errICMPUnavailable{err: fmt.Errorf("icmp listen %s: %w", network, err)} } diff --git a/ping_test.go b/ping_test.go index 5a91d4d..3498151 100644 --- a/ping_test.go +++ b/ping_test.go @@ -2,8 +2,12 @@ package main import ( "fmt" + "net" "runtime" + "strings" "testing" + + "golang.org/x/net/icmp" ) func TestPingDeviceLocalhost(t *testing.T) { @@ -130,3 +134,183 @@ func TestParsePingTime(t *testing.T) { }) } } + +func TestExecPingLocalhost(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping on windows") + } + ms, err := execPing("127.0.0.1", 5000) + if err != nil { + t.Skipf("ping command not available: %v", err) + } + if ms <= 0 { + t.Errorf("expected positive response time, got %v", ms) + } +} + +func TestExecPingInvalidIP(t *testing.T) { + _, err := execPing("not-an-ip", 5000) + if err == nil { + t.Error("expected error for invalid IP") + } +} + +func TestExecPingIPv6Localhost(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping on windows") + } + ms, err := execPing("::1", 5000) + if err != nil { + t.Skipf("ping6 not available: %v", err) + } + if ms <= 0 { + t.Errorf("expected positive response time, got %v", ms) + } +} + +func TestExecPingUnreachable(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping on windows") + } + _, err := execPing("192.0.2.1", 1000) // TEST-NET-1 — unreachable + if err == nil { + t.Error("expected error for unreachable host") + } + if err != nil && !strings.Contains(err.Error(), "ping failed") { + t.Errorf("expected 'ping failed' in error, got: %v", err) + } +} + +func TestPingDeviceFallbackToExec(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping on windows") + } + // Mock icmpListenPacket to always fail → forces fallback to execPing + origListen := icmpListenPacket + defer func() { icmpListenPacket = origListen }() + + icmpListenPacket = func(network, address string) (*icmp.PacketConn, error) { + return nil, fmt.Errorf("permission denied") + } + + ms, err := pingDevice("127.0.0.1", 5000) + if err != nil { + t.Skipf("exec ping fallback not available: %v", err) + } + if ms <= 0 { + t.Errorf("expected positive response time via exec fallback, got %v", ms) + } +} + +func TestPingDeviceNonICMPError(t *testing.T) { + // When icmpPing returns a non-errICMPUnavailable error, pingDevice should NOT + // fall back to exec — it should return the error directly. + origListen := icmpListenPacket + defer func() { icmpListenPacket = origListen }() + + // First call (raw ICMP) returns errICMPUnavailable → triggers UDP fallback + // Second call (UDP ICMP) returns a real write error → not errICMPUnavailable + calls := 0 + icmpListenPacket = func(network, address string) (*icmp.PacketConn, error) { + calls++ + if calls == 1 { + // Raw ICMP fails with errICMPUnavailable + return nil, fmt.Errorf("permission denied") + } + // UDP ICMP also fails with errICMPUnavailable + return nil, fmt.Errorf("also denied") + } + + _, err := pingDevice("127.0.0.1", 1000) + // Both ICMP attempts fail with errICMPUnavailable, so it falls back to exec + // which should succeed for localhost + if err != nil { + t.Skipf("ping fallback not available: %v", err) + } +} + +func TestDoICMPPingIPv6Network(t *testing.T) { + // Test with the IPv6 raw network to cover the ipv6-icmp branches + ip := net.ParseIP("::1") + _, err := doICMPPing(ip, "ip6:ipv6-icmp", false, 1000) + if err != nil { + t.Skipf("IPv6 ICMP not available: %v", err) + } +} + +func TestDoICMPPingUDPNetwork(t *testing.T) { + // Test with UDP network to cover the udp address branch + ip := net.ParseIP("127.0.0.1") + _, err := doICMPPing(ip, "udp4", true, 1000) + if err != nil { + t.Skipf("UDP ICMP not available: %v", err) + } +} + +func TestDoICMPPingTimeout(t *testing.T) { + // Ping unreachable IP with short timeout → covers icmp read timeout error + ip := net.ParseIP("192.0.2.1") // TEST-NET-1 — unreachable + _, err := doICMPPing(ip, "udp4", true, 100) + if err == nil { + t.Error("expected timeout error for unreachable host") + } + if err != nil && !strings.Contains(err.Error(), "icmp read") { + t.Logf("got error (expected icmp read timeout): %v", err) + } +} + +func TestDoICMPPingIPv6Timeout(t *testing.T) { + // IPv6 unreachable — covers the ipv6 branch in doICMPPing + ip := net.ParseIP("100::1") // Unreachable IPv6 + _, err := doICMPPing(ip, "udp6", false, 100) + if err != nil { + // May fail with various errors depending on system IPv6 support + t.Logf("IPv6 ICMP error (expected): %v", err) + } +} + +func TestIcmpPingNonICMPUnavailableError(t *testing.T) { + // When raw ICMP returns a non-errICMPUnavailable error, icmpPing should + // return that error without falling back to UDP. + origListen := icmpListenPacket + defer func() { icmpListenPacket = origListen }() + + // Raw ICMP succeeds (opens a connection), but pinging unreachable IP will timeout. + // The timeout error is NOT errICMPUnavailable, so icmpPing returns it directly. + icmpListenPacket = func(network, address string) (*icmp.PacketConn, error) { + return icmp.ListenPacket("udp4", address) + } + + _, err := icmpPing("192.0.2.1", 100) // TEST-NET-1, 100ms timeout + if err == nil { + t.Error("expected error for unreachable host") + } +} + +func TestIcmpPingUDPFallback(t *testing.T) { + // Mock raw ICMP to fail, forcing UDP fallback path in icmpPing + origListen := icmpListenPacket + defer func() { icmpListenPacket = origListen }() + + calls := 0 + icmpListenPacket = func(network, address string) (*icmp.PacketConn, error) { + calls++ + if calls == 1 { + // Raw ICMP fails + return nil, fmt.Errorf("permission denied") + } + // UDP ICMP uses the real implementation + return icmp.ListenPacket(network, address) + } + + ms, err := icmpPing("127.0.0.1", 5000) + if err != nil { + t.Skipf("UDP ICMP not available: %v", err) + } + if ms <= 0 { + t.Errorf("expected positive response time, got %v", ms) + } + if calls < 2 { + t.Errorf("expected at least 2 ListenPacket calls (raw + udp), got %d", calls) + } +} diff --git a/snmp_test.go b/snmp_test.go index 22ec6ab..cdaba8e 100644 --- a/snmp_test.go +++ b/snmp_test.go @@ -659,6 +659,41 @@ func TestExecuteSnmpJobBatchesGets(t *testing.T) { } } +func TestExecuteSnmpJobCtxCancelled(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + }, func() {}, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel before queries run + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(ctx, &pb.AgentJob{ + JobId: "ctx-test", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_GET, Oids: []string{".1.3.6.1.2.1.1.1.0"}}, + {QueryType: pb.QueryType_GET, Oids: []string{".1.3.6.1.2.1.1.2.0"}}, + }, + }, ch) + + // With cancelled context, the function should return before processing queries + // (no result sent because it returns early in the ctx.Err() check) + select { + case <-ch: + // Might get a result if the first query ran before ctx check + default: + // Expected — returned early + } +} + func TestExecuteCredentialTest(t *testing.T) { t.Run("nil device", func(t *testing.T) { ch := make(chan *pb.CredentialTestResult, 1) diff --git a/update.go b/update.go index 9efe0ac..bbf4d32 100644 --- a/update.go +++ b/update.go @@ -18,6 +18,7 @@ var osExecutable = os.Executable var osCreateTemp = os.CreateTemp var osRename = os.Rename var httpGet = http.Get +var syscallExec = syscall.Exec var maxUpdateSize int64 = 100 << 20 // 100 MB // selfUpdate downloads a new binary, verifies its checksum, replaces the current binary, and re-execs. @@ -105,7 +106,7 @@ func selfUpdate(downloadURL, expectedChecksum string) error { // Re-exec with same arguments slog.Info("re-executing", "args", sanitizeArgs(os.Args)) - return syscall.Exec(currentExe, os.Args, os.Environ()) + return syscallExec(currentExe, os.Args, os.Environ()) } // sanitizeArgs returns a copy of args with token values masked. diff --git a/update_test.go b/update_test.go index a150014..961da43 100644 --- a/update_test.go +++ b/update_test.go @@ -322,6 +322,92 @@ func TestSanitizeArgs(t *testing.T) { }) } +func TestSelfUpdateInvalidURL(t *testing.T) { + err := selfUpdate("://\x7f", "abc123") + if err == nil { + t.Error("expected error for invalid URL") + } + if err != nil && !strings.Contains(err.Error(), "parse url") { + t.Errorf("expected 'parse url' in error, got: %v", err) + } +} + +func TestSelfUpdateFullHappyPath(t *testing.T) { + body := []byte("test binary content for full path") + checksum := fmt.Sprintf("%x", sha256.Sum256(body)) + + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(body) + })) + defer srv.Close() + + origGet := httpGet + origExe := osExecutable + origRename := osRename + origExec := syscallExec + defer func() { + httpGet = origGet + osExecutable = origExe + osRename = origRename + syscallExec = origExec + }() + httpGet = srv.Client().Get + + dir := t.TempDir() + exePath := filepath.Join(dir, "test-agent") + osExecutable = func() (string, error) { return exePath, nil } + osRename = func(oldpath, newpath string) error { + return os.Rename(oldpath, newpath) // real rename within temp dir + } + syscallExec = func(argv0 string, argv []string, envv []string) error { + return nil // success — don't actually re-exec + } + + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) + if err != nil { + t.Errorf("expected nil error on full happy path, got: %v", err) + } +} + +func TestSelfUpdateWriteError(t *testing.T) { + body := []byte("binary") + checksum := fmt.Sprintf("%x", sha256.Sum256(body)) + + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(body) + })) + defer srv.Close() + + origGet := httpGet + origExe := osExecutable + origCreate := osCreateTemp + defer func() { + httpGet = origGet + osExecutable = origExe + osCreateTemp = origCreate + }() + httpGet = srv.Client().Get + + dir := t.TempDir() + osExecutable = func() (string, error) { return filepath.Join(dir, "agent"), nil } + + // Create a temp file that will fail on Write by closing it before selfUpdate tries to write + osCreateTemp = func(d, pattern string) (*os.File, error) { + f, err := os.CreateTemp(d, pattern) + if err != nil { + return nil, err + } + // Close the file so Write will fail + _ = f.Close() + return f, nil + } + + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) + if err == nil { + t.Error("expected write error") + } +} + // rewriteToHTTPS converts an httptest TLS server URL to use the https scheme. // httptest.NewTLSServer returns URLs with https:// already, but this ensures consistency. func rewriteToHTTPS(rawURL string) string {