chore: u

ory · Mar 25, 2024 · 151f092 · 151f092
1 parent 8537f71
commit 151f092
Show file tree

Hide file tree

Showing 4 changed files with 138 additions and 51 deletions.
diff --git a/selfservice/flow/login/flow.go b/selfservice/flow/login/flow.go
@@ -152,6 +152,9 @@ type Flow struct {
 	// It can, for example, contain a reference to the verification flow, created as part of the user's
 	// registration.
 	ContinueWithItems []flow.ContinueWith `json:"-" db:"-" faker:"-" `
+
+	// ReturnToVerification contains the redirect URL for the verification flow.
+	ReturnToVerification string `json:"-" db:"-"`
 }
 
 var _ flow.Flow = new(Flow)
@@ -320,3 +323,7 @@ func (f *Flow) AddContinueWith(c flow.ContinueWith) {
 func (f *Flow) ContinueWith() []flow.ContinueWith {
 	return f.ContinueWithItems
 }
+
+func (f *Flow) SetReturnToVerification(to string) {
+	f.ReturnToVerification = to
+}
diff --git a/selfservice/flow/registration/flow.go b/selfservice/flow/registration/flow.go
@@ -275,3 +275,7 @@ func (f *Flow) SetState(state State) {
 func (t *Flow) GetTransientPayload() json.RawMessage {
 	return t.TransientPayload
 }
+
+func (f *Flow) SetReturnToVerification(to string) {
+	f.ReturnToVerification = to
+}
diff --git a/selfservice/hook/show_verification_ui.go b/selfservice/hook/show_verification_ui.go
@@ -9,13 +9,18 @@ import (
 
 	"github.com/ory/kratos/driver/config"
 	"github.com/ory/kratos/selfservice/flow"
+	"github.com/ory/kratos/selfservice/flow/login"
 	"github.com/ory/kratos/selfservice/flow/registration"
 	"github.com/ory/kratos/session"
+	"github.com/ory/kratos/ui/node"
 	"github.com/ory/kratos/x"
 	"github.com/ory/x/otelx"
 )
 
-var _ registration.PostHookPostPersistExecutor = new(ShowVerificationUIHook)
+var (
+	_ registration.PostHookPostPersistExecutor = new(ShowVerificationUIHook)
+	_ login.PostHookExecutor                   = new(ShowVerificationUIHook)
+)
 
 type (
 	showVerificationUIDependencies interface {
@@ -45,15 +50,28 @@ func (e *ShowVerificationUIHook) ExecutePostRegistrationPostPersistHook(_ http.R
 	})
 }
 
-func (e *ShowVerificationUIHook) execute(r *http.Request, f *registration.Flow) error {
+// ExecutePostRegistrationPostPersistHook adds redirect headers and status code if the request is a browser request.
+// If the request is not a browser request, this hook does nothing.
+func (e *ShowVerificationUIHook) ExecuteLoginPostHook(_ http.ResponseWriter, r *http.Request, _ node.UiNodeGroup, f *login.Flow, _ *session.Session) error {
+	return otelx.WithSpan(r.Context(), "selfservice.hook.ShowVerificationUIHook.ExecutePostRegistrationPostPersistHook", func(ctx context.Context) error {
+		return e.execute(r.WithContext(ctx), f)
+	})
+}
+
+type loginOrRegistrationFlow interface {
+	ContinueWith() []flow.ContinueWith
+	SetReturnToVerification(string)
+}
+
+func (e *ShowVerificationUIHook) execute(r *http.Request, f loginOrRegistrationFlow) error {
 	if !x.IsBrowserRequest(r) {
 		// this hook is only intended to be used by browsers, as it redirects to the verification ui
 		// JSON API clients should use the `continue_with` field to continue the flow
 		return nil
 	}
 
 	var vf *flow.ContinueWithVerificationUI
-	for _, c := range f.ContinueWithItems {
+	for _, c := range f.ContinueWith() {
 		if item, ok := c.(*flow.ContinueWithVerificationUI); ok {
 			vf = item
 		}
@@ -62,7 +80,7 @@ func (e *ShowVerificationUIHook) execute(r *http.Request, f *registration.Flow)
 	ctx := r.Context()
 	if vf != nil {
 		redirURL := e.d.Config().SelfServiceFlowVerificationUI(ctx)
-		f.ReturnToVerification = vf.AppendTo(redirURL).String()
+		f.SetReturnToVerification(vf.AppendTo(redirURL).String())
 	}
 
 	return nil

diff --git a/selfservice/hook/show_verification_ui_test.go b/selfservice/hook/show_verification_ui_test.go
@@ -15,62 +15,120 @@ import (
 	"github.com/ory/kratos/driver/config"
 	"github.com/ory/kratos/internal"
 	"github.com/ory/kratos/selfservice/flow"
+	"github.com/ory/kratos/selfservice/flow/login"
 	"github.com/ory/kratos/selfservice/flow/registration"
 	"github.com/ory/kratos/selfservice/flow/verification"
 	"github.com/ory/kratos/selfservice/hook"
 )
 
 func TestExecutePostRegistrationPostPersistHook(t *testing.T) {
-	t.Run("case=no continue with items returns 200 OK", func(t *testing.T) {
-		_, reg := internal.NewVeryFastRegistryWithoutDB(t)
-		h := hook.NewShowVerificationUIHook(reg)
-		browserRequest := httptest.NewRequest("GET", "/", nil)
-		f := &registration.Flow{}
-		rec := httptest.NewRecorder()
-		require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, f, nil))
-		require.Equal(t, 200, rec.Code)
-	})
+	t.Run("flow=registration", func(t *testing.T) {
+		t.Run("case=no continue with items returns 200 OK", func(t *testing.T) {
+			_, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			f := &registration.Flow{}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, f, nil))
+			require.Equal(t, 200, rec.Code)
+		})
 
-	t.Run("case=not a browser request returns 200 OK", func(t *testing.T) {
-		_, reg := internal.NewVeryFastRegistryWithoutDB(t)
-		h := hook.NewShowVerificationUIHook(reg)
-		browserRequest := httptest.NewRequest("GET", "/", nil)
-		browserRequest.Header.Add("Accept", "application/json")
-		f := &registration.Flow{}
-		rec := httptest.NewRecorder()
-		require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, f, nil))
-		require.Equal(t, 200, rec.Code)
-	})
+		t.Run("case=not a browser request returns 200 OK", func(t *testing.T) {
+			_, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			browserRequest.Header.Add("Accept", "application/json")
+			f := &registration.Flow{}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, f, nil))
+			require.Equal(t, 200, rec.Code)
+		})
+
+		t.Run("case=verification ui in continue with item returns redirect", func(t *testing.T) {
+			conf, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			conf.Set(context.Background(), config.ViperKeySelfServiceVerificationUI, "/verification")
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			vf := &verification.Flow{
+				ID: uuid.Must(uuid.NewV4()),
+			}
+			rf := &registration.Flow{}
+			rf.ContinueWithItems = []flow.ContinueWith{
+				flow.NewContinueWithVerificationUI(vf, "[email protected]", ""),
+			}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, rf, nil))
+			assert.Equal(t, 200, rec.Code)
+			assert.Equal(t, "/verification?flow="+vf.ID.String(), rf.ReturnToVerification)
+		})
 
