diff --git a/workspace-server/internal/handlers/workspace_display_control.go b/workspace-server/internal/handlers/workspace_display_control.go new file mode 100644 index 000000000..54fc1fa1c --- /dev/null +++ b/workspace-server/internal/handlers/workspace_display_control.go @@ -0,0 +1,360 @@ +package handlers + +import ( + "context" + "crypto/subtle" + "database/sql" + "encoding/json" + "fmt" + "log" + "net/http" + "os" + "time" + + "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/models" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth" + "github.com/gin-gonic/gin" +) + +const ( + displayControlDefaultTTLSeconds = 300 + displayControlMinTTLSeconds = 30 + displayControlMaxTTLSeconds = 3600 +) + +type workspaceDisplayControlResponse struct { + Controller string `json:"controller"` + ControlledBy string `json:"controlled_by,omitempty"` + ExpiresAt time.Time `json:"expires_at"` +} + +type workspaceDisplayControlNoneResponse struct { + Controller string `json:"controller"` +} + +type acquireDisplayControlRequest struct { + Controller string `json:"controller"` + TTLSeconds int `json:"ttl_seconds"` +} + +type releaseDisplayControlRequest struct { + Force bool `json:"force"` +} + +// DisplayControl handles GET /workspaces/:id/display/control. +func (h *WorkspaceHandler) DisplayControl(c *gin.Context) { + lock, found, err := h.loadActiveDisplayControl(c, c.Param("id")) + if err != nil { + log.Printf("DisplayControl: load lock for %s failed: %v", c.Param("id"), err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load display control"}) + return + } + if !found { + c.JSON(http.StatusOK, workspaceDisplayControlNoneResponse{Controller: "none"}) + return + } + c.JSON(http.StatusOK, lock) +} + +// AcquireDisplayControl handles POST /workspaces/:id/display/control/acquire. +func (h *WorkspaceHandler) AcquireDisplayControl(c *gin.Context) { + var req acquireDisplayControlRequest + if c.Request.Body != nil && c.Request.ContentLength != 0 { + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid display control request"}) + return + } + } + if req.Controller == "" { + req.Controller = "user" + } + if req.Controller != "user" { + c.JSON(http.StatusBadRequest, gin.H{"error": "browser callers may only acquire user display control"}) + return + } + if req.TTLSeconds == 0 { + req.TTLSeconds = displayControlDefaultTTLSeconds + } + if req.TTLSeconds < displayControlMinTTLSeconds || req.TTLSeconds > displayControlMaxTTLSeconds { + c.JSON(http.StatusBadRequest, gin.H{"error": "ttl_seconds must be between 30 and 3600"}) + return + } + if ok := h.displayControlEnabled(c, c.Param("id")); !ok { + return + } + + controlledBy, ok := displayControlActor(c) + if !ok { + c.JSON(http.StatusForbidden, gin.H{"error": "display control requires admin-token or org-token auth"}) + return + } + workspaceID := c.Param("id") + startedAt := time.Now() + emitDisplayControlEvent(c.Request.Context(), "display.control.acquire.started", workspaceID, map[string]any{ + "controller": req.Controller, + "controlled_by": controlledBy, + "ttl_seconds": req.TTLSeconds, + }) + var lock workspaceDisplayControlResponse + err := db.DB.QueryRowContext(c.Request.Context(), ` +INSERT INTO workspace_display_control_locks + (workspace_id, controller, controlled_by, expires_at) +VALUES + ($1, $2, $3, now() + ($4 * interval '1 second')) +ON CONFLICT (workspace_id) DO UPDATE +SET controller = EXCLUDED.controller, + controlled_by = EXCLUDED.controlled_by, + expires_at = EXCLUDED.expires_at, + updated_at = now() +WHERE workspace_display_control_locks.expires_at <= now() + OR workspace_display_control_locks.controlled_by = EXCLUDED.controlled_by +RETURNING controller, controlled_by, expires_at`, + workspaceID, req.Controller, controlledBy, req.TTLSeconds, + ).Scan(&lock.Controller, &lock.ControlledBy, &lock.ExpiresAt) + if err == nil { + emitDisplayControlEvent(c.Request.Context(), "display.control.acquire.completed", workspaceID, map[string]any{ + "controller": lock.Controller, + "controlled_by": lock.ControlledBy, + "ttl_seconds": req.TTLSeconds, + "duration_ms": time.Since(startedAt).Milliseconds(), + }) + c.JSON(http.StatusOK, lock) + return + } + if err == sql.ErrNoRows { + current, found, loadErr := h.loadActiveDisplayControl(c, workspaceID) + if loadErr != nil { + log.Printf("AcquireDisplayControl: load active lock for %s failed: %v", workspaceID, loadErr) + emitDisplayControlEvent(c.Request.Context(), "display.control.acquire.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": loadErr.Error(), + }) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load display control"}) + return + } + emitDisplayControlEvent(c.Request.Context(), "display.control.acquire.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": "display control already held", + }) + if !found { + c.JSON(http.StatusConflict, gin.H{ + "error": "display control already held", + "current": workspaceDisplayControlNoneResponse{Controller: "none"}, + }) + return + } + c.JSON(http.StatusConflict, gin.H{ + "error": "display control already held", + "current": current, + }) + return + } + log.Printf("AcquireDisplayControl: acquire lock for %s failed: %v", workspaceID, err) + emitDisplayControlEvent(c.Request.Context(), "display.control.acquire.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": err.Error(), + }) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to acquire display control"}) +} + +// ReleaseDisplayControl handles POST /workspaces/:id/display/control/release. +func (h *WorkspaceHandler) ReleaseDisplayControl(c *gin.Context) { + var req releaseDisplayControlRequest + if c.Request.Body != nil && c.Request.ContentLength != 0 { + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid display control release request"}) + return + } + } + if req.Force { + if !displayControlIsAdminToken(c) { + c.JSON(http.StatusForbidden, gin.H{"error": "force release requires admin-token auth"}) + return + } + } + + controlledBy, ok := displayControlActor(c) + if !ok { + c.JSON(http.StatusForbidden, gin.H{"error": "display control requires admin-token or org-token auth"}) + return + } + workspaceID := c.Param("id") + startedAt := time.Now() + emitDisplayControlEvent(c.Request.Context(), "display.control.release.started", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "force": req.Force, + }) + query := `DELETE FROM workspace_display_control_locks WHERE workspace_id = $1 AND controlled_by = $2` + args := []interface{}{workspaceID, controlledBy} + if req.Force { + query = `DELETE FROM workspace_display_control_locks WHERE workspace_id = $1` + args = []interface{}{workspaceID} + } + result, err := db.DB.ExecContext(c.Request.Context(), query, args...) + if err != nil { + log.Printf("ReleaseDisplayControl: release lock for %s failed: %v", workspaceID, err) + emitDisplayControlEvent(c.Request.Context(), "display.control.release.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": err.Error(), + "force": req.Force, + }) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to release display control"}) + return + } + rowsAffected, err := result.RowsAffected() + if err != nil { + log.Printf("ReleaseDisplayControl: rows affected for %s failed: %v", workspaceID, err) + emitDisplayControlEvent(c.Request.Context(), "display.control.release.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": err.Error(), + "force": req.Force, + }) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to release display control"}) + return + } + if rowsAffected == 0 { + current, found, loadErr := h.loadActiveDisplayControl(c, workspaceID) + if loadErr != nil { + log.Printf("ReleaseDisplayControl: load active lock for %s failed: %v", workspaceID, loadErr) + emitDisplayControlEvent(c.Request.Context(), "display.control.release.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": loadErr.Error(), + "force": req.Force, + }) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load display control"}) + return + } + if !found { + emitDisplayControlEvent(c.Request.Context(), "display.control.release.completed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "force": req.Force, + "rows_affected": rowsAffected, + }) + c.JSON(http.StatusOK, workspaceDisplayControlNoneResponse{Controller: "none"}) + return + } + emitDisplayControlEvent(c.Request.Context(), "display.control.release.failed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "error": "display control held by another caller", + "force": req.Force, + }) + c.JSON(http.StatusConflict, gin.H{ + "error": "display control held by another caller", + "current": current, + }) + return + } + emitDisplayControlEvent(c.Request.Context(), "display.control.release.completed", workspaceID, map[string]any{ + "controlled_by": controlledBy, + "duration_ms": time.Since(startedAt).Milliseconds(), + "force": req.Force, + "rows_affected": rowsAffected, + }) + c.JSON(http.StatusOK, workspaceDisplayControlNoneResponse{Controller: "none"}) +} + +func (h *WorkspaceHandler) loadActiveDisplayControl(c *gin.Context, workspaceID string) (workspaceDisplayControlResponse, bool, error) { + var lock workspaceDisplayControlResponse + err := db.DB.QueryRowContext(c.Request.Context(), + `SELECT controller, controlled_by, expires_at FROM workspace_display_control_locks WHERE workspace_id = $1 AND expires_at > now()`, + workspaceID, + ).Scan(&lock.Controller, &lock.ControlledBy, &lock.ExpiresAt) + if err == nil { + return lock, true, nil + } + if err == sql.ErrNoRows { + return workspaceDisplayControlResponse{}, false, nil + } + return workspaceDisplayControlResponse{}, false, err +} + +func (h *WorkspaceHandler) displayControlEnabled(c *gin.Context, workspaceID string) bool { + var raw string + err := db.DB.QueryRowContext(c.Request.Context(), + `SELECT COALESCE(compute, '{}'::jsonb) FROM workspaces WHERE id = $1`, + workspaceID, + ).Scan(&raw) + if err != nil { + if err == sql.ErrNoRows { + c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"}) + return false + } + log.Printf("displayControlEnabled: load compute for %s failed: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load display config"}) + return false + } + compute, err := parseWorkspaceDisplayCompute(workspaceID, raw) + if err != nil { + log.Printf("displayControlEnabled: invalid display config for %s: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "invalid display config"}) + return false + } + if compute.Display.Mode == "" || compute.Display.Mode == "none" { + c.JSON(http.StatusBadRequest, gin.H{"error": "display not enabled"}) + return false + } + return true +} + +func parseWorkspaceDisplayCompute(workspaceID, raw string) (models.WorkspaceCompute, error) { + var compute models.WorkspaceCompute + if raw == "" || raw == "{}" { + return compute, nil + } + if err := json.Unmarshal([]byte(raw), &compute); err != nil { + return models.WorkspaceCompute{}, fmt.Errorf("invalid compute JSON for %s: %w", workspaceID, err) + } + if err := validateWorkspaceDisplayConfig(compute.Display); err != nil { + return models.WorkspaceCompute{}, err + } + return compute, nil +} + +func displayControlActor(c *gin.Context) (string, bool) { + if v, ok := c.Get("org_token_prefix"); ok { + if s, ok := v.(string); ok && s != "" { + return actorOrgTokenPrefix + s, true + } + } + if displayControlIsAdminToken(c) { + return actorAdminToken, true + } + // Browser session auth is intentionally observe-only until AdminAuth + // exposes a stable per-user or per-session identity in gin.Context. + return "", false +} + +func displayControlIsAdminToken(c *gin.Context) bool { + adminSecret := os.Getenv("ADMIN_TOKEN") + if adminSecret == "" { + return false + } + tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization")) + return subtle.ConstantTimeCompare([]byte(tok), []byte(adminSecret)) == 1 +} + +func emitDisplayControlEvent(ctx context.Context, eventType string, workspaceID string, payload map[string]any) { + if payload == nil { + payload = map[string]any{} + } + payloadJSON, err := json.Marshal(payload) + if err != nil { + log.Printf("emitDisplayControlEvent: marshal %s payload failed: %v", eventType, err) + return + } + if _, err := db.DB.ExecContext(ctx, ` + INSERT INTO structure_events (event_type, workspace_id, payload, created_at) + VALUES ($1, $2, $3::jsonb, now()) + `, eventType, workspaceID, string(payloadJSON)); err != nil { + log.Printf("emitDisplayControlEvent: insert %s failed: %v", eventType, err) + } +} diff --git a/workspace-server/internal/handlers/workspace_display_control_test.go b/workspace-server/internal/handlers/workspace_display_control_test.go new file mode 100644 index 000000000..910e8ef80 --- /dev/null +++ b/workspace-server/internal/handlers/workspace_display_control_test.go @@ -0,0 +1,321 @@ +package handlers + +import ( + "bytes" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/gin-gonic/gin" +) + +func attachDisplayControlAdminToken(t *testing.T, c *gin.Context) { + t.Helper() + t.Setenv("ADMIN_TOKEN", "test-admin-secret") + c.Request.Header.Set("Authorization", "Bearer test-admin-secret") +} + +func TestWorkspaceDisplayControl_NoActiveLockReturnsNone(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectQuery(`SELECT controller, controlled_by, expires_at FROM workspace_display_control_locks WHERE workspace_id = \$1 AND expires_at > now\(\)`). + WithArgs("ws-display"). + WillReturnError(sql.ErrNoRows) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("GET", "/workspaces/ws-display/display/control", nil) + + handler.DisplayControl(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["controller"] != "none" { + t.Fatalf("controller = %v, want none", resp["controller"]) + } + if _, ok := resp["expires_at"]; ok { + t.Fatalf("none response included expires_at: %#v", resp) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlAcquire_ClaimsUnlockedDisplay(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + expiresAt := time.Date(2026, 5, 23, 18, 30, 0, 0, time.UTC) + + mock.ExpectQuery(`SELECT COALESCE\(compute, '\{\}'::jsonb\) FROM workspaces WHERE id = \$1`). + WithArgs("ws-display"). + WillReturnRows(sqlmock.NewRows([]string{"compute"}).AddRow(`{"display":{"mode":"desktop-control","protocol":"dcv","width":1920,"height":1080}}`)) + mock.ExpectQuery(`INSERT INTO workspace_display_control_locks`). + WithArgs("ws-display", "user", "admin-token", 300). + WillReturnRows(sqlmock.NewRows([]string{"controller", "controlled_by", "expires_at"}). + AddRow("user", "admin-token", expiresAt)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/acquire", bytes.NewBufferString(`{"controller":"user","ttl_seconds":300}`)) + c.Request.Header.Set("Content-Type", "application/json") + attachDisplayControlAdminToken(t, c) + + handler.AcquireDisplayControl(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["controller"] != "user" || resp["controlled_by"] != "admin-token" { + t.Fatalf("lock response = %#v, want user/admin-token", resp) + } + if resp["expires_at"] == "" { + t.Fatalf("expires_at missing in response: %#v", resp) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlAcquire_ActiveLockReturnsConflict(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + expiresAt := time.Date(2026, 5, 23, 18, 30, 0, 0, time.UTC) + + mock.ExpectQuery(`SELECT COALESCE\(compute, '\{\}'::jsonb\) FROM workspaces WHERE id = \$1`). + WithArgs("ws-display"). + WillReturnRows(sqlmock.NewRows([]string{"compute"}).AddRow(`{"display":{"mode":"desktop-control","protocol":"dcv","width":1920,"height":1080}}`)) + mock.ExpectQuery(`INSERT INTO workspace_display_control_locks`). + WithArgs("ws-display", "user", "admin-token", 300). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(`SELECT controller, controlled_by, expires_at FROM workspace_display_control_locks WHERE workspace_id = \$1 AND expires_at > now\(\)`). + WithArgs("ws-display"). + WillReturnRows(sqlmock.NewRows([]string{"controller", "controlled_by", "expires_at"}). + AddRow("agent", "sidecar", expiresAt)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/acquire", bytes.NewBufferString(`{"controller":"user","ttl_seconds":300}`)) + c.Request.Header.Set("Content-Type", "application/json") + attachDisplayControlAdminToken(t, c) + + handler.AcquireDisplayControl(c) + + if w.Code != http.StatusConflict { + t.Fatalf("expected status 409, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["error"] != "display control already held" { + t.Fatalf("error = %v, want display control already held", resp["error"]) + } + current, ok := resp["current"].(map[string]interface{}) + if !ok || current["controller"] != "agent" || current["controlled_by"] != "sidecar" { + t.Fatalf("current lock = %#v, want agent/sidecar", resp["current"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlAcquire_RejectsDisplayDisabledWorkspace(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectQuery(`SELECT COALESCE\(compute, '\{\}'::jsonb\) FROM workspaces WHERE id = \$1`). + WithArgs("ws-no-display"). + WillReturnRows(sqlmock.NewRows([]string{"compute"}).AddRow(`{}`)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-no-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-no-display/display/control/acquire", bytes.NewBufferString(`{"controller":"user","ttl_seconds":300}`)) + c.Request.Header.Set("Content-Type", "application/json") + attachDisplayControlAdminToken(t, c) + + handler.AcquireDisplayControl(c) + + if w.Code != http.StatusBadRequest { + t.Fatalf("expected status 400, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["error"] != "display not enabled" { + t.Fatalf("error = %v, want display not enabled", resp["error"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlAcquire_RejectsCoarseSessionActor(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectQuery(`SELECT COALESCE\(compute, '\{\}'::jsonb\) FROM workspaces WHERE id = \$1`). + WithArgs("ws-display"). + WillReturnRows(sqlmock.NewRows([]string{"compute"}).AddRow(`{"display":{"mode":"desktop-control","protocol":"dcv","width":1920,"height":1080}}`)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/acquire", bytes.NewBufferString(`{"controller":"user","ttl_seconds":300}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("Cookie", "molecule_session=present") + + handler.AcquireDisplayControl(c) + + if w.Code != http.StatusForbidden { + t.Fatalf("expected status 403, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["error"] != "display control requires admin-token or org-token auth" { + t.Fatalf("error = %v, want display control requires admin-token or org-token auth", resp["error"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlRelease_RemovesCallerLock(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectExec(`DELETE FROM workspace_display_control_locks WHERE workspace_id = \$1 AND controlled_by = \$2`). + WithArgs("ws-display", "admin-token"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/release", nil) + attachDisplayControlAdminToken(t, c) + + handler.ReleaseDisplayControl(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["controller"] != "none" { + t.Fatalf("controller = %v, want none", resp["controller"]) + } + if _, ok := resp["expires_at"]; ok { + t.Fatalf("none response included expires_at: %#v", resp) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlRelease_ConflictWhenCallerDoesNotOwnLock(t *testing.T) { + mock := setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + expiresAt := time.Date(2026, 5, 23, 18, 30, 0, 0, time.UTC) + + mock.ExpectExec(`DELETE FROM workspace_display_control_locks WHERE workspace_id = \$1 AND controlled_by = \$2`). + WithArgs("ws-display", "admin-token"). + WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectQuery(`SELECT controller, controlled_by, expires_at FROM workspace_display_control_locks WHERE workspace_id = \$1 AND expires_at > now\(\)`). + WithArgs("ws-display"). + WillReturnRows(sqlmock.NewRows([]string{"controller", "controlled_by", "expires_at"}). + AddRow("user", "org-token:abcd1234", expiresAt)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/release", nil) + attachDisplayControlAdminToken(t, c) + + handler.ReleaseDisplayControl(c) + + if w.Code != http.StatusConflict { + t.Fatalf("expected status 409, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["error"] != "display control held by another caller" { + t.Fatalf("error = %v, want display control held by another caller", resp["error"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceDisplayControlRelease_RejectsOrgTokenForceRelease(t *testing.T) { + setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Set("org_token_prefix", "abcd1234") + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/release", bytes.NewBufferString(`{"force":true}`)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.ReleaseDisplayControl(c) + + if w.Code != http.StatusForbidden { + t.Fatalf("expected status 403, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["error"] != "force release requires admin-token auth" { + t.Fatalf("error = %v, want force release requires admin-token auth", resp["error"]) + } +} + +func TestWorkspaceDisplayControlAcquire_RejectsAgentImpersonation(t *testing.T) { + setupTestDB(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-display"}} + c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/acquire", bytes.NewBufferString(`{"controller":"agent","ttl_seconds":300}`)) + c.Request.Header.Set("Content-Type", "application/json") + attachDisplayControlAdminToken(t, c) + + handler.AcquireDisplayControl(c) + + if w.Code != http.StatusBadRequest { + t.Fatalf("expected status 400, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to parse response: %v", err) + } + if resp["error"] != "browser callers may only acquire user display control" { + t.Fatalf("error = %v, want browser callers may only acquire user display control", resp["error"]) + } +} diff --git a/workspace-server/internal/router/router.go b/workspace-server/internal/router/router.go index 5d2bad31e..d32eba5ed 100644 --- a/workspace-server/internal/router/router.go +++ b/workspace-server/internal/router/router.go @@ -182,6 +182,9 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi // URLs, so keep the endpoint admin-gated from the first unavailable // state rather than widening it later. wsAdmin.GET("/workspaces/:id/display", wh.Display) + wsAdmin.GET("/workspaces/:id/display/control", wh.DisplayControl) + wsAdmin.POST("/workspaces/:id/display/control/acquire", wh.AcquireDisplayControl) + wsAdmin.POST("/workspaces/:id/display/control/release", wh.ReleaseDisplayControl) // Admin memory backup/restore (#1051) — bulk export/import of agent // memories for safe Docker rebuilds. Matches workspaces by name on import. diff --git a/workspace-server/internal/router/workspace_display_route_test.go b/workspace-server/internal/router/workspace_display_route_test.go index 072df6513..f89dab002 100644 --- a/workspace-server/internal/router/workspace_display_route_test.go +++ b/workspace-server/internal/router/workspace_display_route_test.go @@ -18,6 +18,7 @@ func buildWorkspaceDisplayEngine(t *testing.T) *gin.Engine { r := gin.New() wh := handlers.NewWorkspaceHandler(nil, nil, "http://localhost:8080", t.TempDir()) r.GET("/workspaces/:id/display", middleware.AdminAuth(db.DB), wh.Display) + r.POST("/workspaces/:id/display/control/acquire", middleware.AdminAuth(db.DB), wh.AcquireDisplayControl) return r } @@ -39,3 +40,22 @@ func TestWorkspaceDisplayRoute_RequiresAdminAuth(t *testing.T) { t.Errorf("sqlmock unmet: %v", err) } } + +func TestWorkspaceDisplayControlRoute_RequiresAdminAuth(t *testing.T) { + t.Setenv("ADMIN_TOKEN", "test-admin-secret-not-presented-by-caller") + mock := setupRouterTestDB(t) + mock.ExpectQuery("SELECT COUNT.*FROM workspace_auth_tokens"). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1)) + + r := buildWorkspaceDisplayEngine(t) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/workspaces/ws-display/display/control/acquire", nil) + r.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Errorf("expected 401 for unauthenticated request, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock unmet: %v", err) + } +} diff --git a/workspace-server/migrations/20260523120000_workspace_display_control_locks.down.sql b/workspace-server/migrations/20260523120000_workspace_display_control_locks.down.sql new file mode 100644 index 000000000..793560fb3 --- /dev/null +++ b/workspace-server/migrations/20260523120000_workspace_display_control_locks.down.sql @@ -0,0 +1,2 @@ +DROP INDEX IF EXISTS idx_workspace_display_control_locks_expires; +DROP TABLE IF EXISTS workspace_display_control_locks; diff --git a/workspace-server/migrations/20260523120000_workspace_display_control_locks.up.sql b/workspace-server/migrations/20260523120000_workspace_display_control_locks.up.sql new file mode 100644 index 000000000..03aac281b --- /dev/null +++ b/workspace-server/migrations/20260523120000_workspace_display_control_locks.up.sql @@ -0,0 +1,11 @@ +CREATE TABLE IF NOT EXISTS workspace_display_control_locks ( + workspace_id uuid PRIMARY KEY REFERENCES workspaces(id) ON DELETE CASCADE, + controller text NOT NULL CHECK (controller IN ('user', 'agent')), + controlled_by text NOT NULL CHECK (length(controlled_by) > 0 AND length(controlled_by) <= 200), + expires_at timestamptz NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_workspace_display_control_locks_expires + ON workspace_display_control_locks (expires_at);