fix: harden trust and transport validation

checks.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"crypto/tls"
+	"crypto/x509"
 	"fmt"
 	"io"
 	"net"
@@ -27,6 +28,7 @@ var (
 		MaxIdleConnsPerHost: 10,
 		IdleConnTimeout:     90 * time.Second,
 	}
+	sslRootCAs = x509.SystemCertPool
 )
 
 // ExecuteCheck runs a service check and returns the result.
@@ -325,9 +327,14 @@ func executeSSLCheck(ctx context.Context, config *pb.SslCheckConfig, timeoutMs u
 	address := net.JoinHostPort(host, strconv.Itoa(int(port)))
 
 	dialer := &net.Dialer{Timeout: timeout}
+	rootCAs, err := sslRootCAs()
+	if err != nil {
+		return 3, fmt.Sprintf("Failed to load system root CAs: %v", err), 0
+	}
 	tlsConfig := &tls.Config{
-		InsecureSkipVerify: true,
+		MinVersion: tls.VersionTLS12,
-		ServerName:         host,
+		RootCAs:    rootCAs,
+		ServerName: host,
 	}
 
 	startTime := time.Now()

checks_test.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"context"
 	"crypto/tls"
+	"crypto/x509"
 	"fmt"
 	"net"
 	"net/http"
@@ -1378,6 +1379,7 @@ func TestExecuteCheck_SSLRouting(t *testing.T) {
 		w.WriteHeader(200)
 	}))
 	defer srv.Close()
+	withTestSSLRootCA(t, srv.Certificate())
 
 	// Parse the port from the test server URL
 	_, portStr, _ := net.SplitHostPort(strings.TrimPrefix(strings.TrimPrefix(srv.URL, "https://"), "http://"))
@@ -1405,6 +1407,7 @@ func TestSSLCheck_ValidCert(t *testing.T) {
 		w.WriteHeader(200)
 	}))
 	defer srv.Close()
+	withTestSSLRootCA(t, srv.Certificate())
 
 	_, portStr, _ := net.SplitHostPort(strings.TrimPrefix(srv.URL, "https://"))
 
@@ -1431,6 +1434,7 @@ func TestSSLCheck_WarningThreshold(t *testing.T) {
 		w.WriteHeader(200)
 	}))
 	defer srv.Close()
+	withTestSSLRootCA(t, srv.Certificate())
 
 	_, portStr, _ := net.SplitHostPort(strings.TrimPrefix(srv.URL, "https://"))
 
@@ -1503,6 +1507,7 @@ func TestSSLCheck_DefaultWarningDays30(t *testing.T) {
 		w.WriteHeader(200)
 	}))
 	defer srv.Close()
+	withTestSSLRootCA(t, srv.Certificate())
 
 	_, portStr, _ := net.SplitHostPort(strings.TrimPrefix(srv.URL, "https://"))
 
@@ -1519,6 +1524,28 @@ func TestSSLCheck_DefaultWarningDays30(t *testing.T) {
 	}
 }
 
