Squashed commit of the following:

commit e4779e1
Author: splaunov <[email protected]>
Date:   Tue Sep 17 13:49:38 2024 +0300

    fix: transient_payload lost in API flow with session token exchange code (PS-482)

Author: github-actions[bot]

selfservice/flow/transient_payload.go

@@ -0,0 +1,38 @@
+package flow
+
+import (
+	"github.com/ory/x/sqlxx"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"
+)
+
+const internalContextTransientPayloadPath = "transient_payload"
+
+func SetTransientPayloadIntoInternalContext(flow InternalContexter, transientPayload sqlxx.JSONRawMessage) error {
+	if flow.GetInternalContext() == nil {
+		flow.EnsureInternalContext()
+	}
+	bytes, err := sjson.SetBytes(
+		flow.GetInternalContext(),
+		internalContextTransientPayloadPath,
+		transientPayload,
+	)
+	if err != nil {
+		return err
+	}
+	flow.SetInternalContext(bytes)
+
+	return nil
+}
+
+func GetTransientPayloadFromInternalContext(flow InternalContexter) (sqlxx.JSONRawMessage, error) {
+	if flow.GetInternalContext() == nil {
+		flow.EnsureInternalContext()
+	}
+	raw := gjson.GetBytes(flow.GetInternalContext(), internalContextTransientPayloadPath)
+	if !raw.IsObject() {
+		return nil, nil
+	}
+
+	return sqlxx.JSONRawMessage(raw.Raw), nil
+}

selfservice/strategy/oidc/strategy.go

@@ -316,6 +316,14 @@ func (s *Strategy) ValidateCallback(w http.ResponseWriter, r *http.Request, ps h
 		}
 		cntnr.State = stateParam
 		cntnr.FlowID = uuid.FromBytesOrNil(state.FlowId).String()
