package handler import ( "bytes" "context" "encoding/json" "net/http" "net/http/httptest" "testing" "github.com/go-chi/chi/v5" "github.com/multica-ai/multica/server/internal/middleware" ) // newDaemonTokenRequest creates an HTTP request with daemon token context set // (simulating DaemonAuth middleware for mdt_ tokens). func newDaemonTokenRequest(method, path string, body any, workspaceID, daemonID string) *http.Request { var buf bytes.Buffer if body != nil { json.NewEncoder(&buf).Encode(body) } req := httptest.NewRequest(method, path, &buf) req.Header.Set("Content-Type", "application/json") // No X-User-ID — daemon tokens don't set it. ctx := middleware.WithDaemonContext(req.Context(), workspaceID, daemonID) return req.WithContext(ctx) } func TestDaemonRegister_WithDaemonToken(t *testing.T) { if testHandler == nil { t.Skip("database not available") } w := httptest.NewRecorder() req := newDaemonTokenRequest("POST", "/api/daemon/register", map[string]any{ "workspace_id": testWorkspaceID, "daemon_id": "test-daemon-mdt", "device_name": "test-device", "runtimes": []map[string]any{ {"name": "test-runtime", "type": "claude", "version": "1.0.0", "status": "online"}, }, }, testWorkspaceID, "test-daemon-mdt") testHandler.DaemonRegister(w, req) if w.Code != http.StatusOK { t.Fatalf("DaemonRegister with daemon token: expected 200, got %d: %s", w.Code, w.Body.String()) } var resp map[string]any json.NewDecoder(w.Body).Decode(&resp) runtimes, ok := resp["runtimes"].([]any) if !ok || len(runtimes) == 0 { t.Fatalf("DaemonRegister: expected runtimes in response, got %v", resp) } // Clean up: deregister the runtime. rt := runtimes[0].(map[string]any) runtimeID := rt["id"].(string) testPool.Exec(context.Background(), `DELETE FROM agent_runtime WHERE id = $1`, runtimeID) } func TestDaemonRegister_WithDaemonToken_WorkspaceMismatch(t *testing.T) { if testHandler == nil { t.Skip("database not available") } w := httptest.NewRecorder() // Daemon token is for a different workspace than the request body. req := newDaemonTokenRequest("POST", "/api/daemon/register", map[string]any{ "workspace_id": testWorkspaceID, "daemon_id": "test-daemon-mdt", "device_name": "test-device", "runtimes": []map[string]any{ {"name": "test-runtime", "type": "claude", "version": "1.0.0", "status": "online"}, }, }, "00000000-0000-0000-0000-000000000000", "test-daemon-mdt") testHandler.DaemonRegister(w, req) if w.Code != http.StatusNotFound { t.Fatalf("DaemonRegister with mismatched workspace: expected 404, got %d: %s", w.Code, w.Body.String()) } } func TestDaemonHeartbeat_WithDaemonToken_CrossWorkspace(t *testing.T) { if testHandler == nil { t.Skip("database not available") } // First, register a runtime using PAT (existing flow). w := httptest.NewRecorder() req := newRequest("POST", "/api/daemon/register", map[string]any{ "workspace_id": testWorkspaceID, "daemon_id": "test-daemon-heartbeat", "device_name": "test-device", "runtimes": []map[string]any{ {"name": "test-runtime-hb", "type": "claude", "version": "1.0.0", "status": "online"}, }, }) testHandler.DaemonRegister(w, req) if w.Code != http.StatusOK { t.Fatalf("Setup: DaemonRegister failed: %d: %s", w.Code, w.Body.String()) } var regResp map[string]any json.NewDecoder(w.Body).Decode(®Resp) runtimes := regResp["runtimes"].([]any) runtimeID := runtimes[0].(map[string]any)["id"].(string) defer testPool.Exec(context.Background(), `DELETE FROM agent_runtime WHERE id = $1`, runtimeID) // Try heartbeat with a daemon token from a DIFFERENT workspace — should fail. w = httptest.NewRecorder() req = newDaemonTokenRequest("POST", "/api/daemon/heartbeat", map[string]any{ "runtime_id": runtimeID, }, "00000000-0000-0000-0000-000000000000", "attacker-daemon") testHandler.DaemonHeartbeat(w, req) if w.Code != http.StatusNotFound { t.Fatalf("DaemonHeartbeat with cross-workspace token: expected 404, got %d: %s", w.Code, w.Body.String()) } } func TestGetTaskStatus_WithDaemonToken_CrossWorkspace(t *testing.T) { if testHandler == nil { t.Skip("database not available") } // Create a task in the test workspace. var issueID, taskID string err := testPool.QueryRow(context.Background(), ` INSERT INTO issue (workspace_id, title, status, priority, creator_id, creator_type) VALUES ($1, 'daemon-auth-test-issue', 'todo', 'medium', $2, 'member') RETURNING id `, testWorkspaceID, testUserID).Scan(&issueID) if err != nil { t.Fatalf("setup: create issue: %v", err) } defer testPool.Exec(context.Background(), `DELETE FROM issue WHERE id = $1`, issueID) // Get an agent and runtime from the test workspace. var agentID, runtimeID string err = testPool.QueryRow(context.Background(), ` SELECT a.id, a.runtime_id FROM agent a WHERE a.workspace_id = $1 LIMIT 1 `, testWorkspaceID).Scan(&agentID, &runtimeID) if err != nil { t.Fatalf("setup: get agent: %v", err) } err = testPool.QueryRow(context.Background(), ` INSERT INTO agent_task_queue (agent_id, issue_id, status, runtime_id) VALUES ($1, $2, 'queued', $3) RETURNING id `, agentID, issueID, runtimeID).Scan(&taskID) if err != nil { t.Fatalf("setup: create task: %v", err) } defer testPool.Exec(context.Background(), `DELETE FROM agent_task_queue WHERE id = $1`, taskID) // Try GetTaskStatus with a daemon token from a DIFFERENT workspace — should fail. w := httptest.NewRecorder() req := newDaemonTokenRequest("GET", "/api/daemon/tasks/"+taskID+"/status", nil, "00000000-0000-0000-0000-000000000000", "attacker-daemon") rctx := chi.NewRouteContext() rctx.URLParams.Add("taskId", taskID) req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx)) testHandler.GetTaskStatus(w, req) if w.Code != http.StatusNotFound { t.Fatalf("GetTaskStatus with cross-workspace token: expected 404, got %d: %s", w.Code, w.Body.String()) } // Same request with the CORRECT workspace should succeed. w = httptest.NewRecorder() req = newDaemonTokenRequest("GET", "/api/daemon/tasks/"+taskID+"/status", nil, testWorkspaceID, "legit-daemon") req = req.WithContext(context.WithValue( middleware.WithDaemonContext(req.Context(), testWorkspaceID, "legit-daemon"), chi.RouteCtxKey, rctx)) testHandler.GetTaskStatus(w, req) if w.Code != http.StatusOK { t.Fatalf("GetTaskStatus with correct workspace token: expected 200, got %d: %s", w.Code, w.Body.String()) } }