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

.forgejo/workflows/release.yml

@@ -0,0 +1,53 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  docker:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: https://code.forgejo.org/actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - uses: https://code.forgejo.org/actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+
+      - name: Set up QEMU
+        uses: https://github.com/docker/setup-qemu-action@v3
+
+      - name: Set up Docker Buildx
+        uses: https://github.com/docker/setup-buildx-action@v3
+
+      - name: Log in to Container Registry
+        uses: https://github.com/docker/login-action@v3
+        with:
+          registry: ${{ secrets.REGISTRY_URL }}
+          username: ${{ secrets.REGISTRY_USER }}
+          password: ${{ secrets.REGISTRY_PASSWORD }}
+
+      - name: Extract version
+        id: version
+        run: echo "tag=${GITHUB_REF_NAME#v}" >> "$GITHUB_OUTPUT"
+
+      - name: Build and push Docker image
+        uses: https://github.com/docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          platforms: linux/amd64,linux/arm64
+          build-args: |
+            VERSION=${{ steps.version.outputs.tag }}
+          tags: |
+            ${{ secrets.REGISTRY_URL }}/${{ github.repository }}:latest
+            ${{ secrets.REGISTRY_URL }}/${{ github.repository }}:${{ steps.version.outputs.tag }}
+
+      - name: Run tests
+        run: go test -race -v -count=1 -timeout 60s ./...
+
+      - name: Vet
+        run: go vet ./...

.forgejo/workflows/renovate.yaml

@@ -1,27 +0,0 @@
-name: Renovate
-
-on:
-  schedule:
-    # Run daily at 6:00 AM UTC
-    - cron: '0 6 * * *'
-  workflow_dispatch:
-
-jobs:
-  renovate:
-    runs-on: ubuntu-latest
-    container:
-      image: renovate/renovate:latest
-    steps:
-      - name: Checkout code
-        uses: https://github.com/actions/checkout@v4
-
-      - name: Run Renovate
-        env:
-          RENOVATE_PLATFORM: forgejo
-          RENOVATE_ENDPOINT: ${{ secrets.FORGEJO_ENDPOINT }}
-          RENOVATE_TOKEN: ${{ secrets.RENOVATE_TOKEN }}
-          RENOVATE_GIT_AUTHOR: Renovate Bot <bot@renovateapp.com>
-          RENOVATE_AUTODISCOVER: false
-          LOG_LEVEL: info
-        run: |
-          renovate graham/towerops-agent

.forgejo/workflows/test.yml

@@ -19,7 +19,7 @@ jobs:
           go-version-file: go.mod
 
       - name: Run tests
-        run: go test -race -v ./...
+        run: go test -race -v -count=1 -timeout 60s ./...
 
       - name: Vet
         run: go vet ./...

.github/dependabot.yml

@@ -1,10 +0,0 @@
-version: 2
-updates:
-  - package-ecosystem: "gomod"
-    directory: "/"
-    schedule:
-      interval: "weekly"
-  - package-ecosystem: "github-actions"
-    directory: "/"
-    schedule:
-      interval: "weekly"

.tool-versions

@@ -0,0 +1 @@
+golang 1.26.4

Makefile

@@ -1,4 +1,4 @@
-.PHONY: proto test build vet lint
+.PHONY: proto test test-fast test-race build vet lint
 
 proto:
 	protoc --go_out=. --go_opt=paths=source_relative proto/agent.proto
@@ -6,7 +6,13 @@ proto:
 	mv proto/agent.pb.go pb/agent.pb.go
 
 test:
-	go test -race -v ./...
+	go test -race -v -count=1 -timeout 60s ./...
+
+test-fast:
+	go test -count=1 -timeout 60s ./...
+
+test-short:
+	go test -count=1 -timeout 30s -short ./...
 
 build:
 	go build -o towerops-agent .

agent.go

@@ -26,6 +26,7 @@ var errChannelReloaded = fmt.Errorf("channel reloaded")
 var joinTimeout = 10 * time.Second
 var heartbeatInterval = 60 * time.Second
 var channelHeartbeatInterval = 25 * time.Second
