ultisuite-backend/internal/authentik/flow_executor.go

package authentik

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

const flowExecutorTimeout = 45 * time.Second

// FlowChallenge is a raw Authentik flow executor challenge payload.
type FlowChallenge map[string]any

// FlowExecutor runs an Authentik flow via the backend flow executor API.
type FlowExecutor struct {
	baseURL string
	slug    string
	client  *http.Client
}

// NewFlowExecutor creates a flow executor with an isolated cookie jar.
func NewFlowExecutor(baseURL, slug string) (*FlowExecutor, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}
	return &FlowExecutor{
		baseURL: APIBaseURL(baseURL),
		slug:    strings.Trim(slug, "/"),
		client: &http.Client{
			Timeout: flowExecutorTimeout,
			Jar:     jar,
			CheckRedirect: func(req *http.Request, via []*http.Request) error {
				if len(via) >= 8 {
					return fmt.Errorf("too many redirects")
				}
				return nil
			},
		},
	}, nil
}

func (fe *FlowExecutor) executorURL(query string) string {
	u := fmt.Sprintf("%s/api/v3/flows/executor/%s/", fe.baseURL, fe.slug)
	if query != "" {
		// Authentik's flow executor reads the original frontend querystring from a `query`
		// parameter (e.g. ?query=next%3D...). Passing the params directly is ignored.
		return u + "?query=" + url.QueryEscape(query)
	}
	return u
}

// GetChallenge starts or resumes a flow and returns the pending challenge.
func (fe *FlowExecutor) GetChallenge(ctx context.Context, query string) (FlowChallenge, error) {
	// Establish the flow plan via the executor GET with the query string so the OAuth `next`
	// continuation is captured in the plan. Calling execute/ first would create a plan without
	// `next`, and the executor GET would then reuse that plan and ignore our continuation.
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fe.executorURL(query), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	challenge, err := fe.doChallenge(req)
	if err == nil {
		return challenge, nil
	}
	// Fallback: some Authentik configs require a warmed session cookie before the executor GET.
	if warmErr := fe.warmSession(ctx, query); warmErr != nil {
		return nil, err
	}
	req2, err2 := http.NewRequestWithContext(ctx, http.MethodGet, fe.executorURL(query), nil)
	if err2 != nil {
		return nil, err2
	}
	req2.Header.Set("Accept", "application/json")
	return fe.doChallenge(req2)
}

