Agent-System-Backend/internal/server/auth.go

package server

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

	"agent-system-backend/internal/store"
)

// ownerCtxKey 携带租户身份（认证中间件注入；未启用认证 = store.OwnerDefault）。
type ownerCtxKey struct{}

// authViaCtxKey 携带凭证通道（AuthViaKey / AuthViaSession；未启用认证与豁免路径为空）。
type authViaCtxKey struct{}

// sessionScopeCtxKey 携带会话 scope（full / set_password_only）。
type sessionScopeCtxKey struct{}

// 凭证通道值：/api/me 的 authVia 字段；前端据此区分「改密门」只对会话登录者放行。
const (
	AuthViaKey     = "key"
	AuthViaSession = "session"
)

// Authenticator 是 API key 认证（阶段 1：owner 即租户）。
// Enabled=false（嵌入式档默认，回环本地单用户）时全部请求回落默认租户。
// SSE 的 EventSource 无法带 Authorization 头，允许 ?key= 查询参数兜底。
type Authenticator struct {
	st      *store.Store
	Enabled bool
}

func NewAuthenticator(st *store.Store, enabled bool) *Authenticator {
	return &Authenticator{st: st, Enabled: enabled}
}

// SessionCookie 会话 cookie 名（HttpOnly；SSE 的 EventSource 同源自动携带）。
const SessionCookie = "as_session"

func (a *Authenticator) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if a == nil || !a.Enabled {
			next.ServeHTTP(w, r.WithContext(withOwner(r, store.OwnerDefault)))
			return
		}
		p := r.URL.Path
		// 豁免：静态前端、登录、一次性 Key 兑换（body 专用，不走 Bearer）。
		if !strings.HasPrefix(p, "/api/") || p == "/api/auth/login" || p == "/api/auth/login-key/redeem" || p == "/api/auth/recovery-grants/redeem" {
			next.ServeHTTP(w, r)
			return
		}
		key := bearerKey(r)
		if key == "" {
			key = r.URL.Query().Get("key")
		}
		if key != "" && sessionToken(r) != "" {
			http.Error(w, "conflicting credentials", http.StatusUnauthorized)
			return
		}
		// 一次性登录 Key 不得经 Bearer / ?key= 使用。
		if strings.HasPrefix(key, "alk_") {
			http.Error(w, "unauthorized: one-time login keys must use /api/auth/login-key/redeem", http.StatusUnauthorized)
			return
		}
		var owner string
		var via string
		var scope string
		ok := false
		if key != "" {
			if o, keyOK := a.st.AuthOwner(key); keyOK {
				owner, via, scope, ok = o, AuthViaKey, store.SessionScopeFull, true
			} else {
				http.Error(w, "unauthorized: invalid or revoked API key", http.StatusUnauthorized)
				return
			}
		} else if o, sc, sessOK := cookieSession(r, a.st); sessOK {
			owner, via, scope, ok = o, AuthViaSession, sc, true
		}
		if !ok {
			http.Error(w, "unauthorized: login required", http.StatusUnauthorized)
			return
		}
		ctx := withAuthVia(withOwner(r, owner), via)
		ctx = context.WithValue(ctx, sessionScopeCtxKey{}, scope)
		if err := a.enforcePasswordGate(w, r, owner, via, scope, p); err != nil {
			return
		}
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func (a *Authenticator) enforcePasswordGate(w http.ResponseWriter, r *http.Request, owner, via, scope, path string) error {
	allow := path == "/api/me" || path == "/api/auth/password" || path == "/api/auth/logout" ||
		path == "/api/health"
	if scope == store.SessionScopeSetPasswordOnly {
		if !allow {
			http.Error(w, "forbidden: set-password session only", http.StatusForbidden)
			return errGate
		}
		return nil
	}
	st, err := a.st.GetUserAuthState(owner)
	if err != nil || !st.Enabled {
		http.Error(w, "unauthorized: account unavailable", http.StatusUnauthorized)
		return errGate
	}
	// 长期 API Key：仅当 must_change（已有密码待改）时拒绝业务，防止绕过改密门；
	// 无密码的纯集成账户仍可走 Key（与历史 apikey 语义兼容）。
	if via == AuthViaKey {
		if st.MustChangePassword && !allow {
			http.Error(w, "forbidden: password change required", http.StatusForbidden)
			return errGate
		}
		return nil
	}
	// 会话：待改密或尚未设密 → 只允许身份/设密/退出。
	if (st.MustChangePassword || !st.HasPassword) && !allow {
		http.Error(w, "forbidden: password setup required", http.StatusForbidden)
		return errGate
	}
	return nil
}

var errGate = errorsGate{}

type errorsGate struct{}

func (errorsGate) Error() string { return "gate" }

func cookieSession(r *http.Request, st *store.Store) (owner, scope string, ok bool) {
	if tok := sessionToken(r); tok != "" {
		return st.SessionInfo(tok)
	}
	return "", "", false
}

// cookieOwner 会话 cookie 有效时返回其 owner。
func cookieOwner(r *http.Request, st *store.Store) (string, bool) {
	if tok := sessionToken(r); tok != "" {
		return st.SessionOwner(tok)
	}
	return "", false
}

// sessionToken 从请求取会话明文（无会话返回空）。
func sessionToken(r *http.Request) string {
	if c, err := r.Cookie(SessionCookie); err == nil {
		return c.Value
	}
	return ""
}

func withOwner(r *http.Request, owner string) context.Context {
	return context.WithValue(r.Context(), ownerCtxKey{}, owner)
}

func withAuthVia(ctx context.Context, via string) context.Context {
	return context.WithValue(ctx, authViaCtxKey{}, via)
}

// AuthVia 返回当前请求的凭证通道（未启用认证/豁免路径为空串）。
func AuthVia(r *http.Request) string {
	if v, ok := r.Context().Value(authViaCtxKey{}).(string); ok {
		return v
	}
	return ""
}

// SessionScope 返回当前会话 scope（非会话为空）。
func SessionScope(r *http.Request) string {
	if v, ok := r.Context().Value(sessionScopeCtxKey{}).(string); ok {
		return v
	}
	return ""
}

func bearerKey(r *http.Request) string {
	h := r.Header.Get("Authorization")
	if strings.HasPrefix(h, "Bearer ") {
		return strings.TrimSpace(strings.TrimPrefix(h, "Bearer "))
	}
	return ""
}

// owner 从请求上下文取租户身份（未注入回落默认）。
func (s *Server) owner(r *http.Request) string {
	if v, ok := r.Context().Value(ownerCtxKey{}).(string); ok && v != "" {
		return v
	}
	return store.OwnerDefault
}