MailUI4Agents/gateway/internal/middleware/auth.go

package middleware

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

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

type contextKey string

const AgentNameKey contextKey = "agent_name"

// bearerToken 从 Authorization: Bearer <token> 取出令牌，缺失时返回空串。
func bearerToken(r *http.Request) string {
	h := r.Header.Get("Authorization")
	if h == "" {
		return ""
	}
	const p = "Bearer "
	if len(h) > len(p) && strings.EqualFold(h[:len(p)], p) {
		return strings.TrimSpace(h[len(p):])
	}
	return ""
}

// BearerToken 导出给 handler 层用（注册接口不过中间件，需要自己取密钥）。
func BearerToken(r *http.Request) string { return bearerToken(r) }

// keyAuthError 把密钥校验错误翻译成对外文案。
// 「不存在」与「已使用/已过期」区分开：前两者是拿错了密钥，后者是密钥生命周期到了，
// 运维需要据此判断该重新签发还是该检查配置。
func keyAuthError(err error) string {
	switch {
	case errors.Is(err, repo.ErrKeyUsed):
		return `{"error":"密钥已使用（一次性密钥只能用一次）"}`
	case errors.Is(err, repo.ErrKeyExpired):
		return `{"error":"密钥已过期"}`
	default:
		return `{"error":"密钥无效"}`
	}
}

// AgentAuth 验证 Agent 身份，支持两种凭证：
//
//	Authorization: Bearer <agent_key_token>   —— 密钥认证（推荐）
//	X-Agent-Name + X-Agent-Secret             —— 旧的 name/secret 方式（兼容保留）
//
// 用户密钥（user_keys）不接受：两类密钥共享 token 命名空间但走各自的验证表，
// 因此用用户密钥调 Agent 接口只会得到「密钥无效」。
func AgentAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if token := bearerToken(r); token != "" {
			agentName, err := repo.VerifyAgentKey(r.Context(), token)
			if err != nil {
				http.Error(w, keyAuthError(err), http.StatusUnauthorized)
				return
			}
			if agentName == "" {
				// 密钥有效但尚未绑定 Agent：注册接口会用请求里的 name 落定它，
				// 其余接口无法确定调用者身份，只能拒。
				http.Error(w, `{"error":"密钥尚未绑定 Agent，请先调用 /agent/register 完成注册"}`, http.StatusForbidden)
				return
			}
			repo.HeartbeatAgent(r.Context(), agentName)
			ctx := context.WithValue(r.Context(), AgentNameKey, agentName)
			next.ServeHTTP(w, r.WithContext(ctx))
			return
		}

		agentName := r.Header.Get("X-Agent-Name")
		agentSecret := r.Header.Get("X-Agent-Secret")
		if agentName == "" || agentSecret == "" {
			http.Error(w, `{"error":"Missing Authorization: Bearer <key> or X-Agent-Name/X-Agent-Secret header"}`, http.StatusUnauthorized)
			return
		}

		agent, err := repo.VerifyAgent(r.Context(), agentName, agentSecret)
		if err != nil {
			http.Error(w, `{"error":"Invalid credentials"}`, http.StatusUnauthorized)
			return
		}

		repo.HeartbeatAgent(r.Context(), agent.Name)
		ctx := context.WithValue(r.Context(), AgentNameKey, agent.Name)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// GetAgentName 从 context 中获取 agent_name
func GetAgentName(r *http.Request) string {
	if v := r.Context().Value(AgentNameKey); v != nil {
		return v.(string)
	}
	return ""
}