MailUI4Agents/gateway/internal/handler/attachments.go

package handler

import (
	"errors"
	"fmt"
	"mime"
	"net/http"
	"path/filepath"
	"strings"
	"unicode/utf8"

	"github.com/agentmail/gateway/internal/blob"
	"github.com/agentmail/gateway/internal/config"
	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/google/uuid"
)

// ---------- 附件 ----------
//
// 上传与发信是两步：
//   1. POST /attachments        （multipart）→ 拿到 attachment_id
//   2. 发信时把 id 放进 attachment_ids
// 之所以不合成一步：Agent 侧的工具接口是 JSON，没法带 multipart；
// 而人类侧若只支持一步，就无法在写信过程中先传文件再改正文。
//
// 未挂载的附件是合法中间态，超时由 GC 清理（repo.SweepOrphanAttachments）。

// Blobs 是附件内容存储，由 main 在启动时注入。
var Blobs *blob.Store

// sanitizeFilename 清理用户提供的文件名。
//
// 文件名只用于展示与下载时的 Content-Disposition，磁盘路径完全由 sha256 派生，
// 因此这里的目的不是防路径穿越（那已由内容寻址杜绝），而是：
//   - 去掉目录成分，避免下载时浏览器按 "a/b/c.txt" 解释
//   - 去掉控制字符与换行，避免污染 HTTP 响应头
//   - 限长，避免超出数据库列宽
func sanitizeFilename(name string) string {
	name = strings.TrimSpace(name)
	// 同时处理 / 与 \：上传方可能是 Windows 客户端
	if i := strings.LastIndexAny(name, `/\`); i >= 0 {
		name = name[i+1:]
	}

	var b strings.Builder
	for _, r := range name {
		if r < 0x20 || r == 0x7f {
			continue // 控制字符一律丢弃
		}
		b.WriteRune(r)
	}
	name = strings.TrimSpace(b.String())

	// "." 与 ".." 作为文件名毫无意义，且容易在各层被特殊解释
	if name == "" || name == "." || name == ".." {
		return "unnamed"
	}

	const maxBytes = 255
	if len(name) > maxBytes {
		cut := name[:maxBytes]
		for len(cut) > 0 && !utf8.ValidString(cut) {
			cut = cut[:len(cut)-1]
		}
		name = cut
	}
	return name
}

// detectContentType 优先用客户端声明的类型，缺失时按扩展名猜，兜底 octet-stream。
// 无论如何都不回显未经处理的客户端值到响应头（下载时统一用 octet-stream，见 DownloadAttachment）。
func detectContentType(declared, filename string) string {
	if ct := strings.TrimSpace(declared); ct != "" && ct != "application/octet-stream" {
		if parsed, _, err := mime.ParseMediaType(ct); err == nil {
			return parsed
		}
	}
	if ext := filepath.Ext(filename); ext != "" {
		if byExt := mime.TypeByExtension(ext); byExt != "" {
			if parsed, _, err := mime.ParseMediaType(byExt); err == nil {
				return parsed
			}
		}
	}
	return "application/octet-stream"
}

// uploadAttachment 是 Agent 与人类两条上传路径的公共实现。
func uploadAttachment(w http.ResponseWriter, r *http.Request, uploader string) {
	if Blobs == nil {
		Error(w, http.StatusServiceUnavailable, "附件存储未初始化")
		return
	}

	max := config.C.MaxAttachmentBytes

	// 双层限制：MaxBytesReader 卡整个请求体（含 multipart 边界与其他字段），
	// blob.Put 的 max 卡单个文件内容。少了外层，攻击者可以用超大 multipart 头拖死内存。
	r.Body = http.MaxBytesReader(w, r.Body, max+1<<20)

	// 32MB 内存缓冲上限，超出部分 multipart 会自动落临时文件
	if err := r.ParseMultipartForm(32 << 20); err != nil {
		Error(w, http.StatusBadRequest, "解析 multipart 失败（是否超过大小上限？）")
		return
	}
	defer func() {
		if r.MultipartForm != nil {
			r.MultipartForm.RemoveAll()
		}
	}()

	file, header, err := r.FormFile("file")
	if err != nil {
		Error(w, http.StatusBadRequest, "缺少 file 字段")
		return
	}
	defer file.Close()

	name := sanitizeFilename(header.Filename)
	ctype := detectContentType(header.Header.Get("Content-Type"), name)

	// 先落盘再入库：反过来会出现「库里有记录、磁盘没文件」的下载 500
	sum, size, err := Blobs.Put(file, max)
	if errors.Is(err, blob.ErrTooLarge) {
		Error(w, http.StatusRequestEntityTooLarge,
			fmt.Sprintf("附件超过上限 %.1f MB", float64(max)/(1<<20)))
		return
	}
	if err != nil {
		Error(w, http.StatusInternalServerError, "保存附件失败")
		return
	}

	a, err := repo.CreateAttachment(r.Context(), uploader, name, ctype, size, sum)
	if err != nil {
		// 落盘成功但入库失败：留下的孤立文件由 GC 回收，不影响正确性
		Error(w, http.StatusInternalServerError, "登记附件失败")
		return
	}

	JSON(w, http.StatusOK, map[string]any{"attachment": a})
}

// POST /api/v1/attachments —— Agent 侧上传
func UploadAttachment(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}
	uploadAttachment(w, r, agentName)
}

// POST /api/v1/me/attachments —— 人类侧上传
func MeUploadAttachment(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	uploadAttachment(w, r, user.Username)
}

// downloadAttachment 是 Agent 与人类两条下载路径的公共实现。
func downloadAttachment(w http.ResponseWriter, r *http.Request, viewer string) {
	if Blobs == nil {
		Error(w, http.StatusServiceUnavailable, "附件存储未初始化")
		return
	}
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}

	a, err := repo.GetAttachment(r.Context(), id)
	if errors.Is(err, repo.ErrAttachmentNotFound) {
		Error(w, http.StatusNotFound, "附件不存在")
		return
	}
	if err != nil {
		Error(w, http.StatusInternalServerError, "读取附件失败")
		return
	}

	allowed, err := repo.AttachmentAccessible(r.Context(), a, viewer)
	if err != nil {
		Error(w, http.StatusInternalServerError, "校验权限失败")
		return
	}
	if !allowed {
		Error(w, http.StatusForbidden, "无权访问该附件")
		return
	}

	f, err := Blobs.Open(a.SHA256)
	if err != nil {
		// 元数据在库但文件不在盘：说明存储被外部改动过，这是运维问题而非用户输入问题
		Error(w, http.StatusInternalServerError, "附件内容缺失")
		return
	}
	defer f.Close()

	// 一律 octet-stream + attachment：绝不按声明的 MIME 内联渲染。
	// 否则一个上传的 .html/.svg 就能在本站域下执行脚本，等于自带 XSS。
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("Content-Length", fmt.Sprintf("%d", a.SizeBytes))
	w.Header().Set("Content-Disposition", contentDisposition(a.Filename))

	http.ServeContent(w, r, a.Filename, a.CreatedAt, f)
}

// contentDisposition 构造下载头。
// filename* 用 RFC 5987 编码承载非 ASCII 名字，filename= 给只认 ASCII 的老客户端兜底；
// 兜底值里的引号与反斜杠必须去掉，否则能截断响应头。
func contentDisposition(name string) string {
	var ascii strings.Builder
	for _, r := range name {
		switch {
		case r == '"' || r == '\\':
			ascii.WriteByte('_')
		case r < 0x20 || r > 0x7e:
			ascii.WriteByte('_')
		default:
			ascii.WriteRune(r)
		}
	}
	fallback := ascii.String()
	if fallback == "" {
		fallback = "attachment"
	}
	return fmt.Sprintf(`attachment; filename="%s"; filename*=UTF-8''%s`,
		fallback, urlEncodeRFC5987(name))
}

// urlEncodeRFC5987 按 RFC 5987 的 attr-char 集合做百分号编码。
func urlEncodeRFC5987(s string) string {
	const safe = "!#$&+-.^_`|~" // attr-char 中除字母数字外允许的字符
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		c := s[i]
		isAlnum := (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')
		if isAlnum || strings.IndexByte(safe, c) >= 0 {
			b.WriteByte(c)
		} else {
			fmt.Fprintf(&b, "%%%02X", c)
		}
	}
	return b.String()
}

// GET /api/v1/attachments/{id} —— Agent 侧下载
func DownloadAttachment(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}
	downloadAttachment(w, r, agentName)
}

