pangolin/server/internal/auth/handler.go

package auth

import (
	"encoding/json"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/wangjia/pangolin/server/internal/apierr"
)

// Handler adapts the auth Service to HTTP. Route shapes follow the OpenAPI
// contract (POST /v1/auth/{code,register,login,refresh}).
type Handler struct {
	svc *Service
}

// NewHandler builds a Handler.
func NewHandler(svc *Service) *Handler { return &Handler{svc: svc} }

// RegisterRoutes mounts the four public auth endpoints onto r. The caller is
// expected to mount this group WITHOUT the bearer-auth middleware.
func (h *Handler) RegisterRoutes(r chi.Router) {
	r.Post("/auth/code", h.SendCode)
	r.Post("/auth/register", h.Register)
	r.Post("/auth/login", h.Login)
	r.Post("/auth/refresh", h.Refresh)
	r.Post("/auth/logout", h.Logout)
}

// Logout handles POST /v1/auth/logout. The refresh token to revoke is taken from
// the X-Refresh-Token header (the access token in Authorization is not enough —
// revocation keys on the refresh JTI). Always 204; the client clears its session
// regardless of the server-side outcome.
func (h *Handler) Logout(w http.ResponseWriter, r *http.Request) {
	h.svc.Logout(r.Context(), r.Header.Get("X-Refresh-Token"))
	w.WriteHeader(http.StatusNoContent)
}

// ---- request/response bodies (mirror openapi.yaml) ----

type sendCodeRequest struct {
	Email string `json:"email"`
}

type registerRequest struct {
	Email    string `json:"email"`
	Code     string `json:"code"`
	Password string `json:"password"`
}

type loginRequest struct {
	Email    string `json:"email"`
	Password string `json:"password"`
}

type refreshRequest struct {
	RefreshToken string `json:"refresh_token"`
}

type tokenPairResponse struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	ExpiresIn    int    `json:"expires_in"`
}

// SendCode handles POST /v1/auth/code.
func (h *Handler) SendCode(w http.ResponseWriter, r *http.Request) {
	var req sendCodeRequest
	if !decodeJSON(w, r, &req) {
		return
	}
	retryAfter, apiErr := h.svc.SendCode(r.Context(), req.Email, clientIP(r))
	if apiErr != nil {
		writeAPIErr(w, apiErr, retryAfter)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// Register handles POST /v1/auth/register.
func (h *Handler) Register(w http.ResponseWriter, r *http.Request) {
	var req registerRequest
	if !decodeJSON(w, r, &req) {
		return
	}
	pair, apiErr := h.svc.Register(r.Context(), req.Email, req.Code, req.Password)
	if apiErr != nil {
		writeAPIErr(w, apiErr, 0)
		return
	}
	writeTokenPair(w, pair)
}

// Login handles POST /v1/auth/login.
func (h *Handler) Login(w http.ResponseWriter, r *http.Request) {
	var req loginRequest
	if !decodeJSON(w, r, &req) {
		return
	}
	out, retryAfter, apiErr := h.svc.Login(r.Context(), req.Email, req.Password, clientIP(r))
	if apiErr != nil {
		writeAPIErr(w, apiErr, retryAfter)
		return
	}
	// TOTP-enabled accounts get a pending token instead of tokens; the web
	// user-center then completes login at /auth/login/totp. The app (no TOTP)
	// always receives the flat token pair below.
	if out.PendingToken != "" {
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]any{
			"totp_required": true,
			"pending_token": out.PendingToken,
		})
		return
	}
	writeTokenPair(w, out.Tokens)
}

// Refresh handles POST /v1/auth/refresh.
func (h *Handler) Refresh(w http.ResponseWriter, r *http.Request) {
	var req refreshRequest
	if !decodeJSON(w, r, &req) {
		return
	}
	pair, apiErr := h.svc.Refresh(r.Context(), req.RefreshToken)
	if apiErr != nil {
		writeAPIErr(w, apiErr, 0)
		return
	}
	writeTokenPair(w, pair)
}

// ---- helpers ----

// decodeJSON decodes the request body, writing a 400 on malformed input.
// Returns false if the caller should stop (error already written).
func decodeJSON(w http.ResponseWriter, r *http.Request, dst interface{}) bool {
	dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<16))
	dec.DisallowUnknownFields()
	if err := dec.Decode(dst); err != nil {
		writeAPIErr(w, ErrInvalidRequest, 0)
		return false
	}
	return true
}

// writeTokenPair writes a 200 TokenPair body.
func writeTokenPair(w http.ResponseWriter, pair *TokenPair) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(tokenPairResponse{
		AccessToken:  pair.AccessToken,
		RefreshToken: pair.RefreshToken,
		ExpiresIn:    pair.ExpiresIn,
	})
}

// writeAPIErr maps an apierr.Error to the right HTTP status, attaching a
// Retry-After header when provided.
func writeAPIErr(w http.ResponseWriter, e *apierr.Error, retryAfter time.Duration) {
	status := statusFor(e)
	if retryAfter > 0 {
		secs := int(retryAfter.Seconds())
		if secs < 1 {
			secs = 1
		}
		w.Header().Set("Retry-After", strconv.Itoa(secs))
	}
	apierr.WriteJSON(w, status, e)
}

// statusFor maps auth error codes to HTTP status codes.
func statusFor(e *apierr.Error) int {
	switch e.Code {
	case ErrInvalidRequest.Code, ErrEmailDisposable.Code, ErrCodeInvalid.Code:
		return http.StatusBadRequest
	case ErrEmailExists.Code:
		return http.StatusConflict
	case ErrRateLimited.Code, ErrAccountLocked.Code:
		return http.StatusTooManyRequests
	case ErrInvalidCredentials.Code, ErrInvalidToken.Code, ErrUnauthorized.Code:
		return http.StatusUnauthorized
	case ErrAccountBanned.Code:
		return http.StatusForbidden
	default:
		return http.StatusInternalServerError
	}
}

// clientIP extracts the best-effort client IP, honouring X-Forwarded-For and
// X-Real-IP set by the edge proxy. Used only as a rate-limit key — never logged.
func clientIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		// First entry is the original client.
		if i := strings.IndexByte(xff, ','); i >= 0 {
			return strings.TrimSpace(xff[:i])
		}
		return strings.TrimSpace(xff)
	}
	if xr := r.Header.Get("X-Real-IP"); xr != "" {
		return strings.TrimSpace(xr)
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}