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/ultisuite/ulti-backend/internal/api/apiresponse"
	"github.com/ultisuite/ulti-backend/internal/authentik"
)

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

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

// Handler exposes public Authentik flow executor endpoints for custom auth UI.
type Handler struct {
	flows *authentik.FlowSessionStore
}

func NewHandler(baseURL string) *Handler {
	return &Handler{
		flows: authentik.NewFlowSessionStore(baseURL),
	}
}

func (h *Handler) Routes() chi.Router {
	r := chi.NewRouter()
	r.Post("/flows/{slug}/start", h.StartFlow)
	r.Post("/flows/{slug}/respond", h.RespondFlow)
	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) {
	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)
	} 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
		}
		apiresponse.WriteError(w, r, http.StatusBadGateway, "flow_respond_failed", err.Error(), nil)
		return
	}

	done, denied := authentik.FlowDone(challenge)
	if done {
		clearFlowSessionCookie(w)
	}
	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)
}