func (fe *FlowExecutor) warmSession(ctx context.Context, query string) error {
	u := fmt.Sprintf("%s/api/v3/flows/instances/%s/execute/", fe.baseURL, fe.slug)
	if query != "" {
		u += "?query=" + url.QueryEscape(query)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	resp, err := fe.client.Do(req)
	if err != nil {
		return err
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()
	if resp.StatusCode >= 500 {
		return fmt.Errorf("warm session: %d", resp.StatusCode)
	}
	return nil
}

// PostResponse submits a stage response and returns the next challenge or completion.
func (fe *FlowExecutor) PostResponse(ctx context.Context, query string, payload map[string]any) (FlowChallenge, error) {
	raw, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fe.executorURL(query), bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Referer", fe.executorURL(query))
	if csrf := fe.csrfToken(); csrf != "" {
		req.Header.Set("X-authentik-CSRF", csrf)
	}
	return fe.doChallenge(req)
}

// FollowFlowRedirect loads the flow completion URL so Authentik can bind an authenticated session.
func (fe *FlowExecutor) FollowFlowRedirect(ctx context.Context, target string) error {
	target = strings.TrimSpace(target)
	if target == "" {
		return nil
	}
	resolved, err := fe.resolveRedirectURL(target)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, resolved, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	resp, err := fe.client.Do(req)
	if err != nil {
		return err
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()
	if resp.StatusCode >= 500 {
		return fmt.Errorf("follow flow redirect: %d", resp.StatusCode)
	}
	return nil
}

// CaptureOAuthCallback follows the post-authentication authorize redirect (`to`) within the
// executor's live, authenticated jar and returns the external callback URL carrying the
// authorization code. This mirrors what a browser does after login and is the supported way to
// bridge an embedded (API-driven) Authentik session to an OIDC authorization code.
func (fe *FlowExecutor) CaptureOAuthCallback(ctx context.Context, to string) (string, error) {
	toURL, err := url.Parse(to)
	if err != nil {
		return "", fmt.Errorf("parse authorize redirect: %w", err)
	}
	redirectURI := toURL.Query().Get("redirect_uri")
	if redirectURI == "" {
		return "", fmt.Errorf("authorize redirect missing redirect_uri")
	}
	resolved, err := fe.resolveRedirectURL(to)
	if err != nil {
		return "", err
	}

	var captured *url.URL
	client := &http.Client{
		Jar:     fe.client.Jar,
		Timeout: fe.client.Timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if strings.HasPrefix(req.URL.String(), redirectURI) {
				captured = req.URL
				return http.ErrUseLastResponse
			}
			if len(via) >= 15 {
				return fmt.Errorf("too many authorize redirects")
			}
			return nil
		},
	}

	current := resolved
	for i := 0; i < 6; i++ {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, current, nil)
		if err != nil {
			return "", err
		}
		req.Header.Set("Accept", "text/html,application/xhtml+xml")
		resp, err := client.Do(req)
		if err != nil {
			return "", err
		}
		loc := resp.Header.Get("Location")
		finalURL := resp.Request.URL
		io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
		resp.Body.Close()

		if captured != nil {
			return captured.String(), nil
		}
		if strings.HasPrefix(loc, redirectURI) {
			return loc, nil
		}

		// Authorize routed through a flow (e.g. implicit-consent authorization flow). Drive it
		// via the executor API to obtain the terminal redirect with the code.
		slug := flowSlugFromURL(finalURL)
		if slug == "" {
			return "", fmt.Errorf("authorize did not yield a callback (landed on %s)", redactURL(finalURL))
		}
		callback, next, err := fe.driveFlowToRedirect(ctx, client, slug, finalURL, redirectURI)
		if err != nil {
			return "", err
		}
		if callback != "" {
			return callback, nil
		}
		if next == "" {
			return "", fmt.Errorf("authorization flow %q did not complete", slug)
		}
		current = next
	}
	return "", fmt.Errorf("authorize redirect chain did not resolve to a callback")
}

// driveFlowToRedirect runs a flow executor (used for the implicit-consent authorization flow)
// until it produces a terminal redirect. Returns either a captured callback URL, or a next URL
// to continue following.
func (fe *FlowExecutor) driveFlowToRedirect(ctx context.Context, client *http.Client, slug string, flowURL *url.URL, redirectURI string) (callback string, next string, err error) {
	q := flowURL.Query().Get("query")
	if q == "" {
		q = flowURL.RawQuery
	}
	execURL := fmt.Sprintf("%s/api/v3/flows/executor/%s/", fe.baseURL, slug)
	if q != "" {
		execURL += "?query=" + url.QueryEscape(q)
	}

	for i := 0; i < 6; i++ {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, execURL, nil)
		if err != nil {
			return "", "", err
		}
		req.Header.Set("Accept", "application/json")
		resp, err := client.Do(req)
		if err != nil {
			return "", "", err
		}
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<18))
		resp.Body.Close()
		var ch FlowChallenge
		_ = json.Unmarshal(body, &ch)
		comp := FlowComponent(ch)
		slog.Info("capture authz flow stage", "slug", slug, "component", comp)

		switch comp {
		case "xak-flow-redirect":
			toStr, _ := ch["to"].(string)
			if strings.HasPrefix(toStr, redirectURI) {
				return toStr, "", nil
			}
			resolved, rErr := fe.resolveRedirectURL(toStr)
			if rErr != nil {
				return "", "", rErr
			}
			return "", resolved, nil
		case "ak-stage-consent":
			payload := map[string]any{"component": "ak-stage-consent"}
			pr, pErr := fe.postFlowStage(ctx, client, execURL, payload)
			if pErr != nil {
				return "", "", pErr
			}
			if to, _ := pr["to"].(string); to != "" {
				if strings.HasPrefix(to, redirectURI) {
					return to, "", nil
				}
				resolved, _ := fe.resolveRedirectURL(to)
				return "", resolved, nil
			}
		default:
			return "", "", fmt.Errorf("authorization flow %q stuck at %q", slug, comp)
		}
	}
	return "", "", fmt.Errorf("authorization flow %q did not complete", slug)
}

