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

package authapi

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

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

type flowCompleteRequest struct {
	ReturnTo string `json:"returnTo"`
}

type flowCompleteResponse struct {
	RedirectURL string `json:"redirectUrl"`
}

func (h *Handler) CompleteFlow(w http.ResponseWriter, r *http.Request) {
	returnTo := strings.TrimSpace(r.URL.Query().Get("returnTo"))
	if returnTo == "" && r.Body != nil {
		var req flowCompleteRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		returnTo = strings.TrimSpace(req.ReturnTo)
	}
	if returnTo == "" {
		returnTo = "/mail/inbox"
	}
	if !strings.HasPrefix(returnTo, "/") || strings.HasPrefix(returnTo, "//") {
		apiresponse.WriteError(w, r, http.StatusBadRequest, "invalid_return_to", "returnTo must be a relative path", nil)
		return
	}

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

	slug := FlowAuthentication
	cookies, capturedCallback, err := h.flows.CompleteSessionOAuth(r.Context(), sessionID, slug)
	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.ErrFlowSessionNotCompleted) {
			apiresponse.WriteError(w, r, http.StatusConflict, "flow_not_completed", "authentication flow is not finished yet", 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_complete_failed", err.Error(), nil)
		return
	}

	clearFlowSessionCookie(w)

	if !authentik.SessionAuthenticated(cookies) {
		msg := "authentik session not authenticated"
		if r.Method == http.MethodGet {
			base := suiteAuthAppOrigin(h.appURL)
			http.Redirect(w, r, base+"/login?error="+url.QueryEscape("oidc_bridge_failed:"+msg), http.StatusFound)
			return
		}
		apiresponse.WriteError(w, r, http.StatusBadGateway, "oidc_bridge_failed", msg, nil)
		return
	}

	// The embedded flow already followed the OIDC authorize continuation within its authenticated
	// session and captured the callback URL (carrying the authorization code). The browser holds the
	// matching PKCE verifier + state cookies (set by the flow context endpoint), so we just redirect
	// it to the Next.js callback which performs the token exchange.
	var redirectURL string
	if capturedCallback != "" {
		redirectURL = capturedCallback
	} else {
		// Fallback: authorize server-side using the authenticated Authentik session.
		callbackURL, pkceVerifier, state, bErr := h.oidc.callbackURL(r.Context(), cookies)
		if bErr != nil {
			if r.Method == http.MethodGet {
				base := suiteAuthAppOrigin(h.appURL)
				http.Redirect(w, r, base+"/login?error="+url.QueryEscape("oidc_bridge_failed:"+bErr.Error()), http.StatusFound)
				return
			}
			apiresponse.WriteError(w, r, http.StatusBadGateway, "oidc_bridge_failed", bErr.Error(), nil)
			return
		}
		setOAuthBridgeCookies(w, r, pkceVerifier, state, returnTo)
		redirectURL = callbackURL
	}

	if r.Method == http.MethodGet {
		http.Redirect(w, r, redirectURL, http.StatusFound)
		return
	}

	apiresponse.WriteJSON(w, http.StatusOK, flowCompleteResponse{RedirectURL: redirectURL})
}

func setOAuthBridgeCookies(w http.ResponseWriter, r *http.Request, verifier, state, returnTo string) {
	secure := requestIsHTTPS(r)
	maxAge := 600
	set := func(name, value string) {
		http.SetCookie(w, &http.Cookie{
			Name:     name,
			Value:    value,
			Path:     "/",
			MaxAge:   maxAge,
			HttpOnly: true,
			SameSite: http.SameSiteLaxMode,
			Secure:   secure,
		})
	}
	set("ulti_pkce_verifier", verifier)
	set("ulti_oauth_state", state)
	set("ulti_auth_return", returnTo)
}

// suiteAuthAppOrigin is the suite root (auth routes live at /api/auth/*, not under /mail).
func suiteAuthAppOrigin(appURL string) string {
	base := strings.TrimRight(strings.TrimSpace(appURL), "/")
	if strings.HasSuffix(base, "/mail") {
		base = strings.TrimSuffix(base, "/mail")
	}
	if base == "" {
		return "http://localhost:3004"
	}
	return base
}

func buildLoginRedirectURL(appURL, returnTo string) string {
	base := suiteAuthAppOrigin(appURL)
	params := url.Values{}
	params.Set("returnTo", returnTo)
	params.Set("bridge", "1")
	return base + "/api/auth/login?" + params.Encode()
}

func setBrowserAuthentikCookies(w http.ResponseWriter, r *http.Request, stored []authentik.SerializedCookie) {
	secure := requestIsHTTPS(r)
	for _, c := range authentik.BrowserAuthentikCookies(stored) {
		if secure {
			c.Secure = true
		}
		http.SetCookie(w, c)
	}
}

func requestIsHTTPS(r *http.Request) bool {
	if r.TLS != nil {
		return true
	}
	if strings.EqualFold(r.Header.Get("X-Forwarded-Proto"), "https") {
		return true
	}
	if strings.EqualFold(r.Header.Get("X-Forwarded-Ssl"), "on") {
		return true
	}
	return false
}

func forwardFlowCookies(w http.ResponseWriter, r *http.Request, stored []authentik.SerializedCookie) {
	setBrowserAuthentikCookies(w, r, stored)
}