Cyrene_v2/internal/services/gateway/chat_documents_test.go

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 位 + 扩展名）在
// `<DataDir>/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("被拒的请求起了回合")
	}
}