func (fe *FlowExecutor) postFlowStage(ctx context.Context, client *http.Client, execURL string, payload map[string]any) (FlowChallenge, error) {
	raw, _ := json.Marshal(payload)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, execURL, bytes.NewReader(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)
	if csrf := fe.csrfToken(); csrf != "" {
		req.Header.Set("X-authentik-CSRF", csrf)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<18))
	var ch FlowChallenge
	_ = json.Unmarshal(body, &ch)
	return ch, nil
}

func flowSlugFromURL(u *url.URL) string {
	if u == nil {
		return ""
	}
	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 ""
			}
			return slug
		}
	}
	return ""
}

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

func (fe *FlowExecutor) resolveRedirectURL(target string) (string, error) {
	if strings.HasPrefix(target, "http://") || strings.HasPrefix(target, "https://") {
		return target, nil
	}
	base, err := url.Parse(fe.baseURL)
	if err != nil {
		return "", err
	}
	rel, err := url.Parse(target)
	if err != nil {
		return "", err
	}
	return base.ResolveReference(rel).String(), nil
}

func (fe *FlowExecutor) cookieURL() *url.URL {
	u, err := url.Parse(fe.baseURL)
	if err != nil || u.Host == "" {
		return &url.URL{Scheme: "http", Host: "localhost", Path: "/auth/"}
	}
	if u.Path == "" || u.Path == "/" {
		u.Path = "/auth/"
	} else if !strings.HasSuffix(u.Path, "/") {
		u.Path += "/"
	}
	u.RawQuery = ""
	u.Fragment = ""
	return u
}

func (fe *FlowExecutor) csrfToken() string {
	for _, c := range fe.client.Jar.Cookies(fe.cookieURL()) {
		if c.Name == "authentik_csrf" {
			return c.Value
		}
	}
	return ""
}

func (fe *FlowExecutor) doChallenge(req *http.Request) (FlowChallenge, error) {
	resp, err := fe.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, err
	}

	if resp.StatusCode == http.StatusFound || resp.StatusCode == http.StatusSeeOther {
		loc := resp.Header.Get("Location")
		if loc == "" {
			return FlowChallenge{"component": "xak-flow-redirect"}, nil
		}
		return FlowChallenge{
			"component": "xak-flow-redirect",
			"to":        loc,
		}, nil
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		msg := strings.TrimSpace(string(body))
		if msg == "" {
			msg = resp.Status
		}
		return nil, fmt.Errorf("flow executor %s: %d %s", fe.slug, resp.StatusCode, msg)
	}

	if len(body) == 0 {
		return FlowChallenge{"component": "xak-flow-redirect"}, nil
	}

	var challenge FlowChallenge
	if err := json.Unmarshal(body, &challenge); err != nil {
		return nil, fmt.Errorf("decode flow challenge: %w", err)
	}
	return challenge, nil
}

// FlowComponent returns the active stage component identifier.
func FlowComponent(challenge FlowChallenge) string {
	if challenge == nil {
		return ""
	}
	v, _ := challenge["component"].(string)
	return v
}

// FlowDone reports whether the flow finished successfully or was denied.
func FlowDone(challenge FlowChallenge) (done bool, denied bool) {
	switch FlowComponent(challenge) {
	case "xak-flow-redirect":
		return true, false
	case "ak-stage-access-denied":
		return true, true
	default:
		return false, false
	}
}