+var initialRetryDelay = time.Second
 
 const maxJobPayloadBytes = 4 << 20 // 4 MB — well above any legitimate job list
 
@@ -40,7 +41,7 @@ type channelMsg struct {
 // runAgent connects to the server and runs the event loop with reconnect.
 func runAgent(ctx context.Context, wsURL, token string) {
 	baseURL := strings.TrimRight(wsURL, "/")
-	retryDelay := time.Second
+	retryDelay := initialRetryDelay
 	maxRetry := 10 * time.Second
 	const successfulConnectionThreshold = 30 * time.Second
 
@@ -60,7 +61,7 @@ func runAgent(ctx context.Context, wsURL, token string) {
 			slog.Debug("resetting reconnect backoff after successful session",
 				"duration", sessionDuration,
 				"previous_delay", retryDelay)
-			retryDelay = time.Second
+			retryDelay = initialRetryDelay
 		}
 
 		if ctx.Err() != nil {
@@ -68,7 +69,7 @@ func runAgent(ctx context.Context, wsURL, token string) {
 		}
 		if errors.Is(err, errRestartRequested) {
 			slog.Info("restart requested, reconnecting immediately")
-			retryDelay = time.Second
+			retryDelay = initialRetryDelay
 			continue
 		}
 		if err != nil {
@@ -164,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)
@@ -195,7 +197,11 @@ func runSession(ctx context.Context, baseURL, token string) error {
 				sessionCancel()
 				return
 			}
-			msgCh <- data
+			select {
+			case msgCh <- data:
+			case <-sessionCtx.Done():
+				return
+			}
 		}
 	}()
 
@@ -225,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:
@@ -307,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)
@@ -318,7 +331,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
 		case data := <-msgCh:
 			var msg channelMsg
 			if err := json.Unmarshal(data, &msg); err != nil {
-				slog.Warn("invalid message", "error", err)
+				slog.Debug("invalid message", "error", err)
 				continue
 			}
 			shouldEnd, endErr := handleMessage(sessionCtx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh)

agent_test.go

@@ -35,7 +35,7 @@ func TestChannelMsgSerialization(t *testing.T) {
 	s := string(data)
 	checks := []string{"agent:123", "phx_join", "token", "test"}
 	for _, c := range checks {
-		if !contains(s, c) {
+		if !strings.Contains(s, c) {
 			t.Errorf("expected %q in JSON output %q", c, s)
 		}
 	}
@@ -69,19 +69,6 @@ func TestChannelMsgNullRef(t *testing.T) {
 	}
 }
 
-func contains(s, substr string) bool {
-	return len(s) >= len(substr) && searchString(s, substr)
-}
-
-func searchString(s, substr string) bool {
-	for i := 0; i <= len(s)-len(substr); i++ {
-		if s[i:i+len(substr)] == substr {
-			return true
-		}
-	}
-	return false
-}
-
 func testPools(t *testing.T) *jobPools {
 	t.Helper()
 	p := &jobPools{
@@ -141,7 +128,7 @@ func TestHandleMessage(t *testing.T) {
 		// Wait for goroutine to finish
 		select {
 		case <-snmpCh:
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out waiting for snmp result")
 		}
 	})
@@ -337,7 +324,7 @@ func TestHandleMessage(t *testing.T) {
 		_, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		select {
 		case <-checkCh:
-		case <-time.After(5 * time.Second):
+		case <-time.After(time.Second):
 			t.Error("timed out waiting for check result")
 		}
 	})
@@ -400,7 +387,7 @@ func TestHandleMessage(t *testing.T) {
 		_, _ = handleMessage(context.Background(), channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		select {
 		case <-snmpCh:
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out waiting for discovery result")
 		}
 	})
@@ -431,7 +418,7 @@ func TestHandleMessage(t *testing.T) {
 			if result.Error != "" {
 				t.Errorf("unexpected error: %s", result.Error)
 			}
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out waiting for backup result")
 		}
 	})
