NixMsg/internal/admin/auth.go

package admin

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"net/http"
	"strings"
	"time"

	"git.asio.asia/nixevol/NixMsg/internal/auth"
	"git.asio.asia/nixevol/NixMsg/internal/httpx"
)

type authKind string

const (
	authCookie authKind = "cookie"
	authToken  authKind = "token"
)

type principal struct {
	Kind      authKind
	TokenName string
	TokenID   string
	Session   string
}

type ctxKey int

const principalKey ctxKey = 1

func principalFrom(ctx context.Context) (principal, bool) {
	p, ok := ctx.Value(principalKey).(principal)
	return p, ok
}

func (h *Handler) auth(next http.HandlerFunc) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ip := httpx.ClientIP(r, h.trusted)

		if locked, retry := h.locks.Check(auth.LockKey{Kind: auth.LockAdminIP, IP: ip}); locked {
			w.Header().Set("Retry-After", formatRetryAfter(retry))
			httpx.WriteError(w, http.StatusTooManyRequests, "rate_limited", "登录已锁定，请稍后再试")
			return
		}

		p, errCode, errMsg, status := h.authenticate(r, ip)
		if status != 0 {
			if status == http.StatusTooManyRequests {
				w.Header().Set("Retry-After", "300")
			}
			httpx.WriteError(w, status, errCode, errMsg)
			return
		}

		if p.Kind == authCookie && isMutating(r.Method) {
			if r.Header.Get(csrfHeader) != csrfValue {
				httpx.WriteError(w, http.StatusForbidden, "forbidden", "缺少 X-Nixmsg-Request 头")
				return
			}
		}

		if p.Kind == authToken && isTokenForbiddenPath(r.URL.Path) {
			httpx.WriteError(w, http.StatusForbidden, "forbidden", "API 令牌无权访问该接口")
			return
		}

		ctx := context.WithValue(r.Context(), principalKey, p)
		next(w, r.WithContext(ctx))
	})
}

func isMutating(method string) bool {
	switch method {
	case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete:
		return true
	default:
		return false
	}
}

func isTokenForbiddenPath(path string) bool {
	if path == "/api/admin/password" {
		return true
	}
	return path == "/api/admin/tokens" || strings.HasPrefix(path, "/api/admin/tokens/")
}

func (h *Handler) authenticate(r *http.Request, ip string) (principal, string, string, int) {
	authz := r.Header.Get("Authorization")
	if strings.HasPrefix(strings.ToLower(authz), "bearer ") {
		raw := strings.TrimSpace(authz[len("Bearer "):])
		if raw == "" || !h.tokens.LooksLikeAPIToken(raw) {
			return authFail(h, ip)
		}
		hash := h.tokens.HashToken(raw)
		info, err := h.lookupAPIToken(r.Context(), hash)
		if err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return authFail(h, ip)
			}
			return principal{}, "internal", "内部错误", http.StatusInternalServerError
		}
		if !info.Enabled {
			return authFail(h, ip)
		}
		h.touchLastUsed(r.Context(), info.ID)
		return principal{Kind: authToken, TokenName: info.Name, TokenID: info.ID}, "", "", 0
	}

	c, err := r.Cookie(cookieName)
	if err != nil || c.Value == "" {
		return principal{}, "unauthorized", "未登录", http.StatusUnauthorized
	}
	hashHex := hashSessionHex(c.Value)
	ok, err := h.sessionValid(r.Context(), hashHex)
	if err != nil {
		return principal{}, "internal", "内部错误", http.StatusInternalServerError
	}
	if !ok {
		return principal{}, "unauthorized", "未登录", http.StatusUnauthorized
	}
	return principal{Kind: authCookie, Session: c.Value}, "", "", 0
}

func authFail(h *Handler, ip string) (principal, string, string, int) {
	if locked, _ := h.locks.Fail(auth.LockKey{Kind: auth.LockAdminIP, IP: ip}); locked {
		return principal{}, "rate_limited", "登录已锁定，请稍后再试", http.StatusTooManyRequests
	}
	return principal{}, "unauthorized", "令牌无效", http.StatusUnauthorized
}

func hashSessionHex(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

func newSessionToken() (plain string, hashHex string, err error) {
	buf := make([]byte, 32)
	if _, err = rand.Read(buf); err != nil {
		return "", "", err
	}
	plain = base64.RawURLEncoding.EncodeToString(buf)
	sum := sha256.Sum256([]byte(plain))
	return plain, hex.EncodeToString(sum[:]), nil
}

func formatRetryAfter(d time.Duration) string {
	sec := int(d.Seconds())
	if sec < 1 {
		sec = 1
	}
	return itoa(sec)
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [16]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}

func actorString(p principal) string {
	if p.Kind == authToken {
		return "token:" + p.TokenName
	}
	return "admin"
}