From 3f33d1ab0d4170d2986e7df23e536582ac4addd7 Mon Sep 17 00:00:00 2001
From: Roger Yang <80478925+RogerHYang@users.noreply.github.com>
Date: Thu, 19 Sep 2024 12:06:18 -0700
Subject: [PATCH] fix: user should be able to initiate password reset again
 before existing token is used or expires (#4674)

---
 integration_tests/auth/test_auth.py    | 19 ++++++++++++++-----
 src/phoenix/db/models.py               |  2 +-
 src/phoenix/server/api/routers/auth.py |  9 +++------
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/integration_tests/auth/test_auth.py b/integration_tests/auth/test_auth.py
index ad9a5aae12..4d3d654d1e 100644
--- a/integration_tests/auth/test_auth.py
+++ b/integration_tests/auth/test_auth.py
@@ -15,7 +15,6 @@
     TypeVar,
 )
 
-import httpx
 import jwt
 import pytest
 import smtpdfix
@@ -149,16 +148,26 @@ def test_initiate_password_reset_does_not_change_existing_password(
         u.log_in()
 
     @pytest.mark.parametrize("role_or_user", [_MEMBER, _ADMIN])
-    def test_password_reset_cannot_be_initiated_again_while_in_progress(
+    def test_password_reset_can_be_initiated_multiple_times(
         self,
         role_or_user: _RoleOrUser,
         _get_user: _GetUser,
+        _passwords: Iterator[_Password],
         _smtpd: smtpdfix.AuthController,
     ) -> None:
         u = _get_user(role_or_user)
-        assert u.initiate_password_reset(_smtpd)
-        with pytest.raises(httpx.HTTPStatusError):
-            assert u.initiate_password_reset(_smtpd)
+        new_password = next(_passwords)
+        assert new_password != u.password
+        tokens = [u.initiate_password_reset(_smtpd) for _ in range(2)]
+        assert sum(map(bool, tokens)) > 1
+        for i, token in enumerate(tokens):
+            assert token
+            if i < len(tokens) - 1:
+                with _EXPECTATION_401:
+                    token.reset(new_password)
+                continue
+            # only the last one works
+            token.reset(new_password)
 
     @pytest.mark.parametrize("role_or_user", [_MEMBER, _ADMIN])
     def test_password_reset_can_be_initiated_immediately_after_password_reset(
diff --git a/src/phoenix/db/models.py b/src/phoenix/db/models.py
index 5a81d57820..7a2ffa4633 100644
--- a/src/phoenix/db/models.py
+++ b/src/phoenix/db/models.py
@@ -648,7 +648,7 @@ class User(Base):
         UtcTimeStamp, server_default=func.now(), onupdate=func.now()
     )
     deleted_at: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp)
-    password_reset_token: Mapped["PasswordResetToken"] = relationship(
+    password_reset_token: Mapped[Optional["PasswordResetToken"]] = relationship(
         "PasswordResetToken",
         back_populates="user",
         uselist=False,
diff --git a/src/phoenix/server/api/routers/auth.py b/src/phoenix/server/api/routers/auth.py
index 0872e192b6..c2d7966f0c 100644
--- a/src/phoenix/server/api/routers/auth.py
+++ b/src/phoenix/server/api/routers/auth.py
@@ -39,6 +39,7 @@
     AccessTokenAttributes,
     AccessTokenClaims,
     PasswordResetTokenClaims,
+    PasswordResetTokenId,
     RefreshTokenAttributes,
     RefreshTokenClaims,
     TokenStore,
@@ -227,14 +228,14 @@ async def initiate_password_reset(request: Request) -> Response:
     if user is None or user.auth_method != enums.AuthMethod.LOCAL.value:
         # Withold privileged information
         return Response(status_code=HTTP_204_NO_CONTENT)
+    token_store: TokenStore = request.app.state.get_token_store()
     if user.password_reset_token:
-        raise RESET_IN_PROGRESS
+        await token_store.revoke(PasswordResetTokenId(user.password_reset_token.id))
     password_reset_token_claims = PasswordResetTokenClaims(
         subject=UserId(user.id),
         issued_at=datetime.now(timezone.utc),
         expiration_time=datetime.now(timezone.utc) + token_expiry,
     )
-    token_store: TokenStore = request.app.state.get_token_store()
     token, _ = await token_store.create_password_reset_token(password_reset_token_claims)
     await sender.send_password_reset_email(email, PasswordResetTemplateBody(token, get_base_url()))
     return Response(status_code=HTTP_204_NO_CONTENT)
@@ -293,10 +294,6 @@ def _select_active_user() -> Select[Tuple[models.User]]:
 UNAVAILABLE = HTTPException(
     status_code=HTTP_503_SERVICE_UNAVAILABLE,
 )
-RESET_IN_PROGRESS = HTTPException(
-    status_code=HTTP_503_SERVICE_UNAVAILABLE,
-    detail="Password reset already in progress",
-)
 INVALID_TOKEN = HTTPException(
     status_code=HTTP_401_UNAUTHORIZED,
     detail="Invalid token",