Merge pull request #22 from towerops-app/security-updates

Security fixes

agent.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log/slog"
 	"math/rand/v2"
@@ -22,6 +23,11 @@ import (
 var osExit = os.Exit
 var doSelfUpdate = selfUpdate
 
+var errRestartRequested = fmt.Errorf("restart requested")
+var joinTimeout = 10 * time.Second
+
+const maxJobPayloadBytes = 4 << 20 // 4 MB — well above any legitimate job list
+
 // channelMsg is the WebSocket channel message format (JSON wrapper around binary protobuf).
 type channelMsg struct {
 	Topic   string          `json:"topic"`
@@ -47,6 +53,10 @@ func runAgent(ctx context.Context, wsURL, token string) {
 		if ctx.Err() != nil {
 			return
 		}
+		if errors.Is(err, errRestartRequested) {
+			osExit(0)
+			return
+		}
 		if err != nil {
 			slog.Error("agent disconnected", "error", err)
 		}
@@ -135,12 +145,28 @@ func runSession(ctx context.Context, baseURL, token string) error {
 			return
 		}
 		encoded := base64.StdEncoding.EncodeToString(bin)
-		*bp = bin[:0]
+		full := (*bp)[:cap(*bp)]
+		zeroBytes(full)
+		*bp = full[:0]
 		bufPool.Put(bp)
 		payload, _ := json.Marshal(map[string]string{"binary": encoded})
 		sendMsg(event, payload)
 	}
 
+	// Reader goroutine — must start before join so we can receive the reply
+	msgCh := make(chan []byte, 100)
+	errCh := make(chan error, 1)
+	go func() {
+		for {
+			data, _, err := ws.ReadMessage()
+			if err != nil {
+				errCh <- err
+				return
+			}
+			msgCh <- data
+		}
+	}()
+
 	// Join channel
 	joinPayload, _ := json.Marshal(map[string]string{"token": token})
 	joinMsg := channelMsg{
@@ -155,33 +181,47 @@ func runSession(ctx context.Context, baseURL, token string) error {
 	}
 	slog.Debug("sent channel join request")
 
+	// Wait for join reply before entering main loop
+	select {
+	case data := <-msgCh:
+		var reply channelMsg
+		if err := json.Unmarshal(data, &reply); err != nil {
+			return fmt.Errorf("join reply unmarshal: %w", err)
+		}
+		if reply.Event == "phx_reply" {
+			var status struct {
+				Status   string `json:"status"`
+				Response any    `json:"response"`
+			}
+			if err := json.Unmarshal(reply.Payload, &status); err == nil && status.Status != "ok" {
+				return fmt.Errorf("join rejected: %s", status.Status)
+			}
+		}
+		slog.Info("channel joined")
+	case err := <-errCh:
+		return fmt.Errorf("read during join: %w", err)
+	case <-time.After(joinTimeout):
+		return fmt.Errorf("join timeout")
+	}
+
 	// Writer goroutine - serializes all writes to the WebSocket
 	var writerWg sync.WaitGroup
+	writeErrCh := make(chan error, 1)
 	writerWg.Add(1)
 	go func() {
 		defer writerWg.Done()
 		for data := range writeCh {
 			if err := ws.WriteText(data); err != nil {
 				slog.Error("websocket write", "error", err)
+				select {
+				case writeErrCh <- err:
+				default:
+				}
 				return
 			}
 		}
 	}()
 
-	// Reader goroutine - reads messages and dispatches
-	msgCh := make(chan []byte, 100)
-	errCh := make(chan error, 1)
-	go func() {
-		for {
-			data, _, err := ws.ReadMessage()
-			if err != nil {
-				errCh <- err
-				return
-			}
-			msgCh <- data
-		}
-	}()
-
 	heartbeatTicker := time.NewTicker(60 * time.Second)
 	defer heartbeatTicker.Stop()
 	channelHeartbeatTicker := time.NewTicker(25 * time.Second)
@@ -222,13 +262,20 @@ func runSession(ctx context.Context, baseURL, token string) error {
 			flushSnmpBatch()
 			return fmt.Errorf("read: %w", err)
 
+		case err := <-writeErrCh:
+			flushSnmpBatch()
+			return fmt.Errorf("write: %w", err)
+
 		case data := <-msgCh:
 			var msg channelMsg
 			if err := json.Unmarshal(data, &msg); err != nil {
 				slog.Warn("invalid message", "error", err)
 				continue
 			}
-			handleMessage(ctx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
+			if handleMessage(ctx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh) {
+				flushSnmpBatch()
+				return errRestartRequested
+			}
 
 		case result := <-snmpResultCh:
 			snmpBatch = append(snmpBatch, result)
@@ -279,6 +326,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
 }
 
 // handleMessage dispatches incoming channel messages.
+// Returns true if the session should end (e.g. restart requested).
 func handleMessage(
 	ctx context.Context,
 	msg channelMsg,
@@ -287,7 +335,7 @@ func handleMessage(
 	mikrotikResultCh chan<- *pb.MikrotikResult,
 	credTestResultCh chan<- *pb.CredentialTestResult,
 	monitoringCheckCh chan<- *pb.MonitoringCheck,
-) {
+) bool {
 	switch msg.Event {
 	case "phx_reply":
 		slog.Debug("channel reply", "topic", msg.Topic)
@@ -298,17 +346,21 @@ func handleMessage(
 		}
 		if err := json.Unmarshal(msg.Payload, &payload); err != nil {
 			slog.Error("decode job payload", "error", err)
-			return
+			return false
+		}
+		if len(payload.Binary) > maxJobPayloadBytes {
+			slog.Error("job payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes)
+			return false
 		}
 		bin, err := base64.StdEncoding.DecodeString(payload.Binary)
 		if err != nil {
 			slog.Error("decode base64", "error", err)
-			return
+			return false
 		}
 		var jobList pb.AgentJobList
 		if err := proto.Unmarshal(bin, &jobList); err != nil {
 			slog.Error("unmarshal job list", "error", err)
-			return
+			return false
 		}
 		slog.Info("received jobs", "count", len(jobList.Jobs))
 		for _, job := range jobList.Jobs {
@@ -316,8 +368,8 @@ func handleMessage(
 		}
 
 	case "restart":
-		slog.Info("restart requested by server, exiting")
+		slog.Info("restart requested by server")
-		osExit(0)
+		return true
 
 	case "update":
 		var payload struct {
@@ -326,7 +378,7 @@ func handleMessage(
 		}
 		if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" || payload.Checksum == "" {
 			slog.Error("invalid update payload")
-			return
+			return false
 		}
 		slog.Info("update requested", "url", payload.URL)
 		if err := doSelfUpdate(payload.URL, payload.Checksum); err != nil {
@@ -336,6 +388,7 @@ func handleMessage(
 	default:
 		slog.Debug("ignoring event", "event", msg.Event)
 	}
+	return false
 }
 
 // jobPools holds the worker pools for each job type.
@@ -357,15 +410,19 @@ func dispatchJob(
 ) {
 	slog.Info("starting job", "job_id", job.JobId, "type", job.JobType)
 
+	var ok bool
 	switch job.JobType {
 	case pb.JobType_MIKROTIK:
-		pools.mikrotik.submit(func() { executeMikrotikJob(ctx, job, mikrotikResultCh) })
+		ok = pools.mikrotik.submit(ctx, func() { executeMikrotikJob(ctx, job, mikrotikResultCh) })
 	case pb.JobType_TEST_CREDENTIALS:
-		pools.snmp.submit(func() { executeCredentialTest(ctx, job, credTestResultCh) })
+		ok = pools.snmp.submit(ctx, func() { executeCredentialTest(ctx, job, credTestResultCh) })
 	case pb.JobType_PING:
-		pools.ping.submit(func() { executePingJob(ctx, job, monitoringCheckCh) })
+		ok = pools.ping.submit(ctx, func() { executePingJob(ctx, job, monitoringCheckCh) })
 	default:
-		pools.snmp.submit(func() { executeSnmpJob(ctx, job, snmpResultCh) })
+		ok = pools.snmp.submit(ctx, func() { executeSnmpJob(ctx, job, snmpResultCh) })
+	}
+	if !ok {
+		slog.Warn("job dropped, pool full", "job_id", job.JobId)
 	}
 }
 

agent_test.go

@@ -5,6 +5,9 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"net"
+	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -165,20 +168,14 @@ func TestHandleMessage(t *testing.T) {
 	})
 
 	t.Run("restart", func(t *testing.T) {
-		origExit := osExit
-		defer func() { osExit = origExit }()
-
-		var exitCode int
-		osExit = func(code int) { exitCode = code }
-
 		snmpCh := make(chan *pb.SnmpResult, 1)
 		mtCh := make(chan *pb.MikrotikResult, 1)
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
-		handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
+		shouldEnd := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
 
-		if exitCode != 0 {
+		if !shouldEnd {
-			t.Errorf("expected exit code 0, got %d", exitCode)
+			t.Error("expected handleMessage to return true for restart")
 		}
 	})
 
@@ -336,6 +333,57 @@ func TestHandleMessage(t *testing.T) {
 	})
 }
 
+func TestHandleMessageRejectsOversizedPayload(t *testing.T) {
+	snmpCh := make(chan *pb.SnmpResult, 1)
+	mtCh := make(chan *pb.MikrotikResult, 1)
+	credCh := make(chan *pb.CredentialTestResult, 1)
+	monCh := make(chan *pb.MonitoringCheck, 1)
+
+	// Create a binary payload larger than maxJobPayloadBytes
+	oversized := make([]byte, maxJobPayloadBytes+1)
+	encoded := base64.StdEncoding.EncodeToString(oversized)
+	payload, _ := json.Marshal(map[string]string{"binary": encoded})
+
+	handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
+
+	// Verify no jobs were dispatched
+	select {
+	case <-snmpCh:
+		t.Error("expected no SNMP result for oversized payload")
+	case <-time.After(100 * time.Millisecond):
+		// Good — nothing dispatched
+	}
+}
+
+func TestBufferPoolZeroesOnReturn(t *testing.T) {
+	pool := &sync.Pool{
+		New: func() any {
+			b := make([]byte, 0, 64)
+			return &b
+		},
+	}
+
+	// Get a buffer, write some data, return it
+	bp := pool.Get().(*[]byte)
+	*bp = append((*bp)[:0], []byte("sensitive credentials data here")...)
+	full := (*bp)[:cap(*bp)]
+
+	// Simulate the zeroing that sendBinaryResult should do
+	zeroBytes(full)
+	*bp = full[:0]
+	pool.Put(bp)
+
+	// Get the buffer back and verify it's zeroed
+	bp2 := pool.Get().(*[]byte)
+	full2 := (*bp2)[:cap(*bp2)]
+	for i, b := range full2 {
+		if b != 0 {
+			t.Errorf("byte[%d] = %d, expected 0 — pool buffer not zeroed", i, b)
+			break
+		}
+	}
+}
+
 func TestZeroBytes(t *testing.T) {
 	b := []byte{1, 2, 3, 4, 5}
 	zeroBytes(b)
@@ -496,3 +544,138 @@ func TestDispatchJob(t *testing.T) {
 		}
 	})
 }
+
+func TestRunSessionRejectsFailedJoin(t *testing.T) {
+	origTimeout := joinTimeout
+	defer func() { joinTimeout = origTimeout }()
+	joinTimeout = 2 * time.Second
+
+	// Start a fake WebSocket server that accepts the upgrade then sends a phx_error join reply
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+
+		// Read HTTP upgrade request
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		reqStr := string(buf[:n])
+
+		// Extract key and compute accept
+		key := extractWSKey(reqStr)
+		accept := computeAcceptKey(key)
+
+		// Send valid 101 upgrade
+		resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+
+		// Read the join message (masked WebSocket frame) — just consume it
+		frameBuf := make([]byte, 4096)
+		_, _ = conn.Read(frameBuf)
+
+		// Send phx_error reply as an unmasked text frame
+		reply, _ := json.Marshal(channelMsg{
+			Topic:   "agent:agent-0",
+			Event:   "phx_reply",
+			Payload: json.RawMessage(`{"status":"error","response":{"reason":"invalid token"}}`),
+			Ref:     strPtr("1"),
+		})
+		frame := makeTextFrame(reply)
+		_, _ = conn.Write(frame)
+
+		// Keep connection open for a bit
+		time.Sleep(time.Second)
+	}()
+
+	addr := ln.Addr().String()
+	err = runSession(context.Background(), "ws://"+addr, "bad-token")
+	if err == nil {
+		t.Fatal("expected error from rejected join")
+	}
+	if !strings.Contains(err.Error(), "join rejected") {
+		t.Errorf("expected 'join rejected' in error, got: %v", err)
+	}
+}
+
+func TestRunSessionJoinTimeout(t *testing.T) {
+	origTimeout := joinTimeout
+	defer func() { joinTimeout = origTimeout }()
+	joinTimeout = 500 * time.Millisecond
+
+	// Server that upgrades but never sends a join reply
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		reqStr := string(buf[:n])
+		key := extractWSKey(reqStr)
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+
+		// Read join frame but never reply
+		frameBuf := make([]byte, 4096)
+		_, _ = conn.Read(frameBuf)
+
+		time.Sleep(5 * time.Second)
+	}()
+
+	addr := ln.Addr().String()
+	start := time.Now()
+	err = runSession(context.Background(), "ws://"+addr, "token")
+	elapsed := time.Since(start)
+
+	if err == nil {
+		t.Fatal("expected timeout error")
+	}
+	if !strings.Contains(err.Error(), "join timeout") {
+		t.Errorf("expected 'join timeout' in error, got: %v", err)
+	}
+	if elapsed > 3*time.Second {
+		t.Errorf("took too long (%v), timeout didn't trigger", elapsed)
+	}
+}
+
+// extractWSKey extracts the Sec-WebSocket-Key from a raw HTTP request.
+func extractWSKey(req string) string {
+	for _, line := range strings.Split(req, "\r\n") {
+		lower := strings.ToLower(line)
+		if strings.HasPrefix(lower, "sec-websocket-key: ") {
+			return strings.TrimSpace(line[len("Sec-WebSocket-Key: "):])
+		}
+	}
+	return ""
+}
+
+// makeTextFrame creates an unmasked WebSocket text frame (server→client).
+func makeTextFrame(payload []byte) []byte {
+	length := len(payload)
+	var frame []byte
+	frame = append(frame, 0x81) // FIN + text
+	if length <= 125 {
+		frame = append(frame, byte(length))
+	} else if length <= 65535 {
+		frame = append(frame, 126, byte(length>>8), byte(length))
+	}
+	frame = append(frame, payload...)
+	return frame
+}

ssh.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"net"
+	"strconv"
 	"time"
 
 	"github.com/towerops-app/towerops-agent/pb"
@@ -11,6 +13,7 @@ import (
 )
 
 var sshBackup = executeMikrotikBackup
+var sshDial = ssh.Dial
 var doPing = pingDevice
 
 // executeMikrotikBackup connects via SSH and runs /export compact.
@@ -26,8 +29,8 @@ func executeMikrotikBackup(ip string, port uint16, username, password string) (s
 		Timeout:         30 * time.Second,
 	}
 
-	addr := fmt.Sprintf("%s:%d", ip, port)
+	addr := net.JoinHostPort(ip, strconv.Itoa(int(port)))
-	conn, err := ssh.Dial("tcp", addr, config)
+	conn, err := sshDial("tcp", addr, config)
 	if err != nil {
 		return "", fmt.Errorf("ssh dial %s: %w", addr, err)
 	}

ssh_test.go

@@ -270,6 +270,40 @@ func TestExecuteMikrotikBackupViaSSH(t *testing.T) {
 	})
 }
 
+func TestSSHBackupIPv6Address(t *testing.T) {
+	origDial := sshDial
+	defer func() { sshDial = origDial }()
+
+	var capturedAddr string
+	sshDial = func(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
+		capturedAddr = addr
+		return nil, fmt.Errorf("mock dial")
+	}
+
+	_, _ = executeMikrotikBackup("::1", 22, "admin", "pass")
+
+	if capturedAddr != "[::1]:22" {
+		t.Errorf("expected [::1]:22, got %q", capturedAddr)
+	}
+}
+
+func TestSSHBackupIPv4Address(t *testing.T) {
+	origDial := sshDial
+	defer func() { sshDial = origDial }()
+
+	var capturedAddr string
+	sshDial = func(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
+		capturedAddr = addr
+		return nil, fmt.Errorf("mock dial")
+	}
+
+	_, _ = executeMikrotikBackup("10.0.0.1", 22, "admin", "pass")
+
+	if capturedAddr != "10.0.0.1:22" {
+		t.Errorf("expected 10.0.0.1:22, got %q", capturedAddr)
+	}
+}
+
 func TestExecuteMikrotikBackupDialError(t *testing.T) {
 	_, err := executeMikrotikBackup("127.0.0.1", 1, "admin", "pass")
 	if err == nil {

update.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strings"
 	"syscall"
 )
 
@@ -78,6 +79,20 @@ func selfUpdate(downloadURL, expectedChecksum string) error {
 	slog.Info("binary replaced", "path", currentExe)
 
 	// Re-exec with same arguments
-	slog.Info("re-executing", "args", os.Args)
+	slog.Info("re-executing", "args", sanitizeArgs(os.Args))
 	return syscall.Exec(currentExe, os.Args, os.Environ())
 }
+
+// sanitizeArgs returns a copy of args with token values masked.
+func sanitizeArgs(args []string) []string {
+	out := make([]string, len(args))
+	copy(out, args)
+	for i, a := range out {
+		if (a == "--token" || a == "-token") && i+1 < len(out) {
+			out[i+1] = "***"
+		} else if strings.HasPrefix(a, "--token=") || strings.HasPrefix(a, "-token=") {
+			out[i] = a[:strings.Index(a, "=")+1] + "***"
+		}
+	}
+	return out
+}

update_test.go

@@ -283,6 +283,45 @@ func TestSelfUpdateTooLarge(t *testing.T) {
 	}
 }
 
+func TestSanitizeArgs(t *testing.T) {
+	tests := []struct {
+		name string
+		in   []string
+		want []string
+	}{
+		{"flag with separate value", []string{"agent", "--token", "secret"}, []string{"agent", "--token", "***"}},
+		{"no sensitive flags", []string{"agent", "--api-url", "wss://x"}, []string{"agent", "--api-url", "wss://x"}},
+		{"equals syntax", []string{"agent", "--token=secret"}, []string{"agent", "--token=***"}},
+		{"empty args", []string{"agent"}, []string{"agent"}},
+		{"trailing flag no value", []string{"agent", "--token"}, []string{"agent", "--token"}},
+		{"short flag with value", []string{"agent", "-token", "secret"}, []string{"agent", "-token", "***"}},
+		{"short equals syntax", []string{"agent", "-token=secret"}, []string{"agent", "-token=***"}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := sanitizeArgs(tt.in)
+			if len(got) != len(tt.want) {
+				t.Fatalf("length: got %d, want %d", len(got), len(tt.want))
+			}
+			for i := range got {
+				if got[i] != tt.want[i] {
+					t.Errorf("arg[%d]: got %q, want %q", i, got[i], tt.want[i])
+				}
+			}
+		})
+	}
+
+	// Verify original slice is not mutated
+	t.Run("does not mutate input", func(t *testing.T) {
+		orig := []string{"agent", "--token", "secret"}
+		_ = sanitizeArgs(orig)
+		if orig[2] != "secret" {
+			t.Error("sanitizeArgs mutated the input slice")
+		}
+	})
+}
+
 // 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 {

