ultisuite-backend/internal/authentik/flow_executor.go

package authentik

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"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 != "" {
		return u + "?" + 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) {
	if err := fe.warmSession(ctx); err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fe.executorURL(query), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	return fe.doChallenge(req)
}

func (fe *FlowExecutor) warmSession(ctx context.Context) error {
	u := fmt.Sprintf("%s/api/v3/flows/instances/%s/execute/", fe.baseURL, fe.slug)
	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)
}

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
	}
}