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

test: push coverage from 92% to 97.6% (249 tests)

Browse source 
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.
This commit is contained in:
Graham McIntire 2026-03-04 16:35:02 -06:00
parent aef770ea83
commit 6bf9584cb2
No known key found for this signature in database
8 changed files with 570 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
agent_test.go
View file

@ -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
}

57
main.go
View file

@ -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
}

180
main_test.go
View file

@ -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)
}
}

4
ping.go
View file

@ -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)}
}

184
ping_test.go
View file

@ -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)
}
}

35
snmp_test.go
View file

@ -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)

3
update.go
View file

@ -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.

86
update_test.go
View file

@ -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 {