feat: Add option for propagating id_token to upstream app

Fixes #315 Co-authored-by: tronghn <trong.huu.nguyen@nav.no>
nais · Jan 20, 2025 · 2feb6a3 · 2feb6a3
1 parent bc30791
commit 2feb6a3
Show file tree

Hide file tree

Showing 7 changed files with 132 additions and 66 deletions.
diff --git a/docs/configuration.md b/docs/configuration.md
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -24,16 +24,17 @@ type Config struct {
 	ShutdownWaitBeforePeriod time.Duration `json:"shutdown-wait-before-period"`
 	Version                  string        `json:"version"`
 
-	AutoLogin            bool     `json:"auto-login"`
-	AutoLoginIgnorePaths []string `json:"auto-login-ignore-paths"`
-	Cookie               Cookie   `json:"cookie"`
-	EncryptionKey        string   `json:"encryption-key"`
-	Ingresses            []string `json:"ingress"`
-	LegacyCookie         bool     `json:"legacy-cookie"`
-	UpstreamAccessLogs   bool     `json:"upstream-access-logs"`
-	UpstreamHost         string   `json:"upstream-host"`
-	UpstreamIP           string   `json:"upstream-ip"`
-	UpstreamPort         int      `json:"upstream-port"`
+	AutoLogin              bool     `json:"auto-login"`
+	AutoLoginIgnorePaths   []string `json:"auto-login-ignore-paths"`
+	Cookie                 Cookie   `json:"cookie"`
+	EncryptionKey          string   `json:"encryption-key"`
+	Ingresses              []string `json:"ingress"`
+	LegacyCookie           bool     `json:"legacy-cookie"`
+	UpstreamAccessLogs     bool     `json:"upstream-access-logs"`
+	UpstreamHost           string   `json:"upstream-host"`
+	UpstreamIP             string   `json:"upstream-ip"`
+	UpstreamPort           int      `json:"upstream-port"`
+	UpstreamIncludeIdToken bool     `json:"upstream-include-id-token"`
 
 	OpenTelemetry OpenTelemetry `json:"otel"`
 	OpenID        OpenID        `json:"openid"`
@@ -50,13 +51,14 @@ const (
 	ShutdownGracefulPeriod   = "shutdown-graceful-period"
 	ShutdownWaitBeforePeriod = "shutdown-wait-before-period"
 
-	AutoLogin            = "auto-login"
-	AutoLoginIgnorePaths = "auto-login-ignore-paths"
-	Ingress              = "ingress"
-	UpstreamAccessLogs   = "upstream-access-logs"
-	UpstreamHost         = "upstream-host"
-	UpstreamIP           = "upstream-ip"
-	UpstreamPort         = "upstream-port"
+	AutoLogin              = "auto-login"
+	AutoLoginIgnorePaths   = "auto-login-ignore-paths"
+	Ingress                = "ingress"
+	UpstreamAccessLogs     = "upstream-access-logs"
+	UpstreamHost           = "upstream-host"
+	UpstreamIP             = "upstream-ip"
+	UpstreamPort           = "upstream-port"
+	UpstreamIncludeIdToken = "upstream-include-id-token"
 )
 
 var logger = log.WithField("logger", "wonderwall.config")
@@ -78,6 +80,7 @@ func Initialize() (*Config, error) {
 	flag.String(UpstreamHost, "127.0.0.1:8080", "Address of upstream host.")
 	flag.String(UpstreamIP, "", "IP of upstream host. Overrides 'upstream-host' if set.")
 	flag.Int(UpstreamPort, 0, "Port of upstream host. Overrides 'upstream-host' if set.")
+	flag.Bool(UpstreamIncludeIdToken, false, "Include ID token in upstream requests in 'X-Wonderwall-Id-Token' header.")
 
 	cookieFlags()
 	openidFlags()

diff --git a/pkg/handler/handler.go b/pkg/handler/handler.go
@@ -82,7 +82,7 @@ func NewStandalone(
 		Ingresses:      ingresses,
 		Redirect:       url.NewStandaloneRedirect(),
 		SessionManager: sessionManager,
-		UpstreamProxy:  NewUpstreamProxy(upstream, cfg.UpstreamAccessLogs),
+		UpstreamProxy:  NewUpstreamProxy(upstream, cfg.UpstreamAccessLogs, cfg.UpstreamIncludeIdToken),
 	}, nil
 }
 

diff --git a/pkg/handler/handler_sso_proxy.go b/pkg/handler/handler_sso_proxy.go
@@ -69,7 +69,7 @@ func NewSSOProxy(cfg *config.Config, crypter crypto.Crypter) (*SSOProxy, error)
 		SSOServerURL:          serverURL,
 		SSOServerReverseProxy: NewReverseProxy(serverURL, false),
 		SessionReader:         sessionReader,
-		UpstreamProxy:         NewUpstreamProxy(upstream, cfg.UpstreamAccessLogs),
+		UpstreamProxy:         NewUpstreamProxy(upstream, cfg.UpstreamAccessLogs, cfg.UpstreamIncludeIdToken),
 	}, nil
 }
 

