From 7a7c4be9fa27aa998db4e8c930653acba23055d1 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Tue, 14 Jan 2025 12:59:48 +0800
Subject: [PATCH] fix OIDC middleware not working with Authentik

---
 internal/api/v1/auth/oidc.go         | 17 +++++++++++------
 internal/net/http/middleware/oidc.go |  2 +-
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/internal/api/v1/auth/oidc.go b/internal/api/v1/auth/oidc.go
index f61c445..bd376a5 100644
--- a/internal/api/v1/auth/oidc.go
+++ b/internal/api/v1/auth/oidc.go
@@ -76,9 +76,6 @@ func (auth *OIDCProvider) TokenCookieName() string {
 
 func (auth *OIDCProvider) SetIsMiddleware(enabled bool) {
 	auth.isMiddleware = enabled
-	if auth.isMiddleware {
-		auth.oauthConfig.RedirectURL = OIDCMiddlewareCallbackPath
-	}
 }
 
 func (auth *OIDCProvider) SetAllowedUsers(users []string) {
@@ -152,13 +149,22 @@ func (auth *OIDCProvider) RedirectLoginPage(w http.ResponseWriter, r *http.Reque
 			return
 		}
 		q := u.Query()
-		q.Set("redirect_uri", "https://"+r.Host+q.Get("redirect_uri"))
+		q.Set("redirect_uri", "https://"+r.Host+OIDCMiddlewareCallbackPath+q.Get("redirect_uri"))
 		u.RawQuery = q.Encode()
 		redirURL = u.String()
 	}
 	http.Redirect(w, r, redirURL, http.StatusTemporaryRedirect)
 }
 
+func (auth *OIDCProvider) exchange(r *http.Request) (*oauth2.Token, error) {
+	if auth.isMiddleware {
+		cfg := *auth.oauthConfig
+		cfg.RedirectURL = "https://" + r.Host + OIDCMiddlewareCallbackPath
+		return cfg.Exchange(r.Context(), r.URL.Query().Get("code"))
+	}
+	return auth.oauthConfig.Exchange(r.Context(), r.URL.Query().Get("code"))
+}
+
 // OIDCCallbackHandler handles the OIDC callback.
 func (auth *OIDCProvider) LoginCallbackHandler(w http.ResponseWriter, r *http.Request) {
 	// For testing purposes, skip provider verification
@@ -179,8 +185,7 @@ func (auth *OIDCProvider) LoginCallbackHandler(w http.ResponseWriter, r *http.Re
 		return
 	}
 
-	code := query.Get("code")
-	oauth2Token, err := auth.oauthConfig.Exchange(r.Context(), code)
+	oauth2Token, err := auth.exchange(r)
 	if err != nil {
 		U.HandleErr(w, r, fmt.Errorf("failed to exchange token: %w", err), http.StatusInternalServerError)
 		return
diff --git a/internal/net/http/middleware/oidc.go b/internal/net/http/middleware/oidc.go
index f732d81..b8ad7cd 100644
--- a/internal/net/http/middleware/oidc.go
+++ b/internal/net/http/middleware/oidc.go
@@ -19,7 +19,7 @@ var OIDC = NewMiddleware[oidcMiddleware]()
 
 func (amw *oidcMiddleware) finalize() error {
 	if !auth.IsOIDCEnabled() {
-		return E.New("OIDC not enabled but Auth middleware is used")
+		return E.New("OIDC not enabled but ODIC middleware is used")
 	}
 	authProvider, err := auth.NewOIDCProviderFromEnv()
 	if err != nil {