Files
multica/server/internal/handler/daemon_test.go
2026-04-15 19:10:54 +08:00

307 lines
11 KiB
Go

package handler
import (
"bytes"
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/go-chi/chi/v5"
"github.com/multica-ai/multica/server/internal/middleware"
)
func setHandlerTestWorkspaceRepos(t *testing.T, repos []map[string]string) {
t.Helper()
data, err := json.Marshal(repos)
if err != nil {
t.Fatalf("marshal repos: %v", err)
}
if _, err := testPool.Exec(context.Background(), `UPDATE workspace SET repos = $1 WHERE id = $2`, data, testWorkspaceID); err != nil {
t.Fatalf("update workspace repos: %v", err)
}
t.Cleanup(func() {
if _, err := testPool.Exec(context.Background(), `UPDATE workspace SET repos = $1 WHERE id = $2`, []byte("[]"), testWorkspaceID); err != nil {
t.Fatalf("reset workspace repos: %v", err)
}
})
}
// 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)
}
if _, ok := resp["repos_version"].(string); !ok {
t.Fatalf("DaemonRegister: expected repos_version 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(&regResp)
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())
}
}
func TestGetDaemonWorkspaceRepos_WithDaemonToken(t *testing.T) {
if testHandler == nil {
t.Skip("database not available")
}
setHandlerTestWorkspaceRepos(t, []map[string]string{
{"url": "git@example.com:team/api.git", "description": "API"},
{"url": " git@example.com:team/web.git ", "description": " Web "},
})
w := httptest.NewRecorder()
req := newDaemonTokenRequest("GET", "/api/daemon/workspaces/"+testWorkspaceID+"/repos", nil, testWorkspaceID, "test-daemon-mdt")
req = withURLParam(req, "workspaceId", testWorkspaceID)
testHandler.GetDaemonWorkspaceRepos(w, req)
if w.Code != http.StatusOK {
t.Fatalf("GetDaemonWorkspaceRepos: expected 200, got %d: %s", w.Code, w.Body.String())
}
var resp struct {
WorkspaceID string `json:"workspace_id"`
Repos []map[string]string `json:"repos"`
ReposVersion string `json:"repos_version"`
}
if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
t.Fatalf("decode response: %v", err)
}
if resp.WorkspaceID != testWorkspaceID {
t.Fatalf("expected workspace_id %s, got %s", testWorkspaceID, resp.WorkspaceID)
}
if len(resp.Repos) != 2 {
t.Fatalf("expected 2 repos, got %d", len(resp.Repos))
}
if resp.Repos[1]["url"] != "git@example.com:team/web.git" {
t.Fatalf("expected trimmed repo URL, got %q", resp.Repos[1]["url"])
}
if resp.ReposVersion == "" {
t.Fatal("expected repos_version to be set")
}
}
func TestGetDaemonWorkspaceRepos_WithDaemonToken_WorkspaceMismatch(t *testing.T) {
if testHandler == nil {
t.Skip("database not available")
}
w := httptest.NewRecorder()
req := newDaemonTokenRequest("GET", "/api/daemon/workspaces/"+testWorkspaceID+"/repos", nil, "00000000-0000-0000-0000-000000000000", "test-daemon-mdt")
req = withURLParam(req, "workspaceId", testWorkspaceID)
testHandler.GetDaemonWorkspaceRepos(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("GetDaemonWorkspaceRepos with mismatched workspace: expected 404, got %d: %s", w.Code, w.Body.String())
}
}
func TestGetDaemonWorkspaceRepos_VersionIgnoresOrderAndDescription(t *testing.T) {
if testHandler == nil {
t.Skip("database not available")
}
setHandlerTestWorkspaceRepos(t, []map[string]string{
{"url": "git@example.com:team/api.git", "description": "API"},
{"url": "git@example.com:team/web.git", "description": "Web"},
})
getReposVersion := func() string {
t.Helper()
w := httptest.NewRecorder()
req := newDaemonTokenRequest("GET", "/api/daemon/workspaces/"+testWorkspaceID+"/repos", nil, testWorkspaceID, "test-daemon-mdt")
req = withURLParam(req, "workspaceId", testWorkspaceID)
testHandler.GetDaemonWorkspaceRepos(w, req)
if w.Code != http.StatusOK {
t.Fatalf("GetDaemonWorkspaceRepos: expected 200, got %d: %s", w.Code, w.Body.String())
}
var resp struct {
ReposVersion string `json:"repos_version"`
}
if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
t.Fatalf("decode response: %v", err)
}
return resp.ReposVersion
}
version1 := getReposVersion()
if _, err := testPool.Exec(context.Background(), `UPDATE workspace SET repos = $1 WHERE id = $2`, []byte(`[{"url":"git@example.com:team/web.git","description":"frontend"},{"url":"git@example.com:team/api.git","description":"backend"}]`), testWorkspaceID); err != nil {
t.Fatalf("update workspace repos: %v", err)
}
version2 := getReposVersion()
if version1 != version2 {
t.Fatalf("expected repos_version to ignore order/description changes, got %s vs %s", version1, version2)
}
if _, err := testPool.Exec(context.Background(), `UPDATE workspace SET repos = $1 WHERE id = $2`, []byte(`[{"url":"git@example.com:team/api.git","description":"backend"},{"url":"git@example.com:team/mobile.git","description":"mobile"}]`), testWorkspaceID); err != nil {
t.Fatalf("update workspace repos: %v", err)
}
version3 := getReposVersion()
if strings.EqualFold(version2, version3) {
t.Fatalf("expected repos_version to change when URL set changes, got %s", version3)
}
}