security: fixes 1-5 — payload limits, non-blocking submit, panic recovery, handshake timeout, write error propagation

Fix 1: Limit inbound job payload size to 4MB to prevent OOM from malicious server
Fix 2: Worker pool submit accepts context and returns bool, prevents blocking event loop
Fix 3: Panic recovery in worker pool goroutines prevents permanent worker death
Fix 4: WebSocket handshake timeout (30s) prevents indefinite blocking on slow server
Fix 5: Writer goroutine propagates errors to main loop via writeErrCh

agent.go

@@ -22,6 +22,8 @@ import (
 var osExit = os.Exit
 var doSelfUpdate = selfUpdate
 
+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"`
@@ -157,12 +159,17 @@ func runSession(ctx context.Context, baseURL, token string) error {
 
 	// 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
 			}
 		}
@@ -222,6 +229,10 @@ 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 {
@@ -300,6 +311,10 @@ func handleMessage(
 			slog.Error("decode job payload", "error", err)
 			return
 		}
+		if len(payload.Binary) > maxJobPayloadBytes {
+			slog.Error("job payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes)
+			return
+		}
 		bin, err := base64.StdEncoding.DecodeString(payload.Binary)
 		if err != nil {
 			slog.Error("decode base64", "error", err)
@@ -357,15 +372,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

@@ -336,6 +336,28 @@ 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 TestZeroBytes(t *testing.T) {
 	b := []byte{1, 2, 3, 4, 5}
 	zeroBytes(b)

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.
-func (p *workerPool) submit(fn func()) {
-	p.tasks <- fn
+// submit enqueues a task. Returns false if the context is cancelled before the task can be queued.
+func (p *workerPool) submit(ctx context.Context, fn func()) bool {
+	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)
+}