-	t.Run("case=verification ui in continue with item returns redirect", func(t *testing.T) {
-		conf, reg := internal.NewVeryFastRegistryWithoutDB(t)
-		conf.Set(context.Background(), config.ViperKeySelfServiceVerificationUI, "/verification")
-		h := hook.NewShowVerificationUIHook(reg)
-		browserRequest := httptest.NewRequest("GET", "/", nil)
-		vf := &verification.Flow{
-			ID: uuid.Must(uuid.NewV4()),
-		}
-		rf := &registration.Flow{}
-		rf.ContinueWithItems = []flow.ContinueWith{
-			flow.NewContinueWithVerificationUI(vf, "[email protected]", ""),
-		}
-		rec := httptest.NewRecorder()
-		require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, rf, nil))
-		assert.Equal(t, 200, rec.Code)
-		assert.Equal(t, "/verification?flow="+vf.ID.String(), rf.ReturnToVerification)
+		t.Run("case=no verification ui in continue with item returns 200 OK", func(t *testing.T) {
+			conf, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			conf.Set(context.Background(), config.ViperKeySelfServiceVerificationUI, "/verification")
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			rf := &registration.Flow{}
+			rf.ContinueWithItems = []flow.ContinueWith{
+				flow.NewContinueWithSetToken("token"),
+			}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, rf, nil))
+			assert.Equal(t, 200, rec.Code)
+		})
 	})
 
