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

package authapi

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"

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

type oidcBridgeConfig struct {
	clientID      string
	clientSecret  string
	redirectURI   string
	authorizeURL  string
	authentikBase string
}

type oidcBridge struct {
	cfg oidcBridgeConfig
}

func newOIDCBridge(cfg oidcBridgeConfig) *oidcBridge {
	return &oidcBridge{cfg: cfg}
}

func buildOIDCBridgeConfig(_issuer, clientID, clientSecret, appURL, authentikAPIURL string) oidcBridgeConfig {
	base := suiteAuthAppOrigin(appURL)
	redirectURI := base + "/api/auth/callback"
	authentikBase := authentik.APIBaseURL(authentikAPIURL)
	return oidcBridgeConfig{
		clientID:      strings.TrimSpace(clientID),
		clientSecret:  strings.TrimSpace(clientSecret),
		redirectURI:   redirectURI,
		authorizeURL:  strings.TrimSuffix(authentikBase, "/") + "/application/o/authorize/",
		authentikBase: authentikBase,
	}
}

// callbackURL performs a server-side OIDC authorize using the Authentik session from the embedded flow.
func (b *oidcBridge) callbackURL(ctx context.Context, cookies []authentik.SerializedCookie) (callbackURL, pkceVerifier, state string, err error) {
	if b == nil || b.cfg.clientID == "" || b.cfg.clientSecret == "" || b.cfg.redirectURI == "" {
		return "", "", "", fmt.Errorf("oidc bridge not configured")
	}
	if !authentik.SessionAuthenticated(cookies) {
		return "", "", "", fmt.Errorf("authentik session not authenticated")
	}

	verifier, challenge, err := newPKCEPair()
	if err != nil {
		return "", "", "", err
	}
	state, err = randomOAuthState()
	if err != nil {
		return "", "", "", err
	}

	jar, err := cookiejar.New(nil)
	if err != nil {
		return "", "", "", err
	}
	cookieURL, err := url.Parse(strings.TrimSuffix(b.cfg.authentikBase, "/") + "/")
	if err != nil {
		return "", "", "", err
	}
	jar.SetCookies(cookieURL, authentik.BrowserAuthentikCookies(cookies))

	var captured *url.URL
	client := &http.Client{
		Jar: jar,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			slog.Info("oidc bridge redirect", "to", redactQuery(req.URL))
			if strings.HasPrefix(req.URL.String(), b.cfg.redirectURI) {
				captured = req.URL
				return http.ErrUseLastResponse
			}
			if len(via) >= 15 {
				return fmt.Errorf("too many oauth redirects")
			}
			return nil
		},
	}

	params := url.Values{}
	params.Set("client_id", b.cfg.clientID)
	params.Set("redirect_uri", b.cfg.redirectURI)
	params.Set("response_type", "code")
	params.Set("scope", "openid profile email offline_access")
	params.Set("state", state)
	params.Set("code_challenge", challenge)
	params.Set("code_challenge_method", "S256")

	authURL := b.cfg.authorizeURL + "?" + params.Encode()
	captured, err = b.runAuthorize(ctx, client, jar, authURL, captured)
	if err != nil {
		return "", "", "", err
	}
	if captured == nil {
		return "", "", "", fmt.Errorf("oidc authorize did not return an authorization code")
	}
	if captured.Query().Get("code") == "" {
		return "", "", "", fmt.Errorf("oidc authorize missing code (session may be unauthenticated)")
	}
	if captured.Query().Get("state") != state {
		return "", "", "", fmt.Errorf("oidc authorize state mismatch")
	}

	return captured.String(), verifier, state, nil
}

// runAuthorize performs the authorize request and, if Authentik routes through the
// authorization flow (implicit consent), drives that flow via the executor API to obtain the code.
func (b *oidcBridge) runAuthorize(ctx context.Context, client *http.Client, jar http.CookieJar, authURL string, captured *url.URL) (*url.URL, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, authURL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	finalURL := resp.Request.URL
	loc := resp.Header.Get("Location")
	status := resp.StatusCode
	io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
	resp.Body.Close()

	if captured != nil {
		return captured, nil
	}
	if strings.HasPrefix(loc, b.cfg.redirectURI) {
		return url.Parse(loc)
	}

	// Landed on the authorization flow UI: drive it via the flow executor (implicit consent).
	if slug := authorizationFlowSlug(finalURL); slug != "" {
		slog.Info("oidc bridge driving authorization flow", "slug", slug)
		code, err := b.driveAuthorizationFlow(ctx, client, slug, finalURL)
		if err != nil {
			return nil, err
		}
		if code != nil {
			return code, nil
		}
	}

	slog.Warn("oidc bridge no code", "status", status, "final", redactQuery(finalURL), "location", loc)
	if loc != "" {
		return nil, fmt.Errorf("oidc authorize unexpected redirect to %s (status %d)", loc, status)
	}
	return nil, fmt.Errorf("oidc authorize did not return an authorization code (status %d)", status)
}