@@ -560,7 +547,7 @@ func TestDispatchJob(t *testing.T) {
 			if result.Error == "" {
 				t.Error("expected error from unreachable device")
 			}
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out")
 		}
 	})
@@ -589,7 +576,7 @@ func TestDispatchJob(t *testing.T) {
 			if result.Success {
 				t.Error("expected failure")
 			}
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out")
 		}
 	})
@@ -618,7 +605,7 @@ func TestDispatchJob(t *testing.T) {
 			if result.Status != "success" {
 				t.Errorf("expected success, got %q", result.Status)
 			}
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out")
 		}
 	})
@@ -648,7 +635,7 @@ func TestDispatchJob(t *testing.T) {
 
 		select {
 		case <-snmpCh:
-		case <-time.After(2 * time.Second):
+		case <-time.After(500 * time.Millisecond):
 			t.Error("timed out")
 		}
 	})
@@ -701,7 +688,7 @@ func TestRunSessionRejectsFailedJoin(t *testing.T) {
 		_, _ = conn.Write(frame)
 
 		// Keep connection open for a bit
-		time.Sleep(time.Second)
+		time.Sleep(200 * time.Millisecond)
 	}()
 
 	addr := ln.Addr().String()
@@ -745,7 +732,7 @@ func TestRunSessionJoinTimeout(t *testing.T) {
 		frameBuf := make([]byte, 4096)
 		_, _ = conn.Read(frameBuf)
 
-		time.Sleep(5 * time.Second)
+		time.Sleep(2 * time.Second)
 	}()
 
 	addr := ln.Addr().String()
@@ -932,8 +919,8 @@ func TestRunSessionCtxCancel(t *testing.T) {
 		if err != nil {
 			t.Errorf("expected nil error on ctx cancel, got: %v", err)
 		}
-	case <-time.After(5 * time.Second):
-		t.Error("runSession did not exit after ctx cancel")
+		case <-time.After(2 * time.Second):
+			t.Error("runSession did not exit after ctx cancel")
 	}
 	srv.close()
 }
@@ -1088,6 +1075,10 @@ func TestRunSessionReadErrorDuringJoin(t *testing.T) {
 }
 
 func TestRunAgentReconnectOnError(t *testing.T) {
+	origRetry := initialRetryDelay
+	defer func() { initialRetryDelay = origRetry }()
+	initialRetryDelay = 50 * time.Millisecond
+
 	// Server that fails first connection then succeeds, then sends restart.
 	// Agent should reconnect (not exit) after both error and restart.
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
@@ -1130,14 +1121,14 @@ func TestRunAgentReconnectOnError(t *testing.T) {
 					Ref:     strPtr("1"),
 				})
 				_, _ = conn.Write(makeTextFrame(reply))
-				time.Sleep(50 * time.Millisecond)
+				time.Sleep(20 * time.Millisecond)
 				restart, _ := json.Marshal(channelMsg{
 					Topic:   "agent:agent-0",
 					Event:   "restart",
 					Payload: json.RawMessage(`{}`),
 				})
 				_, _ = conn.Write(makeTextFrame(restart))
-				time.Sleep(time.Second)
+				time.Sleep(50 * time.Millisecond)
 				_ = conn.Close()
 			default:
 				// Third connection: agent successfully reconnected after restart
@@ -1146,7 +1137,7 @@ func TestRunAgentReconnectOnError(t *testing.T) {
 		}
 	}()
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
 	done := make(chan struct{})
@@ -1156,7 +1147,7 @@ func TestRunAgentReconnectOnError(t *testing.T) {
 	}()
 
 	// Wait for 3rd connection attempt (proves agent reconnected after both error and restart)
