A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
10 kB · 309 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310package controller
import ( "context" "net/http" "net/http/httptest" "strings" "testing" "time"
"connectrpc.com/connect" agentv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1" controllerv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/controller/v1" "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/controller/v1/controllerv1connect" "forgejo.alexma.top/alexma233/composia/internal/core/config" "forgejo.alexma.top/alexma233/composia/internal/platform/rpcutil" "forgejo.alexma.top/alexma233/composia/internal/platform/store" "github.com/gorilla/websocket")
func TestExecSessionDeliverAndCloseAreConcurrentSafe(t *testing.T) { t.Parallel() for range 100 { manager := newExecTunnelManager() session := &execSession{id: "session", nodeID: "main", incoming: make(chan *agentv1.OpenExecTunnelRequest, 1)} manager.sessions[session.id] = session done := make(chan struct{}, 2) go func() { manager.deliverFromAgent(&agentv1.OpenExecTunnelRequest{SessionId: session.id, Kind: execKindOutput}) done <- struct{}{} }() go func() { manager.closeSession(session.id); done <- struct{}{} }() <-done <-done }}
func TestContainerExecRequiresWebOriginHeader(t *testing.T) { t.Parallel()
db := openControllerTestDB(t) defer func() { _ = db.Close() }()
ctx := context.Background() if err := db.SyncConfiguredNodes(ctx, []string{"main"}); err != nil { t.Fatalf("sync configured nodes: %v", err) } if err := db.RecordHeartbeat(ctx, controllerTestHeartbeat("main")); err != nil { t.Fatalf("record heartbeat: %v", err) }
execManager := newExecTunnelManager() execManager.registerTunnel("main")
interceptor := rpcutil.NewServerBearerAuthInterceptor(func(token string) (string, error) { if token != testAccessToken { return "", assertError("unexpected token") } return "web-admin", nil })
path, handler := controllerv1connect.NewDockerCommandServiceHandler( &dockerCommandServer{ db: db, cfg: &config.ControllerConfig{Nodes: []config.NodeConfig{{ID: "main"}}}, execManager: execManager, }, connect.WithInterceptors(interceptor), ) mux := http.NewServeMux() mux.Handle(path, handler) httpServer := httptest.NewServer(mux) defer httpServer.Close()
client := controllerv1connect.NewDockerCommandServiceClient( httpServer.Client(), httpServer.URL, connect.WithInterceptors(rpcutil.NewStaticBearerAuthInterceptor(testAccessToken)), )
_, err := client.OpenContainerExec(ctx, connect.NewRequest(&controllerv1.OpenContainerExecRequest{NodeId: "main", ContainerId: "ctr"})) if err == nil { t.Fatal("expected missing origin header to fail") } if connect.CodeOf(err) != connect.CodeInvalidArgument { t.Fatalf("expected invalid argument, got %v", err) } if !strings.Contains(err.Error(), execWebOriginHeader) { t.Fatalf("expected error to mention %s, got %v", execWebOriginHeader, err) }}
func TestContainerExecWebsocketRequiresAllowedOriginAndOneTimeToken(t *testing.T) { t.Parallel()
db := openControllerTestDB(t) defer func() { _ = db.Close() }()
ctx := context.Background() if err := db.SyncConfiguredNodes(ctx, []string{"main"}); err != nil { t.Fatalf("sync configured nodes: %v", err) } if err := db.RecordHeartbeat(ctx, controllerTestHeartbeat("main")); err != nil { t.Fatalf("record heartbeat: %v", err) }
execManager := newExecTunnelManager() execManager.registerTunnel("main")
interceptor := rpcutil.NewServerBearerAuthInterceptor(func(token string) (string, error) { if token != testAccessToken { return "", assertError("unexpected token") } return "web-admin", nil })
path, handler := controllerv1connect.NewDockerCommandServiceHandler( &dockerCommandServer{ db: db, cfg: &config.ControllerConfig{Nodes: []config.NodeConfig{{ID: "main"}}}, execManager: execManager, }, connect.WithInterceptors(interceptor), ) mux := http.NewServeMux() mux.Handle(path, handler) mux.HandleFunc(rpcutil.ControllerExecWSPath, execManager.handleWebsocket) httpServer := httptest.NewServer(mux) defer httpServer.Close()
client := controllerv1connect.NewDockerCommandServiceClient( httpServer.Client(), httpServer.URL, connect.WithInterceptors(rpcutil.NewStaticBearerAuthInterceptor(testAccessToken)), )
req := connect.NewRequest(&controllerv1.OpenContainerExecRequest{NodeId: "main", ContainerId: "ctr"}) req.Header().Set(execWebOriginHeader, "https://web.example.test") response, err := client.OpenContainerExec(ctx, req) if err != nil { t.Fatalf("open exec session: %v", err) } if strings.Contains(response.Msg.GetWebsocketPath(), response.Msg.GetSessionId()) { t.Fatalf("websocket path must not expose session id: %q", response.Msg.GetWebsocketPath()) } if !strings.HasPrefix(response.Msg.GetWebsocketPath(), rpcutil.ControllerExecWSPath) { t.Fatalf("websocket path = %q, want prefix %q", response.Msg.GetWebsocketPath(), rpcutil.ControllerExecWSPath) }
wsURL := websocketURL(httpServer.URL, response.Msg.GetWebsocketPath()) biddenHeader := http.Header{} biddenHeader.Set("Origin", "https://evil.example.test") if conn, dialResponse, err := websocket.DefaultDialer.Dial(wsURL, biddenHeader); err == nil { _ = conn.Close() t.Fatal("expected websocket origin mismatch to fail") } else if dialResponse != nil && dialResponse.Body != nil { _ = dialResponse.Body.Close() }
allowedHeader := http.Header{} allowedHeader.Set("Origin", "https://web.example.test") conn, dialResponse, err := websocket.DefaultDialer.Dial(wsURL, allowedHeader) if dialResponse != nil && dialResponse.Body != nil { defer func() { _ = dialResponse.Body.Close() }() } if err != nil { t.Fatalf("dial websocket with allowed origin: %v", err) } defer func() { _ = conn.Close() }()
var ready execWSEvent if err := conn.ReadJSON(&ready); err != nil { t.Fatalf("read ready event: %v", err) } if ready.Type != execKindReady { t.Fatalf("expected ready event, got %+v", ready) } if ready.Session != response.Msg.GetSessionId() { t.Fatalf("expected ready session %q, got %q", response.Msg.GetSessionId(), ready.Session) }
if extraConn, dialResponse, err := websocket.DefaultDialer.Dial(wsURL, allowedHeader); err == nil { _ = extraConn.Close() t.Fatal("expected attach token to be single-use") } else if dialResponse != nil && dialResponse.Body != nil { _ = dialResponse.Body.Close() } _ = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done"))}
func TestExecTunnelManagerExpiresUnusedAttachToken(t *testing.T) { t.Parallel()
manager := newExecTunnelManager() manager.registerTunnel("main")
session, err := manager.openSession("main", "ctr", []string{defaultExecShell}, 24, 80, "https://web.example.test", "web-admin") if err != nil { t.Fatalf("open session: %v", err) }
manager.mu.Lock() session.attachExpiresAt = time.Now().UTC().Add(-time.Second) manager.mu.Unlock()
if _, err := manager.reserveSessionAttach(session.attachToken); err == nil { t.Fatal("expected expired attach token to fail") } if got := manager.lookupSession(session.id); got != nil { t.Fatalf("expected expired session to be removed, got %+v", got) }}
func TestExecTunnelManagerActivelyExpiresUnusedAttachToken(t *testing.T) { t.Parallel()
manager := newExecTunnelManager() manager.registerTunnel("main") ctx, cancel := context.WithCancel(context.Background()) defer cancel() go manager.runSessionSweeper(ctx, 10*time.Millisecond)
session, err := manager.openSession("main", "ctr", []string{defaultExecShell}, 24, 80, "https://web.example.test", "web-admin") if err != nil { t.Fatalf("open session: %v", err) } manager.mu.Lock() session.attachExpiresAt = time.Now().UTC().Add(-time.Second) manager.mu.Unlock()
deadline := time.Now().Add(time.Second) for time.Now().Before(deadline) { manager.mu.Lock() _, ok := manager.sessions[session.id] manager.mu.Unlock() if !ok { return } time.Sleep(time.Millisecond) } t.Fatal("expected sweeper to remove expired session")}
func TestExecTunnelManagerDoesNotBlockWhenAgentQueueIsFull(t *testing.T) { t.Parallel()
manager := newExecTunnelManager() tunnel := manager.registerTunnel("main") for range cap(tunnel.sendCh) { tunnel.sendCh <- &agentv1.OpenExecTunnelResponse{} }
done := make(chan error, 1) go func() { _, err := manager.openSession("main", "ctr", []string{defaultExecShell}, 24, 80, "https://web.example.test", "web-admin") done <- err }()
select { case err := <-done: if err == nil { t.Fatal("expected full tunnel queue to fail") } case <-time.After(time.Second): t.Fatal("openSession blocked on full tunnel queue") }}
func TestExecTunnelManagerOldTunnelUnregisterDoesNotCloseNewSessions(t *testing.T) { t.Parallel()
manager := newExecTunnelManager() oldTunnel := manager.registerTunnel("main") newTunnel := manager.registerTunnel("main") select { case _, ok := <-oldTunnel.sendCh: if ok { t.Fatal("expected old tunnel send channel to be closed") } default: t.Fatal("expected old tunnel send channel to be closed") }
session, err := manager.openSession("main", "ctr", []string{defaultExecShell}, 24, 80, "https://web.example.test", "web-admin") if err != nil { t.Fatalf("open session on new tunnel: %v", err) } select { case message := <-newTunnel.sendCh: if message.GetSessionId() != session.id { t.Fatalf("expected start message for session %q, got %+v", session.id, message) } default: t.Fatal("expected new tunnel to receive session start") }
manager.unregisterTunnel("main", oldTunnel) if got := manager.lookupSession(session.id); got == nil { t.Fatal("expected new tunnel session to survive old tunnel unregister") }}
func websocketURL(baseURL, path string) string { return "ws" + strings.TrimPrefix(baseURL, "http") + path}
func controllerTestHeartbeat(nodeID string) store.NodeHeartbeat { return store.NodeHeartbeat{NodeID: nodeID, HeartbeatAt: time.Date(2026, 4, 12, 12, 0, 0, 0, time.UTC)}}