fix(ui): show the MFA lockout and keep the code step on a 401

openapi/spec/paths/api@user@mfa@auth.yaml

@@ -4,6 +4,12 @@ post:
   description: |
     Authenticate a user who has MFA enabled. This endpoint should be called after the default authUser endpoint,
     which generates an `X-MFA-Token` indicating that the user has already authenticated with a password.
+
+    The token logs in once only. An unknown, expired or already used token answers `401`, the same
+    as a wrong code, and so does a code sent while another code for the same user is still being
+    checked. A token whose user no longer exists answers `404`. Wrong codes count toward an account
+    lockout of their own, which a correct password does not clear: check the `X-Account-Lockout`
+    header, which is 0 when there is no active lockout.
   tags:
     - cloud
     - mfa
@@ -17,6 +23,12 @@ post:
   responses:
     '200':
       description: Success to authenticate.
+      headers:
+        X-Account-Lockout:
+          schema:
+            type: integer
+            description: Indicates the Unix timestamp in seconds of when the account lockout period ends.
+            example: 0
       content:
         application/json:
           schema:
@@ -29,5 +41,17 @@ post:
       $ref: ../components/responses/403.yaml
     '404':
       $ref: ../components/responses/404.yaml
+    '429':
+      description: Too Many Requests
+      headers:
+        X-Account-Lockout:
+          schema:
+            type: integer
+            description: Indicates the Unix timestamp in seconds of when the account lockout period ends.
+            example: 1711575401
+      content:
+        application/json:
+          schema:
+            $ref: ../components/schemas/apiError.yaml
     '500':
       $ref: ../components/responses/500.yaml

pkg/cache/cache.go