-	deadline := time.After(10 * time.Second)
+	deadline := time.After(5 * time.Second)
 	for {
 		if connCount.Load() >= 3 {
 			cancel()
@@ -1166,7 +1157,7 @@ func TestRunAgentReconnectOnError(t *testing.T) {
 		case <-deadline:
 			t.Fatalf("expected at least 3 connections, got %d", connCount.Load())
 		default:
-			time.Sleep(50 * time.Millisecond)
+			time.Sleep(10 * time.Millisecond)
 		}
 	}
 
@@ -1192,6 +1183,10 @@ func TestRunAgentContextCancellation(t *testing.T) {
 }
 
 func TestRunAgentRestart(t *testing.T) {
+	origRetry := initialRetryDelay
+	defer func() { initialRetryDelay = origRetry }()
+	initialRetryDelay = 50 * time.Millisecond
+
 	// After receiving a restart event, the agent should reconnect (not exit).
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
@@ -1228,20 +1223,20 @@ func TestRunAgentRestart(t *testing.T) {
 
 			if count == 1 {
 				// First connection: send restart event
-				time.Sleep(50 * time.Millisecond)
+				time.Sleep(20 * time.Millisecond)
 				restart, _ := json.Marshal(channelMsg{
 					Topic:   "agent:agent-0",
 					Event:   "restart",
 					Payload: json.RawMessage(`{}`),
 				})
 				_, _ = conn.Write(makeTextFrame(restart))
-				time.Sleep(time.Second)
+				time.Sleep(200 * time.Millisecond)
 			}
 			_ = conn.Close()
 		}
 	}()
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
 	done := make(chan struct{})
@@ -1261,7 +1256,7 @@ func TestRunAgentRestart(t *testing.T) {
 		case <-deadline:
 			t.Fatalf("expected at least 2 connections, got %d", connCount.Load())
 		default:
-			time.Sleep(50 * time.Millisecond)
+			time.Sleep(10 * time.Millisecond)
 		}
 	}
 
@@ -1401,7 +1396,7 @@ func TestRunSessionProcessesJobResults(t *testing.T) {
 		srv.sendEvent("check_jobs", checkPayload)
 
 		// Wait for results to flow through all channels
-		time.Sleep(2 * time.Second)
+		time.Sleep(500 * time.Millisecond)
 		close(stopDrain)
 		srv.close()
 	}()
@@ -1447,7 +1442,7 @@ func TestRunSessionSnmpBatchThreshold(t *testing.T) {
 		}
 		srv.sendEvent("jobs", makeJobPayload(jobs...))
 
-		time.Sleep(3 * time.Second)
+		time.Sleep(time.Second)
 		close(stopDrain)
 		srv.close()
 	}()
@@ -1509,10 +1504,10 @@ func TestExecuteCheckCtxDoneInClosure(t *testing.T) {
 	executeCheck(ctx, check, p, checkCh)
 
 	// Wait for the check to complete (TCP to port 1 fails fast)
-	time.Sleep(500 * time.Millisecond)
+	time.Sleep(20 * time.Millisecond)
 	// Cancel ctx so the closure's select picks ctx.Done instead of blocked channel send
 	cancel()
-	time.Sleep(100 * time.Millisecond)
+	time.Sleep(10 * time.Millisecond)
 }
 
 func TestRunSessionRestartInMainLoop(t *testing.T) {
@@ -1523,10 +1518,10 @@ func TestRunSessionRestartInMainLoop(t *testing.T) {
 		stopDrain := make(chan struct{})
 		go drainFrames(srv.conn, stopDrain)
 
-		time.Sleep(100 * time.Millisecond)
+		time.Sleep(50 * time.Millisecond)
 		// Send restart event — exercised in the main loop select
 		srv.sendEvent("restart", json.RawMessage(`{}`))
-		time.Sleep(time.Second)
+		time.Sleep(200 * time.Millisecond)
 		close(stopDrain)
 		srv.close()
 	}()
@@ -1554,7 +1549,7 @@ func TestRunSessionHeartbeats(t *testing.T) {
 		stopDrain := make(chan struct{})
 		go drainFrames(srv.conn, stopDrain)
 		// Let heartbeats fire a few times
-		time.Sleep(500 * time.Millisecond)
+		time.Sleep(200 * time.Millisecond)
 		close(stopDrain)
 		srv.close()
 	}()
@@ -1566,6 +1561,10 @@ func TestRunSessionHeartbeats(t *testing.T) {
 }
 
 func TestRunAgentCancelDuringRetry(t *testing.T) {
+	origRetry := initialRetryDelay
+	defer func() { initialRetryDelay = origRetry }()
+	initialRetryDelay = 100 * time.Millisecond
+
 	// Connects to a port nothing listens on, then cancel during retry delay
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -1576,13 +1575,13 @@ func TestRunAgentCancelDuringRetry(t *testing.T) {
 	}()
 
 	// Wait for first connection attempt to fail and retry delay to start
-	time.Sleep(1500 * time.Millisecond)
+	time.Sleep(150 * time.Millisecond)
 	cancel()
 
 	select {
 	case <-done:
 		// runAgent returned after cancel during retry
-	case <-time.After(5 * time.Second):
+	case <-time.After(time.Second):
 		t.Error("runAgent did not return after cancel during retry")
 	}
 }

checks.go

@@ -11,6 +11,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"codeberg.org/towerops-agent/towerops-agent/pb"
@@ -28,7 +29,26 @@ var (
 		MaxIdleConnsPerHost: 10,
 		IdleConnTimeout:     90 * time.Second,
 	}
-	sslRootCAs = x509.SystemCertPool
+	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.
+	// 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
+	}
 )
 
 // ExecuteCheck runs a service check and returns the result.
