ultisuite-backend/internal/api/auth/handlers.go

package authapi

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

	"github.com/go-chi/chi/v5"
	"github.com/redis/go-redis/v9"

	"github.com/ultisuite/ulti-backend/internal/api/apiresponse"
	"github.com/ultisuite/ulti-backend/internal/authentik"
)

const (
	FlowEnrollment        = "ulti-enrollment"
	FlowRecovery          = "ulti-recovery"
	FlowAuthentication    = "default-authentication-flow"
	flowSessionCookie     = "ulti_flow_session"
	flowSessionMaxAge     = 20 * time.Minute
)

var allowedFlowSlugs = map[string]struct{}{
	FlowEnrollment:     {},
	FlowRecovery:       {},
	FlowAuthentication: {},
}

// Handler exposes public Authentik flow executor endpoints for custom auth UI.
type Handler struct {
	flows        *authentik.FlowSessionStore
	limiter      *FlowRateLimiter
	appURL       string
	oidcClientID string
	oidc         *oidcBridge
}

func NewHandler(baseURL, appURL, oidcIssuer, oidcClientID, oidcClientSecret string, rdb *redis.Client) *Handler {
	return &Handler{
		flows:        authentik.NewFlowSessionStore(baseURL, rdb),
		limiter:      NewFlowRateLimiter(rdb),
		appURL:       appURL,
		oidcClientID: strings.TrimSpace(oidcClientID),
		oidc:         newOIDCBridge(buildOIDCBridgeConfig(oidcIssuer, oidcClientID, oidcClientSecret, appURL, baseURL)),
	}
}

func (h *Handler) Routes() chi.Router {
	r := chi.NewRouter()
	r.Get("/flows/authentication/context", h.PrepareOAuthFlowContext)
	r.Post("/flows/{slug}/start", h.StartFlow)
	r.Post("/flows/{slug}/respond", h.RespondFlow)
	r.Get("/flows/complete", h.CompleteFlow)
	r.Post("/flows/complete", h.CompleteFlow)
	return r
}

type flowStartResponse struct {
	SessionID string                  `json:"sessionId"`
	Challenge authentik.FlowChallenge `json:"challenge"`
	Done      bool                    `json:"done"`
	Denied    bool                    `json:"denied"`
}

type flowRespondRequest struct {
	Payload map[string]any `json:"payload"`
}

func (h *Handler) StartFlow(w http.ResponseWriter, r *http.Request) {
	if !h.limiter.Allow(r) {
		apiresponse.WriteError(w, r, http.StatusTooManyRequests, apiresponse.CodeRateLimited, "too many flow start attempts", nil)
		return
	}

	slug, ok := validateFlowSlug(w, r, chi.URLParam(r, "slug"))
	if !ok {
		return
	}
	query := strings.TrimSpace(r.URL.Query().Get("query"))

	sessionID, challenge, err := h.flows.Start(r.Context(), slug, query)
	if err != nil {
		apiresponse.WriteError(w, r, http.StatusBadGateway, "flow_start_failed", err.Error(), nil)
		return
	}
	done, denied := authentik.FlowDone(challenge)
	if !done {
		setFlowSessionCookie(w, sessionID)
		if cookies, err := h.flows.SessionCookies(r.Context(), sessionID); err == nil {
			forwardFlowCookies(w, r, cookies)
		}
	} else {
		clearFlowSessionCookie(w)
	}
	writeFlowJSON(w, http.StatusOK, flowStartResponse{
		SessionID: sessionID,
		Challenge: challenge,
		Done:      done,
		Denied:    denied,
	})
}

func (h *Handler) RespondFlow(w http.ResponseWriter, r *http.Request) {
	slug, ok := validateFlowSlug(w, r, chi.URLParam(r, "slug"))
	if !ok {
		return
	}

	sessionID := readFlowSessionCookie(r)
	if sessionID == "" {
		apiresponse.WriteError(w, r, http.StatusUnauthorized, "flow_session_missing", "flow session cookie required", nil)
		return
	}

	var req flowRespondRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		apiresponse.WriteError(w, r, http.StatusBadRequest, "invalid_json", "invalid request body", nil)
		return
	}
	if req.Payload == nil {
		apiresponse.WriteError(w, r, http.StatusBadRequest, "invalid_payload", "payload required", nil)
		return
	}

	query := strings.TrimSpace(r.URL.Query().Get("query"))
	challenge, err := h.flows.Respond(r.Context(), sessionID, slug, query, req.Payload)
	if err != nil {
		if errors.Is(err, authentik.ErrFlowSessionNotFound) {
			clearFlowSessionCookie(w)
			apiresponse.WriteError(w, r, http.StatusGone, "flow_session_expired", "flow session expired; restart the flow", nil)
			return
		}
		if errors.Is(err, authentik.ErrFlowSessionSlugMismatch) {
			apiresponse.WriteError(w, r, http.StatusConflict, "flow_session_mismatch", "flow slug does not match active session", nil)
			return
		}
		if errors.Is(err, authentik.ErrFlowSessionAlreadyCompleted) {
			apiresponse.WriteError(w, r, http.StatusConflict, "flow_already_completed", "flow already completed", nil)
			return
		}
		apiresponse.WriteError(w, r, http.StatusBadGateway, "flow_respond_failed", err.Error(), nil)
		return
	}

	done, denied := authentik.FlowDone(challenge)
	if done {
		if slug == FlowAuthentication {
			if denied {
				clearFlowSessionCookie(w)
				h.flows.Delete(r.Context(), sessionID)
			} else if cookies, err := h.flows.SessionCookies(r.Context(), sessionID); err == nil {
				// Push authenticated Authentik session to browser before OIDC bridge.
				forwardFlowCookies(w, r, cookies)
			}
		} else {
			clearFlowSessionCookie(w)
			h.flows.Delete(r.Context(), sessionID)
		}
	} else if cookies, err := h.flows.SessionCookies(r.Context(), sessionID); err == nil {
		forwardFlowCookies(w, r, cookies)
	}

	writeFlowJSON(w, http.StatusOK, flowStartResponse{
		SessionID: sessionID,
		Challenge: challenge,
		Done:      done,
		Denied:    denied,
	})
}

func validateFlowSlug(w http.ResponseWriter, r *http.Request, slug string) (string, bool) {
	slug = strings.TrimSpace(slug)
	if slug == "" {
		apiresponse.WriteError(w, r, http.StatusBadRequest, "invalid_flow", "flow slug required", nil)
		return "", false
	}
	if _, ok := allowedFlowSlugs[slug]; !ok {
		apiresponse.WriteError(w, r, http.StatusNotFound, "flow_not_allowed", "flow slug not allowed", nil)
		return "", false
	}
	return slug, true
}

func setFlowSessionCookie(w http.ResponseWriter, sessionID string) {
	http.SetCookie(w, &http.Cookie{
		Name:     flowSessionCookie,
		Value:    sessionID,
		Path:     "/api/v1/auth/flows",
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   int(flowSessionMaxAge.Seconds()),
	})
}

func clearFlowSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     flowSessionCookie,
		Value:    "",
		Path:     "/api/v1/auth/flows",
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   -1,
	})
}

func readFlowSessionCookie(r *http.Request) string {
	c, err := r.Cookie(flowSessionCookie)
	if err != nil || c == nil {
		return ""
	}
	return strings.TrimSpace(c.Value)
}

func writeFlowJSON(w http.ResponseWriter, status int, payload flowStartResponse) {
	apiresponse.WriteJSON(w, status, payload)
}