From 0a15d488c88aa62b0ab64a2099574b1619718fb7 Mon Sep 17 00:00:00 2001 From: Ryan Richard Date: Tue, 21 May 2024 14:50:35 -0700 Subject: [PATCH] Merge callback_handler_github_test.go into callback_handler_test.go Co-authored-by: Joshua Casey --- hack/prepare-supervisor-on-kind.sh | 7 +- .../callback/callback_handler_github_test.go | 272 ------------------ .../callback/callback_handler_test.go | 131 ++++++++- 3 files changed, 132 insertions(+), 278 deletions(-) delete mode 100644 internal/federationdomain/endpoints/callback/callback_handler_github_test.go diff --git a/hack/prepare-supervisor-on-kind.sh b/hack/prepare-supervisor-on-kind.sh index 0a8a971a3..0bd24c43f 100755 --- a/hack/prepare-supervisor-on-kind.sh +++ b/hack/prepare-supervisor-on-kind.sh @@ -323,10 +323,9 @@ stringData: EOF # Grant the test user some RBAC permissions so we can play with kubectl as that user. - # TODO -# kubectl create clusterrolebinding github-test-user-can-view --clusterrole view \ -# --user "$PINNIPED_TEST_GITHUB_TODO_WE_DONT_HAVE_THIS_VARIABLE_YET" \ -# --dry-run=client --output yaml | kubectl apply -f - + kubectl create clusterrolebinding github-test-user-can-view --clusterrole view \ + --user "$PINNIPED_TEST_GITHUB_USER_USERNAME:$PINNIPED_TEST_GITHUB_USERID" \ + --dry-run=client --output yaml | kubectl apply -f - fi # Create a CA and TLS serving certificates for the Supervisor's FederationDomain. diff --git a/internal/federationdomain/endpoints/callback/callback_handler_github_test.go b/internal/federationdomain/endpoints/callback/callback_handler_github_test.go deleted file mode 100644 index 1d506436d..000000000 --- a/internal/federationdomain/endpoints/callback/callback_handler_github_test.go +++ /dev/null @@ -1,272 +0,0 @@ -// Copyright 2020-2024 the Pinniped contributors. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package callback - -import ( - "context" - "fmt" - "net/http" - "net/http/httptest" - "regexp" - "strings" - "testing" - - "github.com/gorilla/securecookie" - "github.com/stretchr/testify/require" - "golang.org/x/crypto/bcrypt" - "k8s.io/apimachinery/pkg/types" - "k8s.io/client-go/kubernetes/fake" - - idpdiscoveryv1alpha1 "go.pinniped.dev/generated/latest/apis/supervisor/idpdiscovery/v1alpha1" - supervisorfake "go.pinniped.dev/generated/latest/client/supervisor/clientset/versioned/fake" - "go.pinniped.dev/internal/federationdomain/endpoints/jwks" - "go.pinniped.dev/internal/federationdomain/oidc" - "go.pinniped.dev/internal/federationdomain/oidcclientvalidator" - "go.pinniped.dev/internal/federationdomain/storage" - "go.pinniped.dev/internal/federationdomain/upstreamprovider" - "go.pinniped.dev/internal/psession" - "go.pinniped.dev/internal/testutil" - "go.pinniped.dev/internal/testutil/oidctestutil" - "go.pinniped.dev/internal/testutil/testidplister" -) - -var ( - githubIDPName = "upstream-github-idp-name" - githubIDPResourceUID = types.UID("upstream-github-idp-resource-uid") - githubUpstreamUsername = "some-github-login" - githubUpstreamGroups = []string{"org1/team1", "org2/team2"} - githubDownstreamSubject = fmt.Sprintf("https://github.com?idpName=%s&sub=%s", githubIDPName, githubUpstreamUsername) - githubUpstreamAccessToken = "some-opaque-access-token-from-github" //nolint:gosec // this is not a credential - - happyDownstreamGitHubCustomSessionData = &psession.CustomSessionData{ - Username: githubUpstreamUsername, - UpstreamUsername: githubUpstreamUsername, - UpstreamGroups: githubUpstreamGroups, - ProviderUID: githubIDPResourceUID, - ProviderName: githubIDPName, - ProviderType: psession.ProviderTypeGitHub, - GitHub: &psession.GitHubSessionData{ - UpstreamAccessToken: githubUpstreamAccessToken, - }, - } -) - -func TestCallbackEndpointWithGitHubIdentityProviders(t *testing.T) { - require.Len(t, happyDownstreamState, 8, "we expect fosite to allow 8 byte state params, so we want to test that boundary case") - - var stateEncoderHashKey = []byte("fake-hash-secret") - var stateEncoderBlockKey = []byte("0123456789ABCDEF") // block encryption requires 16/24/32 bytes for AES - var cookieEncoderHashKey = []byte("fake-hash-secret2") - var cookieEncoderBlockKey = []byte("0123456789ABCDE2") // block encryption requires 16/24/32 bytes for AES - require.NotEqual(t, stateEncoderHashKey, cookieEncoderHashKey) - require.NotEqual(t, stateEncoderBlockKey, cookieEncoderBlockKey) - - var happyStateCodec = securecookie.New(stateEncoderHashKey, stateEncoderBlockKey) - happyStateCodec.SetSerializer(securecookie.JSONEncoder{}) - var happyCookieCodec = securecookie.New(cookieEncoderHashKey, cookieEncoderBlockKey) - happyCookieCodec.SetSerializer(securecookie.JSONEncoder{}) - - encodedIncomingCookieCSRFValue, err := happyCookieCodec.Encode("csrf", happyDownstreamCSRF) - require.NoError(t, err) - happyCSRFCookie := "__Host-pinniped-csrf=" + encodedIncomingCookieCSRFValue - - happyExchangeAndValidateTokensArgs := &oidctestutil.ExchangeAuthcodeArgs{ - Authcode: happyUpstreamAuthcode, - RedirectURI: happyUpstreamRedirectURI, - } - - // TODO: when we merge this file back into callback_handler_test.go, we do not need to copy this function - // because it is already in callback_handler_test.go - addFullyCapableDynamicClientAndSecretToKubeResources := func(t *testing.T, supervisorClient *supervisorfake.Clientset, kubeClient *fake.Clientset) { - oidcClient, secret := testutil.FullyCapableOIDCClientAndStorageSecret(t, - "some-namespace", downstreamDynamicClientID, downstreamDynamicClientUID, downstreamRedirectURI, nil, - []string{testutil.HashedPassword1AtGoMinCost}, oidcclientvalidator.Validate) - require.NoError(t, supervisorClient.Tracker().Add(oidcClient)) - require.NoError(t, kubeClient.Tracker().Add(secret)) - } - - tests := []struct { - name string - - idps *testidplister.UpstreamIDPListerBuilder - kubeResources func(t *testing.T, supervisorClient *supervisorfake.Clientset, kubeClient *fake.Clientset) - method string - path string - csrfCookie string - - wantRedirectLocationRegexp string - wantDownstreamGrantedScopes []string - wantDownstreamIDTokenSubject string - wantDownstreamIDTokenUsername string - wantDownstreamIDTokenGroups []string - wantDownstreamRequestedScopes []string - wantDownstreamNonce string - wantDownstreamClientID string - wantDownstreamPKCEChallenge string - wantDownstreamPKCEChallengeMethod string - wantDownstreamCustomSessionData *psession.CustomSessionData - wantDownstreamAdditionalClaims map[string]interface{} - wantGitHubAuthcodeExchangeCall *expectedGitHubAuthcodeExchange - }{ - { - name: "GitHub IDP: GET with good state and cookie and successful upstream token exchange returns 303 to downstream client callback", - idps: testidplister.NewUpstreamIDPListerBuilder().WithGitHub( - happyGitHubUpstream(). - WithAccessToken(githubUpstreamAccessToken). - WithUser(&upstreamprovider.GitHubUser{ - Username: githubUpstreamUsername, - Groups: githubUpstreamGroups, - DownstreamSubject: githubDownstreamSubject, - }). - Build()), - method: http.MethodGet, - path: newRequestPath().WithState( - happyUpstreamStateParam(). - WithUpstreamIDPName(githubIDPName). - WithUpstreamIDPType(idpdiscoveryv1alpha1.IDPTypeGitHub). - WithAuthorizeRequestParams( - happyDownstreamRequestParamsQuery.Encode(), - ).Build(t, happyStateCodec), - ).String(), - csrfCookie: happyCSRFCookie, - wantRedirectLocationRegexp: downstreamRedirectURI + `\?code=([^&]+)&scope=` + regexp.QuoteMeta(strings.Join(happyDownstreamScopesGranted, "+")) + `&state=` + happyDownstreamState, - wantDownstreamIDTokenSubject: githubDownstreamSubject, - wantDownstreamIDTokenUsername: githubUpstreamUsername, - wantDownstreamIDTokenGroups: githubUpstreamGroups, - wantDownstreamRequestedScopes: happyDownstreamScopesRequested, - wantDownstreamGrantedScopes: happyDownstreamScopesGranted, - wantDownstreamNonce: downstreamNonce, - wantDownstreamClientID: downstreamPinnipedClientID, - wantDownstreamPKCEChallenge: downstreamPKCEChallenge, - wantDownstreamPKCEChallengeMethod: downstreamPKCEChallengeMethod, - wantDownstreamCustomSessionData: happyDownstreamGitHubCustomSessionData, - wantGitHubAuthcodeExchangeCall: &expectedGitHubAuthcodeExchange{ - performedByUpstreamName: githubIDPName, - args: happyExchangeAndValidateTokensArgs, - }, - }, - { - name: "GitHub IDP: GET with good state and cookie and successful upstream token exchange with dynamic client returns 303 to downstream client callback, with dynamic client", - idps: testidplister.NewUpstreamIDPListerBuilder().WithGitHub( - happyGitHubUpstream(). - WithAccessToken(githubUpstreamAccessToken). - WithUser(&upstreamprovider.GitHubUser{ - Username: githubUpstreamUsername, - Groups: githubUpstreamGroups, - DownstreamSubject: githubDownstreamSubject, - }). - Build()), - method: http.MethodGet, - kubeResources: addFullyCapableDynamicClientAndSecretToKubeResources, - path: newRequestPath().WithState( - happyUpstreamStateParam(). - WithUpstreamIDPName(githubIDPName). - WithUpstreamIDPType(idpdiscoveryv1alpha1.IDPTypeGitHub). - WithAuthorizeRequestParams( - shallowCopyAndModifyQuery( - happyDownstreamRequestParamsQuery, - map[string]string{ - "client_id": downstreamDynamicClientID, - }, - ).Encode(), - ).Build(t, happyStateCodec), - ).String(), - csrfCookie: happyCSRFCookie, - wantRedirectLocationRegexp: downstreamRedirectURI + `\?code=([^&]+)&scope=` + regexp.QuoteMeta(strings.Join(happyDownstreamScopesGranted, "+")) + `&state=` + happyDownstreamState, - wantDownstreamIDTokenSubject: githubDownstreamSubject, - wantDownstreamIDTokenUsername: githubUpstreamUsername, - wantDownstreamIDTokenGroups: githubUpstreamGroups, - wantDownstreamRequestedScopes: happyDownstreamScopesRequested, - wantDownstreamGrantedScopes: happyDownstreamScopesGranted, - wantDownstreamNonce: downstreamNonce, - wantDownstreamClientID: downstreamDynamicClientID, - wantDownstreamPKCEChallenge: downstreamPKCEChallenge, - wantDownstreamPKCEChallengeMethod: downstreamPKCEChallengeMethod, - wantDownstreamCustomSessionData: happyDownstreamGitHubCustomSessionData, - wantGitHubAuthcodeExchangeCall: &expectedGitHubAuthcodeExchange{ - performedByUpstreamName: githubIDPName, - args: happyExchangeAndValidateTokensArgs, - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - kubeClient := fake.NewSimpleClientset() - supervisorClient := supervisorfake.NewSimpleClientset() - secrets := kubeClient.CoreV1().Secrets("some-namespace") - oidcClientsClient := supervisorClient.ConfigV1alpha1().OIDCClients("some-namespace") - - if test.kubeResources != nil { - test.kubeResources(t, supervisorClient, kubeClient) - } - - // Configure fosite the same way that the production code would. - // Inject this into our test subject at the last second, so we get a fresh storage for every test. - timeoutsConfiguration := oidc.DefaultOIDCTimeoutsConfiguration() - // Use lower minimum required bcrypt cost than we would use in production to keep unit the tests fast. - oauthStore := storage.NewKubeStorage(secrets, oidcClientsClient, timeoutsConfiguration, bcrypt.MinCost) - hmacSecretFunc := func() []byte { return []byte("some secret - must have at least 32 bytes") } - require.GreaterOrEqual(t, len(hmacSecretFunc()), 32, "fosite requires that hmac secrets have at least 32 bytes") - jwksProviderIsUnused := jwks.NewDynamicJWKSProvider() - oauthHelper := oidc.FositeOauth2Helper(oauthStore, downstreamIssuer, hmacSecretFunc, jwksProviderIsUnused, timeoutsConfiguration) - - subject := NewHandler(test.idps.BuildFederationDomainIdentityProvidersListerFinder(), oauthHelper, happyStateCodec, happyCookieCodec, happyUpstreamRedirectURI) - reqContext := context.WithValue(context.Background(), struct{ name string }{name: "test"}, "request-context") - req := httptest.NewRequest(test.method, test.path, nil).WithContext(reqContext) - if test.csrfCookie != "" { - req.Header.Set("Cookie", test.csrfCookie) - } - rsp := httptest.NewRecorder() - subject.ServeHTTP(rsp, req) - t.Logf("response: %#v", rsp) - t.Logf("response body: %q", rsp.Body.String()) - - testutil.RequireSecurityHeadersWithFormPostPageCSPs(t, rsp) - - require.NotNil(t, test.wantGitHubAuthcodeExchangeCall, "wantOIDCAuthcodeExchangeCall is required for testing purposes") - - test.wantGitHubAuthcodeExchangeCall.args.Ctx = reqContext - test.idps.RequireExactlyOneGitHubAuthcodeExchange(t, - test.wantGitHubAuthcodeExchangeCall.performedByUpstreamName, - test.wantGitHubAuthcodeExchangeCall.args, - ) - - require.Equal(t, http.StatusSeeOther, rsp.Code) - testutil.RequireEqualContentType(t, rsp.Header().Get("Content-Type"), "") - require.Empty(t, rsp.Body.String()) - - require.Len(t, rsp.Header().Values("Location"), 1) - require.NotEmpty(t, test.wantRedirectLocationRegexp, "wantRedirectLocationRegexp is required for testing purposes") - oidctestutil.RequireAuthCodeRegexpMatch( - t, - rsp.Header().Get("Location"), - test.wantRedirectLocationRegexp, - kubeClient, - secrets, - oauthStore, - test.wantDownstreamGrantedScopes, - test.wantDownstreamIDTokenSubject, - test.wantDownstreamIDTokenUsername, - test.wantDownstreamIDTokenGroups, - test.wantDownstreamRequestedScopes, - test.wantDownstreamPKCEChallenge, - test.wantDownstreamPKCEChallengeMethod, - test.wantDownstreamNonce, - test.wantDownstreamClientID, - downstreamRedirectURI, - test.wantDownstreamCustomSessionData, - test.wantDownstreamAdditionalClaims, - ) - }) - } -} - -func happyGitHubUpstream() *oidctestutil.TestUpstreamGitHubIdentityProviderBuilder { - return oidctestutil.NewTestUpstreamGitHubIdentityProviderBuilder(). - WithName(githubIDPName). - WithResourceUID(githubIDPResourceUID). - WithClientID("some-client-id"). - WithScopes([]string{"these", "scopes", "appear", "unused"}) -} diff --git a/internal/federationdomain/endpoints/callback/callback_handler_test.go b/internal/federationdomain/endpoints/callback/callback_handler_test.go index 7818e7f0a..4968f41bd 100644 --- a/internal/federationdomain/endpoints/callback/callback_handler_test.go +++ b/internal/federationdomain/endpoints/callback/callback_handler_test.go @@ -6,6 +6,7 @@ package callback import ( "context" "errors" + "fmt" "net/http" "net/http/httptest" "net/url" @@ -17,14 +18,17 @@ import ( "github.com/stretchr/testify/require" "golang.org/x/crypto/bcrypt" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/kubernetes/fake" configv1alpha1 "go.pinniped.dev/generated/latest/apis/supervisor/config/v1alpha1" + idpdiscoveryv1alpha1 "go.pinniped.dev/generated/latest/apis/supervisor/idpdiscovery/v1alpha1" supervisorfake "go.pinniped.dev/generated/latest/client/supervisor/clientset/versioned/fake" "go.pinniped.dev/internal/federationdomain/endpoints/jwks" "go.pinniped.dev/internal/federationdomain/oidc" "go.pinniped.dev/internal/federationdomain/oidcclientvalidator" "go.pinniped.dev/internal/federationdomain/storage" + "go.pinniped.dev/internal/federationdomain/upstreamprovider" "go.pinniped.dev/internal/psession" "go.pinniped.dev/internal/testutil" "go.pinniped.dev/internal/testutil/oidctestutil" @@ -73,6 +77,13 @@ const ( ) var ( + githubIDPName = "upstream-github-idp-name" + githubIDPResourceUID = types.UID("upstream-github-idp-resource-uid") + githubUpstreamUsername = "some-github-login" + githubUpstreamGroups = []string{"org1/team1", "org2/team2"} + githubDownstreamSubject = fmt.Sprintf("https://github.com?idpName=%s&sub=%s", githubIDPName, githubUpstreamUsername) + githubUpstreamAccessToken = "some-opaque-access-token-from-github" //nolint:gosec // this is not a credential + oidcUpstreamGroupMembership = []string{"test-pinniped-group-0", "test-pinniped-group-1"} happyDownstreamScopesRequested = []string{"openid", "username", "groups"} happyDownstreamScopesGranted = []string{"openid", "username", "groups"} @@ -129,8 +140,27 @@ var ( UpstreamSubject: oidcUpstreamSubject, }, } + happyDownstreamGitHubCustomSessionData = &psession.CustomSessionData{ + Username: githubUpstreamUsername, + UpstreamUsername: githubUpstreamUsername, + UpstreamGroups: githubUpstreamGroups, + ProviderUID: githubIDPResourceUID, + ProviderName: githubIDPName, + ProviderType: psession.ProviderTypeGitHub, + GitHub: &psession.GitHubSessionData{ + UpstreamAccessToken: githubUpstreamAccessToken, + }, + } ) +func happyGitHubUpstream() *oidctestutil.TestUpstreamGitHubIdentityProviderBuilder { + return oidctestutil.NewTestUpstreamGitHubIdentityProviderBuilder(). + WithName(githubIDPName). + WithResourceUID(githubIDPResourceUID). + WithClientID("some-client-id"). + WithScopes([]string{"these", "scopes", "appear", "unused"}) +} + func TestCallbackEndpoint(t *testing.T) { require.Len(t, happyDownstreamState, 8, "we expect fosite to allow 8 byte state params, so we want to test that boundary case") @@ -166,6 +196,11 @@ func TestCallbackEndpoint(t *testing.T) { RedirectURI: happyUpstreamRedirectURI, } + happyGitHubExchangeAuthcodeArgs := &oidctestutil.ExchangeAuthcodeArgs{ + Authcode: happyUpstreamAuthcode, + RedirectURI: happyUpstreamRedirectURI, + } + // Note that fosite puts the granted scopes as a param in the redirect URI even though the spec doesn't seem to require it happyDownstreamRedirectLocationRegexp := downstreamRedirectURI + `\?code=([^&]+)&scope=openid\+username\+groups&state=` + happyDownstreamState @@ -206,6 +241,7 @@ func TestCallbackEndpoint(t *testing.T) { wantDownstreamCustomSessionData *psession.CustomSessionData wantDownstreamAdditionalClaims map[string]interface{} wantOIDCAuthcodeExchangeCall *expectedOIDCAuthcodeExchange + wantGitHubAuthcodeExchangeCall *expectedGitHubAuthcodeExchange }{ { name: "GET with good state and cookie and successful upstream token exchange with response_mode=form_post returns 200 with HTML+JS form", @@ -1231,6 +1267,90 @@ func TestCallbackEndpoint(t *testing.T) { args: happyExchangeAndValidateTokensArgs, }, }, + { + name: "GitHub IDP: GET with good state and cookie and successful upstream token exchange returns 303 to downstream client callback", + idps: testidplister.NewUpstreamIDPListerBuilder().WithGitHub( + happyGitHubUpstream(). + WithAccessToken(githubUpstreamAccessToken). + WithUser(&upstreamprovider.GitHubUser{ + Username: githubUpstreamUsername, + Groups: githubUpstreamGroups, + DownstreamSubject: githubDownstreamSubject, + }). + Build()), + method: http.MethodGet, + path: newRequestPath().WithState( + happyUpstreamStateParam(). + WithUpstreamIDPName(githubIDPName). + WithUpstreamIDPType(idpdiscoveryv1alpha1.IDPTypeGitHub). + WithAuthorizeRequestParams( + happyDownstreamRequestParamsQuery.Encode(), + ).Build(t, happyStateCodec), + ).String(), + csrfCookie: happyCSRFCookie, + wantStatus: http.StatusSeeOther, + wantRedirectLocationRegexp: happyDownstreamRedirectLocationRegexp, + wantBody: "", + wantDownstreamIDTokenSubject: githubDownstreamSubject, + wantDownstreamIDTokenUsername: githubUpstreamUsername, + wantDownstreamIDTokenGroups: githubUpstreamGroups, + wantDownstreamRequestedScopes: happyDownstreamScopesRequested, + wantDownstreamGrantedScopes: happyDownstreamScopesGranted, + wantDownstreamNonce: downstreamNonce, + wantDownstreamClientID: downstreamPinnipedClientID, + wantDownstreamPKCEChallenge: downstreamPKCEChallenge, + wantDownstreamPKCEChallengeMethod: downstreamPKCEChallengeMethod, + wantDownstreamCustomSessionData: happyDownstreamGitHubCustomSessionData, + wantGitHubAuthcodeExchangeCall: &expectedGitHubAuthcodeExchange{ + performedByUpstreamName: githubIDPName, + args: happyGitHubExchangeAuthcodeArgs, + }, + }, + { + name: "GitHub IDP: GET with good state and cookie and successful upstream token exchange with dynamic client returns 303 to downstream client callback, with dynamic client", + idps: testidplister.NewUpstreamIDPListerBuilder().WithGitHub( + happyGitHubUpstream(). + WithAccessToken(githubUpstreamAccessToken). + WithUser(&upstreamprovider.GitHubUser{ + Username: githubUpstreamUsername, + Groups: githubUpstreamGroups, + DownstreamSubject: githubDownstreamSubject, + }). + Build()), + method: http.MethodGet, + kubeResources: addFullyCapableDynamicClientAndSecretToKubeResources, + path: newRequestPath().WithState( + happyUpstreamStateParam(). + WithUpstreamIDPName(githubIDPName). + WithUpstreamIDPType(idpdiscoveryv1alpha1.IDPTypeGitHub). + WithAuthorizeRequestParams( + shallowCopyAndModifyQuery( + happyDownstreamRequestParamsQuery, + map[string]string{ + "client_id": downstreamDynamicClientID, + }, + ).Encode(), + ).Build(t, happyStateCodec), + ).String(), + csrfCookie: happyCSRFCookie, + wantStatus: http.StatusSeeOther, + wantRedirectLocationRegexp: happyDownstreamRedirectLocationRegexp, + wantBody: "", + wantDownstreamIDTokenSubject: githubDownstreamSubject, + wantDownstreamIDTokenUsername: githubUpstreamUsername, + wantDownstreamIDTokenGroups: githubUpstreamGroups, + wantDownstreamRequestedScopes: happyDownstreamScopesRequested, + wantDownstreamGrantedScopes: happyDownstreamScopesGranted, + wantDownstreamNonce: downstreamNonce, + wantDownstreamClientID: downstreamDynamicClientID, + wantDownstreamPKCEChallenge: downstreamPKCEChallenge, + wantDownstreamPKCEChallengeMethod: downstreamPKCEChallengeMethod, + wantDownstreamCustomSessionData: happyDownstreamGitHubCustomSessionData, + wantGitHubAuthcodeExchangeCall: &expectedGitHubAuthcodeExchange{ + performedByUpstreamName: githubIDPName, + args: happyGitHubExchangeAuthcodeArgs, + }, + }, { name: "the OIDCIdentityProvider resource has been deleted", idps: testidplister.NewUpstreamIDPListerBuilder().WithOIDC(otherUpstreamOIDCIdentityProvider), @@ -1561,13 +1681,20 @@ func TestCallbackEndpoint(t *testing.T) { testutil.RequireSecurityHeadersWithFormPostPageCSPs(t, rsp) - if test.wantOIDCAuthcodeExchangeCall != nil { + switch { + case test.wantOIDCAuthcodeExchangeCall != nil: test.wantOIDCAuthcodeExchangeCall.args.Ctx = reqContext test.idps.RequireExactlyOneOIDCAuthcodeExchange(t, test.wantOIDCAuthcodeExchangeCall.performedByUpstreamName, test.wantOIDCAuthcodeExchangeCall.args, ) - } else { + case test.wantGitHubAuthcodeExchangeCall != nil: + test.wantGitHubAuthcodeExchangeCall.args.Ctx = reqContext + test.idps.RequireExactlyOneGitHubAuthcodeExchange(t, + test.wantGitHubAuthcodeExchangeCall.performedByUpstreamName, + test.wantGitHubAuthcodeExchangeCall.args, + ) + default: test.idps.RequireExactlyZeroAuthcodeExchanges(t) }