internal: optimise outpost's flow executor to use less requests

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>

internal/outpost/flow/executor.go

@@ -7,7 +7,6 @@ import (
 	"net/http"
 	"net/http/cookiejar"
 	"net/url"
-	"strconv"
 	"strings"
 
 	"github.com/getsentry/sentry-go"
@@ -31,11 +30,15 @@ var (
 	}, []string{"stage", "flow"})
 )
 
+type SolverFunction func(*api.ChallengeTypes, api.ApiFlowsExecutorSolveRequest) (api.FlowChallengeResponseRequest, error)
+
 type FlowExecutor struct {
 	Params  url.Values
 	Answers map[StageComponent]string
 	Context context.Context
 
+	solvers map[StageComponent]SolverFunction
+
 	cip       string
 	api       *api.APIClient
 	flowSlug  string
@@ -68,6 +71,11 @@ func NewFlowExecutor(ctx context.Context, flowSlug string, refConfig *api.Config
 		cip:       "",
 		transport: transport,
 	}
+	fe.solvers = map[StageComponent]SolverFunction{
+		StageIdentification:        fe.solveChallenge_Identification,
+		StagePassword:              fe.solveChallenge_Password,
+		StageAuthenticatorValidate: fe.solveChallenge_AuthenticatorValidate,
+	}
 	// Create new http client that also sets the correct ip
 	config := api.NewConfiguration()
 	config.Host = refConfig.Host
@@ -98,7 +106,7 @@ func (fe *FlowExecutor) ApiClient() *api.APIClient {
 	return fe.api
 }
 
-type ChallengeInt interface {
+type challengeInt interface {
 	GetComponent() string
 	GetType() api.ChallengeChoices
 	GetResponseErrors() map[string][]api.ErrorDetail
@@ -145,20 +153,23 @@ func (fe *FlowExecutor) WarmUp() error {
 }
 
 func (fe *FlowExecutor) Execute() (bool, error) {
-	return fe.solveFlowChallenge(1)
+	initial, err := fe.getInitialChallenge()
+	if err != nil {
+		return false, err
+	}
+	defer fe.sp.Finish()
+	return fe.solveFlowChallenge(initial, 1)
 }
 
-func (fe *FlowExecutor) solveFlowChallenge(depth int) (bool, error) {
+func (fe *FlowExecutor) getInitialChallenge() (*api.ChallengeTypes, error) {
-	defer fe.sp.Finish()
-
 	// Get challenge
 	gcsp := sentry.StartSpan(fe.Context, "authentik.outposts.flow_executor.get_challenge")
 	req := fe.api.FlowsApi.FlowsExecutorGet(gcsp.Context(), fe.flowSlug).Query(fe.Params.Encode())
 	challenge, _, err := req.Execute()
 	if err != nil {
-		return false, errors.New("failed to get challenge")
+		return nil, err
 	}
-	ch := challenge.GetActualInstance().(ChallengeInt)
+	ch := challenge.GetActualInstance().(challengeInt)
 	fe.log.WithField("component", ch.GetComponent()).WithField("type", ch.GetType()).Debug("Got challenge")
 	gcsp.SetTag("authentik.flow.challenge", string(ch.GetType()))
 	gcsp.SetTag("authentik.flow.component", ch.GetComponent())
@@ -167,60 +178,16 @@ func (fe *FlowExecutor) solveFlowChallenge(depth int) (bool, error) {
 		"stage": ch.GetComponent(),
 		"flow":  fe.flowSlug,
 	}).Observe(float64(gcsp.EndTime.Sub(gcsp.StartTime)))
+	return challenge, nil
+}
 
+func (fe *FlowExecutor) solveFlowChallenge(challenge *api.ChallengeTypes, depth int) (bool, error) {
 	// Resole challenge
 	scsp := sentry.StartSpan(fe.Context, "authentik.outposts.flow_executor.solve_challenge")
 	responseReq := fe.api.FlowsApi.FlowsExecutorSolve(scsp.Context(), fe.flowSlug).Query(fe.Params.Encode())
-	switch ch.GetComponent() {
+	ch := challenge.GetActualInstance().(challengeInt)
-	case string(StageIdentification):
-		r := api.NewIdentificationChallengeResponseRequest(fe.getAnswer(StageIdentification))
-		r.SetPassword(fe.getAnswer(StagePassword))
-		responseReq = responseReq.FlowChallengeResponseRequest(api.IdentificationChallengeResponseRequestAsFlowChallengeResponseRequest(r))
-	case string(StagePassword):
-		responseReq = responseReq.FlowChallengeResponseRequest(api.PasswordChallengeResponseRequestAsFlowChallengeResponseRequest(api.NewPasswordChallengeResponseRequest(fe.getAnswer(StagePassword))))
-	case string(StageAuthenticatorValidate):
-		// We only support duo as authenticator, check if that's allowed
-		var deviceChallenge *api.DeviceChallenge
-		for _, devCh := range challenge.AuthenticatorValidationChallenge.DeviceChallenges {
-			if devCh.DeviceClass == string(api.DEVICECLASSESENUM_DUO) {
-				deviceChallenge = &devCh
-			}
-		}
-		if deviceChallenge == nil {
-			return false, errors.New("no compatible authenticator class found")
-		}
-		devId, err := strconv.Atoi(deviceChallenge.DeviceUid)
-		if err != nil {
-			return false, errors.New("failed to convert duo device id to int")
-		}
-		devId32 := int32(devId)
-		inner := api.NewAuthenticatorValidationChallengeResponseRequest()
-		inner.SelectedChallenge = (*api.DeviceChallengeRequest)(deviceChallenge)
-		inner.Duo = &devId32
-		responseReq = responseReq.FlowChallengeResponseRequest(api.AuthenticatorValidationChallengeResponseRequestAsFlowChallengeResponseRequest(inner))
-	case string(StageAccessDenied):
-		return false, errors.New("got ak-stage-access-denied")
-	default:
-		return false, fmt.Errorf("unsupported challenge type %s", ch.GetComponent())
-	}
 
-	response, _, err := responseReq.Execute()
+	// Check for any validation errors that we might've gotten
-	ch = response.GetActualInstance().(ChallengeInt)
-	fe.log.WithField("component", ch.GetComponent()).WithField("type", ch.GetType()).Debug("Got response")
-	scsp.SetTag("authentik.flow.challenge", string(ch.GetType()))
-	scsp.SetTag("authentik.flow.component", ch.GetComponent())
-	scsp.Finish()
-
-	switch ch.GetComponent() {
-	case string(StageAccessDenied):
-		return false, nil
-	}
-	if ch.GetType() == "redirect" {
-		return true, nil
-	}
-	if err != nil {
-		return false, fmt.Errorf("failed to submit challenge %w", err)
-	}
 	if len(ch.GetResponseErrors()) > 0 {
 		for key, errs := range ch.GetResponseErrors() {
 			for _, err := range errs {
@@ -228,6 +195,34 @@ func (fe *FlowExecutor) solveFlowChallenge(depth int) (bool, error) {
 			}
 		}
 	}
+
+	switch ch.GetType() {
+	case api.CHALLENGECHOICES_REDIRECT:
+		return true, nil
+	case api.CHALLENGECHOICES_NATIVE:
+		if ch.GetComponent() == string(StageAccessDenied) {
+			return false, nil
+		}
+		solver, ok := fe.solvers[StageComponent(ch.GetComponent())]
+		if !ok {
+			return false, fmt.Errorf("unsupported challenge type %s", ch.GetComponent())
+		}
+		rr, err := solver(challenge, responseReq)
+		if err != nil {
+			return false, err
+		}
+		responseReq = responseReq.FlowChallengeResponseRequest(rr)
+	}
+
+	response, _, err := responseReq.Execute()
+	if err != nil {
+		return false, fmt.Errorf("failed to submit challenge %w", err)
+	}
+	ch = response.GetActualInstance().(challengeInt)
+	fe.log.WithField("component", ch.GetComponent()).WithField("type", ch.GetType()).Debug("Got response")
+	scsp.SetTag("authentik.flow.challenge", string(ch.GetType()))
+	scsp.SetTag("authentik.flow.component", ch.GetComponent())
+	scsp.Finish()
 	FlowTimingPost.With(prometheus.Labels{
 		"stage": ch.GetComponent(),
 		"flow":  fe.flowSlug,
@@ -236,5 +231,5 @@ func (fe *FlowExecutor) solveFlowChallenge(depth int) (bool, error) {
 	if depth >= 10 {
 		return false, errors.New("exceeded stage recursion depth")
 	}
-	return fe.solveFlowChallenge(depth + 1)
+	return fe.solveFlowChallenge(response, depth+1)
 }

internal/outpost/flow/solvers.go

@@ -0,0 +1,41 @@
+package flow
+
+import (
+	"errors"
+	"strconv"
+
+	"goauthentik.io/api/v3"
+)
+
+func (fe *FlowExecutor) solveChallenge_Identification(challenge *api.ChallengeTypes, req api.ApiFlowsExecutorSolveRequest) (api.FlowChallengeResponseRequest, error) {
+	r := api.NewIdentificationChallengeResponseRequest(fe.getAnswer(StageIdentification))
+	r.SetPassword(fe.getAnswer(StagePassword))
+	return api.IdentificationChallengeResponseRequestAsFlowChallengeResponseRequest(r), nil
+}
+
+func (fe *FlowExecutor) solveChallenge_Password(challenge *api.ChallengeTypes, req api.ApiFlowsExecutorSolveRequest) (api.FlowChallengeResponseRequest, error) {
+	r := api.NewPasswordChallengeResponseRequest(fe.getAnswer(StagePassword))
+	return api.PasswordChallengeResponseRequestAsFlowChallengeResponseRequest(r), nil
+}
+
+func (fe *FlowExecutor) solveChallenge_AuthenticatorValidate(challenge *api.ChallengeTypes, req api.ApiFlowsExecutorSolveRequest) (api.FlowChallengeResponseRequest, error) {
+	// We only support duo as authenticator, check if that's allowed
+	var deviceChallenge *api.DeviceChallenge
+	for _, devCh := range challenge.AuthenticatorValidationChallenge.DeviceChallenges {
+		if devCh.DeviceClass == string(api.DEVICECLASSESENUM_DUO) {
+			deviceChallenge = &devCh
+		}
+	}
+	if deviceChallenge == nil {
+		return api.FlowChallengeResponseRequest{}, errors.New("no compatible authenticator class found")
+	}
+	devId, err := strconv.Atoi(deviceChallenge.DeviceUid)
+	if err != nil {
+		return api.FlowChallengeResponseRequest{}, errors.New("failed to convert duo device id to int")
+	}
+	devId32 := int32(devId)
+	inner := api.NewAuthenticatorValidationChallengeResponseRequest()
+	inner.SelectedChallenge = (*api.DeviceChallengeRequest)(deviceChallenge)
+	inner.Duo = &devId32
+	return api.AuthenticatorValidationChallengeResponseRequestAsFlowChallengeResponseRequest(inner), nil
+}