websocket.go

@@ -13,6 +13,7 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"time"
 )
 
 // websocketGUID is the magic GUID from RFC 6455 Section 4.2.2.
@@ -42,6 +43,8 @@ type WSConn struct {
 	mu     sync.Mutex // serializes writes
 }
 
+var wsHandshakeTimeout = 30 * time.Second
+
 var randRead = rand.Read
 var netDial = net.Dial
 var tlsDial = func(network, addr string) (net.Conn, error) {
@@ -75,6 +78,9 @@ func WSDial(rawURL string) (*WSConn, error) {
 		return nil, fmt.Errorf("dial %s: %w", host, err)
 	}
 
+	// Set handshake deadline — prevents a slow/malicious server from blocking indefinitely
+	_ = conn.SetDeadline(time.Now().Add(wsHandshakeTimeout))
+
 	// Generate random key for Sec-WebSocket-Key
 	keyBytes := make([]byte, 16)
 	if _, err := randRead(keyBytes); err != nil {
@@ -124,6 +130,9 @@ func WSDial(rawURL string) (*WSConn, error) {
 		return nil, fmt.Errorf("missing Sec-WebSocket-Accept header")
 	}
 
+	// Clear handshake deadline — normal operation uses no deadline
+	_ = conn.SetDeadline(time.Time{})
+
 	return &WSConn{conn: conn, reader: bufio.NewReaderSize(conn, 8192)}, nil
 }
 

websocket_test.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestWriteFrameMasked(t *testing.T) {
@@ -558,6 +559,44 @@ func TestWSDialBadURL(t *testing.T) {
 	}
 }
 
+func TestWSDialHandshakeTimeout(t *testing.T) {
+	// Server accepts connection but never sends data — should trigger timeout
+	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 the request but never respond
+		buf := make([]byte, 4096)
+		_, _ = conn.Read(buf)
+		// Hold connection open until test ends
+		<-make(chan struct{})
+		_ = conn.Close()
+	}()
+
+	origTimeout := wsHandshakeTimeout
+	defer func() { wsHandshakeTimeout = origTimeout }()
+	wsHandshakeTimeout = 500 * time.Millisecond // Short timeout for test
+
+	addr := ln.Addr().String()
+	start := time.Now()
+	_, err = WSDial("ws://" + addr + "/socket")
+	elapsed := time.Since(start)
+
+	if err == nil {
+		t.Error("expected timeout error")
+	}
+	if elapsed > 5*time.Second {
+		t.Errorf("took too long (%v), timeout didn't work", elapsed)
+	}
+}
+
 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")