+		internalContexter, ok := f.(flow.InternalContexter)
+		if ok {
+			transientPayload, err := flow.GetTransientPayloadFromInternalContext(internalContexter)
+			if err != nil {
+				return nil, state, &cntnr, err
+			}
+			cntnr.TransientPayload = json.RawMessage(transientPayload)
+		}
 	}
 
 	if errorParam != "" {

selfservice/strategy/oidc/strategy_login.go

@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"github.com/ory/x/sqlxx"
 	"net/http"
 	"strings"
 	"time"
@@ -198,6 +199,12 @@ func (s *Strategy) Login(w http.ResponseWriter, r *http.Request, f *login.Flow,
 	f.IDToken = p.IDToken
 	f.RawIDTokenNonce = p.IDTokenNonce
 	f.TransientPayload = p.TransientPayload
+	if err := flow.SetTransientPayloadIntoInternalContext(f, sqlxx.JSONRawMessage(p.TransientPayload)); err != nil {
+		return nil, err
+	}
+	if err := s.d.LoginFlowPersister().UpdateLoginFlow(ctx, f); err != nil {
+		return nil, err
+	}
 
 	pid := p.Provider // this can come from both url query and post body
 	if pid == "" {

selfservice/strategy/oidc/strategy_registration.go

@@ -168,6 +168,12 @@ func (s *Strategy) Register(w http.ResponseWriter, r *http.Request, f *registrat
 	f.TransientPayload = p.TransientPayload
 	f.IDToken = p.IDToken
 	f.RawIDTokenNonce = p.IDTokenNonce
+	if err := flow.SetTransientPayloadIntoInternalContext(f, sqlxx.JSONRawMessage(p.TransientPayload)); err != nil {
+		return s.handleError(ctx, w, r, f, pid, nil, err)
+	}
+	if err := s.d.RegistrationFlowPersister().UpdateRegistrationFlow(ctx, f); err != nil {
+		return s.handleError(ctx, w, r, f, pid, nil, err)
+	}
 
 	if !strings.EqualFold(strings.ToLower(p.Method), s.SettingsStrategyID()) && p.Method != "" {
 		// the user is sending a method that is not oidc, but the payload includes a provider

selfservice/strategy/oidc/strategy_test.go

@@ -199,11 +199,13 @@ func TestStrategy(t *testing.T) {
 		return res, body
 	}
 
-	makeAPICodeFlowRequest := func(t *testing.T, provider, action string) (returnToURL *url.URL) {
-		res, err := testhelpers.NewDebugClient(t).Post(action, "application/json", strings.NewReader(fmt.Sprintf(`{
-	"method": "oidc",
-	"provider": %q
-}`, provider)))
+	makeAPICodeFlowRequest := func(t *testing.T, provider, action string, transientPayload string) (returnToURL *url.URL) {
+		res, err := testhelpers.NewDebugClient(t).Post(action, "application/json",
+			strings.NewReader(fmt.Sprintf(`{
+												"method": "oidc",
+												"provider": %q,
+												"transient_payload": %q
+											}`, provider, transientPayload)))
 		require.NoError(t, err)
 		require.Equal(t, http.StatusUnprocessableEntity, res.StatusCode)
 		var changeLocation flow.BrowserLocationChangeRequiredError
@@ -834,14 +836,25 @@ func TestStrategy(t *testing.T) {
 	})
 
 	t.Run("suite=API with session token exchange code", func(t *testing.T) {
+		postRegistrationWebhook := hooktest.NewServer()
+		t.Cleanup(postRegistrationWebhook.Close)
+		postRegistrationWebhook.SetConfig(t, conf.GetProvider(ctx),
+			config.HookStrategyKey(config.ViperKeySelfServiceRegistrationAfter, identity.CredentialsTypeOIDC.String()))
+
+		postLoginWebhook := hooktest.NewServer()
+		t.Cleanup(postLoginWebhook.Close)
+		postLoginWebhook.SetConfig(t, conf.GetProvider(ctx),
+			config.HookStrategyKey(config.ViperKeySelfServiceLoginAfter, config.HookGlobal))
+
 		scope = []string{"openid"}
+		transientPayload := `{"data": "registration"}`
 
 		loginOrRegister := func(t *testing.T, flowID uuid.UUID, code string) {
 			_, err := exchangeCodeForToken(t, sessiontokenexchange.Codes{InitCode: code})
 			require.Error(t, err)
 
 			action := assertFormValues(t, flowID, "valid")
-			returnToURL := makeAPICodeFlowRequest(t, "valid", action)
+			returnToURL := makeAPICodeFlowRequest(t, "valid", action, transientPayload)
 			returnToCode := returnToURL.Query().Get("code")
 			assert.NotEmpty(t, code, "code query param was empty in the return_to URL")
 
@@ -857,27 +870,39 @@ func TestStrategy(t *testing.T) {
 		performRegistration := func(t *testing.T) {
 			f := newAPIRegistrationFlow(t, returnTS.URL+"?return_session_token_exchange_code=true&return_to=/app_code", 1*time.Minute)
 			loginOrRegister(t, f.ID, f.SessionTokenExchangeCode)
+			postRegistrationWebhook.AssertTransientPayload(t, transientPayload)
+		}
+		startRegistrationButLogin := func(t *testing.T) {
+			f := newAPIRegistrationFlow(t, returnTS.URL+"?return_session_token_exchange_code=true&return_to=/app_code", 1*time.Minute)
+			loginOrRegister(t, f.ID, f.SessionTokenExchangeCode)
+			postLoginWebhook.AssertTransientPayload(t, transientPayload)
 		}
 		performLogin := func(t *testing.T) {
 			f := newAPILoginFlow(t, returnTS.URL+"?return_session_token_exchange_code=true&return_to=/app_code", 1*time.Minute)
 			loginOrRegister(t, f.ID, f.SessionTokenExchangeCode)
+			postLoginWebhook.AssertTransientPayload(t, transientPayload)
+		}
+		startLoginButRegister := func(t *testing.T) {
+			f := newAPILoginFlow(t, returnTS.URL+"?return_session_token_exchange_code=true&return_to=/app_code", 1*time.Minute)
+			loginOrRegister(t, f.ID, f.SessionTokenExchangeCode)
+			postRegistrationWebhook.AssertTransientPayload(t, transientPayload)
 		}
 
 		for _, tc := range []struct {
 			name        string
 			first, then func(*testing.T)
 		}{{
 			name:  "login-twice",
-			first: performLogin, then: performLogin,
+			first: startLoginButRegister, then: performLogin,
 		}, {
 			name:  "login-then-register",
-			first: performLogin, then: performRegistration,
+			first: startLoginButRegister, then: startRegistrationButLogin,
 		}, {
 			name:  "register-then-login",
 			first: performRegistration, then: performLogin,
 		}, {
 			name:  "register-twice",
-			first: performRegistration, then: performRegistration,
+			first: performRegistration, then: startRegistrationButLogin,
 		}} {
 			t.Run("case="+tc.name, func(t *testing.T) {
 				subject = tc.name + "[email protected]"
@@ -902,7 +927,7 @@ func TestStrategy(t *testing.T) {
 			require.Error(t, err)
 
 			action := assertFormValues(t, f.ID, "valid")
-			returnToURL := makeAPICodeFlowRequest(t, "valid", action)
+			returnToURL := makeAPICodeFlowRequest(t, "valid", action, "{}")
 			returnedFlow := returnToURL.Query().Get("flow")
 
 			require.NotEmpty(t, returnedFlow, "flow query param was empty in the return_to URL")