// driveAuthorizationFlow runs the provider authorization flow executor and follows its
// terminal redirect, returning the captured callback URL with the authorization code.
func (b *oidcBridge) driveAuthorizationFlow(ctx context.Context, client *http.Client, slug string, flowURL *url.URL) (*url.URL, error) {
	next := flowURL.Query().Get("next")
	execURL := strings.TrimSuffix(b.cfg.authentikBase, "/") + "/api/v3/flows/executor/" + slug + "/"
	if next != "" {
		execURL += "?query=" + url.QueryEscape(strings.TrimPrefix(next, "/auth"))
	}

	for i := 0; i < 6; i++ {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, execURL, nil)
		if err != nil {
			return nil, err
		}
		req.Header.Set("Accept", "application/json")
		resp, err := client.Do(req)
		if err != nil {
			return nil, err
		}
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<18))
		resp.Body.Close()

		var challenge map[string]any
		_ = json.Unmarshal(body, &challenge)
		comp, _ := challenge["component"].(string)
		slog.Info("oidc bridge authz flow stage", "component", comp)

		switch comp {
		case "xak-flow-redirect":
			to, _ := challenge["to"].(string)
			if strings.HasPrefix(to, b.cfg.redirectURI) {
				return url.Parse(to)
			}
			if u := b.followToCallback(ctx, client, to); u != nil {
				return u, nil
			}
			return nil, fmt.Errorf("authorization flow redirected to %s (no code)", to)
		case "ak-stage-consent":
			token, _ := challenge["token"].(string)
			postReq, err := postAuthorizeConsent(ctx, execURL, token, client)
			if err != nil {
				return nil, err
			}
			if postReq != nil {
				return postReq, nil
			}
		default:
			return nil, fmt.Errorf("authorization flow stuck at stage %q", comp)
		}
	}
	return nil, fmt.Errorf("authorization flow did not complete")
}

func postAuthorizeConsent(ctx context.Context, execURL, token string, client *http.Client) (*url.URL, error) {
	payload := map[string]any{"component": "ak-stage-consent"}
	if token != "" {
		payload["token"] = token
	}
	raw, _ := json.Marshal(payload)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, execURL, strings.NewReader(string(raw)))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Referer", execURL)
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<18))
	resp.Body.Close()
	var challenge map[string]any
	_ = json.Unmarshal(body, &challenge)
	if to, _ := challenge["to"].(string); to != "" {
		return url.Parse(to)
	}
	return nil, nil
}

func (b *oidcBridge) followToCallback(ctx context.Context, client *http.Client, to string) *url.URL {
	if to == "" {
		return nil
	}
	target := to
	if strings.HasPrefix(to, "/") {
		target = strings.TrimSuffix(b.cfg.authentikBase, "/") + to
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
	if err != nil {
		return nil
	}
	var captured *url.URL
	client.CheckRedirect = func(r *http.Request, via []*http.Request) error {
		if strings.HasPrefix(r.URL.String(), b.cfg.redirectURI) {
			captured = r.URL
			return http.ErrUseLastResponse
		}
		if len(via) >= 10 {
			return fmt.Errorf("too many redirects")
		}
		return nil
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil
	}
	loc := resp.Header.Get("Location")
	io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
	resp.Body.Close()
	if captured != nil {
		return captured
	}
	if strings.HasPrefix(loc, b.cfg.redirectURI) {
		u, _ := url.Parse(loc)
		return u
	}
	return nil
}

func authorizationFlowSlug(u *url.URL) string {
	if u == nil {
		return ""
	}
	// Matches /auth/if/flow/<slug>/ and /if/flow/<slug>/
	parts := strings.Split(strings.Trim(u.Path, "/"), "/")
	for i := 0; i+1 < len(parts); i++ {
		if parts[i] == "flow" && i >= 1 && parts[i-1] == "if" {
			slug := parts[i+1]
			if strings.Contains(slug, "authentication") {
				return "" // login flow → not authenticated, don't loop
			}
			return slug
		}
	}
	return ""
}

func redactQuery(u *url.URL) string {
	if u == nil {
		return ""
	}
	c := *u
	c.RawQuery = ""
	return c.String()
}