workerpool.go

@@ -1,6 +1,10 @@
 package main
 
-import "sync"
+import (
+	"context"
+	"log/slog"
+	"sync"
+)
 
 // workerPool is a fixed-size goroutine pool for executing tasks.
 type workerPool struct {
@@ -19,16 +23,28 @@ func newWorkerPool(n int) *workerPool {
 		go func() {
 			defer p.wg.Done()
 			for fn := range p.tasks {
-				fn()
+				func() {
+					defer func() {
+						if r := recover(); r != nil {
+							slog.Error("worker panic recovered", "error", r)
+						}
+					}()
+					fn()
+				}()
 			}
 		}()
 	}
 	return p
 }
 
-// submit enqueues a task. Blocks if all workers are busy and the queue is full.
+// submit enqueues a task. Returns false if the context is cancelled before the task can be queued.
-func (p *workerPool) submit(fn func()) {
+func (p *workerPool) submit(ctx context.Context, fn func()) bool {
-	p.tasks <- fn
+	select {
+	case p.tasks <- fn:
+		return true
+	case <-ctx.Done():
+		return false
+	}
 }
 
 // stop closes the task channel and waits for all workers to finish.

workerpool_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -13,7 +14,7 @@ func TestWorkerPool(t *testing.T) {
 
 		var count atomic.Int32
 		for i := 0; i < 100; i++ {
-			pool.submit(func() {
+			pool.submit(context.Background(), func() {
 				count.Add(1)
 			})
 		}
@@ -32,7 +33,7 @@ func TestWorkerPool(t *testing.T) {
 		var maxConcurrent atomic.Int32
 
 		for i := 0; i < 20; i++ {
-			pool.submit(func() {
+			pool.submit(context.Background(), func() {
 				cur := concurrent.Add(1)
 				for {
 					old := maxConcurrent.Load()
@@ -57,3 +58,56 @@ func TestWorkerPool(t *testing.T) {
 		pool.stop() // should not panic
 	})
 }
+
+func TestWorkerPoolRecoversPanic(t *testing.T) {
+	pool := newWorkerPool(1)
+	defer pool.stop()
+
+	// Submit a function that panics
+	pool.submit(context.Background(), func() { panic("boom") })
+
+	// Give the panic time to be processed
+	time.Sleep(50 * time.Millisecond)
+
+	// Submit a normal function — the worker should still be alive
+	done := make(chan struct{})
+	ok := pool.submit(context.Background(), func() { close(done) })
+	if !ok {
+		t.Fatal("expected submit to succeed after panic recovery")
+	}
+
+	select {
+	case <-done:
+		// Worker survived the panic
+	case <-time.After(2 * time.Second):
+		t.Error("timed out — worker did not survive panic")
+	}
+}
+
+func TestWorkerPoolSubmitRespectsContext(t *testing.T) {
+	pool := newWorkerPool(1) // 1 worker, queue capacity 4
+	defer pool.stop()
+
+	blocker := make(chan struct{})
+
+	// Occupy the single worker
+	pool.submit(context.Background(), func() { <-blocker })
+
+	// Fill the buffered queue (capacity = n*4 = 4)
+	for i := 0; i < 4; i++ {
+		pool.submit(context.Background(), func() { <-blocker })
+	}
+
+	// Now the queue is full and the worker is busy.
+	// Submit with a cancelled context should return false immediately.
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	ok := pool.submit(ctx, func() { t.Error("should not execute") })
+	if ok {
+		t.Error("expected submit to return false with cancelled context")
+	}
+
+	// Unblock everything for cleanup
+	close(blocker)
+}