@@ -38,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"
 		}
@@ -73,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,
@@ -146,17 +163,16 @@ func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs
 
 	// Check content regex if provided
 	if config.Regex != "" {
+		re, err := regexp.Compile(config.Regex)
+		if err != nil {
+			return 3, fmt.Sprintf("Invalid regex: %v", err), responseTime
+		}
 		body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
 		if err != nil {
 			return 2, fmt.Sprintf("Failed to read body: %v", err), responseTime
 		}
 
-		matched, err := regexp.MatchString(config.Regex, string(body))
-		if err != nil {
-			return 2, fmt.Sprintf("Invalid regex: %v", err), responseTime
-		}
-
-		if !matched {
+		if !re.Match(body) {
 			return 2, fmt.Sprintf("Content does not match pattern: %s", config.Regex), responseTime
 		}
 	}
@@ -173,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)

go.mod

@@ -4,9 +4,9 @@ go 1.25.6
 
 require (
 	github.com/gosnmp/gosnmp v1.43.2
-	golang.org/x/crypto v0.49.0
-	golang.org/x/net v0.52.0
+	golang.org/x/crypto v0.52.0
+	golang.org/x/net v0.55.0
 	google.golang.org/protobuf v1.36.11
 )
 
-require golang.org/x/sys v0.42.0 // indirect
+require golang.org/x/sys v0.45.0 // indirect

go.sum

@@ -8,14 +8,14 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
-golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4=
-golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA=
-golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0=
-golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw=
-golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
-golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
-golang.org/x/term v0.41.0 h1:QCgPso/Q3RTJx2Th4bDLqML4W6iJiaXFq2/ftQF13YU=
-golang.org/x/term v0.41.0/go.mod h1:3pfBgksrReYfZ5lvYM0kSO0LIkAl4Yl2bXOkKP7Ec2A=
+golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
+golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
+golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8=
+golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww=
+golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
+golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
+golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4=
+golang.org/x/term v0.43.0/go.mod h1:lrhlHNdQJHO+1qVYiHfFKVuVioJIheAc3fBSMFYEIsk=
 google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
 google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

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
 }

lldp.go

