From f7dc76c6e532895805323ae1929b6815111f258c Mon Sep 17 00:00:00 2001
From: Kenneth Jenkins <51246568+kenjenkins@users.noreply.github.com>
Date: Wed, 6 Dec 2023 14:08:35 -0800
Subject: [PATCH] support both stateful and stateless authenticate

Update the initialization logic for the authenticate, authorize, and
proxy services to automatically select between the stateful
authentication flow and the stateless authentication flow, depending on
whether Pomerium is configured to use the hosted authenticate service.
---
 authenticate/handlers.go |  6 ++++--
 authenticate/state.go    | 18 +++++++++++-------
 authorize/state.go       |  6 +++++-
 config/options.go        | 10 ++++++++++
 config/options_test.go   | 15 +++++++++++++++
 proxy/state.go           |  8 ++++++--
 6 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/authenticate/handlers.go b/authenticate/handlers.go
index 351ea5192..ca5afe92b 100644
--- a/authenticate/handlers.go
+++ b/authenticate/handlers.go
@@ -84,6 +84,7 @@ func (a *Authenticate) mountDashboard(r *mux.Router) {
 		AllowedHeaders:   []string{"*"},
 	})
 	sr.Use(c.Handler)
+	sr.Use(a.RetrieveSession)
 
 	// routes that don't need a session:
 	sr.Path("/sign_out").Handler(httputil.HandlerFunc(a.SignOut))
@@ -91,7 +92,6 @@ func (a *Authenticate) mountDashboard(r *mux.Router) {
 
 	// routes that need a session:
 	sr = sr.NewRoute().Subrouter()
-	sr.Use(a.RetrieveSession)
 	sr.Use(a.VerifySession)
 	sr.Path("/").Handler(a.requireValidSignatureOnRedirect(a.userInfo))
 	sr.Path("/sign_in").Handler(httputil.HandlerFunc(a.SignIn))
@@ -475,7 +475,9 @@ func (a *Authenticate) revokeSession(ctx context.Context, w http.ResponseWriter,
 		return ""
 	}
 
-	return state.flow.RevokeSession(ctx, r, authenticator, nil)
+	sessionState, _ := a.getSessionFromCtx(ctx)
+
+	return state.flow.RevokeSession(ctx, r, authenticator, sessionState)
 }
 
 // Callback handles the result of a successful call to the authenticate service
diff --git a/authenticate/state.go b/authenticate/state.go
index 213060e11..f8db3960d 100644
--- a/authenticate/state.go
+++ b/authenticate/state.go
@@ -144,13 +144,17 @@ func newAuthenticateStateFromConfig(
 		}
 	}
 
-	state.flow, err = authenticateflow.NewStateless(
-		cfg,
-		cookieStore,
-		authenticateConfig.getIdentityProvider,
-		authenticateConfig.profileTrimFn,
-		authenticateConfig.authEventFn,
-	)
+	if cfg.Options.UseStatelessAuthenticateFlow() {
+		state.flow, err = authenticateflow.NewStateless(
+			cfg,
+			cookieStore,
+			authenticateConfig.getIdentityProvider,
+			authenticateConfig.profileTrimFn,
+			authenticateConfig.authEventFn,
+		)
+	} else {
+		state.flow, err = authenticateflow.NewStateful(cfg, cookieStore)
+	}
 	if err != nil {
 		return nil, err
 	}
diff --git a/authorize/state.go b/authorize/state.go
index bb2d9e277..115b0bea2 100644
--- a/authorize/state.go
+++ b/authorize/state.go
@@ -83,7 +83,11 @@ func newAuthorizeStateFromConfig(
 		return nil, fmt.Errorf("authorize: invalid session store: %w", err)
 	}
 
-	state.authenticateFlow, err = authenticateflow.NewStateless(cfg, nil, nil, nil, nil)
+	if cfg.Options.UseStatelessAuthenticateFlow() {
+		state.authenticateFlow, err = authenticateflow.NewStateless(cfg, nil, nil, nil, nil)
+	} else {
+		state.authenticateFlow, err = authenticateflow.NewStateful(cfg, nil)
+	}
 	if err != nil {
 		return nil, err
 	}
diff --git a/config/options.go b/config/options.go
index f329ad39b..0a82caa16 100644
--- a/config/options.go
+++ b/config/options.go
@@ -827,6 +827,16 @@ func (o *Options) GetInternalAuthenticateURL() (*url.URL, error) {
 	return urlutil.ParseAndValidateURL(o.AuthenticateInternalURLString)
 }
 
+// UseStatelessAuthenticateFlow returns true if the stateless authentication
+// flow should be used (i.e. for hosted authenticate).
+func (o *Options) UseStatelessAuthenticateFlow() bool {
+	u, err := o.GetInternalAuthenticateURL()
+	if err != nil {
+		return false
+	}
+	return urlutil.IsHostedAuthenticateDomain(u.Hostname())
+}
+
 // GetAuthorizeURLs returns the AuthorizeURLs in the options or 127.0.0.1:5443.
 func (o *Options) GetAuthorizeURLs() ([]*url.URL, error) {
 	if IsAll(o.Services) && o.AuthorizeURLString == "" && len(o.AuthorizeURLStrings) == 0 {
diff --git a/config/options_test.go b/config/options_test.go
index 39e9b753a..4801e14c8 100644
--- a/config/options_test.go
+++ b/config/options_test.go
@@ -856,6 +856,21 @@ func TestOptions_DefaultURL(t *testing.T) {
 	}
 }
 
+func TestOptions_UseStatelessAuthenticateFlow(t *testing.T) {
+	t.Run("enabled by default", func(t *testing.T) {
+		options := &Options{}
+		assert.True(t, options.UseStatelessAuthenticateFlow())
+	})
+	t.Run("enabled explicitly", func(t *testing.T) {
+		options := &Options{AuthenticateURLString: "https://authenticate.pomerium.app"}
+		assert.True(t, options.UseStatelessAuthenticateFlow())
+	})
+	t.Run("disabled", func(t *testing.T) {
+		options := &Options{AuthenticateURLString: "https://authenticate.example.com"}
+		assert.False(t, options.UseStatelessAuthenticateFlow())
+	})
+}
+
 func TestOptions_GetOauthOptions(t *testing.T) {
 	opts := &Options{AuthenticateURLString: "https://authenticate.example.com"}
 	oauthOptions, err := opts.GetOauthOptions()
diff --git a/proxy/state.go b/proxy/state.go
index bdb4b1a64..42f8ebb44 100644
--- a/proxy/state.go
+++ b/proxy/state.go
@@ -114,8 +114,12 @@ func newProxyStateFromConfig(cfg *config.Config) (*proxyState, error) {
 
 	state.programmaticRedirectDomainWhitelist = cfg.Options.ProgrammaticRedirectDomainWhitelist
 
-	state.authenticateFlow, err = authenticateflow.NewStateless(
-		cfg, state.sessionStore, nil, nil, nil)
+	if cfg.Options.UseStatelessAuthenticateFlow() {
+		state.authenticateFlow, err = authenticateflow.NewStateless(
+			cfg, state.sessionStore, nil, nil, nil)
+	} else {
+		state.authenticateFlow, err = authenticateflow.NewStateful(cfg, state.sessionStore)
+	}
 	if err != nil {
 		return nil, err
 	}