3 files changed, 42 insertions, 17 deletions
diff --git a/app/middleware/id_token.go b/app/middleware/id_token.go
index 8084af0..bc6a05e 100644
--- a/app/middleware/id_token.go
+++ b/app/middleware/id_token.go
@@ -14,6 +14,8 @@ import (
 func IDToken(provider *oidc.Provider, config *oidc.Config, parsers ...TokenParser) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			log.WithFields(r.Context(), log.Fields{"jwt": r.Header.Get("x-jwt-payload")})
+
 			for _, parser := range parsers {
 				rawIDToken := parser(r)
 				if x.IsPresent(rawIDToken) {
diff --git a/app/middleware/user.go b/app/middleware/user.go
index 9a88f8e..6c018f4 100644
--- a/app/middleware/user.go
+++ b/app/middleware/user.go
@@ -14,20 +14,26 @@ import (
 func User(db domain.Repository[*domain.User]) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			idToken := cfg.IDToken.From(r.Context())
-			if x.IsZero(idToken) {
-				next.ServeHTTP(w, r)
-				return
-			}
+			subject := r.Header.Get("x-jwt-claim-sub")
+			user := db.Find(r.Context(), domain.ID(subject))
 
-			user := db.Find(r.Context(), domain.ID(idToken.Subject))
 			if !x.IsPresent(user) {
-				user = mapper.MapFrom[*oidc.IDToken, *domain.User](idToken)
-				if err := db.Save(r.Context(), user); err != nil {
-					pls.LogError(r.Context(), err)
+				idToken := cfg.IDToken.From(r.Context())
+
+				if x.IsZero(idToken) {
 					next.ServeHTTP(w, r)
 					return
 				}
+
+				user = db.Find(r.Context(), domain.ID(idToken.Subject))
+				if !x.IsPresent(user) {
+					user = mapper.MapFrom[*oidc.IDToken, *domain.User](idToken)
+					if err := db.Save(r.Context(), user); err != nil {
+						pls.LogError(r.Context(), err)
+						next.ServeHTTP(w, r)
+						return
+					}
+				}
 			}
 
 			next.ServeHTTP(w, r.WithContext(cfg.CurrentUser.With(r.Context(), user)))
diff --git a/app/middleware/user_test.go b/app/middleware/user_test.go
index aed3582..7653684 100644
--- a/app/middleware/user_test.go
+++ b/app/middleware/user_test.go
@@ -61,16 +61,33 @@ func TestUser(t *testing.T) {
 	})
 
 	t.Run("when ID Token is not provided", func(t *testing.T) {
-		server := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			user := cfg.CurrentUser.From(r.Context())
-			require.Nil(t, user)
+		t.Run("without custom headers", func(t *testing.T) {
+			server := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				user := cfg.CurrentUser.From(r.Context())
+				require.Nil(t, user)
+
+				w.WriteHeader(http.StatusTeapot)
+			}))
 
-			w.WriteHeader(http.StatusTeapot)
-		}))
+			r, w := test.RequestResponse("GET", "/example")
+			server.ServeHTTP(w, r)
 
-		r, w := test.RequestResponse("GET", "/example")
-		server.ServeHTTP(w, r)
+			assert.Equal(t, http.StatusTeapot, w.Code)
+		})
 
-		assert.Equal(t, http.StatusTeapot, w.Code)
+		t.Run("with x-jwt-claim-sub header", func(t *testing.T) {
+			server := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				user := cfg.CurrentUser.From(r.Context())
+				require.NotNil(t, user)
+				require.Equal(t, knownUser.ID, user.ID)
+
+				w.WriteHeader(http.StatusTeapot)
+			}))
+
+			r, w := test.RequestResponse("GET", "/example", test.WithRequestHeader("x-jwt-claim-sub", knownUser.ID.String()))
+			server.ServeHTTP(w, r)
+
+			assert.Equal(t, http.StatusTeapot, w.Code)
+		})
 	})
 }