Add option to remove traits or roles from the app JWT header

api/proto/teleport/legacy/types/types.proto

@@ -1003,6 +1003,8 @@ message Rewrite {
   // Headers is a list of headers to inject when passing the request over
   // to the application.
   repeated Header Headers = 2 [(gogoproto.jsontag) = "headers,omitempty"];
+  // JWTClaims configures whether roles/traits are included in the JWT token.
+  string JWTClaims = 3 [(gogoproto.jsontag) = "jwt_claims,omitempty"];
 }
 
 // Header represents a single http header passed over to the proxied application.

api/types/app.go

@@ -368,6 +368,15 @@ func (a *AppV3) CheckAndSetDefaults() error {
 			constants.KubeTeleportProxyALPNPrefix, a.Spec.PublicAddr)
 	}
 
+	if a.Spec.Rewrite != nil {
+		switch a.Spec.Rewrite.JWTClaims {
+		case "", JWTClaimsRewriteRolesAndTraits, JWTClaimsRewriteRoles, JWTClaimsRewriteNone:
+		default:
+			return trace.BadParameter("app %q has unexpected JWT rewrite value %q", a.GetName(), a.Spec.Rewrite.JWTClaims)
+
+		}
+	}
+
 	return nil
 }
 

api/types/constants.go

@@ -1057,3 +1057,12 @@ const (
 	// already existing users are not deleted
 	TeleportServiceGroup = "teleport-system"
 )
+
+const (
+	// JWTClaimsRewriteRolesAndTraits includes both roles and traits in the JWT token.
+	JWTClaimsRewriteRolesAndTraits = "roles-and-traits"
+	// JWTClaimsRewriteRoles includes only the roles in the JWT token.
+	JWTClaimsRewriteRoles = "roles"
+	// JWTClaimsRewriteNone include neither traits nor roles in the JWT token.
+	JWTClaimsRewriteNone = "none"
+)

api/types/jwt.go