@@ -106,7 +106,6 @@ func discoverLldpNeighbors(client *gosnmp.GoSNMP, deviceID, jobID string) (*pb.L
 
 	// Walk remote port descriptions
 	remotePorts := make(map[string]string)
-	var walkErrors []string
 	if err := client.Walk(oidRemPortDesc, func(pdu gosnmp.SnmpPDU) error {
 		key := parseRemoteKey(pdu.Name, oidRemPortDesc)
 		if key != "" {
@@ -114,7 +113,6 @@ func discoverLldpNeighbors(client *gosnmp.GoSNMP, deviceID, jobID string) (*pb.L
 		}
 		return nil
 	}); err != nil {
-		walkErrors = append(walkErrors, fmt.Sprintf("walk remote port descriptions: %v", err))
 		slog.Warn("failed to walk remote port descriptions", "error", err)
 	}
 
@@ -127,7 +125,6 @@ func discoverLldpNeighbors(client *gosnmp.GoSNMP, deviceID, jobID string) (*pb.L
 		}
 		return nil
 	}); err != nil {
-		walkErrors = append(walkErrors, fmt.Sprintf("walk remote port IDs: %v", err))
 		slog.Warn("failed to walk remote port IDs", "error", err)
 	}
 
@@ -140,7 +137,6 @@ func discoverLldpNeighbors(client *gosnmp.GoSNMP, deviceID, jobID string) (*pb.L
 		}
 		return nil
 	}); err != nil {
-		walkErrors = append(walkErrors, fmt.Sprintf("walk management addresses: %v", err))
 		slog.Warn("failed to walk management addresses", "error", err)
 	}
 
@@ -172,9 +168,6 @@ func discoverLldpNeighbors(client *gosnmp.GoSNMP, deviceID, jobID string) (*pb.L
 		result.Neighbors = append(result.Neighbors, neighbor)
 	}
 
-	// Walk errors are already logged; return result with whatever data was collected
-	_ = walkErrors
-
 	return result, nil
 }
 

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

mikrotik_test.go