+func TestSSLCheck_UntrustedCertFails(t *testing.T) {
+	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	_, portStr, _ := net.SplitHostPort(strings.TrimPrefix(srv.URL, "https://"))
+
+	status, output, _ := executeSSLCheck(context.Background(), &pb.SslCheckConfig{
+		Host:        "127.0.0.1",
+		Port:        parsePort(portStr),
+		WarningDays: 7,
+	}, 5000)
+
+	if status != 3 {
+		t.Fatalf("expected status 3 for untrusted certificate, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Connection failed") {
+		t.Fatalf("expected connection failure output, got %s", output)
+	}
+}
+
 // ---------------------------------------------------------------------------
 // Helper functions
 // ---------------------------------------------------------------------------
@@ -1535,3 +1562,14 @@ func parsePort(s string) uint32 {
 	_, _ = fmt.Sscanf(s, "%d", &port)
 	return port
 }
+
+func withTestSSLRootCA(t *testing.T, cert *x509.Certificate) {
+	t.Helper()
+	orig := sslRootCAs
+	sslRootCAs = func() (*x509.CertPool, error) {
+		pool := x509.NewCertPool()
+		pool.AddCert(cert)
+		return pool, nil
+	}
+	t.Cleanup(func() { sslRootCAs = orig })
+}

hostkeys.go

@@ -65,7 +65,8 @@ func (s *hostKeyStore) verify(host, fingerprint string) error {
 		slog.Warn("TOFU: first connection, trusting host key", "host", host, "fingerprint", fingerprint)
 		s.keys[host] = fingerprint
 		if err := s.save(); err != nil {
-			slog.Error("failed to save host keys", "error", err)
+			delete(s.keys, host)
+			return fmt.Errorf("failed to persist trusted host key for %s: %w", host, err)
 		}
 		return nil
 	}

hostkeys_test.go

@@ -52,9 +52,8 @@ func TestHostKeyStorePersistence(t *testing.T) {
 
 func TestHostKeyStoreMissingFile(t *testing.T) {
 	s := newHostKeyStore("/nonexistent/path/known_hosts.json")
-	// Should work in memory even if file can't be read
+	if err := s.verify("host:22", "fp"); err == nil {
-	if err := s.verify("host:22", "fp"); err != nil {
+		t.Fatal("expected error when trusted key cannot be persisted")
-		t.Fatalf("should work without file: %v", err)
 	}
 }
 
@@ -109,9 +108,11 @@ func TestSSHHostKeyCallback(t *testing.T) {
 func TestHostKeyStoreSaveError(t *testing.T) {
 	// Use a path in a non-existent directory
 	s := newHostKeyStore("/nonexistent/dir/hosts.json")
-	// verify should not error even if save fails (it logs instead)
+	if err := s.verify("host:22", "fp"); err == nil {
-	if err := s.verify("host:22", "fp"); err != nil {
+		t.Fatal("expected persistence error")
-		t.Fatalf("should succeed even if save fails: %v", err)
+	}
+	if _, ok := s.keys["host:22"]; ok {
+		t.Fatal("failed first-use trust should not remain cached in memory")
 	}
 }
 

ssh.go

@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"net"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/towerops-app/towerops-agent/pb"
@@ -42,9 +43,8 @@ func executeMikrotikBackup(ip string, port uint16, username, password string) (s
 
 	output, err := session.CombinedOutput("/export compact")
 	if err != nil {
-		// MikroTik SSH doesn't use exit codes the same way - check if we got output
+		if trimmed := strings.TrimSpace(string(output)); trimmed != "" {
-		if len(output) > 0 {
+			return "", fmt.Errorf("ssh command: %w: %s", err, trimmed)
-			return string(output), nil
 		}
 		return "", fmt.Errorf("ssh command: %w", err)
 	}

ssh_test.go

@@ -360,7 +360,6 @@ func TestExecuteMikrotikBackupCommandError(t *testing.T) {
 }
 
 func TestExecuteMikrotikBackupWithOutput(t *testing.T) {
-	// MikroTik SSH returns output even with non-zero exit code
 	addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) {
 		_, _ = ch.Write([]byte("# partial config\n"))
 		_ = ch.CloseWrite()
@@ -373,12 +372,12 @@ func TestExecuteMikrotikBackupWithOutput(t *testing.T) {
 	var portNum uint16
 	_, _ = fmt.Sscanf(port, "%d", &portNum)
 
-	config, err := executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass")
+	_, err := executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass")
-	if err != nil {
+	if err == nil {
-		t.Fatalf("expected success when output present despite exit code, got: %v", err)
+		t.Fatal("expected command failure when output is present with non-zero exit status")
 	}
-	if config == "" {
+	if !strings.Contains(err.Error(), "# partial config") {
-		t.Error("expected non-empty config")
+		t.Fatalf("expected output to be included in error, got: %v", err)
 	}
 }
 

update.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"crypto/sha256"
 	"crypto/subtle"
 	"fmt"
@@ -12,14 +13,18 @@ import (
 	"path/filepath"
 	"strings"
 	"syscall"
+	"time"
 )
 
 var osExecutable = os.Executable
 var osCreateTemp = os.CreateTemp
 var osRename = os.Rename
-var httpGet = http.Get
+var httpDo = func(req *http.Request) (*http.Response, error) {
+	return http.DefaultClient.Do(req)
+}
 var syscallExec = syscall.Exec
 var maxUpdateSize int64 = 100 << 20 // 100 MB
+var selfUpdateTimeout = 30 * time.Second
 
 // selfUpdate downloads a new binary, verifies its checksum, replaces the current binary, and re-execs.
 func selfUpdate(downloadURL, expectedChecksum string) error {
@@ -36,7 +41,14 @@ func selfUpdate(downloadURL, expectedChecksum string) error {
 
 	slog.Info("downloading update", "url", sanitizeURL(downloadURL))
 
-	resp, err := httpGet(downloadURL)
+	reqCtx, cancel := context.WithTimeout(context.Background(), selfUpdateTimeout)
+	defer cancel()
+	req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, downloadURL, nil)
+	if err != nil {
+		return fmt.Errorf("build request: %w", err)
+	}
+
+	resp, err := httpDo(req)
 	if err != nil {
 		return fmt.Errorf("download: %w", err)
 	}

update_test.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestSelfUpdateRejectsHTTP(t *testing.T) {
@@ -44,9 +45,9 @@ func TestSelfUpdateChecksumMismatch(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL), "0000000000000000000000000000000000000000000000000000000000000000")
 	if err == nil {
@@ -60,9 +61,9 @@ func TestSelfUpdate404(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL)+"/missing", "abc123")
 	if err == nil {
@@ -86,9 +87,9 @@ func TestSelfUpdateReadBodyError(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL), "abc123")
 	// This may or may not error depending on io.ReadAll behavior with truncated body
@@ -111,9 +112,9 @@ func TestSelfUpdateOsExecutableError(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
 	if err == nil {
@@ -144,9 +145,9 @@ func TestSelfUpdateCreateTempError(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
 	if err == nil {
@@ -178,9 +179,9 @@ func TestSelfUpdateRenameError(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
 	if err == nil {
@@ -200,9 +201,9 @@ func TestSelfUpdateChecksumMatch(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	// This will fail at the rename step (writing to os.Executable path),
 	// but the checksum verification should pass
@@ -241,9 +242,9 @@ func TestSelfUpdateFilePermissions(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	_ = selfUpdate(rewriteToHTTPS(srv.URL), checksum)
 
@@ -270,9 +271,9 @@ func TestSelfUpdateTooLarge(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
-	defer func() { httpGet = origGet }()
+	defer func() { httpDo = origDo }()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
 	if err == nil {
@@ -341,17 +342,17 @@ func TestSelfUpdateFullHappyPath(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
 	origExe := osExecutable
 	origRename := osRename
 	origExec := syscallExec
 	defer func() {
-		httpGet = origGet
+		httpDo = origDo
 		osExecutable = origExe
 		osRename = origRename
 		syscallExec = origExec
 	}()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	dir := t.TempDir()
 	exePath := filepath.Join(dir, "test-agent")
@@ -378,15 +379,15 @@ func TestSelfUpdateWriteError(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	origGet := httpGet
+	origDo := httpDo
 	origExe := osExecutable
 	origCreate := osCreateTemp
 	defer func() {
-		httpGet = origGet
+		httpDo = origDo
 		osExecutable = origExe
 		osCreateTemp = origCreate
 	}()
-	httpGet = srv.Client().Get
+	httpDo = srv.Client().Do
 
 	dir := t.TempDir()
 	osExecutable = func() (string, error) { return filepath.Join(dir, "agent"), nil }
@@ -408,6 +409,29 @@ func TestSelfUpdateWriteError(t *testing.T) {
 	}
 }
 
+func TestSelfUpdateDownloadTimeout(t *testing.T) {
+	origDo := httpDo
+	origTimeout := selfUpdateTimeout
+	defer func() {
+		httpDo = origDo
+		selfUpdateTimeout = origTimeout
+	}()
+
+	httpDo = func(req *http.Request) (*http.Response, error) {
+		<-req.Context().Done()
+		return nil, req.Context().Err()
+	}
+	selfUpdateTimeout = 10 * time.Millisecond
+
+	err := selfUpdate("https://example.com/agent", strings.Repeat("0", 64))
+	if err == nil {
+		t.Fatal("expected timeout error")
+	}
+	if !strings.Contains(err.Error(), "context deadline exceeded") {
+		t.Fatalf("expected context deadline exceeded, got: %v", err)
+	}
+}
+
 // rewriteToHTTPS converts an httptest TLS server URL to use the https scheme.
 // httptest.NewTLSServer returns URLs with https:// already, but this ensures consistency.
 func rewriteToHTTPS(rawURL string) string {

websocket.go

@@ -135,7 +135,8 @@ func wsConnect(u *url.URL, host, network string) (*WSConn, error) {
 		_ = conn.Close()
 		return nil, fmt.Errorf("read handshake: %w", err)
 	}
-	if !strings.Contains(statusLine, "101") {
+	statusParts := strings.SplitN(strings.TrimSpace(statusLine), " ", 3)
+	if len(statusParts) < 2 || statusParts[1] != "101" {
 		_ = conn.Close()
 		return nil, fmt.Errorf("handshake failed: %s", strings.TrimSpace(statusLine))
 	}
@@ -143,6 +144,8 @@ func wsConnect(u *url.URL, host, network string) (*WSConn, error) {
 	// Read remaining headers, verify Sec-WebSocket-Accept per RFC 6455
 	expectedAccept := computeAcceptKey(key)
 	acceptFound := false
+	upgradeFound := false
+	connectionFound := false
 	for {
 		line, err := br.ReadString('\n')
 		if err != nil {
@@ -160,12 +163,32 @@ func wsConnect(u *url.URL, host, network string) (*WSConn, error) {
 				return nil, fmt.Errorf("invalid accept key: got %q, want %q", actual, expectedAccept)
 			}
 			acceptFound = true
+			continue
+		}
+		if strings.HasPrefix(strings.ToLower(line), "upgrade: ") {
+			if strings.EqualFold(strings.TrimSpace(line[len("Upgrade: "):]), "websocket") {
+				upgradeFound = true
+			}
+			continue
+		}
+		if strings.HasPrefix(strings.ToLower(line), "connection: ") {
+			if headerHasToken(line[len("Connection: "):], "upgrade") {
+				connectionFound = true
+			}
 		}
 	}
 	if !acceptFound {
 		_ = conn.Close()
 		return nil, fmt.Errorf("missing Sec-WebSocket-Accept header")
 	}
+	if !upgradeFound {
+		_ = conn.Close()
+		return nil, fmt.Errorf("missing or invalid Upgrade header")
+	}
+	if !connectionFound {
+		_ = conn.Close()
+		return nil, fmt.Errorf("missing or invalid Connection header")
+	}
 
 	// Clear handshake deadline — normal operation uses no deadline
 	if err := conn.SetDeadline(time.Time{}); err != nil {
@@ -177,6 +200,15 @@ func wsConnect(u *url.URL, host, network string) (*WSConn, error) {
 	return &WSConn{conn: conn, reader: br}, nil
 }
 
+func headerHasToken(value, token string) bool {
+	for _, part := range strings.Split(value, ",") {
+		if strings.EqualFold(strings.TrimSpace(part), token) {
+			return true
+		}
+	}
+	return false
+}
+
 // ReadMessage reads the next text or binary message, handling control frames internally.
 func (ws *WSConn) ReadMessage() ([]byte, int, error) {
 	for {

websocket_test.go

@@ -711,6 +711,68 @@ func TestWSDialMissingAcceptHeader(t *testing.T) {
 	}
 }
 
+func TestWSDialMissingUpgradeHeader(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		key := extractHeader(string(buf[:n]), "Sec-WebSocket-Key")
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+	}()
+
+	_, err = WSDial("ws://" + ln.Addr().String() + "/socket")
+	if err == nil {
+		t.Fatal("expected missing Upgrade header error")
+	}
+	if !strings.Contains(err.Error(), "missing or invalid Upgrade header") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
+func TestWSDialRejectsNon101Status(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		key := extractHeader(string(buf[:n]), "Sec-WebSocket-Key")
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 1010 Not Switching\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+	}()
+
+	_, err = WSDial("ws://" + ln.Addr().String() + "/socket")
+	if err == nil {
+		t.Fatal("expected non-101 status to fail")
+	}
+	if !strings.Contains(err.Error(), "handshake failed") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
 func TestWSDialDefaultPorts(t *testing.T) {
 	// Test that ws:// defaults to port 80 — will fail to connect but verifies URL parsing
 	_, err := WSDial("ws://127.0.0.1/path")