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

@@ -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)

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 {

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)

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
 }

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

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)

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()
 }