From 1c33225b080f67323e136e756f2347b739a0e746 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Sun, 21 Jun 2026 14:34:24 -0500 Subject: [PATCH] fix: 14 bug fixes across agent, websocket, mikrotik, checks, and hostkeys - agent.go: fix buffer pool leak on protobuf marshal error - agent.go: reject malformed/phx_error join replies as failures - agent.go: add sessionCtx.Done() to main loop select - agent.go/websocket.go: prevent double-close race in WSConn - websocket.go: fix case-insensitive header parsing truncation - websocket.go: fix IPv6 literal host parsing missing default port - mikrotik.go: fix per-word read deadline reset (DoS vector) - mikrotik.go: use context with deadline for TLS dial - checks.go: fix sslRootCAs caching errors forever via sync.Once - checks.go: use context-aware dial for TCP checks - checks.go: fix responseTimeMs=0 sentinel ambiguity - checks.go: report regex compile error as UNKNOWN not CRITICAL - hostkeys.go: log JSON unmarshal errors in known_hosts.json - ssh_test.go: isolate tests with resetHostKeyStore to prevent TOFU contamination --- agent.go | 24 ++++++++++++++++-------- checks.go | 32 ++++++++++++++++++++------------ checks_test.go | 4 ++-- hostkeys.go | 6 +++++- mikrotik.go | 14 ++++++++------ ssh_test.go | 26 ++++++++++++++++++-------- websocket.go | 41 +++++++++++++++++++++++++++++------------ 7 files changed, 98 insertions(+), 49 deletions(-) diff --git a/agent.go b/agent.go index 7026ebb..67fe945 100644 --- a/agent.go +++ b/agent.go @@ -165,6 +165,7 @@ func runSession(ctx context.Context, baseURL, token string) error { bin, err := proto.MarshalOptions{}.MarshalAppend(buf, msg) if err != nil { slog.Error("marshal protobuf", "error", err) + bufPool.Put(bp) return } encoded := base64.StdEncoding.EncodeToString(bin) @@ -230,14 +231,17 @@ func runSession(ctx context.Context, baseURL, token string) error { 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) - } + if reply.Event != "phx_reply" { + return fmt.Errorf("expected phx_reply, got %s", reply.Event) + } + var status struct { + Status string `json:"status"` + } + if err := json.Unmarshal(reply.Payload, &status); err != nil { + return fmt.Errorf("join reply payload: %w", err) + } + if status.Status != "ok" { + return fmt.Errorf("join rejected: %s", status.Status) } slog.Info("channel joined") case err := <-errCh: @@ -312,6 +316,10 @@ func runSession(ctx context.Context, baseURL, token string) error { flushSnmpBatch() return nil + case <-sessionCtx.Done(): + flushSnmpBatch() + return fmt.Errorf("session cancelled") + case err := <-errCh: flushSnmpBatch() return fmt.Errorf("read: %w", err) diff --git a/checks.go b/checks.go index 7e45b01..0935134 100644 --- a/checks.go +++ b/checks.go @@ -30,15 +30,23 @@ var ( IdleConnTimeout: 90 * time.Second, } sslRootCAsOnce sync.Once + sslRootCAsMu sync.Mutex sslRootCAsPool *x509.CertPool sslRootCAsErr error - // sslRootCAs returns the system cert pool, cached after first load. + // sslRootCAs returns the system cert pool, cached after first successful load. + // Errors are not cached — subsequent calls will retry loading. // Overridable for tests. 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{} + } return sslRootCAsPool, sslRootCAsErr } ) @@ -50,33 +58,32 @@ func ExecuteCheck(ctx context.Context, check *pb.Check) *pb.CheckResult { var status uint32 var output string - var responseTimeMs float64 switch check.CheckType { case "http": if httpConfig := check.GetHttp(); httpConfig != nil { - status, output, responseTimeMs = executeHTTPCheck(ctx, httpConfig, check.TimeoutMs) + status, output, _ = executeHTTPCheck(ctx, httpConfig, check.TimeoutMs) } else { status, output = 3, "Missing HTTP config" } case "tcp": if tcpConfig := check.GetTcp(); tcpConfig != nil { - status, output, responseTimeMs = executeTCPCheck(ctx, tcpConfig, check.TimeoutMs) + status, output, _ = executeTCPCheck(ctx, tcpConfig, check.TimeoutMs) } else { status, output = 3, "Missing TCP config" } case "dns": if dnsConfig := check.GetDns(); dnsConfig != nil { - status, output, responseTimeMs = executeDNSCheck(ctx, dnsConfig, check.TimeoutMs) + status, output, _ = executeDNSCheck(ctx, dnsConfig, check.TimeoutMs) } else { status, output = 3, "Missing DNS config" } case "ssl": if sslConfig := check.GetSsl(); sslConfig != nil { - status, output, responseTimeMs = executeSSLCheck(ctx, sslConfig, check.TimeoutMs) + status, output, _ = executeSSLCheck(ctx, sslConfig, check.TimeoutMs) } else { status, output = 3, "Missing SSL config" } @@ -85,10 +92,8 @@ func ExecuteCheck(ctx context.Context, check *pb.Check) *pb.CheckResult { status, output = 3, fmt.Sprintf("Unknown check type: %s", check.CheckType) } - // If responseTimeMs wasn't set by executor, calculate from start time - if responseTimeMs == 0 { - responseTimeMs = float64(time.Since(startTime).Milliseconds()) - } + // Always calculate elapsed from the outer startTime for consistency. + responseTimeMs := float64(time.Since(startTime).Milliseconds()) return &pb.CheckResult{ CheckId: check.Id, @@ -160,7 +165,7 @@ func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs if config.Regex != "" { re, err := regexp.Compile(config.Regex) if err != nil { - return 2, fmt.Sprintf("Invalid regex: %v", err), responseTime + return 3, fmt.Sprintf("Invalid regex: %v", err), responseTime } body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) if err != nil { @@ -184,7 +189,10 @@ func executeTCPCheck(ctx context.Context, config *pb.TcpCheckConfig, timeoutMs u address := net.JoinHostPort(config.Host, strconv.Itoa(int(config.Port))) startTime := time.Now() - conn, err := net.DialTimeout("tcp", address, timeout) + dialCtx, dialCancel := context.WithTimeout(ctx, timeout) + defer dialCancel() + var d net.Dialer + conn, err := d.DialContext(dialCtx, "tcp", address) responseTime := float64(time.Since(startTime).Milliseconds()) if err != nil { diff --git a/checks_test.go b/checks_test.go index 55b3da5..1e2e7ee 100644 --- a/checks_test.go +++ b/checks_test.go @@ -429,8 +429,8 @@ func TestHTTPCheck_InvalidRegex(t *testing.T) { Regex: `[invalid`, }, 5000) - if status != 2 { - t.Fatalf("expected status 2 for invalid regex, got %d: %s", status, output) + if status != 3 { + t.Fatalf("expected status 3 for invalid regex, got %d: %s", status, output) } if !strings.Contains(output, "Invalid regex") { t.Fatalf("expected 'Invalid regex' in output, got %s", output) diff --git a/hostkeys.go b/hostkeys.go index cd8f6ee..e25c4d8 100644 --- a/hostkeys.go +++ b/hostkeys.go @@ -38,7 +38,11 @@ func newHostKeyStore(path string) *hostKeyStore { s := &hostKeyStore{path: path, keys: make(map[string]string)} data, err := os.ReadFile(path) if err == nil { - _ = json.Unmarshal(data, &s.keys) + if err := json.Unmarshal(data, &s.keys); err != nil { + slog.Warn("failed to parse known_hosts.json, starting with empty key store", + "path", path, + "error", err) + } } return s } diff --git a/mikrotik.go b/mikrotik.go index 93278bb..609c279 100644 --- a/mikrotik.go +++ b/mikrotik.go @@ -49,7 +49,9 @@ func mikrotikConnect(ip string, port uint32, username, password string, useSSL b NetDialer: &net.Dialer{Timeout: mikrotikConnTimeout}, Config: &tls.Config{InsecureSkipVerify: true, MinVersion: tls.VersionTLS12}, } - conn, err = dialer.DialContext(context.Background(), "tcp", addr) + dialCtx, dialCancel := context.WithTimeout(context.Background(), mikrotikConnTimeout) + defer dialCancel() + conn, err = dialer.DialContext(dialCtx, "tcp", addr) if err == nil { // Verify TLS cert fingerprint via TOFU tlsConn, ok := conn.(*tls.Conn) @@ -160,13 +162,13 @@ func (c *mikrotikClient) readResponse() (*mikrotikResponse, error) { } func (c *mikrotikClient) readSentence() ([]string, error) { + if tc, ok := c.conn.(net.Conn); ok { + if err := tc.SetReadDeadline(time.Now().Add(mikrotikReadTimeout)); err != nil { + return nil, fmt.Errorf("set read deadline: %w", err) + } + } var words []string for { - if tc, ok := c.conn.(net.Conn); ok { - if err := tc.SetReadDeadline(time.Now().Add(mikrotikReadTimeout)); err != nil { - return nil, fmt.Errorf("set read deadline: %w", err) - } - } word, err := c.readWord() if err != nil { return nil, err diff --git a/ssh_test.go b/ssh_test.go index 4accd2d..455f900 100644 --- a/ssh_test.go +++ b/ssh_test.go @@ -320,7 +320,22 @@ func TestExecuteMikrotikBackupDialError(t *testing.T) { } } +// resetHostKeyStore resets the global host key store for SSH tests +// to prevent cross-test TOFU contamination from different server keys. +func resetHostKeyStore(t *testing.T) { + t.Helper() + origStore := globalHostKeys + t.Cleanup(func() { + hostKeysOnce = sync.Once{} + globalHostKeys = origStore + }) + hostKeysOnce = sync.Once{} + t.Setenv("TOWEROPS_HOST_KEYS_FILE", filepath.Join(t.TempDir(), "hosts.json")) +} + func TestExecuteMikrotikBackupSuccess(t *testing.T) { + resetHostKeyStore(t) + addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) { _, _ = ch.Write([]byte("# RouterOS config\n/ip address\nadd address=10.0.0.1/24\n")) _ = ch.CloseWrite() @@ -344,6 +359,7 @@ func TestExecuteMikrotikBackupSuccess(t *testing.T) { } func TestExecuteMikrotikBackupCommandError(t *testing.T) { + resetHostKeyStore(t) addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) { // Send exit-status 1 with no output (simulates command failure) _, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{1})) @@ -362,6 +378,7 @@ func TestExecuteMikrotikBackupCommandError(t *testing.T) { } func TestExecuteMikrotikBackupWithOutput(t *testing.T) { + resetHostKeyStore(t) addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) { _, _ = ch.Write([]byte("# partial config\n")) _ = ch.CloseWrite() @@ -384,14 +401,7 @@ func TestExecuteMikrotikBackupWithOutput(t *testing.T) { } func TestExecuteMikrotikBackupSessionError(t *testing.T) { - // Reset global host key store to avoid TOFU collisions from other tests - origStore := globalHostKeys - defer func() { - hostKeysOnce = sync.Once{} - globalHostKeys = origStore - }() - hostKeysOnce = sync.Once{} - t.Setenv("TOWEROPS_HOST_KEYS_FILE", filepath.Join(t.TempDir(), "hosts.json")) + resetHostKeyStore(t) // SSH server that accepts connection but rejects all channel requests key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) diff --git a/websocket.go b/websocket.go index 51665eb..52e97ee 100644 --- a/websocket.go +++ b/websocket.go @@ -42,6 +42,7 @@ type WSConn struct { conn io.ReadWriteCloser reader *bufio.Reader mu sync.Mutex // serializes writes + closed bool // prevents double-close } var wsHandshakeTimeout = 30 * time.Second @@ -66,14 +67,16 @@ func WSDial(rawURL string) (*WSConn, error) { slog.Warn("plaintext websocket connection - credentials sent unencrypted", "url", sanitizeURL(rawURL)) } - host := u.Host - if !strings.Contains(host, ":") { + hostname := u.Hostname() + port := u.Port() + if port == "" { if u.Scheme == "wss" { - host += ":443" + port = "443" } else { - host += ":80" + port = "80" } } + host := net.JoinHostPort(hostname, port) ws, err := wsConnect(u, host, "tcp") if err != nil { @@ -156,23 +159,29 @@ func wsConnect(u *url.URL, host, network string) (*WSConn, error) { if line == "" { break // end of headers } - if strings.HasPrefix(strings.ToLower(line), "sec-websocket-accept: ") { - actual := strings.TrimSpace(line[len("Sec-WebSocket-Accept: "):]) - if actual != expectedAccept { + colon := strings.Index(line, ":") + if colon == -1 { + continue + } + name := line[:colon] + value := strings.TrimSpace(line[colon+1:]) + + if strings.EqualFold(name, "sec-websocket-accept") { + if value != expectedAccept { _ = conn.Close() - return nil, fmt.Errorf("invalid accept key: got %q, want %q", actual, expectedAccept) + return nil, fmt.Errorf("invalid accept key: got %q, want %q", value, expectedAccept) } acceptFound = true continue } - if strings.HasPrefix(strings.ToLower(line), "upgrade: ") { - if strings.EqualFold(strings.TrimSpace(line[len("Upgrade: "):]), "websocket") { + if strings.EqualFold(name, "upgrade") { + if strings.EqualFold(value, "websocket") { upgradeFound = true } continue } - if strings.HasPrefix(strings.ToLower(line), "connection: ") { - if headerHasToken(line[len("Connection: "):], "upgrade") { + if strings.EqualFold(name, "connection") { + if headerHasToken(line[colon+1:], "upgrade") { connectionFound = true } } @@ -237,6 +246,14 @@ func (ws *WSConn) WriteText(data []byte) error { // Close sends a close frame and closes the underlying connection. func (ws *WSConn) Close() error { + ws.mu.Lock() + if ws.closed { + ws.mu.Unlock() + return nil + } + ws.closed = true + ws.mu.Unlock() + _ = ws.writeFrame(opClose, nil) // best-effort return ws.conn.Close() }