test(oidc): test sourcing secondary claims from access token

Emyrk · Emyrk · commit 895bb00ec22e · 2025-02-10T15:16:53.000-06:00
diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -105,6 +105,7 @@ type FakeIDP struct {
 	// "Authorized Redirect URLs". This can be used to emulate that.
 	hookValidRedirectURL func(redirectURL string) error
 	hookUserInfo         func(email string) (jwt.MapClaims, error)
+	hookAccessTokenJWT   func(email string, exp time.Time) jwt.MapClaims
 	// defaultIDClaims is if a new client connects and we didn't preset
 	// some claims.
 	defaultIDClaims jwt.MapClaims
@@ -154,6 +155,12 @@ func WithMiddlewares(mws ...func(http.Handler) http.Handler) func(*FakeIDP) {
 	}
 }
 
+func WithAccessTokenJWTHook(hook func(email string, exp time.Time) jwt.MapClaims) func(*FakeIDP) {
+	return func(f *FakeIDP) {
+		f.hookAccessTokenJWT = hook
+	}
+}
+
 func WithHookWellKnown(hook func(r *http.Request, j *ProviderJSON) error) func(*FakeIDP) {
 	return func(f *FakeIDP) {
 		f.hookWellKnown = hook
@@ -316,8 +323,7 @@ const (
 func NewFakeIDP(t testing.TB, opts ...FakeIDPOpt) *FakeIDP {
 	t.Helper()
 
-	block, _ := pem.Decode([]byte(testRSAPrivateKey))
-	pkey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
+	pkey, err := FakeIDPKey()
 	require.NoError(t, err)
 
 	idp := &FakeIDP{
@@ -676,8 +682,13 @@ func (f *FakeIDP) newCode(state string) string {
 
 // newToken enforces the access token exchanged is actually a valid access token
 // created by the IDP.
-func (f *FakeIDP) newToken(email string, expires time.Time) string {
+func (f *FakeIDP) newToken(t testing.TB, email string, expires time.Time) string {
 	accessToken := uuid.NewString()
+	if f.hookAccessTokenJWT != nil {
+		claims := f.hookAccessTokenJWT(email, expires)
+		accessToken = f.encodeClaims(t, claims)
+	}
+
 	f.accessTokens.Store(accessToken, token{
 		issued: time.Now(),
 		email:  email,
@@ -963,7 +974,7 @@ func (f *FakeIDP) httpHandler(t testing.TB) http.Handler {
 		email := getEmail(claims)
 		refreshToken := f.newRefreshTokens(email)
 		token := map[string]interface{}{
-			"access_token":  f.newToken(email, exp),
+			"access_token":  f.newToken(t, email, exp),
 			"refresh_token": refreshToken,
 			"token_type":    "Bearer",
 			"expires_in":    int64((f.defaultExpire).Seconds()),
@@ -1553,3 +1564,8 @@ d8h4Ht09E+f3nhTEc87mODkl7WJZpHL6V2sORfeq/eIkds+H6CJ4hy5w/bSw8tjf
 sz9Di8sGIaUbLZI2rd0CQQCzlVwEtRtoNCyMJTTrkgUuNufLP19RZ5FpyXxBO5/u
 QastnN77KfUwdj3SJt44U/uh1jAIv4oSLBr8HYUkbnI8
 -----END RSA PRIVATE KEY-----`
+
+func FakeIDPKey() (*rsa.PrivateKey, error) {
+	block, _ := pem.Decode([]byte(testRSAPrivateKey))
+	return x509.ParsePKCS1PrivateKey(block.Bytes)
+}
diff --git a/coderd/userauth_test.go b/coderd/userauth_test.go
@@ -892,13 +892,15 @@ func TestUserOIDC(t *testing.T) {
 		Name                string
 		IDTokenClaims       jwt.MapClaims
 		UserInfoClaims      jwt.MapClaims
+		AccessTokenClaims   jwt.MapClaims
 		AllowSignups        bool
 		EmailDomain         []string
 		AssertUser          func(t testing.TB, u codersdk.User)
 		StatusCode          int
 		AssertResponse      func(t testing.TB, resp *http.Response)
 		IgnoreEmailVerified bool
 		IgnoreUserInfo      bool
+		UseAccessToken      bool
 	}{
 		{
 			Name: "NoSub",
@@ -908,6 +910,32 @@ func TestUserOIDC(t *testing.T) {
 			AllowSignups: true,
 			StatusCode:   http.StatusBadRequest,
 		},
+		{
+			Name: "AccessTokenMerge",
+			IDTokenClaims: jwt.MapClaims{
+				"sub": uuid.NewString(),
+			},
+			AccessTokenClaims: jwt.MapClaims{
+				"email": "kyle@kwc.io",
+			},
+			IgnoreUserInfo: true,
+			AllowSignups:   true,
+			UseAccessToken: true,
+			StatusCode:     http.StatusOK,
+			AssertUser: func(t testing.TB, u codersdk.User) {
+				assert.Equal(t, "kyle@kwc.io", u.Email)
+			},
+		},
+		{
+			Name: "AccessTokenMergeNotJWT",
+			IDTokenClaims: jwt.MapClaims{
+				"sub": uuid.NewString(),
+			},
+			IgnoreUserInfo: true,
+			AllowSignups:   true,
+			UseAccessToken: true,
+			StatusCode:     http.StatusBadRequest,
+		},
 		{
 			Name: "EmailOnly",
 			IDTokenClaims: jwt.MapClaims{
@@ -1290,20 +1318,31 @@ func TestUserOIDC(t *testing.T) {
 		tc := tc
 		t.Run(tc.Name, func(t *testing.T) {
 			t.Parallel()
-			fake := oidctest.NewFakeIDP(t,
+			opts := []oidctest.FakeIDPOpt{
 				oidctest.WithRefresh(func(_ string) error {
 					return xerrors.New("refreshing token should never occur")
 				}),
 				oidctest.WithServing(),
 				oidctest.WithStaticUserInfo(tc.UserInfoClaims),
-			)
+			}
+
+			if tc.AccessTokenClaims != nil && len(tc.AccessTokenClaims) > 0 {
+				opts = append(opts, oidctest.WithAccessTokenJWTHook(func(email string, exp time.Time) jwt.MapClaims {
+					return tc.AccessTokenClaims
+				}))
+			}
+
+			fake := oidctest.NewFakeIDP(t, opts...)
 			cfg := fake.OIDCConfig(t, nil, func(cfg *coderd.OIDCConfig) {
 				cfg.AllowSignups = tc.AllowSignups
 				cfg.EmailDomain = tc.EmailDomain
 				cfg.IgnoreEmailVerified = tc.IgnoreEmailVerified
 				if tc.IgnoreUserInfo {
 					cfg.SecondaryClaims = coderd.MergedClaimsSourceUserInfo
 				}
+				if tc.UseAccessToken {
+					cfg.SecondaryClaims = coderd.MergedClaimsSourceAccessToken
+				}
 				cfg.NameField = "name"
 			})
 
