diff --git a/lib/services/app.go b/lib/services/app.go index 1c2bd0b91f1..fc43307d422 100644 --- a/lib/services/app.go +++ b/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, diff --git a/lib/services/app_test.go b/lib/services/app_test.go index 290b9071a9b..e30321cf07c 100644 --- a/lib/services/app_test.go +++ b/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) diff --git a/lib/srv/app/common/jwt.go b/lib/srv/app/common/jwt.go index 65bad26dc24..1816654b3bc 100644 --- a/lib/srv/app/common/jwt.go +++ b/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 diff --git a/lib/srv/app/common/jwt_test.go b/lib/srv/app/common/jwt_test.go index 2ba8f30e7a0..56754162877 100644 --- a/lib/srv/app/common/jwt_test.go +++ b/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") +} diff --git a/lib/srv/app/session.go b/lib/srv/app/session.go index 40935f4d94d..107d35ab0c7 100644 --- a/lib/srv/app/session.go +++ b/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) diff --git a/lib/srv/app/transport.go b/lib/srv/app/transport.go index 1e2da267e0f..26a63e61b8f 100644 --- a/lib/srv/app/transport.go +++ b/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 } diff --git a/lib/srv/mcp/auth.go b/lib/srv/mcp/auth.go index bb7969f0041..3f8e01cc72c 100644 --- a/lib/srv/mcp/auth.go +++ b/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 { diff --git a/lib/srv/mcp/auth_test.go b/lib/srv/mcp/auth_test.go index 9c00dcc03c5..0882739aea3 100644 --- a/lib/srv/mcp/auth_test.go +++ b/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() diff --git a/lib/srv/mcp/session.go b/lib/srv/mcp/session.go index eb84fc44533..a4c38c2bfb5 100644 --- a/lib/srv/mcp/session.go +++ b/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 }