Merge branch 'main' into feature/branch_file_selector

Author: alex289

backend/internal/agent/agent.go

@@ -14,6 +14,7 @@ import (
 	"github.com/OrcaCD/orca-cd/internal/agent/docker"
 	messages "github.com/OrcaCD/orca-cd/internal/proto"
 	"github.com/OrcaCD/orca-cd/internal/version"
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
 	"google.golang.org/protobuf/proto"
@@ -24,6 +25,7 @@ type Config struct {
 	LogJSON   bool
 	HubUrl    string
 	AuthToken string
+	AgentID   string
 }
 
 func DefaultConfig() (Config, error) {
@@ -41,6 +43,11 @@ func DefaultConfig() (Config, error) {
 		return Config{}, errors.New("AUTH_TOKEN is required")
 	}
 
+	agentID, err := parseAgentID(authToken)
+	if err != nil {
+		return Config{}, fmt.Errorf("AUTH_TOKEN: %w", err)
+	}
+
 	logLevel, err := zerolog.ParseLevel(logLevelStr)
 	if err != nil || logLevelStr == "" {
 		logLevel = zerolog.InfoLevel
@@ -53,9 +60,23 @@ func DefaultConfig() (Config, error) {
 		LogJSON:   logJSON,
 		HubUrl:    hubUrl,
 		AuthToken: authToken,
+		AgentID:   agentID,
 	}, nil
 }
 
+func parseAgentID(authToken string) (string, error) {
+	p := jwt.NewParser()
+	token, _, err := p.ParseUnverified(authToken, &jwt.RegisteredClaims{})
+	if err != nil {
+		return "", fmt.Errorf("could not parse token to extract agent ID: %w", err)
+	}
+	claims, ok := token.Claims.(*jwt.RegisteredClaims)
+	if !ok || claims.Subject == "" {
+		return "", errors.New("token is missing subject claim")
+	}
+	return claims.Subject, nil
+}
+
 func newLogger(logJSON bool) zerolog.Logger {
 	if logJSON {
 		return zerolog.New(os.Stderr).With().Timestamp().Str("service", "agent").Logger()
@@ -120,6 +141,13 @@ func Run(cfg Config) error {
 		return nil
 	}
 
+	session, err := performHandshake(conn, cfg.AgentID)
+	if err != nil {
+		Log.Error().Err(err).Msg("handshake failed")
+		_ = conn.Close()
+		return fmt.Errorf("handshake: %w", err)
+	}
+
 	for {
 		_, data, readErr := conn.ReadMessage()
 		if readErr != nil {
@@ -136,13 +164,19 @@ func Run(cfg Config) error {
 				Log.Info().Msg("agent stopped")
 				return nil
 			}
+			session, err = performHandshake(conn, cfg.AgentID)
+			if err != nil {
+				Log.Error().Err(err).Msg("handshake failed on reconnect")
+				_ = conn.Close()
+				return fmt.Errorf("handshake on reconnect: %w", err)
+			}
 			continue
 		}
 		msg := &messages.ServerMessage{}
 		if err := proto.Unmarshal(data, msg); err != nil {
 			Log.Error().Err(err).Msg("unmarshal error")
 			continue
 		}
-		handleServerMessage(msg, conn)
+		handleServerMessage(msg, conn, session)
 	}
 }

backend/internal/agent/agent_config_test.go

@@ -1,14 +1,27 @@
 package agent
 
 import (
+	"context"
+	"encoding/base64"
 	"testing"
+	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
 )
 
+// testAgentToken is a minimal JWT-format string with a known subject.
+// The signature is intentionally fake; DefaultConfig only calls jwt.ParseUnverified,
+// which does not check the signature.
+var testAgentToken = func() string {
+	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"EdDSA","typ":"JWT"}`))
+	payload := base64.RawURLEncoding.EncodeToString([]byte(`{"sub":"test-agent-id"}`))
+	return header + "." + payload + ".fakesig"
+}()
+
 func TestDefaultConfig_Valid(t *testing.T) {
 	t.Setenv("HUB_URL", "https://hub.example.com")
-	t.Setenv("AUTH_TOKEN", "test-token")
+	t.Setenv("AUTH_TOKEN", testAgentToken)
 	t.Setenv("LOG_LEVEL", "debug")
 	t.Setenv("LOG_JSON", "true")
 
@@ -26,14 +39,17 @@ func TestDefaultConfig_Valid(t *testing.T) {
 	if cfg.HubUrl != "wss://hub.example.com/api/v1/ws" {
 		t.Errorf("HubUrl = %q, want %q", cfg.HubUrl, "wss://hub.example.com/api/v1/ws")
 	}
-	if cfg.AuthToken != "test-token" {
-		t.Errorf("AuthToken = %q, want %q", cfg.AuthToken, "test-token")
+	if cfg.AuthToken != testAgentToken {
+		t.Errorf("AuthToken = %q, want %q", cfg.AuthToken, testAgentToken)
+	}
+	if cfg.AgentID != "test-agent-id" {
+		t.Errorf("AgentID = %q, want %q", cfg.AgentID, "test-agent-id")
 	}
 }
 
 func TestDefaultConfig_Defaults(t *testing.T) {
 	t.Setenv("HUB_URL", "https://hub.example.com")
-	t.Setenv("AUTH_TOKEN", "test-token")
+	t.Setenv("AUTH_TOKEN", testAgentToken)
 	t.Setenv("LOG_LEVEL", "")
 	t.Setenv("LOG_JSON", "")
 
@@ -66,7 +82,7 @@ func TestDefaultConfig_LogLevels(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.input, func(t *testing.T) {
 			t.Setenv("HUB_URL", "https://hub.example.com")
-			t.Setenv("AUTH_TOKEN", "test-token")
+			t.Setenv("AUTH_TOKEN", testAgentToken)
 			t.Setenv("LOG_LEVEL", tt.input)
 
 			cfg, err := DefaultConfig()
@@ -95,7 +111,7 @@ func TestDefaultConfig_LogJSON(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.input, func(t *testing.T) {
 			t.Setenv("HUB_URL", "https://hub.example.com")
-			t.Setenv("AUTH_TOKEN", "test-token")
+			t.Setenv("AUTH_TOKEN", testAgentToken)
 			t.Setenv("LOG_JSON", tt.input)
 
 			cfg, err := DefaultConfig()
@@ -109,14 +125,50 @@ func TestDefaultConfig_LogJSON(t *testing.T) {
 	}
 }
 
+func TestParseAgentID_MissingSubject(t *testing.T) {
+	// JWT with empty subject field.
+	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"EdDSA","typ":"JWT"}`))
+	payload := base64.RawURLEncoding.EncodeToString([]byte(`{"sub":""}`))
+	token := header + "." + payload + ".fakesig"
+
+	_, err := parseAgentID(token)
+	if err == nil {
+		t.Fatal("expected error for token with empty subject")
+	}
+}
+
+func TestParseAgentID_InvalidToken(t *testing.T) {
+	_, err := parseAgentID("not.a.jwt")
+	if err == nil {
+		t.Fatal("expected error for malformed JWT")
+	}
+}
+
+func TestConnTracker_SetAndCancelled_CancelledCtx(t *testing.T) {
+	srv := newTestServer(t, func(serverConn *websocket.Conn) {
+		serverConn.SetReadDeadline(time.Now().Add(500 * time.Millisecond)) //nolint:errcheck,gosec
+		_, _, _ = serverConn.ReadMessage()
+	})
+
+	clientConn := dialServer(t, srv)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // pre-cancel so ctx.Err() != nil
+
+	var tracker connTracker
+	if cancelled := tracker.setAndCancelled(ctx, clientConn); !cancelled {
+		t.Error("expected setAndCancelled to return true for a pre-cancelled context")
+	}
+}
+
 func TestDefaultConfig_Errors(t *testing.T) {
 	tests := []struct {
 		name      string
 		hubURL    string
 		authToken string
 	}{
 		{"missing auth token", "https://hub.example.com", ""},
-		{"invalid hub url", "not-a-url", "test-token"},
+		{"invalid hub url", "not-a-url", testAgentToken},
 	}
 
 	for _, tt := range tests {

backend/internal/agent/websocket.go

@@ -2,38 +2,129 @@ package agent
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"strings"
 	"time"
 
 	messages "github.com/OrcaCD/orca-cd/internal/proto"
+	"github.com/OrcaCD/orca-cd/internal/shared/wscrypto"
 	"github.com/gorilla/websocket"
 	"google.golang.org/protobuf/proto"
 )
 
-func handleServerMessage(msg *messages.ServerMessage, conn *websocket.Conn) {
+const handshakeTimeout = 15 * time.Second
+
+func performHandshake(conn *websocket.Conn, agentID string) (*wscrypto.Session, error) {
+	if err := conn.SetReadDeadline(time.Now().Add(handshakeTimeout)); err != nil {
+		return nil, fmt.Errorf("set read deadline: %w", err)
+	}
+
+	_, data, err := conn.ReadMessage()
+	if err != nil {
+		return nil, fmt.Errorf("read KeyExchangeInit: %w", err)
+	}
+
+	serverMsg := &messages.ServerMessage{}
+	if err := proto.Unmarshal(data, serverMsg); err != nil {
+		return nil, fmt.Errorf("unmarshal KeyExchangeInit: %w", err)
+	}
+
+	init, ok := serverMsg.Payload.(*messages.ServerMessage_KeyExchangeInit)
+	if !ok {
+		return nil, fmt.Errorf("expected KeyExchangeInit, got %T", serverMsg.Payload)
+	}
+
+	mlkemCiphertext, agentX25519Pub, sessionKey, err := wscrypto.AgentHandshake(
+		init.KeyExchangeInit.MlkemEncapsulationKey,
+		init.KeyExchangeInit.X25519PublicKey,
+		agentID,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("agent handshake: %w", err)
+	}
+
+	resp := &messages.ClientMessage{
+		Payload: &messages.ClientMessage_KeyExchangeResponse{
+			KeyExchangeResponse: &messages.KeyExchangeResponse{
+				MlkemCiphertext:      mlkemCiphertext,
+				AgentX25519PublicKey: agentX25519Pub,
+			},
+		},
+	}
+	respData, err := proto.Marshal(resp)
+	if err != nil {
+		return nil, fmt.Errorf("marshal KeyExchangeResponse: %w", err)
+	}
+	if err := conn.WriteMessage(websocket.BinaryMessage, respData); err != nil {
+		return nil, fmt.Errorf("send KeyExchangeResponse: %w", err)
+	}
+
+	// Clear the handshake deadline — the hub's read loop will set its own deadline.
+	if err := conn.SetReadDeadline(time.Time{}); err != nil {
+		return nil, fmt.Errorf("clear read deadline: %w", err)
+	}
+
+	Log.Debug().Str("agent_id", agentID).Msg("WS handshake complete")
+	return wscrypto.NewSession(sessionKey)
+}
+
+func sendMessage(conn *websocket.Conn, session *wscrypto.Session, msg *messages.ClientMessage) error {
+	outMsg := msg
+	if !wscrypto.AllowedUnencrypted(msg) {
+		env, err := session.Encrypt(msg)
+		if err != nil {
+			return fmt.Errorf("encrypt: %w", err)
+		}
+		outMsg = &messages.ClientMessage{
+			Payload: &messages.ClientMessage_EncryptedPayload{
+				EncryptedPayload: env,
+			},
+		}
+	}
+	data, err := proto.Marshal(outMsg)
+	if err != nil {
+		return fmt.Errorf("marshal: %w", err)
+	}
+	return conn.WriteMessage(websocket.BinaryMessage, data)
+}
+
+func handleServerMessage(msg *messages.ServerMessage, conn *websocket.Conn, session *wscrypto.Session) {
+	_, isEncrypted := msg.Payload.(*messages.ServerMessage_EncryptedPayload)
+	if !isEncrypted && !wscrypto.AllowedUnencrypted(msg) {
+		Log.Warn().Msgf("dropping unencrypted message of type %T", msg.Payload)
+		return
+	}
+
+	// Unwrap at most one layer of encryption to prevent a malicious server from
+	// crafting a chain of nested EncryptedPayloads that causes unbounded recursion.
+	if p, ok := msg.Payload.(*messages.ServerMessage_EncryptedPayload); ok {
+		inner := &messages.ServerMessage{}
+		if err := session.Decrypt(p.EncryptedPayload, inner); err != nil {
+			Log.Error().Err(err).Msg("failed to decrypt message")
+			return
+		}
+		if _, stillEncrypted := inner.Payload.(*messages.ServerMessage_EncryptedPayload); stillEncrypted {
+			Log.Warn().Msg("dropping doubly-encrypted message")
+			return
+		}
+		msg = inner
+	}
+
 	switch p := msg.Payload.(type) {
 	case *messages.ServerMessage_Ping:
 		latency := time.Now().UnixMilli() - p.Ping.Timestamp
 		Log.Debug().Msgf("ping received, latency: %dms", latency)
 
-		// Respond with Pong
 		pong := &messages.ClientMessage{
 			Payload: &messages.ClientMessage_Pong{
 				Pong: &messages.PongResponse{
 					Timestamp: time.Now().UnixMilli(),
 				},
 			},
 		}
-		data, err := proto.Marshal(pong)
-
-		if err != nil {
-			Log.Error().Err(err).Msg("failed to marshal Pong response")
-			return
-		}
-		if err := conn.WriteMessage(websocket.BinaryMessage, data); err != nil {
+		if err := sendMessage(conn, session, pong); err != nil {
 			Log.Error().Err(err).Msg("failed to send Pong response")
-			return
 		}
 	default:
 		Log.Warn().Msg("unknown message type received")