@@ -39,6 +39,11 @@ type Cache interface {
 	// exactly once under concurrent requests, which Get+Set cannot do race-free.
 	SetNX(ctx context.Context, key string, value any, ttl time.Duration) (bool, error)
 
+	// CompareAndDelete atomically deletes key only while it still holds value, and reports whether
+	// it did. It releases a reservation taken with SetNX under a value unique to its owner, so an
+	// owner whose reservation expired cannot delete the one another caller took after it.
+	CompareAndDelete(ctx context.Context, key, value string) (bool, error)
+
 	// HasAccountLockout reports whether the source is currently blocked from attempting to
 	// log in to a user with the specified userID. It returns the absolute Unix timestamp
 	// in seconds representing the end of the lockout while clock.Now() is before it, or 0

pkg/cache/cache_null.go

@@ -31,6 +31,10 @@ func (*nullCache) Delete(_ context.Context, _ string) error {
 	return nil
 }
 
+func (*nullCache) CompareAndDelete(_ context.Context, _, _ string) (bool, error) {
+	return true, nil
+}
+
 func (*nullCache) HasAccountLockout(_ context.Context, _, _ string) (int64, int, error) {
 	return 0, 0, nil
 }

pkg/cache/cache_redis.go

@@ -74,6 +74,22 @@ func (c *redisCache) SetNX(ctx context.Context, key string, value any, ttl time.
 	return c.client.SetNX(ctx, key, value, ttl).Result()
 }
 
+var compareAndDeleteScript = redis.NewScript(`
+if redis.call("GET", KEYS[1]) == ARGV[1] then
+	return redis.call("DEL", KEYS[1])
+end
+return 0
+`)
+
+func (c *redisCache) CompareAndDelete(ctx context.Context, key, value string) (bool, error) {
+	deleted, err := compareAndDeleteScript.Run(ctx, c.client, []string{key}, value).Int()
+	if err != nil {
+		return false, err
+	}
+
+	return deleted == 1, nil
+}
+
 // Delete deletes cached value by given key.
 func (c *redisCache) Delete(ctx context.Context, key string) error {
 	if err := c.cache.Get(ctx, key, nil); errors.Is(err, rediscache.ErrCacheMiss) {

pkg/cache/cache_redis_test.go

@@ -66,3 +66,40 @@ func TestRedisCacheLockoutEndsAtAnnouncedDeadline(t *testing.T) {
 	require.NoError(t, c.Get(ctx, "account-lockout=10.0.0.1:user", &storedDeadline))
 	assert.Equal(t, "1700000060", storedDeadline)
 }
+
+func TestRedisCacheCompareAndDeleteKeepsAnotherOwnersKey(t *testing.T) {
+	ctx := context.Background()
+
+	redisContainer, err := redis.Run(ctx, "docker.io/valkey/valkey:9.1-alpine")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		require.NoError(t, redisContainer.Terminate(ctx))
+	})
+
+	uri, err := redisContainer.ConnectionString(ctx)
+	require.NoError(t, err)
+
+	c, err := cache.NewRedisCache(uri, 0)
+	require.NoError(t, err)
+
+	taken, err := c.SetNX(ctx, "guard", "owner", time.Minute)
+	require.NoError(t, err)
+	require.True(t, taken)
+
+	deleted, err := c.CompareAndDelete(ctx, "guard", "someone-else")
+	require.NoError(t, err)
+	assert.False(t, deleted)
+
+	taken, err = c.SetNX(ctx, "guard", "intruder", time.Minute)
+	require.NoError(t, err)
+	assert.False(t, taken)
+
+	deleted, err = c.CompareAndDelete(ctx, "guard", "owner")
+	require.NoError(t, err)
+	assert.True(t, deleted)
+
+	taken, err = c.SetNX(ctx, "guard", "next", time.Minute)
+	require.NoError(t, err)
+	assert.True(t, taken)
+}

ui/apps/console/src/api/fetchInterceptors.ts

@@ -70,11 +70,13 @@ client.interceptors.response.use((response) => {
     if (mfaToken) {
       useAuthStore.getState().setMfaToken(mfaToken);
     } else {
-      const isLoginRequest = response.url.includes("/api/login");
+      const isSignInStep =
+        response.url.includes("/api/login") ||
+        response.url.includes("/api/user/mfa/auth");
       const isTokenLogin = new URLSearchParams(window.location.search).has(
         "token",
       );
-      if (!isLoginRequest && !isTokenLogin) {
+      if (!isSignInStep && !isTokenLogin) {
         useAuthStore.getState().logout();
         window.location.href = "/login";
       }

ui/apps/console/src/components/auth/SignInForm.tsx

@@ -1,9 +1,13 @@
-import { useState, useEffect, FormEvent } from "react";
+import { useState, FormEvent } from "react";
 import { useForm } from "react-hook-form";
 import { useNavigate } from "react-router-dom";
 import { Button, Callout } from "@shellhub/design-system/primitives";
 import { isSdkError } from "@/api/errors";
 import { useAuthStore } from "@/stores/authStore";
+import {
+  lockoutEndFrom,
+  useLockoutCountdown,
+} from "@/hooks/useLockoutCountdown";
 import { landingAfterSignIn } from "@/utils/navigation";
 import {
   FormInputField,
@@ -14,52 +18,6 @@ import MfaCodeForm from "@/components/mfa/MfaCodeForm";
 import AuthActions, { type AuthLink } from "@/components/auth/AuthActions";
 import type { LoginFormValues } from "@/pages/setup/loginResolver";
 
-interface CountdownState {
-  display: string;
-  expired: boolean;
-  epoch: number | null;
-}
-
-function useLoginCountdown(lockoutEndEpoch: number | null) {
-  const [state, setState] = useState<CountdownState>({
-    display: "",
-    expired: false,
-    epoch: null,
-  });
-
-  useEffect(() => {
-    if (lockoutEndEpoch === null) return;
-
-    const interval = setInterval(() => {
-      const diff = lockoutEndEpoch - Date.now() / 1000;
-      if (diff <= 0) {
-        clearInterval(interval);
-        setState({ display: "", expired: true, epoch: lockoutEndEpoch });
-      } else if (diff < 60) {
-        const s = Math.floor(diff);
-        setState({
-          display: `${s} ${s === 1 ? "second" : "seconds"}`,
-          expired: false,
-          epoch: lockoutEndEpoch,
-        });
-      } else {
-        const m = Math.floor(diff / 60);
-        setState({
-          display: `${m} ${m === 1 ? "minute" : "minutes"}`,
-          expired: false,
-          epoch: lockoutEndEpoch,
-        });
-      }
-    }, 1000);
-
-    return () => clearInterval(interval);
-  }, [lockoutEndEpoch]);
-
-  if (state.epoch !== lockoutEndEpoch) return { display: "", expired: false };
-
-  return { display: state.display, expired: state.expired };
-}
-
 interface SignInFormProps {
   redirect: string;
   submitLabel?: string;
@@ -88,7 +46,7 @@ export default function SignInForm({
   const [secondFactor, setSecondFactor] = useState(false);
   const [lockoutEndEpoch, setLockoutEndEpoch] = useState<number | null>(null);
   const { display: countdownDisplay, expired: lockoutExpired } =
-    useLoginCountdown(lockoutEndEpoch);
+    useLockoutCountdown(lockoutEndEpoch);
 
   const { control, handleSubmit, formState } = useForm<LoginFormValues>({
     resolver: loginResolver,
@@ -143,8 +101,7 @@ export default function SignInForm({
           );
           break;
         case 429: {
-          const epoch = Number(err.headers.get("x-account-lockout"));
-          setLockoutEndEpoch(isNaN(epoch) ? null : epoch);
+          setLockoutEndEpoch(lockoutEndFrom(err.headers));
           setError(
             "Too many failed login attempts. Please wait before trying again.",
           );

ui/apps/console/src/components/mfa/MfaCodeForm.tsx

@@ -1,6 +1,11 @@
-import type { FormEvent } from "react";
+import { useState, type FormEvent } from "react";
 import { Button, Callout } from "@shellhub/design-system/primitives";
+import { isSdkError } from "@/api/errors";
 import { useAuthStore } from "@/stores/authStore";
+import {
+  lockoutEndFrom,
+  useLockoutCountdown,
+} from "@/hooks/useLockoutCountdown";
 import { useOtpInput } from "@/hooks/useOtpInput";
 import OtpCells from "@/components/mfa/OtpCells";
 import AuthActions, { type AuthLink } from "@/components/auth/AuthActions";
@@ -14,7 +19,9 @@ interface MfaCodeFormProps {
 /**
  * The second step of a sign-in: the six cells for the authenticator code, exchanged for a full
  * session. A wrong code clears the cells and leaves the partial session in place, so the user
- * can try again without signing in from scratch. `onVerified` runs once the session is full.
+ * can try again without signing in from scratch. Too many wrong codes lock the form out, with a
+ * countdown to when the user may try again, still without signing in from scratch. `onVerified`
+ * runs once the session is full.
  */
 export default function MfaCodeForm({
   onVerified,
@@ -23,22 +30,43 @@ export default function MfaCodeForm({
 }: MfaCodeFormProps) {
   const otp = useOtpInput(6);
   const { loginWithMfa, loading, error } = useAuthStore();
+  const [lockoutEndEpoch, setLockoutEndEpoch] = useState<number | null>(null);
+  const { display: countdownDisplay, expired: lockoutExpired } =
+    useLockoutCountdown(lockoutEndEpoch);
 
   const handleSubmit = async (e: FormEvent) => {
     e.preventDefault();
     if (!otp.isComplete) return;
 
+    setLockoutEndEpoch(null);
     try {
       await loginWithMfa(otp.getValue());
       onVerified();
-    } catch {
+    } catch (err) {
+      if (isSdkError(err) && err.status === 429) {
+        setLockoutEndEpoch(lockoutEndFrom(err.headers));
+      }
       otp.reset();
     }
   };
 
   return (
     <form onSubmit={(e) => void handleSubmit(e)} className="space-y-6">
-      {error && <Callout variant="error">{error}</Callout>}
+      {lockoutExpired && (
+        <Callout variant="success">
+          Your timeout has finished. Please enter a new code.
+        </Callout>
+      )}
+      {error && !lockoutExpired && (
+        <Callout variant="error">
+          <span>
+            <span>{error}</span>
+            {countdownDisplay ? (
+              <span className="font-semibold"> ({countdownDisplay})</span>
+            ) : null}
+          </span>
+        </Callout>
+      )}
 
       <OtpCells otp={otp} label="Verification Code" size="lg" numeric />