From 58910eeb073423028327d25b3eb99173d626f52a Mon Sep 17 00:00:00 2001
From: p53 <pavol.ipoth@protonmail.com>
Date: Fri, 7 Jun 2024 01:55:36 +0200
Subject: [PATCH] Fix idp session check and token refreshment

Fix idp session check and token refreshment
---
 pkg/proxy/middleware/oauth.go    | 24 ++++++++++++------------
 pkg/testsuite/middleware_test.go | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/pkg/proxy/middleware/oauth.go b/pkg/proxy/middleware/oauth.go
index cc69d6e3..9a900f75 100644
--- a/pkg/proxy/middleware/oauth.go
+++ b/pkg/proxy/middleware/oauth.go
@@ -85,18 +85,6 @@ func AuthenticationMiddleware(
 			// https://github.com/coreos/go-oidc/issues/402
 			oidcLibCtx := context.WithValue(ctx, oauth2.HTTPClient, httpClient)
 
-			if enableIDPSessionCheck {
-				tokenSource := oauth2.StaticTokenSource(
-					&oauth2.Token{AccessToken: user.RawToken},
-				)
-				_, err := provider.UserInfo(oidcLibCtx, tokenSource)
-				if err != nil {
-					scope.Logger.Error(err.Error())
-					redirectToAuthorization(wrt, req)
-					return
-				}
-			}
-
 			// step: skip if we are running skip-token-verification
 			if skipTokenVerification {
 				scope.Logger.Warn(
@@ -314,6 +302,18 @@ func AuthenticationMiddleware(
 				}
 			}
 
+			if enableIDPSessionCheck {
+				tokenSource := oauth2.StaticTokenSource(
+					&oauth2.Token{AccessToken: scope.Identity.RawToken},
+				)
+				_, err := provider.UserInfo(oidcLibCtx, tokenSource)
+				if err != nil {
+					scope.Logger.Error(err.Error())
+					redirectToAuthorization(wrt, req)
+					return
+				}
+			}
+
 			*req = *(req.WithContext(ctx))
 			next.ServeHTTP(wrt, req)
 		})
diff --git a/pkg/testsuite/middleware_test.go b/pkg/testsuite/middleware_test.go
index 97b32f02..1e9544dc 100644
--- a/pkg/testsuite/middleware_test.go
+++ b/pkg/testsuite/middleware_test.go
@@ -1529,6 +1529,35 @@ func TestRefreshToken(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name: "TestRefreshTokenWithIdpSessionCheck",
+			ProxySettings: func(conf *config.Config) {
+				conf.EnableIDPSessionCheck = true
+				conf.EnableRefreshTokens = true
+				conf.EnableEncryptedToken = true
+				conf.Verbose = true
+				conf.EnableLogging = true
+				conf.EncryptionKey = testEncryptionKey
+			},
+			ExecutionSettings: []fakeRequest{
+				{
+					URI:                           FakeAuthAllURL,
+					HasLogin:                      true,
+					Redirects:                     true,
+					OnResponse:                    delay,
+					ExpectedProxy:                 true,
+					ExpectedCode:                  http.StatusOK,
+					ExpectedLoginCookiesValidator: map[string]func(*testing.T, *config.Config, string) bool{cfg.CookieRefreshName: checkRefreshTokenEncryption},
+				},
+				{
+					URI:           FakeAuthAllURL,
+					Redirects:     false,
+					HasLogin:      false,
+					ExpectedProxy: true,
+					ExpectedCode:  http.StatusOK,
+				},
+			},
+		},
 		{
 			Name: "TestRefreshTokenEncryptionWithClientIDAndIssuerCheckOn",
 			ProxySettings: func(conf *config.Config) {
@@ -2888,7 +2917,7 @@ func TestAuthenticationMiddleware(t *testing.T) {
 					RawToken:          badlySignedToken,
 					HasCookieToken:    true,
 					ExpectedProxy:     false,
-					ExpectedCode:      http.StatusSeeOther,
+					ExpectedCode:      http.StatusForbidden,
 				},
 			},
 		},