graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: harden trust and transport validation

Browse source
This commit is contained in:
Graham McIntire 2026-03-26 12:21:39 -05:00
parent 724870d29c
commit e99f681c87
No known key found for this signature in database
10 changed files with 230 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
checks.go
View file

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

38
checks_test.go
View file

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

3
hostkeys.go
View file

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

13
hostkeys_test.go
View file

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

6
ssh.go
View file

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

11
ssh_test.go
View file

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

16
update.go
View file

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

90
update_test.go
View file

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

34
websocket.go
View file

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

62
websocket_test.go
View file

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