Merge commit from fork

agent/client_test.go

@@ -4,6 +4,16 @@ package agent
 
 import (
 	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"math/big"
+	"net"
+	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -16,6 +26,7 @@ import (
 	"github.com/henrygd/beszel/internal/common"
 
 	"github.com/fxamacker/cbor/v2"
+	"github.com/lxzan/gws"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
@@ -165,6 +176,155 @@ func TestWebSocketClient_GetOptions(t *testing.T) {
 	}
 }
 
+func TestWebSocketClient_TLSVerification(t *testing.T) {
+	agent := createTestAgent(t)
+	serverCert, serverCertPEM := newSelfSignedServerCertificate(t)
+	upgrader := gws.NewUpgrader(&gws.BuiltinEventHandler{}, nil)
+	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r)
+		if err == nil {
+			go conn.ReadLoop()
+		}
+	}))
+	server.TLS = &tls.Config{Certificates: []tls.Certificate{serverCert}}
+	server.StartTLS()
+	t.Cleanup(server.Close)
+
+	caCertFile := filepath.Join(t.TempDir(), "hub-ca.crt")
+	require.NoError(t, os.WriteFile(caCertFile, serverCertPEM, 0600))
+
+	newClient := func(t *testing.T, caCertFile string) *WebSocketClient {
+		t.Helper()
+		t.Setenv("BESZEL_AGENT_HUB_URL", server.URL)
+		t.Setenv("BESZEL_AGENT_TOKEN", "test-token")
+		t.Setenv("BESZEL_AGENT_CA_CERT_FILE", caCertFile)
+		client, err := newWebSocketClient(agent)
+		require.NoError(t, err)
+		return client
+	}
+
+	t.Run("system roots are used by default", func(t *testing.T) {
+		client := newClient(t, "")
+		assert.Nil(t, client.getOptions().TlsConfig)
+		_, _, err := gws.NewClient(&gws.BuiltinEventHandler{}, client.getOptions())
+		require.Error(t, err)
+	})
+
+	t.Run("custom CA trusts self-signed certificate", func(t *testing.T) {
+		systemRoots, err := x509.SystemCertPool()
+		require.NoError(t, err)
+		client := newClient(t, caCertFile)
+		assert.Greater(t, len(client.getOptions().TlsConfig.RootCAs.Subjects()), len(systemRoots.Subjects()))
+		conn, _, err := gws.NewClient(&gws.BuiltinEventHandler{}, client.getOptions())
+		require.NoError(t, err)
+		require.NoError(t, conn.NetConn().Close())
+	})
+
+	t.Run("custom CA does not bypass hostname verification", func(t *testing.T) {
+		client := newClient(t, caCertFile)
+		client.getOptions().TlsConfig.ServerName = "wrong.example.com"
+		_, _, err := gws.NewClient(&gws.BuiltinEventHandler{}, client.getOptions())
+		require.Error(t, err)
+	})
+}
+
+func TestWebSocketClient_NonTLSConnection(t *testing.T) {
+	agent := createTestAgent(t)
+	upgrader := gws.NewUpgrader(&gws.BuiltinEventHandler{}, nil)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r)
+		if err == nil {
+			go conn.ReadLoop()
+		}
+	}))
+	t.Cleanup(server.Close)
+
+	t.Setenv("BESZEL_AGENT_HUB_URL", server.URL)
+	t.Setenv("BESZEL_AGENT_TOKEN", "test-token")
+	t.Setenv("BESZEL_AGENT_CA_CERT_FILE", "")
+	client, err := newWebSocketClient(agent)
+	require.NoError(t, err)
+	assert.Nil(t, client.getOptions().TlsConfig)
+
+	conn, _, err := gws.NewClient(&gws.BuiltinEventHandler{}, client.getOptions())
+	require.NoError(t, err)
+	require.NoError(t, conn.NetConn().Close())
+}
+
+func TestGetTLSConfigErrors(t *testing.T) {
+	tempDir := t.TempDir()
+	testCases := []struct {
+		name       string
+		path       string
+		contents   []byte
+		errorMatch string
+	}{
+		{
+			name:       "missing file",
+			path:       filepath.Join(tempDir, "missing.pem"),
+			errorMatch: "read CA_CERT_FILE",
+		},
+		{
+			name:       "unreadable path",
+			path:       tempDir,
+			errorMatch: "read CA_CERT_FILE",
+		},
+		{
+			name:       "empty file",
+			path:       filepath.Join(tempDir, "empty.pem"),
+			contents:   []byte{},
+			errorMatch: "does not contain any valid PEM certificates",
+		},
+		{
+			name:       "malformed file",
+			path:       filepath.Join(tempDir, "malformed.pem"),
+			contents:   []byte("not a PEM certificate"),
+			errorMatch: "does not contain any valid PEM certificates",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.contents != nil {
+				require.NoError(t, os.WriteFile(tc.path, tc.contents, 0600))
+			}
+			t.Setenv("BESZEL_AGENT_CA_CERT_FILE", tc.path)
+
+			tlsConfig, err := getTLSConfig()
+			require.Error(t, err)
+			assert.Nil(t, tlsConfig)
+			assert.Contains(t, err.Error(), tc.errorMatch)
+			assert.Contains(t, err.Error(), tc.path)
+		})
+	}
+}
+
+func newSelfSignedServerCertificate(t *testing.T) (tls.Certificate, []byte) {
+	t.Helper()
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+
+	template := &x509.Certificate{
+		SerialNumber:          big.NewInt(1),
+		Subject:               pkix.Name{CommonName: "127.0.0.1"},
+		NotBefore:             time.Now().Add(-time.Hour),
+		NotAfter:              time.Now().Add(time.Hour),
+		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
+		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageCertSign,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+	}
+	certDER, err := x509.CreateCertificate(rand.Reader, template, template, &privateKey.PublicKey, privateKey)
+	require.NoError(t, err)
+
+	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
+	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)})
+	certificate, err := tls.X509KeyPair(certPEM, keyPEM)
+	require.NoError(t, err)
+	return certificate, certPEM
+}
+
 // TestWebSocketClient_VerifySignature tests signature verification
 func TestWebSocketClient_VerifySignature(t *testing.T) {
 	agent := createTestAgent(t)