ultisuite-backend/internal/authentik/flow_session_store.go

package authentik

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

const (
	defaultFlowSessionTTL = 20 * time.Minute
	flowSessionKeyPrefix  = "auth_flow_session:"
)

type storedFlowSession struct {
	Slug          string             `json:"slug"`
	Cookies       []SerializedCookie `json:"cookies"`
	CreatedAt     time.Time          `json:"createdAt"`
	Completed     bool               `json:"completed,omitempty"`
	OAuthCallback string             `json:"oauthCallback,omitempty"`
}

type flowSessionEntry struct {
	slug          string
	cookies       []SerializedCookie
	createdAt     time.Time
	completed     bool
	oauthCallback string
}

// FlowSessionStore keeps Authentik flow executor sessions (memory + optional KeyDB).
type FlowSessionStore struct {
	mu      sync.Mutex
	baseURL string
	ttl     time.Duration
	rdb     *redis.Client
	items   map[string]*flowSessionEntry
}

func NewFlowSessionStore(baseURL string, rdb *redis.Client) *FlowSessionStore {
	return &FlowSessionStore{
		baseURL: baseURL,
		ttl:     defaultFlowSessionTTL,
		rdb:     rdb,
		items:   make(map[string]*flowSessionEntry),
	}
}

func (s *FlowSessionStore) Start(ctx context.Context, slug, query string) (sessionID string, challenge FlowChallenge, err error) {
	executor, err := NewFlowExecutor(s.baseURL, slug)
	if err != nil {
		return "", nil, err
	}
	challenge, err = executor.GetChallenge(ctx, query)
	if err != nil {
		return "", nil, err
	}
	id, err := randomSessionID()
	if err != nil {
		return "", nil, err
	}
	done, _ := FlowDone(challenge)
	if !done {
		entry := &flowSessionEntry{
			slug:      slug,
			cookies:   executor.ExportCookies(),
			createdAt: time.Now(),
		}
		if err := s.save(ctx, id, entry); err != nil {
			return "", nil, err
		}
	}
	return id, challenge, nil
}

func (s *FlowSessionStore) Respond(ctx context.Context, sessionID, slug, query string, payload map[string]any) (FlowChallenge, error) {
	entry, err := s.load(ctx, sessionID)
	if err != nil {
		return nil, err
	}
	if entry.slug != slug {
		return nil, ErrFlowSessionSlugMismatch
	}
	if entry.completed {
		return nil, ErrFlowSessionAlreadyCompleted
	}
	executor, err := RestoreFlowExecutor(s.baseURL, slug, entry.cookies)
	if err != nil {
		return nil, err
	}
	challenge, err := executor.PostResponse(ctx, query, payload)
	if err != nil {
		return nil, err
	}
	entry.cookies = executor.ExportCookies()
	challenge, err = s.advanceHeadlessStages(ctx, executor, query, challenge)
	if err != nil {
		return nil, err
	}
	entry.cookies = executor.ExportCookies()

	if SessionAuthenticated(entry.cookies) {
		entry.completed = true
		// If the flow's terminal redirect is an OIDC authorize continuation, follow it within the
		// live authenticated jar to capture the authorization code (the only reliable way to bridge
		// an API-driven session to OIDC; replaying the session cookie elsewhere is rejected).
		if to, ok := challenge["to"].(string); ok && isOAuthAuthorizeRedirect(to) {
			callback, capErr := executor.CaptureOAuthCallback(ctx, to)
			if capErr != nil {
				slog.Warn("capture oauth callback failed", "err", capErr.Error())
			} else {
				entry.oauthCallback = callback
				slog.Info("captured oauth callback", "ok", callback != "")
			}
		}
		if err := s.save(ctx, sessionID, entry); err != nil {
			return nil, err
		}
		if FlowComponent(challenge) != "xak-flow-redirect" {
			challenge = FlowChallenge{"component": "xak-flow-redirect"}
		}
		return challenge, nil
	}

	done, denied := FlowDone(challenge)
	if done && !denied {
		return nil, fmt.Errorf("authentik flow finished without authenticated session")
	}
	if err := s.save(ctx, sessionID, entry); err != nil {
		return nil, err
	}
	return challenge, nil
}

func (s *FlowSessionStore) advanceHeadlessStages(ctx context.Context, executor *FlowExecutor, query string, challenge FlowChallenge) (FlowChallenge, error) {
	const maxSteps = 8
	for step := 0; step < maxSteps; step++ {
		if SessionAuthenticated(executor.ExportCookies()) {
			return challenge, nil
		}
		if _, denied := FlowDone(challenge); denied {
			return challenge, nil
		}

		switch FlowComponent(challenge) {
		case "ak-stage-user-login":
			var err error
			challenge, err = executor.PostResponse(ctx, query, map[string]any{"component": "ak-stage-user-login"})
			if err != nil {
				return nil, err
			}
		case "xak-flow-redirect":
			// Resume pending stages before following terminal redirects (/auth/ resets session).
			next, err := executor.GetChallenge(ctx, query)
			if err != nil {
				return nil, err
			}
			if comp := FlowComponent(next); comp != "" && comp != "xak-flow-redirect" {
				challenge = next
				continue
			}
			if to, ok := challenge["to"].(string); ok && isTerminalFlowRedirect(to) {
				loginChallenge, loginErr := executor.PostResponse(ctx, query, map[string]any{"component": "ak-stage-user-login"})
				if loginErr == nil {
					challenge = loginChallenge
					continue
				}
				return challenge, nil
			}
			if to, ok := challenge["to"].(string); ok && strings.TrimSpace(to) != "" {
				if err := executor.FollowFlowRedirect(ctx, to); err != nil {
					return nil, err
				}
			}
			next, err = executor.GetChallenge(ctx, query)
			if err != nil {
				return nil, err
			}
			challenge = next
		default:
			return challenge, nil
		}
	}
	return challenge, nil
}

