From b2453bab38531d3bf5b25372ff9ba86736dc6de5 Mon Sep 17 00:00:00 2001
From: Carsten Dietrich <3203968+carstendietrich@users.noreply.github.com>
Date: Fri, 19 May 2023 10:39:33 +0200
Subject: [PATCH] feat(core/auth/oidc): Add support of `AuthCodeOption` in the
 token exchange during callback (#338)

* feat(core/auth/oidc): Add support of `AuthCodeOption` in the token exchange during callback

Currently, it is only possible to modify the authorization URL using the `AuthCodeOptioner` interface. Since the Exchange function also supports this pattern we should also pass the registered `AuthCodeOptioner` to it.

For example, both my authorization and my token retrieval endpoint need some custom URL param to work. Currently, only the one for the auth endpoint would be possible.

---------

Co-authored-by: Karol Nowak <karol.nowak@omnevo.net>
---
 core/auth/oauth/oidc.go      | 10 ++++-
 core/auth/oauth/oidc_test.go | 73 ++++++++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/core/auth/oauth/oidc.go b/core/auth/oauth/oidc.go
index 5e837e270..8a3aae155 100644
--- a/core/auth/oauth/oidc.go
+++ b/core/auth/oauth/oidc.go
@@ -467,7 +467,15 @@ func (i *openIDIdentifier) Callback(ctx context.Context, request *web.Request, r
 		return i.responder.ServerError(err)
 	}
 
-	oauth2Token, err := oauthConfig.Exchange(ctx, code)
+	options := make([]oauth2.AuthCodeOption, 0)
+
+	if i.authCodeOptionerProvider != nil {
+		for _, o := range i.authCodeOptionerProvider() {
+			options = append(options, o.Options(ctx, i.Broker(), request)...)
+		}
+	}
+
+	oauth2Token, err := oauthConfig.Exchange(ctx, code, options...)
 	if err != nil {
 		return i.responder.ServerError(err)
 	}
diff --git a/core/auth/oauth/oidc_test.go b/core/auth/oauth/oidc_test.go
index 8b782023d..3447e01b1 100644
--- a/core/auth/oauth/oidc_test.go
+++ b/core/auth/oauth/oidc_test.go
@@ -11,12 +11,13 @@ import (
 	"testing"
 	"time"
 
-	"flamingo.me/flamingo/v3/core/auth"
-	"flamingo.me/flamingo/v3/framework/flamingo"
-	"flamingo.me/flamingo/v3/framework/web"
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/oauth2"
+
+	"flamingo.me/flamingo/v3/core/auth"
+	"flamingo.me/flamingo/v3/framework/flamingo"
+	"flamingo.me/flamingo/v3/framework/web"
 )
 
 type mockRouter struct {
@@ -138,6 +139,12 @@ func (p *testOidcProvider) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 	}
 }
 
+type testAuthCodeOptioner struct{}
+
+func (*testAuthCodeOptioner) Options(_ context.Context, _ string, req *web.Request) []oauth2.AuthCodeOption {
+	return []oauth2.AuthCodeOption{oauth2.SetAuthURLParam("redirect_uri", "foobar123test")}
+}
+
 func TestOidcCallback(t *testing.T) {
 	t.Run("Test Callback", func(t *testing.T) {
 		t.Parallel()
@@ -233,6 +240,66 @@ func TestOidcCallback(t *testing.T) {
 		expectedURL, _ := url.Parse("https://example.com/callback-error-handler")
 		assert.Equal(t, result, &web.URLRedirectResponse{URL: expectedURL}, "Result of callback handler was ignored")
 	})
+
+	t.Run("Test add auth code option to exchange token call", func(t *testing.T) {
+		t.Parallel()
+
+		testServer := httptest.NewUnstartedServer(nil)
+		tokenCalled := false
+
+		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch strings.Trim(r.URL.Path, "/") {
+			case "token":
+				{
+					_ = r.ParseForm()
+					redirectURI := r.PostForm.Get("redirect_uri")
+
+					assert.Equal(t, redirectURI, "foobar123test")
+					tokenCalled = true
+
+					w.Header().Set("Content-type", "application/json")
+					return
+				}
+			case ".well-known/openid-configuration":
+				_, _ = fmt.Fprintf(w, `{
+				"issuer": "%s",
+				"token_endpoint": "%s/token",
+				"jwks_uri": "%s/certs"
+			}`, testServer.URL, testServer.URL, testServer.URL)
+			}
+		})
+
+		testServer.Config.Handler = handler
+		testServer.Start()
+		defer testServer.Close()
+
+		identifier := new(openIDIdentifier)
+		identifier.reverseRouter = new(mockRouter)
+		identifier.responder = new(web.Responder)
+
+		identifier.authCodeOptionerProvider = func() []AuthCodeOptioner {
+			return []AuthCodeOptioner{new(testAuthCodeOptioner)}
+		}
+
+		var err error
+		identifier.provider, err = oidc.NewProvider(context.Background(), testServer.URL)
+		assert.NoError(t, err)
+		identifier.oauth2Config = &oauth2.Config{
+			Endpoint: oauth2.Endpoint{AuthURL: "", TokenURL: testServer.URL + "/token"},
+		}
+
+		session := web.EmptySession()
+		request := web.CreateRequest(nil, session)
+
+		identifier.createSessionCode(request, "test-callback-state")
+
+		request.Request().URL.RawQuery = "state=test-callback-state&code=test-callback-code"
+		identifier.Callback(context.Background(), request, func(request *web.Request) *url.URL {
+			return new(url.URL)
+		})
+
+		assert.True(t, tokenCalled)
+	})
 }
 
 func Test_openIDIdentifier_RefreshIdentity(t *testing.T) {