diff --git a/agent.go b/agent.go index 8eedb2d..1899259 100644 --- a/agent.go +++ b/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) } } diff --git a/agent_test.go b/agent_test.go index 852a49c..bbdf277 100644 --- a/agent_test.go +++ b/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) diff --git a/websocket.go b/websocket.go index 20d21e5..a9f5bdc 100644 --- a/websocket.go +++ 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 } diff --git a/websocket_test.go b/websocket_test.go index 9543286..ec51ea3 100644 --- a/websocket_test.go +++ b/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") diff --git a/workerpool.go b/workerpool.go index 2c565b9..adfbe71 100644 --- a/workerpool.go +++ b/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. diff --git a/workerpool_test.go b/workerpool_test.go index b83d4d6..9deb518 100644 --- a/workerpool_test.go +++ b/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) +}