func isOAuthAuthorizeRedirect(to string) bool {
	return strings.Contains(to, "/application/o/authorize")
}

func isTerminalFlowRedirect(to string) bool {
	to = strings.TrimSpace(to)
	if to == "" || to == "/auth/" || to == "/auth" {
		return true
	}
	return strings.HasSuffix(strings.TrimSuffix(to, "/"), "/auth")
}

// CompleteSession returns cookies from a completed flow and removes the session.
func (s *FlowSessionStore) CompleteSession(ctx context.Context, sessionID, slug string) ([]SerializedCookie, error) {
	cookies, _, err := s.CompleteSessionOAuth(ctx, sessionID, slug)
	return cookies, err
}

// CompleteSessionOAuth returns the completed flow's cookies plus any captured OIDC callback URL
// (carrying the authorization code) and removes the session.
func (s *FlowSessionStore) CompleteSessionOAuth(ctx context.Context, sessionID, slug string) ([]SerializedCookie, string, error) {
	entry, err := s.load(ctx, sessionID)
	if err != nil {
		return nil, "", err
	}
	if entry.slug != slug {
		return nil, "", ErrFlowSessionSlugMismatch
	}
	if !entry.completed {
		return nil, "", ErrFlowSessionNotCompleted
	}
	cookies := entry.cookies
	callback := entry.oauthCallback
	s.delete(ctx, sessionID)
	return cookies, callback, nil
}

// SessionCookies returns current Authentik cookies for an active session.
func (s *FlowSessionStore) SessionCookies(ctx context.Context, sessionID string) ([]SerializedCookie, error) {
	entry, err := s.load(ctx, sessionID)
	if err != nil {
		return nil, err
	}
	return entry.cookies, nil
}

func (s *FlowSessionStore) Delete(ctx context.Context, sessionID string) {
	s.delete(ctx, sessionID)
}

func (s *FlowSessionStore) save(ctx context.Context, sessionID string, entry *flowSessionEntry) error {
	s.mu.Lock()
	s.items[sessionID] = entry
	s.mu.Unlock()
	if s.rdb == nil {
		return nil
	}
	stored := storedFlowSession{
		Slug:          entry.slug,
		Cookies:       entry.cookies,
		CreatedAt:     entry.createdAt,
		Completed:     entry.completed,
		OAuthCallback: entry.oauthCallback,
	}
	raw, err := json.Marshal(stored)
	if err != nil {
		return err
	}
	return s.rdb.Set(ctx, flowSessionKeyPrefix+sessionID, raw, s.ttl).Err()
}

func (s *FlowSessionStore) load(ctx context.Context, sessionID string) (*flowSessionEntry, error) {
	if s.rdb != nil {
		raw, err := s.rdb.Get(ctx, flowSessionKeyPrefix+sessionID).Bytes()
		if err != nil {
			if errors.Is(err, redis.Nil) {
				return nil, ErrFlowSessionNotFound
			}
			return nil, err
		}
		var stored storedFlowSession
		if err := json.Unmarshal(raw, &stored); err != nil {
			return nil, err
		}
		if time.Since(stored.CreatedAt) > s.ttl {
			s.delete(ctx, sessionID)
			return nil, ErrFlowSessionNotFound
		}
		entry := &flowSessionEntry{
			slug:          stored.Slug,
			cookies:       stored.Cookies,
			createdAt:     stored.CreatedAt,
			completed:     stored.Completed,
			oauthCallback: stored.OAuthCallback,
		}
		s.mu.Lock()
		s.items[sessionID] = entry
		s.mu.Unlock()
		return entry, nil
	}

	s.mu.Lock()
	entry, ok := s.items[sessionID]
	s.mu.Unlock()
	if !ok {
		return nil, ErrFlowSessionNotFound
	}
	if time.Since(entry.createdAt) > s.ttl {
		s.delete(ctx, sessionID)
		return nil, ErrFlowSessionNotFound
	}
	return entry, nil
}

func (s *FlowSessionStore) delete(ctx context.Context, sessionID string) {
	s.mu.Lock()
	delete(s.items, sessionID)
	s.mu.Unlock()
	if s.rdb != nil {
		_ = s.rdb.Del(ctx, flowSessionKeyPrefix+sessionID).Err()
	}
}

func randomSessionID() (string, error) {
	buf := make([]byte, 24)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf), nil
}