// GET /api/v1/me/attachments/{id} —— 人类侧下载
func MeDownloadAttachment(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	downloadAttachment(w, r, user.Username)
}

// DELETE /api/v1/me/attachments/{id} —— 删除自己上传且尚未挂载的附件
//
// 已挂载的不允许删：邮件是不可篡改的历史记录，附件是它的一部分。
func MeDeleteAttachment(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}

	a, err := repo.GetAttachment(r.Context(), id)
	if errors.Is(err, repo.ErrAttachmentNotFound) {
		Error(w, http.StatusNotFound, "附件不存在")
		return
	}
	if err != nil {
		Error(w, http.StatusInternalServerError, "读取附件失败")
		return
	}
	if a.Uploader != user.Username {
		Error(w, http.StatusForbidden, "只能删除自己上传的附件")
		return
	}
	if a.MailID != nil {
		Error(w, http.StatusConflict, "附件已随邮件发出，不能删除")
		return
	}

	sum, orphaned, err := repo.DeleteAttachment(r.Context(), id)
	if err != nil {
		Error(w, http.StatusInternalServerError, "删除附件失败")
		return
	}
	// 内容寻址下多条记录可能共享同一文件，只有最后一条引用消失才删磁盘
	if orphaned && Blobs != nil {
		_ = Blobs.Remove(sum)
	}
	JSON(w, http.StatusOK, map[string]string{"status": "deleted"})
}

