weichen
02e2036fc6
Details (original message): Add dual-tier storage (embedded SQLite + optional PostgreSQL), tenant owner scoping, API key auth, versioned dual-dialect migrations, integrity hardening, and batched event persistence
261 lines
8.3 KiB
Go
261 lines
8.3 KiB
Go
package server
|
||
|
||
// 认证与多租户 e2e（阶段 1）：API key 中间件、租户隔离（404 负例）、
|
||
// 会话在途互斥（409）、备份端点、SSE 越权拒绝。
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"agent-system-backend/internal/bus"
|
||
"agent-system-backend/internal/plugins"
|
||
"agent-system-backend/internal/profiles"
|
||
"agent-system-backend/internal/runmgr"
|
||
"agent-system-backend/internal/runtime"
|
||
"agent-system-backend/internal/settings"
|
||
"agent-system-backend/internal/store"
|
||
"agent-system-backend/internal/tools"
|
||
|
||
core "agent-system-core"
|
||
)
|
||
|
||
// newAuthTestServer：开启 API key 认证的组合根（测试版）。
|
||
func newAuthTestServer(t *testing.T) (*Server, *store.Store) {
|
||
t.Helper()
|
||
dir := t.TempDir()
|
||
st, err := store.OpenSQLite(filepath.Join(dir, "test.db"))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() { st.Close() })
|
||
|
||
b := bus.New()
|
||
reg := tools.NewRegistry()
|
||
reg.RegisterBuiltin()
|
||
setSvc := settings.NewService(st)
|
||
profSvc := profiles.New(st)
|
||
profSvc.EnsureDefault()
|
||
mgr := runmgr.New(st, b, reg, profSvc, setSvc, &runtime.LLMRunner{}, &runtime.SubLLMRunner{}, &runtime.MockRunner{}, &runtime.MockSubagent{})
|
||
ph := plugins.New(st, reg, setSvc, filepath.Join(dir, "plugins"))
|
||
srv := New(st, mgr, reg, profSvc, setSvc, ph)
|
||
srv.Auth = NewAuthenticator(st, true)
|
||
srv.Subscribe = func(cid string) (<-chan core.StreamEvent, func()) { return b.Subscribe(cid) }
|
||
return srv, st
|
||
}
|
||
|
||
func TestAuthMiddleware(t *testing.T) {
|
||
srv, st := newAuthTestServer(t)
|
||
h := srv.Handler()
|
||
|
||
// 未带 key：401。
|
||
req := httptest.NewRequest("GET", "/api/projects", nil)
|
||
rec := httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
if rec.Code != http.StatusUnauthorized {
|
||
t.Fatalf("no key: %d", rec.Code)
|
||
}
|
||
|
||
// 创建密钥。
|
||
key, err := st.CreateAPIKey("alice", "e2e")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
// Bearer 头：200。
|
||
req = httptest.NewRequest("GET", "/api/projects", nil)
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
rec = httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
if rec.Code != http.StatusOK {
|
||
t.Fatalf("bearer: %d", rec.Code)
|
||
}
|
||
|
||
// ?key= 查询参数（EventSource 路径）：200。
|
||
req = httptest.NewRequest("GET", "/api/projects?key="+key, nil)
|
||
rec = httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
if rec.Code != http.StatusOK {
|
||
t.Fatalf("query key: %d", rec.Code)
|
||
}
|
||
|
||
// 撤销后 401。
|
||
if ok, err := st.RevokeAPIKeyByName("alice", "e2e"); err != nil || !ok {
|
||
t.Fatalf("revoke: %v %v", ok, err)
|
||
}
|
||
req = httptest.NewRequest("GET", "/api/projects", nil)
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
rec = httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
if rec.Code != http.StatusUnauthorized {
|
||
t.Fatalf("revoked: %d", rec.Code)
|
||
}
|
||
}
|
||
|
||
// 租户隔离（HTTP 负例）：alice 的资源对 bob 一律 404。
|
||
func TestTenantIsolationHTTP(t *testing.T) {
|
||
srv, st := newAuthTestServer(t)
|
||
h := srv.Handler()
|
||
keyA, _ := st.CreateAPIKey("alice", "a")
|
||
keyB, _ := st.CreateAPIKey("bob", "b")
|
||
|
||
do := func(key, method, path string, body any) int {
|
||
var buf bytes.Buffer
|
||
if body != nil {
|
||
json.NewEncoder(&buf).Encode(body)
|
||
}
|
||
req := httptest.NewRequest(method, path, &buf)
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
rec := httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
return rec.Code
|
||
}
|
||
|
||
if c := do(keyA, "POST", "/api/projects", map[string]any{"name": "A"}); c != 200 {
|
||
t.Fatalf("alice create: %d", c)
|
||
}
|
||
var projects []map[string]any
|
||
req := httptest.NewRequest("GET", "/api/projects", nil)
|
||
req.Header.Set("Authorization", "Bearer "+keyA)
|
||
rec := httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
json.Unmarshal(rec.Body.Bytes(), &projects)
|
||
if len(projects) != 1 {
|
||
t.Fatalf("alice projects = %v", projects)
|
||
}
|
||
projectID := projects[0]["id"].(string)
|
||
|
||
// bob 看不到 alice 的项目/会话/工件。
|
||
if c := do(keyB, "PATCH", "/api/projects/"+projectID, map[string]any{"name": "劫持"}); c != 404 {
|
||
t.Fatalf("bob patch alice project: %d", c)
|
||
}
|
||
if c := do(keyB, "GET", "/api/nonexistent", nil); c != 404 {
|
||
t.Fatalf("unknown api path must 404, got %d", c)
|
||
}
|
||
if c := do(keyB, "GET", "/api/projects/"+projectID+"/conversations", nil); c != 404 {
|
||
t.Fatalf("bob list alice conversations: %d", c)
|
||
}
|
||
if c := do(keyB, "DELETE", "/api/projects/"+projectID, nil); c != 404 {
|
||
t.Fatalf("bob delete alice project: %d", c)
|
||
}
|
||
// bob 项目列表为空（不是 404，是空数组）。
|
||
req = httptest.NewRequest("GET", "/api/projects", nil)
|
||
req.Header.Set("Authorization", "Bearer "+keyB)
|
||
rec = httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
var bobProjects []map[string]any
|
||
json.Unmarshal(rec.Body.Bytes(), &bobProjects)
|
||
if len(bobProjects) != 0 {
|
||
t.Fatalf("bob projects = %v", bobProjects)
|
||
}
|
||
// bob 不能在 alice 项目下建会话。
|
||
if c := do(keyB, "POST", "/api/projects/"+projectID+"/conversations", map[string]any{"title": "x"}); c != 404 {
|
||
t.Fatalf("bob create conversation in alice project: %d", c)
|
||
}
|
||
// 设置按租户隔离。
|
||
if c := do(keyA, "PUT", "/api/settings/user", map[string]any{"ui": map[string]any{"theme": "light"}}); c != 200 {
|
||
t.Fatalf("alice settings: %d", c)
|
||
}
|
||
req = httptest.NewRequest("GET", "/api/settings", nil)
|
||
req.Header.Set("Authorization", "Bearer "+keyB)
|
||
rec = httptest.NewRecorder()
|
||
h.ServeHTTP(rec, req)
|
||
var s map[string]any
|
||
json.Unmarshal(rec.Body.Bytes(), &s)
|
||
if u := s["user"].(map[string]any); len(u) != 0 {
|
||
t.Fatalf("bob user settings must be empty, got %v", u)
|
||
}
|
||
}
|
||
|
||
// 会话在途互斥：mock run 进行中（阻塞 runner）第二个消息 409。
|
||
func TestConversationBusyConflict(t *testing.T) {
|
||
dir := t.TempDir()
|
||
st, err := store.OpenSQLite(filepath.Join(dir, "test.db"))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() { st.Close() })
|
||
|
||
b := bus.New()
|
||
reg := tools.NewRegistry()
|
||
reg.RegisterBuiltin()
|
||
setSvc := settings.NewService(st)
|
||
profSvc := profiles.New(st)
|
||
profSvc.EnsureDefault()
|
||
|
||
block := make(chan struct{}, 1) // 缓冲：通知不阻塞执行
|
||
release := make(chan struct{})
|
||
// 阻塞 runner：第一个 run 挂住直到 release。
|
||
blocking := blockingRunner{block: block, release: release}
|
||
mgr := runmgr.New(st, b, reg, profSvc, setSvc, blocking, blocking, &runtime.MockRunner{}, &runtime.MockSubagent{})
|
||
ph := plugins.New(st, reg, setSvc, filepath.Join(dir, "plugins"))
|
||
srv := New(st, mgr, reg, profSvc, setSvc, ph)
|
||
h := srv.Handler()
|
||
|
||
// 建项目与会话（认证关闭，默认租户）。
|
||
_, out := doJSON(t, h, "POST", "/api/projects", map[string]any{"name": "p"})
|
||
projectID := out["id"].(string)
|
||
_, out = doJSON(t, h, "POST", "/api/projects/"+projectID+"/conversations", map[string]any{"title": "c"})
|
||
convID := out["id"].(string)
|
||
|
||
msg := core.Message{ConversationID: convID, Role: core.RoleUser, Parts: []core.Part{{Type: core.PartText, Text: "hi"}}}
|
||
conv, _ := st.GetConversation(store.OwnerDefault, convID)
|
||
if _, err := mgr.StartConversationRun(conv, msg, nil); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// 第二个 run：409。
|
||
if _, err := mgr.StartConversationRun(conv, msg, nil); err != runmgr.ErrConversationBusy {
|
||
t.Fatalf("second run must be busy, got %v", err)
|
||
}
|
||
// 释放第一个 run，槽位归还后可再启动。
|
||
close(release)
|
||
deadline := time.Now().Add(2 * time.Second)
|
||
for time.Now().Before(deadline) {
|
||
_, err := mgr.StartConversationRun(conv, msg, nil)
|
||
if err == nil {
|
||
return // 槽位已归还：通过
|
||
}
|
||
if err != runmgr.ErrConversationBusy {
|
||
t.Fatalf("unexpected: %v", err)
|
||
}
|
||
time.Sleep(10 * time.Millisecond)
|
||
}
|
||
t.Fatal("conversation slot not released after run finished")
|
||
}
|
||
|
||
type blockingRunner struct {
|
||
block chan struct{}
|
||
release chan struct{}
|
||
}
|
||
|
||
func (r blockingRunner) Name() string { return "blocking" }
|
||
func (r blockingRunner) Execute(rc runtime.RunCtx) error {
|
||
rc.SetAssistantParts([]core.Part{{Type: core.PartText, Text: "done"}})
|
||
r.block <- struct{}{} // 通知已到位
|
||
<-r.release // 等待释放
|
||
return nil
|
||
}
|
||
|
||
// 备份端点：产出一致性副本文件。
|
||
func TestBackupEndpoint(t *testing.T) {
|
||
srv, _ := newTestServer(t)
|
||
h := srv.Handler()
|
||
code, out := doJSON(t, h, "POST", "/api/admin/backup", nil)
|
||
if code != 200 {
|
||
t.Fatalf("backup: %d %v", code, out)
|
||
}
|
||
path := out["path"].(string)
|
||
if !strings.HasSuffix(path, ".db") {
|
||
t.Fatalf("backup path: %s", path)
|
||
}
|
||
if _, err := os.Stat(path); err != nil {
|
||
t.Fatalf("backup file: %v", err)
|
||
}
|
||
}
|