diff --git a/checks.go b/checks.go index 89f4980..f95a606 100644 --- a/checks.go +++ b/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, - ServerName: host, + MinVersion: tls.VersionTLS12, + RootCAs: rootCAs, + ServerName: host, } startTime := time.Now() diff --git a/checks_test.go b/checks_test.go index 01feaef..b1bd442 100644 --- a/checks_test.go +++ b/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 }) +} diff --git a/hostkeys.go b/hostkeys.go index 4470022..cd8f6ee 100644 --- a/hostkeys.go +++ b/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 } diff --git a/hostkeys_test.go b/hostkeys_test.go index 7ecc781..e7dd801 100644 --- a/hostkeys_test.go +++ b/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 { - t.Fatalf("should work without file: %v", err) + if err := s.verify("host:22", "fp"); err == nil { + t.Fatal("expected error when trusted key cannot be persisted") } } @@ -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 { - t.Fatalf("should succeed even if save fails: %v", err) + if err := s.verify("host:22", "fp"); err == nil { + t.Fatal("expected persistence error") + } + if _, ok := s.keys["host:22"]; ok { + t.Fatal("failed first-use trust should not remain cached in memory") } } diff --git a/ssh.go b/ssh.go index ecffe87..ea9b53d 100644 --- a/ssh.go +++ b/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 len(output) > 0 { - return string(output), nil + if trimmed := strings.TrimSpace(string(output)); trimmed != "" { + return "", fmt.Errorf("ssh command: %w: %s", err, trimmed) } return "", fmt.Errorf("ssh command: %w", err) } diff --git a/ssh_test.go b/ssh_test.go index 0d1f620..795ebbc 100644 --- a/ssh_test.go +++ b/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") - if err != nil { - t.Fatalf("expected success when output present despite exit code, got: %v", err) + _, err := executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass") + if err == nil { + t.Fatal("expected command failure when output is present with non-zero exit status") } - if config == "" { - t.Error("expected non-empty config") + if !strings.Contains(err.Error(), "# partial config") { + t.Fatalf("expected output to be included in error, got: %v", err) } } diff --git a/update.go b/update.go index e8a176b..00658b9 100644 --- a/update.go +++ b/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) } diff --git a/update_test.go b/update_test.go index 961da43..07439ab 100644 --- a/update_test.go +++ b/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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + httpDo = srv.Client().Do _ = selfUpdate(rewriteToHTTPS(srv.URL), checksum) @@ -270,9 +271,9 @@ func TestSelfUpdateTooLarge(t *testing.T) { })) defer srv.Close() - origGet := httpGet - defer func() { httpGet = origGet }() - httpGet = srv.Client().Get + origDo := httpDo + defer func() { httpDo = origDo }() + 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 { diff --git a/websocket.go b/websocket.go index 5eeab3e..51665eb 100644 --- a/websocket.go +++ b/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 { diff --git a/websocket_test.go b/websocket_test.go index 5679479..c2850b5 100644 --- a/websocket_test.go +++ b/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")