package gateway // 本文件锁定切片 S3 的 B 路(网关侧,design/24 §4.2 / §6.1 / §9.5): // - `chat.documents` 解析为**用户回合**的文档部件(红线:只允许 RoleUser); // - 单回合条数上限(too_many_documents)与存储名形态 / 存在性校验 // (document_not_found),且**绝不接受 URL**(不重开 SSRF 面); // - 历史重建带回文档引用、纯附件回合不被判空(S2 移交的两处「静默吞掉」)。 // // 断言一律走哨兵错误 + errors.Is / errors.As(机器可读的 code),不按错误文本 // 子串断言 —— 文案会变,契约不变。文档落盘用临时目录,绝不触碰 data/cyrene.db。 import ( "context" "crypto/sha256" "encoding/hex" "errors" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "strconv" "sync" "sync/atomic" "testing" "time" "git.yeij.top/AskaEth/Cyrene_v2/internal/auth" "git.yeij.top/AskaEth/Cyrene_v2/internal/config" apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors" "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop" "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model" ) // chatDocRunner 记录每次 Run 看到的完整消息(含文档部件)。流式回合跑在后台 // goroutine 上,因此用互斥锁守住记录。 type chatDocRunner struct { mu sync.Mutex calls [][]model.Message } func (r *chatDocRunner) Run(_ context.Context, _ string, msgs []model.Message) (*loop.Result, error) { r.mu.Lock() r.calls = append(r.calls, msgs) r.mu.Unlock() return &loop.Result{Text: "收到了", Steps: 1, StopReason: "stop"}, nil } func (r *chatDocRunner) snapshot() [][]model.Message { r.mu.Lock() defer r.mu.Unlock() return append([][]model.Message(nil), r.calls...) } // waitCalls 等到至少 n 次 Run 被记录(后台流式回合),超时即失败。 func (r *chatDocRunner) waitCalls(t *testing.T, n int) [][]model.Message { t.Helper() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if calls := r.snapshot(); len(calls) >= n { return calls } time.Sleep(time.Millisecond) } t.Fatalf("等待 %d 次回合超时(实际 %d 次)", n, len(r.snapshot())) return nil } // chatDocumentsEnv 打包文档聊天测试所需的依赖:真实 JWT、临时数据目录(文档 // 落盘在它下面)、接线的录制 runner。 type chatDocumentsEnv struct { srv *Server ts *httptest.Server dataDir string auth *auth.JWTService runner *chatDocRunner } func newChatDocumentsEnv(t *testing.T) *chatDocumentsEnv { t.Helper() dataDir := t.TempDir() jwt := auth.New(config.AuthConfig{JWTSecret: testSecret}) cfg := &config.Config{Auth: config.AuthConfig{AdminUser: "admin", AdminPass: "secret", JWTSecret: testSecret}} runner := &chatDocRunner{} srv, err := New(cfg, Deps{Auth: jwt, DataDir: dataDir, Loop: runner}) if err != nil { t.Fatalf("构造网关失败: %v", err) } ts := httptest.NewServer(srv.Handler()) t.Cleanup(ts.Close) t.Cleanup(srv.Close) return &chatDocumentsEnv{srv: srv, ts: ts, dataDir: dataDir, auth: jwt, runner: runner} } func (e *chatDocumentsEnv) token(t *testing.T) string { t.Helper() tok, err := e.auth.Sign("acct", auth.RoleOwner, time.Hour) if err != nil { t.Fatalf("签发 token 失败: %v", err) } return tok } // storeDocumentFile 按 S1 的内容寻址规则(sha256 前 32 位 + 扩展名)在 // `/uploads/documents` 下落一份文档,并返回它的存储名。 func (e *chatDocumentsEnv) storeDocumentFile(t *testing.T, content, ext string) string { t.Helper() sum := sha256.Sum256([]byte(content)) name := hex.EncodeToString(sum[:])[:documentHashHexLen] + ext if !validDocumentName(name) { t.Fatalf("构造的存储名 %q 不是合法形态(与 S1 口径不一致)", name) } dir := filepath.Join(e.dataDir, "uploads", documentSubDir) if err := os.MkdirAll(dir, 0o700); err != nil { t.Fatalf("创建文档目录失败: %v", err) } if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o600); err != nil { t.Fatalf("写入文档失败: %v", err) } return name } // validFormDocumentName 返回一个**形态合法但不存在**的存储名(用于条数上限与 // 存在性两条判定的分辨:先撞上限还是先撞不存在)。 func validFormDocumentName(seed string) string { sum := sha256.Sum256([]byte(seed)) return hex.EncodeToString(sum[:])[:documentHashHexLen] + ".txt" } // chatDocumentsOf 取出这些消息里的全部文档部件。 func chatDocumentsOf(msgs []model.Message) []model.Content { out := make([]model.Content, 0, len(msgs)) for _, m := range msgs { for _, c := range m.Content { if c.Type == model.ContentTypeDocument { out = append(out, c) } } } return out } // assertOnlyUserCarriesDocuments 断言文档部件只挂在用户回合上(design/24 §2.3-③ // 的红线:编码器会对非 user 的文档块 fail-loud,因此这一条必须在这里就成立)。 func assertOnlyUserCarriesDocuments(t *testing.T, msgs []model.Message) { t.Helper() for _, m := range msgs { for _, c := range m.Content { if c.Type == model.ContentTypeDocument && m.Role != model.RoleUser { t.Fatalf("文档部件挂在了 %s 回合上(只允许 user)", m.Role) } } } } // TestChatDocumentsAttachAsUserParts 钉住主路径:REST chat 带 documents → // 该回合的用户消息上多出一个文档部件,且它的 name / MIME / URL / 大小全部由 // **服务端**解析(客户端只给了 name)。 func TestChatDocumentsAttachAsUserParts(t *testing.T) { t.Parallel() env := newChatDocumentsEnv(t) const content = "行1\n行2\n" name := env.storeDocumentFile(t, content, ".txt") resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{ "sessionId": "web", "text": "读一下这份文件", // 客户端只给 name;MIME / URL / 大小就算谎报也必须被服务端自己的解析覆盖 // (design/24 §3.2:字节嗅探是唯一裁决)。 "documents": []map[string]any{{ "name": name, "mime": "text/html", "url": "https://evil.example/steal.txt", "size": 999999, }}, }, env.token(t)) if resp.StatusCode != http.StatusOK { t.Fatalf("status = %d, want 200 (body=%s)", resp.StatusCode, body) } calls := env.runner.snapshot() if len(calls) != 1 { t.Fatalf("驱动了 %d 个回合,want 1", len(calls)) } msgs := calls[0] assertOnlyUserCarriesDocuments(t, msgs) docs := chatDocumentsOf(msgs) if len(docs) != 1 { t.Fatalf("文档部件数 = %d, want 1 (msgs=%+v)", len(docs), msgs) } got := docs[0] if got.DocumentName != name { t.Fatalf("DocumentName = %q, want %q", got.DocumentName, name) } if got.DocumentMIME != documentSniffOKMIME { t.Fatalf("DocumentMIME = %q, want %q(嗅探所得,绝不是客户端声明)", got.DocumentMIME, documentSniffOKMIME) } if want := documentFilePrefix + name; got.DocumentURL != want { t.Fatalf("DocumentURL = %q, want %q(服务端本地引用)", got.DocumentURL, want) } if got.DocumentSize != int64(len(content)) { t.Fatalf("DocumentSize = %d, want %d", got.DocumentSize, len(content)) } // 文本仍在同一条用户消息上,顺序是「文本在前、附件在后」。 last := msgs[len(msgs)-1] if last.Role != model.RoleUser || last.Content[0].Type != model.ContentTypeText { t.Fatalf("末条消息 = %+v, want 用户文本消息在前", last) } } // TestChatDocumentsPureAttachmentTurn 钉住「只带附件」的回合:既不是空回合, // 也不会塞一个空文本块(空文本块会被 provider 拒)。 func TestChatDocumentsPureAttachmentTurn(t *testing.T) { t.Parallel() env := newChatDocumentsEnv(t) name := env.storeDocumentFile(t, "只有附件", ".md") resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{ "sessionId": "web", "text": "", "documents": []chatDocumentRef{{Name: name}}, }, env.token(t)) if resp.StatusCode != http.StatusOK { t.Fatalf("纯附件回合被拒:status = %d (body=%s)", resp.StatusCode, body) } calls := env.runner.snapshot() if len(calls) != 1 { t.Fatalf("驱动了 %d 个回合,want 1(纯附件回合是真实回合)", len(calls)) } msgs := calls[0] if len(chatDocumentsOf(msgs)) != 1 { t.Fatalf("文档部件丢失:%+v", msgs) } for _, m := range msgs { for _, c := range m.Content { if c.Type == model.ContentTypeText && c.Text == "" { t.Fatalf("纯附件回合塞了空文本部件:%+v", m) } } } } // TestChatDocumentsTooManyRejected 钉住单回合条数上限:恰好 5 条通过,第 6 条 // 起 400 + too_many_documents,且条数判定**先于**任何文件系统访问(6 条全不存在 // 也报条数超限,而不是 document_not_found)。回合一次都不许起。 func TestChatDocumentsTooManyRejected(t *testing.T) { t.Parallel() env := newChatDocumentsEnv(t) refs := make([]chatDocumentRef, 0, maxChatDocuments+1) for i := 0; i < maxChatDocuments; i++ { refs = append(refs, chatDocumentRef{Name: env.storeDocumentFile(t, "doc-"+strconv.Itoa(i), ".txt")}) } if docs, err := env.srv.resolveChatDocuments(refs); err != nil || len(docs) != maxChatDocuments { t.Fatalf("恰好 %d 条必须通过:docs=%d err=%v", maxChatDocuments, len(docs), err) } over := append(append([]chatDocumentRef(nil), refs...), chatDocumentRef{Name: validFormDocumentName("sixth")}) if _, err := env.srv.resolveChatDocuments(over); !errors.Is(err, apperr.ErrBadRequest) { t.Fatalf("第 %d 条应被拒为 ErrBadRequest,得到 %v", maxChatDocuments+1, err) } var appErr *apperr.AppError if _, err := env.srv.resolveChatDocuments(over); !errors.As(err, &appErr) || appErr.Code != tooManyDocumentsCode { t.Fatalf("error.code = %q, want %q", appErr.Code, tooManyDocumentsCode) } tok := env.token(t) resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{ "sessionId": "web", "text": "太多了", "documents": over, }, tok) if resp.StatusCode != http.StatusBadRequest { t.Fatalf("status = %d, want 400 (body=%s)", resp.StatusCode, body) } if code := decodeDocumentEnvelope(t, body).Error.Code; code != tooManyDocumentsCode { t.Fatalf("error.code = %q, want %q", code, tooManyDocumentsCode) } if len(env.runner.snapshot()) != 0 { t.Fatalf("被拒的请求起了回合 —— 附件条数必须在起回合之前判定") } // 恰好 5 条(都存在)照常放行:上限不得误伤边界。 resp, body = doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{ "sessionId": "web", "text": "这五份", "documents": refs, }, tok) if resp.StatusCode != http.StatusOK { t.Fatalf("恰好 %d 条必须通过:status = %d (body=%s)", maxChatDocuments, resp.StatusCode, body) } } // TestChatDocumentsRejectMalformedOrMissingName 钉住存储名校验:形态非法(路径 // 穿越、非 32 位十六进制、白名单外扩展名、空名)与形态合法但不存在,一律 // 404 document_not_found(与读端点同一口径),且都不起回合。 func TestChatDocumentsRejectMalformedOrMissingName(t *testing.T) { t.Parallel() env := newChatDocumentsEnv(t) // 形态校验必须真的在把关:目录里真有一份**非内容寻址**的同名文件时也不 // 能因为「它存在」就被放行 —— 客户端能引用的只有 S1 签发过的存储名。 dir := filepath.Join(env.dataDir, "uploads", documentSubDir) if err := os.MkdirAll(dir, 0o700); err != nil { t.Fatalf("创建文档目录失败: %v", err) } if err := os.WriteFile(filepath.Join(dir, "notes.txt"), []byte("不是内容寻址名"), 0o600); err != nil { t.Fatalf("写入非内容寻址文件失败: %v", err) } cases := []struct{ label, name string }{ {"路径穿越", "../../../etc/passwd"}, {"反斜杠穿越", `..\..\cyrene.exe`}, {"主名不是 32 位十六进制", "deadbeef.txt"}, {"扩展名不在白名单", fmt.Sprintf("%032d.html", 1)}, {"空名", ""}, {"形态合法但不存在", validFormDocumentName("missing")}, {"非内容寻址但确实存在", "notes.txt"}, } tok := env.token(t) for _, tc := range cases { _, err := env.srv.resolveChatDocuments([]chatDocumentRef{{Name: tc.name}}) if !errors.Is(err, apperr.ErrNotFound) { t.Fatalf("%s: err = %v, want ErrNotFound", tc.label, err) } var appErr *apperr.AppError if _, err := env.srv.resolveChatDocuments([]chatDocumentRef{{Name: tc.name}}); !errors.As(err, &appErr) || appErr.Code != documentNotFoundCode { t.Fatalf("%s: error.code = %q, want %q", tc.label, appErr.Code, documentNotFoundCode) } resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{ "sessionId": "web", "text": "读一下", "documents": []chatDocumentRef{{Name: tc.name}}, }, tok) if resp.StatusCode != http.StatusNotFound { t.Fatalf("%s: status = %d, want 404 (body=%s)", tc.label, resp.StatusCode, body) } if code := decodeDocumentEnvelope(t, body).Error.Code; code != documentNotFoundCode { t.Fatalf("%s: error.code = %q, want %q", tc.label, code, documentNotFoundCode) } } if len(env.runner.snapshot()) != 0 { t.Fatalf("被拒的请求起了回合 —— 存储名校验必须在起回合之前完成") } } // TestChatDocumentsNeverFetchURLs 是 SSRF 红线的正向反证(design/24 §5):把 // 一个真实可命中的 URL 当作 name 发进来,网关**一次都不许去取它** —— 文档只认 // 服务端签发的存储名。 func TestChatDocumentsNeverFetchURLs(t *testing.T) { t.Parallel() env := newChatDocumentsEnv(t) var hits atomic.Int64 evil := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hits.Add(1) w.WriteHeader(http.StatusOK) })) defer evil.Close() tok := env.token(t) names := []string{evil.URL + "/report.txt", "http://127.0.0.1:8080/api/v1/documents/whatever.txt"} for _, name := range names { resp, body := doReq(t, http.MethodPost, env.ts.URL+"/api/v1/chat", map[string]any{ "sessionId": "web", "text": "读这个链接", "documents": []chatDocumentRef{{Name: name}}, }, tok) if resp.StatusCode != http.StatusNotFound { t.Fatalf("URL 形态的 name 必须被拒:status = %d (body=%s)", resp.StatusCode, body) } } if got := hits.Load(); got != 0 { t.Fatalf("网关替客户端取了 URL %d 次 —— SSRF 面被重开", got) } if len(env.runner.snapshot()) != 0 { t.Fatalf("被拒的请求起了回合") } }