@@ -47,9 +47,6 @@ func (p *GenerateAppTokenRequest) Check() error {
 	if p.Username == "" {
 		return trace.BadParameter("username missing")
 	}
-	if len(p.Roles) == 0 {
-		return trace.BadParameter("roles missing")
-	}
 	if p.Expires.IsZero() {
 		return trace.BadParameter("expires missing")
 	}

lib/auth/auth.go

@@ -1746,6 +1746,8 @@ type AppTestCertRequest struct {
 	GCPServiceAccount string
 	// PinnedIP is optional IP to pin certificate to.
 	PinnedIP string
+	// LoginTrait is the login to include in the cert
+	LoginTrait string
 }
 
 // GenerateUserAppTestCert generates an application specific certificate, used
@@ -1768,6 +1770,12 @@ func (a *Server) GenerateUserAppTestCert(req AppTestCertRequest) ([]byte, error)
 	if sessionID == "" {
 		sessionID = uuid.New().String()
 	}
+
+	login := req.LoginTrait
+	if login == "" {
+		login = uuid.New().String()
+	}
+
 	certs, err := a.generateUserCert(certRequest{
 		user:      user,
 		publicKey: req.PublicKey,
@@ -1777,7 +1785,7 @@ func (a *Server) GenerateUserAppTestCert(req AppTestCertRequest) ([]byte, error)
 		// used to log into servers but SSH certificate generation code requires a
 		// principal be in the certificate.
 		traits: wrappers.Traits(map[string][]string{
-			constants.TraitLogins: {uuid.New().String()},
+			constants.TraitLogins: {login},
 		}),
 		// Only allow this certificate to be used for applications.
 		usage: []string{teleport.UsageAppsOnly},

lib/config/configuration.go

@@ -1613,8 +1613,9 @@ func applyAppsConfig(fc *FileConfig, cfg *servicecfg.Config) error {
 					application.Name)
 			}
 			app.Rewrite = &servicecfg.Rewrite{
-				Redirect: application.Rewrite.Redirect,
-				Headers:  headers,
+				Redirect:  application.Rewrite.Redirect,
+				Headers:   headers,
+				JWTClaims: application.Rewrite.JWTClaims,
 			}
 		}
 		if application.AWS != nil {

lib/config/fileconf.go

@@ -2086,6 +2086,8 @@ type Rewrite struct {
 	Redirect []string `yaml:"redirect"`
 	// Headers is a list of extra headers to inject in the request.
 	Headers []string `yaml:"headers,omitempty"`
+	// JWTClaims configures whether roles/traits are included in the JWT token
+	JWTClaims string `yaml:"jwt_claims,omitempty"`
 }
 
 // AppAWS contains additional options for AWS applications.

lib/jwt/jwt.go

@@ -132,9 +132,6 @@ func (p *SignParams) Check() error {
 	if p.Username == "" {
 		return trace.BadParameter("username missing")
 	}
-	if len(p.Roles) == 0 {
-		return trace.BadParameter("roles missing")
-	}
 	if p.Expires.IsZero() {
 		return trace.BadParameter("expires missing")
 	}

lib/service/service.go

@@ -4968,7 +4968,8 @@ func (process *TeleportProcess) initApps() {
 			var rewrite *types.Rewrite
 			if app.Rewrite != nil {
 				rewrite = &types.Rewrite{
-					Redirect: app.Rewrite.Redirect,
+					Redirect:  app.Rewrite.Redirect,
+					JWTClaims: app.Rewrite.JWTClaims,
 				}
 				for _, header := range app.Rewrite.Headers {
 					rewrite.Headers = append(rewrite.Headers,

lib/service/servicecfg/app.go

@@ -148,6 +148,8 @@ type Rewrite struct {
 	Redirect []string
 	// Headers is a list of extra headers to inject in the request.
 	Headers []Header
+	// JWTClaims configures whether roles/traits are included in the JWT token.
+	JWTClaims string
 }
 
 // Header represents a single http header passed over to the proxied application.

lib/srv/app/server_test.go

@@ -40,18 +40,21 @@ import (
 	"github.com/gorilla/websocket"
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/require"
+	"gopkg.in/square/go-jose.v2/jwt"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/constants"
 	"github.com/gravitational/teleport/api/defaults"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
+	"github.com/gravitational/teleport/api/types/wrappers"
 	"github.com/gravitational/teleport/lib/auth"
 	"github.com/gravitational/teleport/lib/auth/native"
 	"github.com/gravitational/teleport/lib/auth/testauthority"
 	"github.com/gravitational/teleport/lib/authz"
 	"github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/httplib/reverseproxy"
+	libjwt "github.com/gravitational/teleport/lib/jwt"
 	"github.com/gravitational/teleport/lib/labels"
 	"github.com/gravitational/teleport/lib/services"
 	"github.com/gravitational/teleport/lib/session"
@@ -89,6 +92,8 @@ type Suite struct {
 	user       types.User
 	role       types.Role
 	serverPort string
+
+	login string
 }
 
 func (s *Suite) TearDown(t *testing.T) {
@@ -128,6 +133,10 @@ type suiteConfig struct {
 	AppLabels map[string]string
 	// RoleAppLabels are the labels set to allow for the user role.
 	RoleAppLabels types.Labels
+	// Rewrite configures the rewrite rules for the app.
+	Rewrite *types.Rewrite
+	// Login is used to specify "login" trait in the jwt token
+	Login string
 }
 
 type fakeConnMonitor struct{}
@@ -146,6 +155,7 @@ func SetUpSuiteWithConfig(t *testing.T, config suiteConfig) *Suite {
 	s.clock = clockwork.NewFakeClock()
 	s.dataDir = t.TempDir()
 	s.hostUUID = uuid.New().String()
+	s.login = config.Login
 
 	var err error
 	// Create Auth Server.
@@ -262,6 +272,7 @@ func SetUpSuiteWithConfig(t *testing.T, config suiteConfig) *Suite {
 		PublicAddr:         "foo.example.com",
 		InsecureSkipVerify: true,
 		DynamicLabels:      types.LabelsToV2(dynamicLabels),
+		Rewrite:            config.Rewrite,
 	})
 	require.NoError(t, err)
 	s.appAWS, err = types.NewAppV3(types.Metadata{
@@ -362,6 +373,7 @@ func (s *Suite) generateCertificate(t *testing.T, user types.User, publicAddr, A
 		TTL:         1 * time.Hour,
 		PublicAddr:  publicAddr,
 		ClusterName: "root.example.com",
+		LoginTrait:  s.login,
 	}
 	if AWSRoleARN != "" {
 		req.AWSRoleARN = AWSRoleARN
@@ -712,6 +724,76 @@ func TestHandleConnectionHTTP2WS(t *testing.T) {
 	})
 }
 
+func TestRewriteJWT(t *testing.T) {
+	login := uuid.New().String()
+	for _, tc := range []struct {
+		name           string
+		expectedRoles  []string
+		expectedTraits wrappers.Traits
+		jwtRewrite     string
+	}{
+		{
+			name:          "test default behavior",
+			expectedRoles: []string{"foo"},
+			expectedTraits: wrappers.Traits{
+				"logins": []string{login},
+			},
+			jwtRewrite: "",
+		},
+		{
+			name:          "test roles-and-traits behavior",
+			expectedRoles: []string{"foo"},
+			expectedTraits: wrappers.Traits{
+				"logins": []string{login},
+			},
+			jwtRewrite: types.JWTClaimsRewriteRolesAndTraits,
+		},
+		{
+			name:           "test roles behavior",
+			expectedRoles:  []string{"foo"},
+			expectedTraits: wrappers.Traits{},
+			jwtRewrite:     types.JWTClaimsRewriteRoles,
+		},
+		{
+			name:           "test none behavior",
+			expectedRoles:  nil,
+			expectedTraits: wrappers.Traits{},
+			jwtRewrite:     types.JWTClaimsRewriteNone,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			s := SetUpSuiteWithConfig(t,
+				suiteConfig{
+					Rewrite: &types.Rewrite{
+						JWTClaims: tc.jwtRewrite,
+						Headers: []*types.Header{
+							{Name: "TestHeader", Value: "{{internal.jwt}}"},
+						},
+					},
+					ValidateRequest: func(s *Suite, r *http.Request) {
+						token, err := jwt.ParseSigned(r.Header.Get("TestHeader"))
+						require.NoError(t, err)
+
+						claims := libjwt.Claims{}
+						err = token.UnsafeClaimsWithoutVerification(&claims)
+						require.NoError(t, err)
+
+						require.Equal(t, tc.expectedTraits, claims.Traits)
+						require.Equal(t, tc.expectedRoles, claims.Roles)
+					},
+					Login: login,
+				})
+			s.checkHTTPResponse(t, s.clientCertificate, func(resp *http.Response) {
+				require.Equal(t, resp.StatusCode, http.StatusOK)
+				buf, err := io.ReadAll(resp.Body)
+				require.NoError(t, err)
+				require.Equal(t, strings.TrimSpace(string(buf)), s.message)
+			})
+		})
+	}
+
+}
+
 // TestAuthorize verifies that only authorized requests are handled.
 func TestAuthorize(t *testing.T) {
 	tests := []struct {

lib/srv/app/session.go

@@ -151,11 +151,25 @@ func (s *Server) newSessionChunk(ctx context.Context, identity *tlsca.Identity,
 // withJWTTokenForwarder is a sessionOpt that creates a forwarder that attaches
 // a generated JWT token to all requests.
 func (s *Server) withJWTTokenForwarder(ctx context.Context, sess *sessionChunk, identity *tlsca.Identity, app types.Application) error {
+	rewrite := app.GetRewrite()
+	traits := identity.Traits
+	roles := identity.Groups
+	if rewrite != nil {
+		switch rewrite.JWTClaims {
+		case types.JWTClaimsRewriteNone:
+			traits = nil
+			roles = nil
+		case types.JWTClaimsRewriteRoles:
+			traits = nil
+		case "", types.JWTClaimsRewriteRolesAndTraits:
+		}
+	}
+
 	// Request a JWT token that will be attached to all requests.
 	jwt, err := s.c.AuthClient.GenerateAppToken(ctx, types.GenerateAppTokenRequest{
 		Username: identity.Username,
-		Roles:    identity.Groups,
-		Traits:   identity.Traits,
+		Roles:    roles,
+		Traits:   traits,
 		URI:      app.GetURI(),
 		Expires:  identity.Expires,
 	})
@@ -164,7 +178,6 @@ func (s *Server) withJWTTokenForwarder(ctx context.Context, sess *sessionChunk,
 	}
 
 	// Add JWT token to the traits so it can be used in headers templating.
-	traits := identity.Traits
 	if traits == nil {
 		traits = make(wrappers.Traits)
 	}