-	t.Run("case=no verification ui in continue with item returns 200 OK", func(t *testing.T) {
-		conf, reg := internal.NewVeryFastRegistryWithoutDB(t)
-		conf.Set(context.Background(), config.ViperKeySelfServiceVerificationUI, "/verification")
-		h := hook.NewShowVerificationUIHook(reg)
-		browserRequest := httptest.NewRequest("GET", "/", nil)
-		rf := &registration.Flow{}
-		rf.ContinueWithItems = []flow.ContinueWith{
-			flow.NewContinueWithSetToken("token"),
-		}
-		rec := httptest.NewRecorder()
-		require.NoError(t, h.ExecutePostRegistrationPostPersistHook(rec, browserRequest, rf, nil))
-		assert.Equal(t, 200, rec.Code)
+	t.Run("flow=login", func(t *testing.T) {
+		t.Run("case=no continue with items returns 200 OK", func(t *testing.T) {
+			_, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			f := &login.Flow{}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecuteLoginPostHook(rec, browserRequest, "", f, nil))
+			require.Equal(t, 200, rec.Code)
+		})
+
+		t.Run("case=not a browser request returns 200 OK", func(t *testing.T) {
+			_, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			browserRequest.Header.Add("Accept", "application/json")
+			f := &login.Flow{}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecuteLoginPostHook(rec, browserRequest, "", f, nil))
+			require.Equal(t, 200, rec.Code)
+		})
+
+		t.Run("case=verification ui in continue with item returns redirect", func(t *testing.T) {
+			conf, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			conf.Set(context.Background(), config.ViperKeySelfServiceVerificationUI, "/verification")
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			vf := &verification.Flow{
+				ID: uuid.Must(uuid.NewV4()),
+			}
+			rf := &login.Flow{}
+			rf.ContinueWithItems = []flow.ContinueWith{
+				flow.NewContinueWithVerificationUI(vf, "[email protected]", ""),
+			}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecuteLoginPostHook(rec, browserRequest, "", rf, nil))
+			assert.Equal(t, 200, rec.Code)
+			assert.Equal(t, "/verification?flow="+vf.ID.String(), rf.ReturnToVerification)
+		})
+
+		t.Run("case=no verification ui in continue with item returns 200 OK", func(t *testing.T) {
+			conf, reg := internal.NewVeryFastRegistryWithoutDB(t)
+			conf.Set(context.Background(), config.ViperKeySelfServiceVerificationUI, "/verification")
+			h := hook.NewShowVerificationUIHook(reg)
+			browserRequest := httptest.NewRequest("GET", "/", nil)
+			rf := &login.Flow{}
+			rf.ContinueWithItems = []flow.ContinueWith{
+				flow.NewContinueWithSetToken("token"),
+			}
+			rec := httptest.NewRecorder()
+			require.NoError(t, h.ExecuteLoginPostHook(rec, browserRequest, "", rf, nil))
+			assert.Equal(t, 200, rec.Code)
+		})
 	})
 }