fix: jwt token leak in audit event user_traits for app and mcp sessions (#66090)

refactor: rename traits to rewriteTraits in JWT handling and header rewriting

chore: minor updates to docs and tests based on pr feedback

lib/services/app.go

@@ -412,9 +412,9 @@ func getClusterDomain() string {
 
 // RewriteHeadersAndApplyValueTraits rewrites the provided request's headers
 // while applying value traits to them.
-func RewriteHeadersAndApplyValueTraits(r *http.Request, rewrites iter.Seq[*types.Header], traits wrappers.Traits, log *slog.Logger) {
+func RewriteHeadersAndApplyValueTraits(r *http.Request, rewrites iter.Seq[*types.Header], rewriteTraits wrappers.Traits, log *slog.Logger) {
 	for header := range rewrites {
-		values, err := ApplyValueTraits(header.Value, traits)
+		values, err := ApplyValueTraits(header.Value, rewriteTraits)
 		if err != nil {
 			log.DebugContext(r.Context(), "Failed to apply traits",
 				"header_value", header.Value,

lib/services/app_test.go

@@ -577,10 +577,10 @@ func TestRewriteHeadersAndApplyValueTraits(t *testing.T) {
 		// Missing traits should log a debug message that this rewrite is skipped.
 		{Name: "x-bad-rewrite", Value: "{{external.bad_rewrite}}"},
 	}
-	traits := map[string][]string{
+	rewriteTraits := map[string][]string{
 		"rewrite": {"value1", "value2"},
 	}
-	RewriteHeadersAndApplyValueTraits(r, slices.Values(rewrites), traits, slog.Default())
+	RewriteHeadersAndApplyValueTraits(r, slices.Values(rewrites), rewriteTraits, slog.Default())
 
 	assert.Equal(t, "1.2.3.4", r.Host)
 	wantHeaders := make(http.Header)

lib/srv/app/common/jwt.go

@@ -20,6 +20,7 @@ package common
 
 import (
 	"context"
+	"maps"
 	"time"
 
 	"github.com/gravitational/trace"
@@ -37,7 +38,8 @@ type AppTokenGenerator interface {
 }
 
 // GenerateJWTAndTraits is helper that generates a JWT for an application and
-// populates the user traits with the result JWT for templating.
+// populates the rewrite traits with the result JWT for templating. On success, the
+// returned rewrite traits map is guaranteed to be non-nil.
 func GenerateJWTAndTraits(
 	ctx context.Context,
 	identity *tlsca.Identity,
@@ -58,11 +60,12 @@ func GenerateJWTAndTraits(
 	if err != nil {
 		return "", nil, trace.Wrap(err)
 	}
-	if traits == nil {
-		traits = make(wrappers.Traits)
+	rewriteTraits := maps.Clone(traits)
+	if rewriteTraits == nil {
+		rewriteTraits = make(wrappers.Traits)
 	}
-	traits[constants.TraitJWT] = []string{jwt}
-	return jwt, traits, trace.Wrap(err)
+	rewriteTraits[constants.TraitJWT] = []string{jwt}
+	return jwt, rewriteTraits, nil
 }
 
 // RolesAndTraitsForAppToken is a helper to populate roles and traits that are

lib/srv/app/common/jwt_test.go

@@ -19,11 +19,15 @@
 package common
 
 import (
+	"context"
+	"maps"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/gravitational/teleport/api/constants"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/types/wrappers"
 	"github.com/gravitational/teleport/lib/tlsca"
@@ -92,3 +96,41 @@ func TestRolesAndTraitsForAppToken(t *testing.T) {
 		})
 	}
 }
+
+type fakeTokenGenerator struct{}
+
+func (f fakeTokenGenerator) GenerateAppToken(_ context.Context, _ types.GenerateAppTokenRequest) (string, error) {
+	return "fake-jwt-token", nil
+}
+
+func TestGenerateJWTAndTraitsDoesNotMutateIdentity(t *testing.T) {
+	identity := &tlsca.Identity{
+		Username: "test",
+		Groups:   []string{"access", "editor"},
+		Traits: wrappers.Traits{
+			"team": []string{"dev"},
+		},
+	}
+	originalTraits := maps.Clone(identity.Traits)
+
+	app, err := types.NewAppV3(
+		types.Metadata{Name: "test-app"},
+		types.AppSpecV3{URI: "http://localhost:12345"},
+	)
+	require.NoError(t, err)
+
+	jwt, rewriteTraits, err := GenerateJWTAndTraits(
+		t.Context(),
+		identity,
+		app,
+		&fakeTokenGenerator{},
+		time.Now().Add(time.Hour).In(time.UTC),
+	)
+	require.NoError(t, err)
+	assert.NotEmpty(t, jwt)
+
+	assert.Equal(t, []string{"fake-jwt-token"}, rewriteTraits[constants.TraitJWT])
+
+	assert.Equal(t, originalTraits, identity.Traits)
+	assert.Empty(t, identity.Traits[constants.TraitJWT], "identity.Traits must not contain the JWT")
+}

lib/srv/app/session.go

@@ -151,7 +151,7 @@ func (c *ConnectionsHandler) newSessionChunk(ctx context.Context, identity *tlsc
 func (c *ConnectionsHandler) withJWTTokenForwarder(ctx context.Context, sess *sessionChunk, identity *tlsca.Identity, app types.Application) error {
 	// TODO(greedy52) consider using a shorter ttl for the token. The chunk is
 	// only 5 minutes anyway.
-	jwt, traits, err := common.GenerateJWTAndTraits(ctx, identity, app, c.cfg.AuthClient, identity.Expires)
+	jwt, rewriteTraits, err := common.GenerateJWTAndTraits(ctx, identity, app, c.cfg.AuthClient, identity.Expires)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -159,14 +159,14 @@ func (c *ConnectionsHandler) withJWTTokenForwarder(ctx context.Context, sess *se
 	// Create a rewriting transport that will be used to forward requests.
 	transport, err := newTransport(c.closeContext,
 		&transportConfig{
-			app:          app,
-			publicPort:   c.proxyPort,
-			cipherSuites: c.cfg.CipherSuites,
-			jwt:          jwt,
-			traits:       traits,
-			log:          c.log,
-			hostID:       c.cfg.HostID,
-			insecureMode: c.cfg.InsecureMode,
+			app:           app,
+			publicPort:    c.proxyPort,
+			cipherSuites:  c.cfg.CipherSuites,
+			jwt:           jwt,
+			rewriteTraits: rewriteTraits,
+			log:           c.log,
+			hostID:        c.cfg.HostID,
+			insecureMode:  c.cfg.InsecureMode,
 		})
 	if err != nil {
 		return trace.Wrap(err)

lib/srv/app/transport.go

@@ -44,12 +44,12 @@ import (
 
 // transportConfig is configuration for a rewriting transport.
 type transportConfig struct {
-	app          types.Application
-	publicPort   string
-	cipherSuites []uint16
-	jwt          string
-	traits       wrappers.Traits
-	log          *slog.Logger
+	app           types.Application
+	publicPort    string
+	cipherSuites  []uint16
+	jwt           string
+	rewriteTraits wrappers.Traits
+	log           *slog.Logger
 	// hostID is purely for troubleshooting purposes (put in the error messages)
 	hostID       string
 	insecureMode bool
@@ -208,7 +208,7 @@ func (t *transport) rewriteRequest(r *http.Request) error {
 	r.Header.Set(teleport.AppJWTHeader, t.jwt)
 	// Add headers from rewrite configuration.
 	rewriteHeaders := common.AppRewriteHeaders(r.Context(), t.app.GetRewrite(), t.log)
-	services.RewriteHeadersAndApplyValueTraits(r, rewriteHeaders, t.traits, t.log)
+	services.RewriteHeadersAndApplyValueTraits(r, rewriteHeaders, t.rewriteTraits, t.log)
 	return nil
 }
 

lib/srv/mcp/auth.go

@@ -48,10 +48,10 @@ type sessionAuth struct {
 	authClient appcommon.AppTokenGenerator
 	clock      clockwork.Clock
 
-	mu         sync.Mutex
-	jwt        string
-	traits     wrappers.Traits
-	lastUpdate time.Time
+	mu            sync.Mutex
+	jwt           string
+	rewriteTraits wrappers.Traits
+	lastUpdate    time.Time
 }
 
 func (a *sessionAuth) generateJWTAndTraits(ctx context.Context) (string, wrappers.Traits, error) {
@@ -59,7 +59,7 @@ func (a *sessionAuth) generateJWTAndTraits(ctx context.Context) (string, wrapper
 	defer a.mu.Unlock()
 
 	if a.clock.Now().Before(a.lastUpdate.Add(maxTokenDuration)) {
-		return a.jwt, a.traits, nil
+		return a.jwt, a.rewriteTraits, nil
 	}
 
 	// Note that token validation on server side usually has some leeway like a
@@ -70,7 +70,7 @@ func (a *sessionAuth) generateJWTAndTraits(ctx context.Context) (string, wrapper
 		expires = maxExpires
 	}
 
-	jwt, traitsForRewriteHeaders, err := appcommon.GenerateJWTAndTraits(ctx, &a.Identity, a.App, a.authClient, expires)
+	jwt, rewriteTraits, err := appcommon.GenerateJWTAndTraits(ctx, &a.Identity, a.App, a.authClient, expires)
 	if err != nil {
 		return "", nil, trace.Wrap(err)
 	}
@@ -80,13 +80,13 @@ func (a *sessionAuth) generateJWTAndTraits(ctx context.Context) (string, wrapper
 		if err != nil {
 			return "", nil, trace.Wrap(err)
 		}
-		traitsForRewriteHeaders[constants.TraitIDToken] = []string{idToken}
+		rewriteTraits[constants.TraitIDToken] = []string{idToken}
 	}
 
 	a.jwt = jwt
-	a.traits = traitsForRewriteHeaders
+	a.rewriteTraits = rewriteTraits
 	a.lastUpdate = now
-	return jwt, traitsForRewriteHeaders, nil
+	return jwt, rewriteTraits, nil
 }
 
 type rewriteAuthDetails struct {

lib/srv/mcp/auth_test.go

@@ -145,12 +145,12 @@ func Test_generateJWTAndTraits(t *testing.T) {
 		clock:      clock,
 	}
 
-	jwt, traits, err := auth.generateJWTAndTraits(t.Context())
+	jwt, rewriteTraits, err := auth.generateJWTAndTraits(t.Context())
 	require.NoError(t, err)
 	require.Equal(t, "app-token-for-ai-by-jwt", jwt)
-	require.NotEmpty(t, traits)
-	require.Equal(t, []string{"app-token-for-ai-by-jwt"}, traits[constants.TraitJWT])
-	require.Equal(t, []string{"app-token-for-ai-by-oidc_idp"}, traits[constants.TraitIDToken])
+	require.NotEmpty(t, rewriteTraits)
+	require.Equal(t, []string{"app-token-for-ai-by-jwt"}, rewriteTraits[constants.TraitJWT])
+	require.Equal(t, []string{"app-token-for-ai-by-oidc_idp"}, rewriteTraits[constants.TraitIDToken])
 
 	// Two calls, one for JWT, and one for ID token.
 	appTokenRequests := authClient.getAppTokenRequests()

lib/srv/mcp/session.go

@@ -313,7 +313,7 @@ func (s *sessionHandler) makeToolsCallResponse(ctx context.Context, resp *mcputi
 }
 
 func (s *sessionHandler) rewriteHTTPRequestHeaders(r *http.Request) error {
-	jwt, traits, err := s.generateJWTAndTraits(r.Context())
+	jwt, rewriteTraits, err := s.generateJWTAndTraits(r.Context())
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -327,7 +327,7 @@ func (s *sessionHandler) rewriteHTTPRequestHeaders(r *http.Request) error {
 	r.Header.Set(teleport.AppJWTHeader, jwt)
 	// Add headers from rewrite configuration.
 	rewriteHeaders := appcommon.AppRewriteHeaders(r.Context(), s.App.GetRewrite(), s.logger)
-	services.RewriteHeadersAndApplyValueTraits(r, rewriteHeaders, traits, s.logger)
+	services.RewriteHeadersAndApplyValueTraits(r, rewriteHeaders, rewriteTraits, s.logger)
 	return nil
 }