diff --git a/pkg/handler/reverseproxy.go b/pkg/handler/reverseproxy.go
@@ -29,11 +29,13 @@ type ReverseProxySource interface {
 type ReverseProxy struct {
 	*httputil.ReverseProxy
 	EnableAccessLogs bool
+	IncludeIdToken   bool
 }
 
-func NewUpstreamProxy(upstream *urllib.URL, enableAccessLogs bool) *ReverseProxy {
+func NewUpstreamProxy(upstream *urllib.URL, enableAccessLogs bool, includeIdToken bool) *ReverseProxy {
 	rp := NewReverseProxy(upstream, true)
 	rp.EnableAccessLogs = enableAccessLogs
+	rp.IncludeIdToken = includeIdToken
 	return rp
 }
 
@@ -68,6 +70,11 @@ func NewReverseProxy(upstream *urllib.URL, preserveInboundHostHeader bool) *Reve
 			if ok {
 				r.Out.Header.Set("authorization", "Bearer "+accessToken)
 			}
+
+			idToken, ok := mw.IdTokenFrom(r.In.Context())
+			if ok {
+				r.Out.Header.Set("X-Wonderwall-Id-Token", idToken)
+			}
 		},
 		Transport: server.DefaultTransport(),
 	}
@@ -117,6 +124,10 @@ func (rp *ReverseProxy) Handler(src ReverseProxySource, w http.ResponseWriter, r
 
 	if isAuthenticated {
 		ctx = mw.WithAccessToken(ctx, accessToken)
+		if rp.IncludeIdToken {
+			idToken := sess.IDToken()
+			ctx = mw.WithIdToken(ctx, idToken)
+		}
 
 		if rp.EnableAccessLogs && isRelevantAccessLog(r) {
 			logger.Info("default: authenticated request")

diff --git a/pkg/handler/reverseproxy_test.go b/pkg/handler/reverseproxy_test.go
@@ -437,4 +437,45 @@ func TestReverseProxy(t *testing.T) {
 		}...)
 		assertUpstreamOKResponse(t, resp)
 	})
+
+	t.Run("request should not include idToken by default", func(t *testing.T) {
+		cfg := mock.Config()
+		cfg.UpstreamHost = up.URL.Host
+		idp := mock.NewIdentityProvider(cfg)
+		defer idp.Close()
+
+		up.SetIdentityProvider(idp)
+		rpClient := idp.RelyingPartyClient()
+
+		// acquire session
+		login(t, rpClient, idp)
+
+		up.requestCallback = func(r *http.Request) {
+			assert.Empty(t, r.Header.Get("x-wonderwall-id-token"))
+		}
+
+		resp := get(t, rpClient, idp.RelyingPartyServer.URL)
+		assertUpstreamOKResponse(t, resp)
+	})
+
+	t.Run("request should include idToken", func(t *testing.T) {
+		cfg := mock.Config()
+		cfg.UpstreamHost = up.URL.Host
+		cfg.UpstreamIncludeIdToken = true
+		idp := mock.NewIdentityProvider(cfg)
+		defer idp.Close()
+
+		up.SetIdentityProvider(idp)
+		rpClient := idp.RelyingPartyClient()
+
+		// acquire session
+		login(t, rpClient, idp)
+
+		up.requestCallback = func(r *http.Request) {
+			assert.NotEmpty(t, r.Header.Get("x-wonderwall-id-token"))
+		}
+
+		resp := get(t, rpClient, idp.RelyingPartyServer.URL)
+		assertUpstreamOKResponse(t, resp)
+	})
 }
diff --git a/pkg/middleware/context.go b/pkg/middleware/context.go
@@ -11,6 +11,7 @@ type contextKey string
 
 const (
 	ctxAccessToken = contextKey("AccessToken")
+	ctxIdToken     = contextKey("IdToken")
 	ctxIngress     = contextKey("Ingress")
 	ctxPath        = contextKey("Path")
 )
@@ -24,6 +25,15 @@ func WithAccessToken(ctx context.Context, accessToken string) context.Context {
 	return context.WithValue(ctx, ctxAccessToken, accessToken)
 }
 
+func IdTokenFrom(ctx context.Context) (string, bool) {
+	idToken, ok := ctx.Value(ctxIdToken).(string)
+	return idToken, ok
+}
+
+func WithIdToken(ctx context.Context, idToken string) context.Context {
+	return context.WithValue(ctx, ctxIdToken, idToken)
+}
+
 func IngressFrom(ctx context.Context) (ingress.Ingress, bool) {
 	i, ok := ctx.Value(ctxIngress).(ingress.Ingress)
 	return i, ok