// parseAttachmentIDs 把请求里的附件 id 列表解析为 UUID。
func parseAttachmentIDs(raw []string) ([]uuid.UUID, error) {
	out := make([]uuid.UUID, 0, len(raw))
	for _, s := range raw {
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}
		id, err := uuid.Parse(s)
		if err != nil {
			return nil, fmt.Errorf("非法的 attachment_id %q", s)
		}
		out = append(out, id)
	}
	return out, nil
}

// attachAll 把附件挂到刚创建的邮件上，并把错误翻译成 HTTP 响应。
// 返回 false 表示已写出错误响应，调用方应立即返回。
func attachAll(w http.ResponseWriter, r *http.Request, mailID uuid.UUID, ids []uuid.UUID, uploader string) bool {
	if len(ids) == 0 {
		return true
	}
	return writeAttachErr(w, repo.AttachToMail(r.Context(), mailID, ids, uploader))
}

// checkAttachable 在**产生任何副作用之前**校验附件可不可挂。
//
// 返回 false 表示已写出错误响应，调用方应立即返回。
//
// # 为什么不能只靠 attachAll
//
// attachAll 在 CreateMail **之后**调用，于是附件不合法时请求返回 403/409，
// 但那封邮件**已经入库、已经通知了收件人、已经扣掉了会话预算**。
// 生产实测：两封探针邮件（一封 403「只能附加自己上传的附件」、一封 409
// 「附件已随其他邮件发出」）都躺在 mails 表里，used_rounds 也涨了。
// 发件方看到 4xx 会重试，收件方于是收到两封。
//
// 纯输入校验必须在副作用之前做完 —— 与「400 之后会话已建好」是同一个教训。
//
// 它**不取代** attachAll：两次调用之间仍有竞态窗口（另一个请求把同一个附件
// 挂走了），那一次由 attachAll 的原子 UPDATE 拦下、并由调用方回滚。
// 双层分工：这里挡住绝大多数（拼错 id、拿别人的附件、重复挂），
// attachAll 挡住真正的并发。
func checkAttachable(w http.ResponseWriter, r *http.Request, ids []uuid.UUID, uploader string) bool {
	if len(ids) == 0 {
		return true
	}
	return writeAttachErr(w, repo.EnsureAttachable(r.Context(), ids, uploader))
}

// writeAttachErr 把 repo 层的附件错误映射成 HTTP 响应。
//
// checkAttachable 与 attachAll 共用一份：同一种错误在两条路径上必须给出同一个
// 状态码与同一句话 —— 分开写早晚会分叉，而调用方无法区分自己碰上的是哪一层。
func writeAttachErr(w http.ResponseWriter, err error) bool {
	switch {
	case err == nil:
		return true
	case errors.Is(err, repo.ErrAttachmentNotFound):
		Error(w, http.StatusNotFound, "附件不存在")
	case errors.Is(err, repo.ErrAttachmentNotOwned):
		Error(w, http.StatusForbidden, "只能附加自己上传的附件")
	case errors.Is(err, repo.ErrAttachmentAlreadyAttached):
		Error(w, http.StatusConflict, "附件已随其他邮件发出")
	default:
		Error(w, http.StatusInternalServerError, "附加附件失败")
	}
	return false
}

// fillAttachments 给邮件填充附件列表（读取单封/线程时用）。
// 读附件失败不该让整封邮件打不开，因此吞错只留空列表。
//
// 一批邮件走一次查询：逐封调 ListAttachmentsFor 是 N+1，
// 一个 200 封的会话打开一次要打 200 次库。
func fillAttachments(r *http.Request, mails ...*models.Mail) {
	ids := make([]uuid.UUID, 0, len(mails))
	for _, m := range mails {
		if m != nil {
			ids = append(ids, m.ID)
		}
	}
	if len(ids) == 0 {
		return
	}

	byMail, err := repo.ListAttachmentsForMails(r.Context(), ids)
	if err != nil {
		return
	}
	for _, m := range mails {
		if m == nil {
			continue
		}
		// 没有附件的邮件保持 nil：Attachments 带 omitempty，
		// 填空切片只会给每封邮件的 JSON 加一个 "attachments":[]
		if as := byMail[m.ID]; len(as) > 0 {
			m.Attachments = as
		}
	}
}