MailUI4Agents/server/internal/middleware/user.go

package middleware

import (
	"context"
	"net/http"
	"strings"

	"github.com/agentmail/gateway/internal/config"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
)

const UserKey contextKey = "auth_user"

// SetSessionCookie 写入登录 Cookie
func SetSessionCookie(w http.ResponseWriter, token string, maxAge int) {
	http.SetCookie(w, &http.Cookie{
		Name:     config.C.CookieName,
		Value:    token,
		Path:     "/",
		HttpOnly: true,
		Secure:   config.C.SecureCookie,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   maxAge,
	})
}

// ClearSessionCookie 清除登录 Cookie
func ClearSessionCookie(w http.ResponseWriter) {
	SetSessionCookie(w, "", -1)
}

// SessionToken 从请求中取出登录令牌
func SessionToken(r *http.Request) string {
	c, err := r.Cookie(config.C.CookieName)
	if err != nil {
		return ""
	}
	return c.Value
}

// QueryToken 从 ?access_token= 取出令牌。
//
// 仅为浏览器 EventSource 存在：它不支持自定义请求头，因此订阅 SSE 时
// 除了 Cookie 就只剩 query 一条路。代价是令牌会进访问日志，
// 所以只在 SSE 端点启用，其余接口一律要求 Authorization 头。
func QueryToken(r *http.Request) string {
	return strings.TrimSpace(r.URL.Query().Get("access_token"))
}

// UserAuth 校验人类用户登录态，把 *models.User 注入 context。
// 支持两种凭证：浏览器 Cookie，或 Authorization: Bearer <user_key_token>（第三方客户端）。
func UserAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u, err := resolve(r)
		if err != nil {
			// 只有 Cookie 路径才清 Cookie；密钥认证失败不应频带浏览器会话
			if bearerToken(r) == "" {
				ClearSessionCookie(w)
			}
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte(`{"error":"not authenticated"}`))
			return
		}
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), UserKey, u)))
	})
}

// AdminOnly 叠在 UserAuth 之后，要求 role = admin
func AdminOnly(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u := GetUser(r)
		if u == nil || !u.IsAdmin() {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte(`{"error":"admin only"}`))
			return
		}
		next.ServeHTTP(w, r)
	})
}

// UserAuthAllowQueryToken 与 UserAuth 相同，但额外接受 ?access_token=。
//
// 只给那些【由浏览器直接发起、无法设置请求头】的端点用（附件下载的 <a download>）。
// URL 里的令牌会进访问日志与 Referer，所以不能全局开启。
func UserAuthAllowQueryToken(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u, err := resolve(r)
		if err != nil {
			if token := QueryToken(r); token != "" {
				if ku, kErr := repo.VerifyUserKey(r.Context(), token); kErr == nil {
					u, err = ku, nil
				}
			}
		}
		if err != nil {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte(`{"error":"not authenticated"}`))
			return
		}
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), UserKey, u)))
	})
}

// GetUser 从 context 取登录用户；未登录返回 nil
func GetUser(r *http.Request) *models.User {
	if v := r.Context().Value(UserKey); v != nil {
		if u, ok := v.(*models.User); ok {
			return u
		}
	}
	return nil
}

// GetUserName 便捷取登录用户名
func GetUserName(r *http.Request) string {
	if u := GetUser(r); u != nil {
		return u.Username
	}
	return ""
}

// OptionalUser 解析登录态但不拦截（SSE 等需要区分匿名/登录的场景）
func OptionalUser(r *http.Request) *models.User {
	u, err := resolve(r)
	if err != nil {
		return nil
	}
	return u
}

// OptionalUserWithQuery 在 resolve 的三种凭证之外额外接受 ?access_token=。
// 仅 SSE 用：浏览器 EventSource 无法带自定义头。
func OptionalUserWithQuery(r *http.Request) *models.User {
	if u := OptionalUser(r); u != nil {
		return u
	}
	if token := QueryToken(r); token != "" {
		if u, err := repo.VerifyUserKey(r.Context(), token); err == nil {
			return u
		}
	}
	return nil
}

// resolve 解析调用者身份：Cookie 优先，其次 Bearer 用户密钥。
//
// 用户密钥只能走到这里（/me/* 与会话级接口），Agent 密钥只能走 AgentAuth，
// 两者各自查自己的表，因此拿 Agent 密钥读人类邮箱会得到 not authenticated。
func resolve(r *http.Request) (*models.User, error) {
	if token := SessionToken(r); token != "" {
		return repo.ResolveUserSession(r.Context(), token)
	}
	if token := bearerToken(r); token != "" {
		return repo.VerifyUserKey(r.Context(), token)
	}
	return nil, repo.ErrSessionInvalid
}