@@ -660,24 +660,11 @@ func TestReadWordExceedsMaxSize(t *testing.T) {
 	if err == nil {
 		t.Error("expected error for word exceeding max size")
 	}
-	if !containsStr(err.Error(), "exceeds max") {
+	if !strings.Contains(err.Error(), "exceeds max") {
 		t.Errorf("expected 'exceeds max' in error, got: %v", err)
 	}
 }
 
-func containsStr(s, sub string) bool {
-	return len(s) >= len(sub) && searchStr(s, sub)
-}
-
-func searchStr(s, sub string) bool {
-	for i := 0; i <= len(s)-len(sub); i++ {
-		if s[i:i+len(sub)] == sub {
-			return true
-		}
-	}
-	return false
-}
-
 func TestReadSentenceError(t *testing.T) {
 	// Buffer with valid length byte but truncated word data
 	buf := bytes.NewBuffer([]byte{0x03, 'a'}) // length=3 but only 1 byte of data

ping_test.go

@@ -14,7 +14,10 @@ func TestPingDeviceLocalhost(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("skipping ping test on windows")
 	}
-	ms, err := pingDevice("127.0.0.1", 5000)
+	if testing.Short() {
+		t.Skip("skipping real ICMP ping in short mode")
+	}
+	ms, err := pingDevice("127.0.0.1", 2000)
 	if err != nil {
 		t.Skipf("ping not available: %v", err)
 	}
@@ -34,7 +37,10 @@ func TestPingDeviceIPv6(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("skipping ping test on windows")
 	}
-	ms, err := pingDevice("::1", 5000)
+	if testing.Short() {
+		t.Skip("skipping real ICMP ping in short mode")
+	}
+	ms, err := pingDevice("::1", 2000)
 	if err != nil {
 		t.Skipf("IPv6 not available: %v", err)
 	}
@@ -44,7 +50,10 @@ func TestPingDeviceIPv6(t *testing.T) {
 }
 
 func TestIcmpPingLocalhost(t *testing.T) {
-	ms, err := icmpPing("127.0.0.1", 5000)
+	if testing.Short() {
+		t.Skip("skipping real ICMP ping in short mode")
+	}
+	ms, err := icmpPing("127.0.0.1", 2000)
 	if err != nil {
 		t.Skipf("ICMP not available: %v", err)
 	}
@@ -54,7 +63,10 @@ func TestIcmpPingLocalhost(t *testing.T) {
 }
 
 func TestIcmpPingIPv6(t *testing.T) {
-	ms, err := icmpPing("::1", 5000)
+	if testing.Short() {
+		t.Skip("skipping real ICMP ping in short mode")
+	}
+	ms, err := icmpPing("::1", 2000)
 	if err != nil {
 		t.Skipf("IPv6 ICMP not available: %v", err)
 	}
@@ -139,7 +151,10 @@ func TestExecPingLocalhost(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("skipping on windows")
 	}
-	ms, err := execPing("127.0.0.1", 5000)
+	if testing.Short() {
+		t.Skip("skipping real exec ping in short mode")
+	}
+	ms, err := execPing("127.0.0.1", 2000)
 	if err != nil {
 		t.Skipf("ping command not available: %v", err)
 	}
@@ -159,7 +174,10 @@ func TestExecPingIPv6Localhost(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("skipping on windows")
 	}
-	ms, err := execPing("::1", 5000)
+	if testing.Short() {
+		t.Skip("skipping real exec ping in short mode")
+	}
+	ms, err := execPing("::1", 2000)
 	if err != nil {
 		t.Skipf("ping6 not available: %v", err)
 	}
@@ -185,6 +203,9 @@ func TestPingDeviceFallbackToExec(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("skipping on windows")
 	}
+	if testing.Short() {
+		t.Skip("skipping real ping in short mode")
+	}
 	// Mock icmpListenPacket to always fail → forces fallback to execPing
 	origListen := icmpListenPacket
 	defer func() { icmpListenPacket = origListen }()
@@ -193,7 +214,7 @@ func TestPingDeviceFallbackToExec(t *testing.T) {
 		return nil, fmt.Errorf("permission denied")
 	}
 
-	ms, err := pingDevice("127.0.0.1", 5000)
+	ms, err := pingDevice("127.0.0.1", 2000)
 	if err != nil {
 		t.Skipf("exec ping fallback not available: %v", err)
 	}
@@ -288,6 +309,9 @@ func TestIcmpPingNonICMPUnavailableError(t *testing.T) {
 }
 
 func TestIcmpPingUDPFallback(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping real ICMP in short mode")
+	}
 	// Mock raw ICMP to fail, forcing UDP fallback path in icmpPing
 	origListen := icmpListenPacket
 	defer func() { icmpListenPacket = origListen }()
@@ -303,7 +327,7 @@ func TestIcmpPingUDPFallback(t *testing.T) {
 		return icmp.ListenPacket(network, address)
 	}
 
-	ms, err := icmpPing("127.0.0.1", 5000)
+	ms, err := icmpPing("127.0.0.1", 2000)
 	if err != nil {
 		t.Skipf("UDP ICMP not available: %v", err)
 	}

renovate.json

@@ -1,64 +0,0 @@
-{
-  "$schema": "https://docs.renovatebot.com/renovate-schema.json",
-  "extends": [
-    "config:recommended"
-  ],
-  "gitAuthor": "Renovate Bot <bot@renovateapp.com>",
-  "branchPrefix": "renovate/",
-  "semanticCommits": "enabled",
-  "dependencyDashboard": true,
-  "labels": [
-    "dependencies"
-  ],
-  "schedule": [
-    "at any time"
-  ],
-  "timezone": "UTC",
-  "prConcurrentLimit": 10,
-  "prCreation": "immediate",
-  "automerge": false,
-  "rangeStrategy": "bump",
-  "separateMajorMinor": true,
-  "separateMinorPatch": false,
-  "packageRules": [
-    {
-      "description": "Go module dependencies",
-      "matchManagers": [
-        "gomod"
-      ],
-      "matchUpdateTypes": [
-        "major",
-        "minor",
-        "patch"
-      ],
-      "groupName": null,
-      "commitMessagePrefix": "chore(deps):"
-    },
-    {
-      "description": "Docker dependencies",
-      "matchManagers": [
-        "dockerfile"
-      ],
-      "matchUpdateTypes": [
-        "major",
-        "minor",
-        "patch"
-      ],
-      "groupName": null,
-      "commitMessagePrefix": "chore(deps):"
-    },
-    {
-      "description": "GitHub Actions",
-      "matchManagers": [
-        "github-actions"
-      ],
-      "matchUpdateTypes": [
-        "major",
-        "minor",
-        "patch"
-      ],
-      "groupName": null,
-      "commitMessagePrefix": "chore(deps):"
-    }
-  ]
-}

ssh_test.go

@@ -7,7 +7,9 @@ import (
 	"crypto/rand"
 	"fmt"
 	"net"
+	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -318,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()
@@ -342,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}))
@@ -360,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()
@@ -382,6 +401,8 @@ func TestExecuteMikrotikBackupWithOutput(t *testing.T) {
 }
 
 func TestExecuteMikrotikBackupSessionError(t *testing.T) {
+	resetHostKeyStore(t)
+
 	// SSH server that accepts connection but rejects all channel requests
 	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	if err != nil {

update.go

@@ -85,7 +85,7 @@ func selfUpdate(downloadURL, expectedChecksum string) error {
 		return fmt.Errorf("create temp: %w", err)
 	}
 	tempPath := tempFile.Name()
-	defer func() { _ = os.Remove(tempPath) }() // cleanup on any failure
+	defer func() { _ = os.Remove(tempPath) }()
 
 	if _, err := tempFile.Write(body); err != nil {
 		_ = tempFile.Close()

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

websocket_test.go

@@ -712,6 +712,10 @@ func TestWSDialMissingAcceptHeader(t *testing.T) {
 }
 
 func TestWSDialMissingUpgradeHeader(t *testing.T) {
+	origTimeout := wsHandshakeTimeout
+	defer func() { wsHandshakeTimeout = origTimeout }()
+	wsHandshakeTimeout = time.Second
+
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
@@ -743,6 +747,10 @@ func TestWSDialMissingUpgradeHeader(t *testing.T) {
 }
 
 func TestWSDialRejectsNon101Status(t *testing.T) {
+	origTimeout := wsHandshakeTimeout
+	defer func() { wsHandshakeTimeout = origTimeout }()
+	wsHandshakeTimeout = time.Second
+
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
@@ -774,17 +782,44 @@ func TestWSDialRejectsNon101Status(t *testing.T) {
 }
 
 func TestWSDialDefaultPorts(t *testing.T) {
-	// Test that ws:// defaults to port 80 — will fail to connect but verifies URL parsing
+	origTimeout := wsHandshakeTimeout
+	origDial := netDial
+	origTLS := tlsDial
+	defer func() {
+		wsHandshakeTimeout = origTimeout
+		netDial = origDial
+		tlsDial = origTLS
+	}()
+	wsHandshakeTimeout = time.Second
+
+	var wsCalls, wssCalls int
+	netDial = func(network, addr string) (net.Conn, error) {
+		wsCalls++
+		if addr == "127.0.0.1:80" {
+			return nil, fmt.Errorf("connection refused (mock)")
+		}
+		return nil, fmt.Errorf("mock dial: %s", addr)
+	}
+	tlsDial = func(network, addr string) (net.Conn, error) {
+		wssCalls++
+		return nil, fmt.Errorf("tls dial refused (mock)")
+	}
+
 	_, err := WSDial("ws://127.0.0.1/path")
 	if err == nil {
 		t.Error("expected connection error (nothing on port 80)")
 	}
+	if wsCalls == 0 {
+		t.Error("expected netDial to be called for ws://")
+	}
 
-	// Test that wss:// defaults to port 443
 	_, err = WSDial("wss://127.0.0.1/path")
 	if err == nil {
 		t.Error("expected connection error (nothing on port 443)")
 	}
+	if wssCalls == 0 {
+		t.Error("expected tlsDial to be called for wss://")
+	}
 }
 
 func TestWSDialRealHTTPServer(t *testing.T) {