Merge pull request #287 from canonical/IAM-1039

Fix redirect on first login
canonical · Sep 27, 2024 · 3883fa3 · 3883fa3
2 parents bb26052 + 4f92c65
commit 3883fa3
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 29 deletions.
diff --git a/pkg/kratos/handlers.go b/pkg/kratos/handlers.go
@@ -75,7 +75,7 @@ func (a *API) handleCreateFlow(w http.ResponseWriter, r *http.Request) {
 		shouldEnforceMfa, err = a.shouldEnforceMFAWithSession(r.Context(), session)
 
 		if shouldEnforceMfa {
-			a.mfaSettingsRedirect(w)
+			a.mfaSettingsRedirect(w, returnTo)
 			return
 		}
 
@@ -217,7 +217,7 @@ func (a *API) handleUpdateFlow(w http.ResponseWriter, r *http.Request) {
 	setCookies(w, cookies)
 
 	if shouldEnforceMfa {
-		a.mfaSettingsRedirect(w)
+		a.mfaSettingsRedirect(w, *loginFlow.ReturnTo)
 		return
 	}
 
@@ -317,8 +317,9 @@ func (a *API) is40xError(err error) bool {
 	return false
 }
 
-func (a *API) mfaSettingsRedirect(w http.ResponseWriter) {
+func (a *API) mfaSettingsRedirect(w http.ResponseWriter, returnTo string) {
 	redirect, err := url.JoinPath("/", a.contextPath, "/ui/setup_secure")
+
 	if err != nil {
 		err = fmt.Errorf("unable to build mfa redirect path, possible misconfiguration, err: %v", err)
 		a.logger.Error(err.Error())
@@ -328,11 +329,19 @@ func (a *API) mfaSettingsRedirect(w http.ResponseWriter) {
 
 	errorId := "session_aal2_required"
 
+	// Set the original login URL as return_to, to continue the flow after mfa
+	// has been set.
+	u, _ := url.Parse(redirect)
+	q := u.Query()
+	q.Set("return_to", returnTo)
+	u.RawQuery = q.Encode()
+	r := u.String()
+
 	w.WriteHeader(http.StatusSeeOther)
 	_ = json.NewEncoder(w).Encode(
 		ErrorBrowserLocationChangeRequired{
 			Error:             &client.GenericError{Id: &errorId},
-			RedirectBrowserTo: &redirect,
+			RedirectBrowserTo: &r,
 		},
 	)
 }
@@ -531,11 +540,7 @@ func (a *API) handleUpdateSettingsFlow(w http.ResponseWriter, r *http.Request) {
 }
 
 func (a *API) handleCreateSettingsFlow(w http.ResponseWriter, r *http.Request) {
-	returnTo, err := url.JoinPath(a.baseURL, "/ui/setup_complete")
-	if err != nil {
-		a.logger.Errorf("Failed to construct returnTo URL: ", err)
-		http.Error(w, "Failed to construct returnTo URL", http.StatusBadRequest)
-	}
+	returnTo := r.URL.Query().Get("return_to")
 
 	flow, response, err := a.service.CreateBrowserSettingsFlow(context.Background(), returnTo, r.Cookies())
 	if err != nil {

diff --git a/pkg/kratos/handlers_test.go b/pkg/kratos/handlers_test.go
@@ -905,6 +905,7 @@ func TestHandleCreateSettingsFlow(t *testing.T) {
 
 	req := httptest.NewRequest(http.MethodGet, HANDLE_CREATE_SETTINGS_FLOW_URL, nil)
 	values := req.URL.Query()
+	values.Set("return_to", redirect)
 	req.URL.RawQuery = values.Encode()
 
 	flow := kClient.NewSettingsFlowWithDefaults()
@@ -941,6 +942,7 @@ func TestHandleCreateSettingsFlowWithRedirect(t *testing.T) {
 
 	req := httptest.NewRequest(http.MethodGet, HANDLE_CREATE_SETTINGS_FLOW_URL, nil)
 	values := req.URL.Query()
+	values.Set("return_to", redirect)
 	req.URL.RawQuery = values.Encode()
 
 	flow := kClient.NewSettingsFlowWithDefaults()
@@ -974,6 +976,7 @@ func TestHandleCreateSettingsFlowFailOnCreateBrowserSettingsFlow(t *testing.T) {
 
 	req := httptest.NewRequest(http.MethodGet, HANDLE_CREATE_SETTINGS_FLOW_URL, nil)
 	values := req.URL.Query()
+	values.Set("return_to", redirect)
 	req.URL.RawQuery = values.Encode()
 
 	mockService.EXPECT().CreateBrowserSettingsFlow(gomock.Any(), redirect, req.Cookies()).Return(nil, nil, fmt.Errorf("error"))

diff --git a/ui/pages/setup_secure.tsx b/ui/pages/setup_secure.tsx
@@ -44,26 +44,27 @@ const SetupSecure: NextPage = () => {
     // Otherwise we initialize it
     kratos
       .createBrowserSettingsFlow({
-        returnTo: returnTo ? String(returnTo) : undefined,
+        returnTo: returnTo
+          ? returnTo.toString()
+          : window.location.pathname.replace("setup_secure", "setup_complete"),
       })
-      .then(({ data }) => {
-        if (data.request_url !== undefined) {
-          const pwParam = pwChanged
-            ? `&pw_changed=${pwChanged.toString()}`
-            : "";
-          window.location.href = `./setup_secure?flow=${data.id}${pwParam}`;
-          return;
-        }
+      .then(async ({ data }) => {
+        const pwParam = pwChanged ? pwChanged.toString() : "false";
+        router.query.flow = data.id;
+        router.query.pw_changed = pwParam;
+
+        await router.replace(
+          {
+            pathname: window.location.pathname,
+            query: router.query,
+          },
+          undefined,
+          { shallow: true },
+        );
         setFlow(data);
       })
       .catch(handleFlowError("settings", setFlow))
       .catch(async (err: AxiosError<string>) => {
-        if (err.response?.data.trim() === "Failed to create settings flow") {
-          setFlow(undefined);
-          window.location.href = "./login";
-          return;
-        }
-
         return Promise.reject(err);
       });
   }, [flowId, router, router.isReady, returnTo, flow]);
@@ -83,15 +84,12 @@ const SetupSecure: NextPage = () => {
           },
         })
         .then(({ data }) => {
-          if (flow?.state === "success") {
-            window.location.href = "./setup_complete";
-          }
           if ("redirect_to" in data) {
             window.location.href = data.redirect_to as string;
             return;
           }
-          if (flow?.return_to) {
-            window.location.href = flow.return_to;
+          if (data?.return_to) {
+            window.location.href = data?.return_to;
             return;
           }
           window.location.href = "./error";