From b33caaa14acfe2d61556720ae0151c0484c3e753 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Sun, 26 Jul 2026 15:54:57 -0500 Subject: [PATCH] fix: handle missing results, topic validation, and fragmented WebSocket frames - LLDP topology and SNMP jobs now always send a result to the server, preventing indefinite waits when errors or cancellation occur - handleMessage validates message topic to reject cross-agent messages - WebSocket client reassembles fragmented (continuation) frames per RFC 6455 - Replace fragile sync.Once reset in sslRootCAs with mutex-guarded nil check - Guard nextBackoff against zero/negative input that would panic - Drain HTTP response body on regex compilation failure for connection reuse --- agent.go | 18 +++++++++++-- agent_test.go | 60 ++++++++++++++++++++++++----------------- checks.go | 18 ++++++------- lldp.go | 28 ++++++++++++++++--- snmp.go | 8 +++++- websocket.go | 68 ++++++++++++++++++++++++++++++++++------------- websocket_test.go | 20 +++++++------- 7 files changed, 152 insertions(+), 68 deletions(-) diff --git a/agent.go b/agent.go index 22eee65..1b73a8a 100644 --- a/agent.go +++ b/agent.go @@ -334,7 +334,7 @@ func runSession(ctx context.Context, baseURL, token string) error { slog.Debug("invalid message", "error", err) continue } - shouldEnd, endErr := handleMessage(sessionCtx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh) + shouldEnd, endErr := handleMessage(sessionCtx, msg, topic, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh) if shouldEnd { flushSnmpBatch() return endErr @@ -401,6 +401,7 @@ func runSession(ctx context.Context, baseURL, token string) error { func handleMessage( ctx context.Context, msg channelMsg, + topic string, pools *jobPools, snmpResultCh chan<- *pb.SnmpResult, mikrotikResultCh chan<- *pb.MikrotikResult, @@ -409,6 +410,12 @@ func handleMessage( checkResultCh chan<- *pb.CheckResult, lldpTopologyResultCh chan<- *pb.LldpTopologyResult, ) (bool, error) { + // Ignore messages not addressed to our topic (except Phoenix control messages) + if msg.Topic != topic && msg.Topic != "phoenix" { + slog.Debug("ignoring message for different topic", "got", msg.Topic, "want", topic) + return false, nil + } + switch msg.Event { case "phx_reply": slog.Debug("channel reply", "topic", msg.Topic) @@ -539,11 +546,18 @@ func dispatchJob( // nextBackoff doubles the current delay (capped at max) and adds up to 25% jitter. func nextBackoff(current, maxDelay time.Duration) time.Duration { + if current <= 0 { + current = initialRetryDelay + } next := current * 2 if next > maxDelay { next = maxDelay } - jitter := time.Duration(rand.Int64N(int64(next / 4))) + jitterRange := int64(next / 4) + if jitterRange <= 0 { + return next + } + jitter := time.Duration(rand.Int64N(jitterRange)) return next + jitter } diff --git a/agent_test.go b/agent_test.go index 15b645d..84bec20 100644 --- a/agent_test.go +++ b/agent_test.go @@ -96,7 +96,7 @@ func TestHandleMessage(t *testing.T) { credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) - _, _ = handleMessage(context.Background(), channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "phx_reply", Payload: json.RawMessage(`{}`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Just verify it doesn't panic }) @@ -124,7 +124,7 @@ func TestHandleMessage(t *testing.T) { SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, }) - _, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Wait for goroutine to finish select { case <-snmpCh: @@ -139,7 +139,7 @@ func TestHandleMessage(t *testing.T) { credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) - _, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: json.RawMessage(`not json`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Should log error but not panic }) @@ -150,7 +150,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) }) t.Run("invalid protobuf", func(t *testing.T) { @@ -160,7 +160,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) }) t.Run("restart", func(t *testing.T) { @@ -169,7 +169,7 @@ func TestHandleMessage(t *testing.T) { credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) - shouldEnd, reason := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + shouldEnd, reason := handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "restart", Payload: json.RawMessage(`{}`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) if !shouldEnd { t.Error("expected handleMessage to return true for restart") @@ -192,6 +192,7 @@ func TestHandleMessage(t *testing.T) { shouldEnd, reason := handleMessage( context.Background(), channelMsg{Topic: "agent:test", Event: event, Payload: json.RawMessage(`{}`)}, + "agent:test", testPools(t), snmpCh, mtCh, @@ -226,7 +227,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) if calledURL != "https://example.com/agent" { t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL) @@ -250,7 +251,7 @@ func TestHandleMessage(t *testing.T) { checkCh := make(chan *pb.CheckResult, 1) // Missing URL field payload, _ := json.Marshal(map[string]string{"checksum": "abc123"}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) if called { t.Error("selfUpdate should not be called with empty URL") @@ -271,7 +272,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Should log error but not panic }) @@ -291,7 +292,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) if called { t.Error("selfUpdate should not be called with empty checksum") @@ -304,7 +305,7 @@ func TestHandleMessage(t *testing.T) { credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) - _, _ = handleMessage(context.Background(), channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Should just log and not panic }) @@ -321,7 +322,7 @@ func TestHandleMessage(t *testing.T) { credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) - _, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) select { case <-checkCh: case <-time.After(time.Second): @@ -335,7 +336,7 @@ func TestHandleMessage(t *testing.T) { credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) - _, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: json.RawMessage(`not json`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Should log error but not panic }) @@ -346,7 +347,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Should log error but not panic }) @@ -357,7 +358,7 @@ func TestHandleMessage(t *testing.T) { monCh := make(chan *pb.MonitoringCheck, 1) checkCh := make(chan *pb.CheckResult, 1) payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})}) - _, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Should log error but not panic }) @@ -384,7 +385,7 @@ func TestHandleMessage(t *testing.T) { JobType: pb.JobType_DISCOVER, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, }) - _, _ = handleMessage(context.Background(), channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "discovery_job", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) select { case <-snmpCh: case <-time.After(500 * time.Millisecond): @@ -412,7 +413,7 @@ func TestHandleMessage(t *testing.T) { JobType: pb.JobType_MIKROTIK, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, }) - _, _ = handleMessage(context.Background(), channelMsg{Event: "backup_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "backup_job", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) select { case result := <-mtCh: if result.Error != "" { @@ -436,7 +437,7 @@ func TestHandleMessageRejectsOversizedPayload(t *testing.T) { 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, checkCh, make(chan *pb.LldpTopologyResult, 1)) + _, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) // Verify no jobs were dispatched select { @@ -833,8 +834,9 @@ func TestDispatchJobCancelledContext(t *testing.T) { // fakeWSServer is a test helper that sets up a WebSocket server for runSession tests. type fakeWSServer struct { - ln net.Listener - conn net.Conn + ln net.Listener + conn net.Conn + topic string } func newFakeWSServer(t *testing.T) *fakeWSServer { @@ -867,13 +869,21 @@ func (s *fakeWSServer) acceptAndJoin(t *testing.T) { 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 - frameBuf := make([]byte, 4096) - _, _ = conn.Read(frameBuf) + // Read join frame and extract topic + joinPayload, err := readMaskedFrame(conn) + if err == nil { + var joinMsg channelMsg + if err := json.Unmarshal(joinPayload, &joinMsg); err == nil && joinMsg.Topic != "" { + s.topic = joinMsg.Topic + } + } + if s.topic == "" { + s.topic = "agent:agent-0" + } // Send join OK reply, _ := json.Marshal(channelMsg{ - Topic: "agent:agent-0", + Topic: s.topic, Event: "phx_reply", Payload: json.RawMessage(`{"status":"ok"}`), Ref: strPtr("1"), @@ -884,7 +894,7 @@ func (s *fakeWSServer) acceptAndJoin(t *testing.T) { // sendEvent sends a channel message to the connected client. func (s *fakeWSServer) sendEvent(event string, payload json.RawMessage) { msg, _ := json.Marshal(channelMsg{ - Topic: "agent:agent-0", + Topic: s.topic, Event: event, Payload: payload, }) diff --git a/checks.go b/checks.go index c118447..2f62c1b 100644 --- a/checks.go +++ b/checks.go @@ -29,10 +29,8 @@ var ( MaxIdleConnsPerHost: 10, IdleConnTimeout: 90 * time.Second, } - sslRootCAsOnce sync.Once sslRootCAsMu sync.Mutex sslRootCAsPool *x509.CertPool - sslRootCAsErr error // sslRootCAs returns the system cert pool, cached after first successful load. // Errors are not cached — subsequent calls will retry loading. @@ -40,14 +38,15 @@ var ( sslRootCAs = func() (*x509.CertPool, error) { sslRootCAsMu.Lock() defer sslRootCAsMu.Unlock() - sslRootCAsOnce.Do(func() { - sslRootCAsPool, sslRootCAsErr = x509.SystemCertPool() - }) - if sslRootCAsErr != nil { - // Reset the once so next call will retry - sslRootCAsOnce = sync.Once{} + if sslRootCAsPool != nil { + return sslRootCAsPool, nil } - return sslRootCAsPool, sslRootCAsErr + pool, err := x509.SystemCertPool() + if err != nil { + return nil, err + } + sslRootCAsPool = pool + return pool, nil } ) @@ -168,6 +167,7 @@ func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs if config.Regex != "" { re, err := regexp.Compile(config.Regex) if err != nil { + _, _ = io.Copy(io.Discard, resp.Body) return 3, fmt.Sprintf("Invalid regex: %v", err), responseTime } body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) diff --git a/lldp.go b/lldp.go index ccbbc3a..1c169d3 100644 --- a/lldp.go +++ b/lldp.go @@ -27,9 +27,15 @@ const ( func executeLldpTopologyJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.LldpTopologyResult) { deviceID := job.DeviceId jobID := job.JobId + timestamp := time.Now().Unix() if job.SnmpDevice == nil { slog.Error("missing SNMP config for LLDP job", "job_id", jobID, "device_id", deviceID) + sendLldpResultWithTimeout(ctx, resultCh, &pb.LldpTopologyResult{ + DeviceId: deviceID, + JobId: jobID, + Timestamp: timestamp, + }, jobID) return } @@ -37,6 +43,11 @@ func executeLldpTopologyJob(ctx context.Context, job *pb.AgentJob, resultCh chan client, err := newSnmpConn(snmpDev) if err != nil { slog.Error("failed to connect SNMP for LLDP", "job_id", jobID, "device_id", deviceID, "error", err) + sendLldpResultWithTimeout(ctx, resultCh, &pb.LldpTopologyResult{ + DeviceId: deviceID, + JobId: jobID, + Timestamp: timestamp, + }, jobID) return } defer func() { @@ -48,14 +59,25 @@ func executeLldpTopologyJob(ctx context.Context, job *pb.AgentJob, resultCh chan result, err := discoverLldpNeighbors(client, deviceID, jobID) if err != nil { slog.Error("LLDP discovery failed", "job_id", jobID, "device_id", deviceID, "error", err) + sendLldpResultWithTimeout(ctx, resultCh, &pb.LldpTopologyResult{ + DeviceId: deviceID, + JobId: jobID, + Timestamp: timestamp, + }, jobID) return } + sendLldpResultWithTimeout(ctx, resultCh, result, jobID) + slog.Info("LLDP topology discovered", "job_id", jobID, "device_id", deviceID, "neighbors", len(result.Neighbors)) +} + +func sendLldpResultWithTimeout(ctx context.Context, resultCh chan<- *pb.LldpTopologyResult, result *pb.LldpTopologyResult, jobID string) { + sendCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() select { case resultCh <- result: - slog.Info("LLDP topology discovered", "job_id", jobID, "device_id", deviceID, "neighbors", len(result.Neighbors)) - case <-ctx.Done(): - slog.Warn("LLDP result send cancelled", "job_id", jobID) + case <-sendCtx.Done(): + slog.Error("LLDP result send timeout - agent overloaded", "job_id", jobID) } } diff --git a/snmp.go b/snmp.go index afed76d..c01793e 100644 --- a/snmp.go +++ b/snmp.go @@ -67,9 +67,11 @@ func executeSnmpJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.S } oidValues := make(map[string]string, totalOIDs) + cancelled := false for _, q := range job.Queries { if ctx.Err() != nil { - return + cancelled = true + break } switch q.QueryType { case pb.QueryType_GET: @@ -115,6 +117,10 @@ func executeSnmpJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.S } } + if cancelled { + slog.Warn("snmp job cancelled, sending partial result", "job_id", job.JobId, "oids", len(oidValues)) + } + result := &pb.SnmpResult{ DeviceId: job.DeviceId, JobType: job.JobType, diff --git a/websocket.go b/websocket.go index 1056790..3e6fbdf 100644 --- a/websocket.go +++ b/websocket.go @@ -28,11 +28,12 @@ func computeAcceptKey(key string) string { } const ( - opText = 1 - opBinary = 2 - opClose = 8 - opPing = 9 - opPong = 10 + opText = 1 + opBinary = 2 + opClose = 8 + opPing = 9 + opPong = 10 + opContinuation = 0 maxFrameSize = 16 << 20 // 16 MB ) @@ -219,22 +220,52 @@ func headerHasToken(value, token string) bool { } // ReadMessage reads the next text or binary message, handling control frames internally. +// Fragmented messages (continuation frames per RFC 6455 Section 5.4) are reassembled. func (ws *WSConn) ReadMessage() ([]byte, int, error) { + var accumulated []byte + var msgOpcode int + for { - opcode, payload, err := ws.readFrame() + fin, opcode, payload, err := ws.readFrame() if err != nil { return nil, 0, err } + switch opcode { - case opText, opBinary: - return payload, opcode, nil case opPing: if err := ws.writeFrame(opPong, payload); err != nil { return nil, 0, fmt.Errorf("pong: %w", err) } + continue case opClose: - _ = ws.writeFrame(opClose, nil) // best-effort close reply + _ = ws.writeFrame(opClose, nil) return nil, opClose, io.EOF + case opPong: + continue + } + + if opcode == opText || opcode == opBinary { + // New message — if we already had an incomplete fragmented message, discard it + // (protocol violation by the server, but be resilient). + if !fin { + msgOpcode = opcode + accumulated = append(accumulated[:0], payload...) + continue + } + // Single unfragmented message + return payload, opcode, nil + } + + // Continuation frame (opcode 0) + if accumulated == nil { + // Unexpected continuation without a start frame — skip + continue + } + accumulated = append(accumulated, payload...) + if fin { + result := accumulated + accumulated = nil + return result, msgOpcode, nil } } } @@ -261,19 +292,20 @@ func (ws *WSConn) Close() error { const wsReadTimeout = 90 * time.Second const wsWriteTimeout = 30 * time.Second -func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) { +func (ws *WSConn) readFrame() (fin bool, opcode int, payload []byte, err error) { // Set read deadline to detect dead connections (MEDIUM-4) if nc, ok := ws.conn.(net.Conn); ok { if err := nc.SetReadDeadline(time.Now().Add(wsReadTimeout)); err != nil { - return 0, nil, fmt.Errorf("set read deadline: %w", err) + return false, 0, nil, fmt.Errorf("set read deadline: %w", err) } } var header [2]byte if _, err = io.ReadFull(ws.reader, header[:]); err != nil { - return 0, nil, err + return false, 0, nil, err } + fin = header[0]&0x80 != 0 opcode = int(header[0] & 0x0F) masked := header[1]&0x80 != 0 length := uint64(header[1] & 0x7F) @@ -282,32 +314,32 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) { case 126: var ext [2]byte if _, err = io.ReadFull(ws.reader, ext[:]); err != nil { - return 0, nil, err + return false, 0, nil, err } length = uint64(binary.BigEndian.Uint16(ext[:])) case 127: var ext [8]byte if _, err = io.ReadFull(ws.reader, ext[:]); err != nil { - return 0, nil, err + return false, 0, nil, err } length = binary.BigEndian.Uint64(ext[:]) } if length > maxFrameSize { - return 0, nil, fmt.Errorf("frame size %d exceeds max %d", length, maxFrameSize) + return false, 0, nil, fmt.Errorf("frame size %d exceeds max %d", length, maxFrameSize) } var maskKey [4]byte if masked { if _, err = io.ReadFull(ws.reader, maskKey[:]); err != nil { - return 0, nil, err + return false, 0, nil, err } } payload = make([]byte, length) if length > 0 { if _, err = io.ReadFull(ws.reader, payload); err != nil { - return 0, nil, err + return false, 0, nil, err } } @@ -317,7 +349,7 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) { } } - return opcode, payload, nil + return fin, opcode, payload, nil } func (ws *WSConn) writeFrame(opcode int, payload []byte) error { diff --git a/websocket_test.go b/websocket_test.go index 599a3f4..69bc0ae 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -120,7 +120,7 @@ func TestReadFrame(t *testing.T) { buf.Write(payload) ws := testWSConn(&nopCloser{readWriter: &buf}) - opcode, data, err := ws.readFrame() + _, opcode, data, err := ws.readFrame() if err != nil { t.Fatal(err) } @@ -146,7 +146,7 @@ func TestReadFrame16BitLength(t *testing.T) { buf.Write(payload) ws := testWSConn(&nopCloser{readWriter: &buf}) - opcode, data, err := ws.readFrame() + _, opcode, data, err := ws.readFrame() if err != nil { t.Fatal(err) } @@ -169,7 +169,7 @@ func TestReadFrame64BitLength(t *testing.T) { buf.Write(payload) ws := testWSConn(&nopCloser{readWriter: &buf}) - opcode, data, err := ws.readFrame() + _, opcode, data, err := ws.readFrame() if err != nil { t.Fatal(err) } @@ -196,7 +196,7 @@ func TestReadFrameMasked(t *testing.T) { buf.Write(masked) ws := testWSConn(&nopCloser{readWriter: &buf}) - opcode, data, err := ws.readFrame() + _, opcode, data, err := ws.readFrame() if err != nil { t.Fatal(err) } @@ -265,7 +265,7 @@ func TestReadFrame16BitLengthError(t *testing.T) { buf.WriteByte(126) // 16-bit length indicator, but no length bytes follow ws := testWSConn(&nopCloser{readWriter: &buf}) - _, _, err := ws.readFrame() + _, _, _, err := ws.readFrame() if err == nil { t.Error("expected error for truncated 16-bit length") } @@ -277,7 +277,7 @@ func TestReadFrame64BitLengthError(t *testing.T) { buf.WriteByte(127) // 64-bit length indicator, but no length bytes follow ws := testWSConn(&nopCloser{readWriter: &buf}) - _, _, err := ws.readFrame() + _, _, _, err := ws.readFrame() if err == nil { t.Error("expected error for truncated 64-bit length") } @@ -289,7 +289,7 @@ func TestReadFrameMaskKeyError(t *testing.T) { buf.WriteByte(0x80 | 0x01) // masked + length 1, but no mask key or payload ws := testWSConn(&nopCloser{readWriter: &buf}) - _, _, err := ws.readFrame() + _, _, _, err := ws.readFrame() if err == nil { t.Error("expected error for missing mask key") } @@ -301,7 +301,7 @@ func TestReadFramePayloadError(t *testing.T) { buf.WriteByte(5) // length 5, but no payload ws := testWSConn(&nopCloser{readWriter: &buf}) - _, _, err := ws.readFrame() + _, _, _, err := ws.readFrame() if err == nil { t.Error("expected error for truncated payload") } @@ -348,7 +348,7 @@ func TestReadFrameExceedsMaxSize(t *testing.T) { // Don't need to write the payload — should reject before reading it ws := testWSConn(&nopCloser{readWriter: &buf}) - _, _, err := ws.readFrame() + _, _, _, err := ws.readFrame() if err == nil { t.Error("expected error for frame exceeding max size") } @@ -363,7 +363,7 @@ func TestReadFrameEmptyPayload(t *testing.T) { buf.WriteByte(0) // zero length ws := testWSConn(&nopCloser{readWriter: &buf}) - opcode, data, err := ws.readFrame() + _, opcode, data, err := ws.readFrame() if err != nil { t.Fatal(err) }