Improve the email confirmation logic using the router and more robust JWT handling

src/services/auth.py

@@ -1,3 +1,4 @@
+from datetime import datetime, UTC, timedelta
 import abc
 from src.config.email import EmailAddress
 from src.config.parse import assert_key_of_type, ParseError
@@ -14,10 +15,12 @@ JWT = str
 T = TypeVar('T', bound='Claim')
 
 
-# TODO: Do these claims need explicit exp attributes? I think the ID fields should actually be sub now that I think
-# about it...
 @dataclass
 class Claim(abc.ABC):
+    sub: int
+    exp: datetime
+    iat: datetime
+
     @classmethod
     @abc.abstractmethod
     def from_dict(cls, claims: dict[str, Any]) -> T:
@@ -25,39 +28,55 @@ class Claim(abc.ABC):
 
     @classmethod
     @abc.abstractmethod
-    def from_user(cls, user: User) -> T:
+    def from_user(cls, user: User, expire_in_secs: int) -> T:
         pass
 
 
-@dataclass
-class UserClaims(Claim):
-    id: int
-
-    @classmethod
-    def from_user(cls, user: User) -> 'UserClaims':
-        return UserClaims(id=user.id)
-
-    @classmethod
-    def from_dict(cls, claims: dict[str, Any]) -> 'UserClaims':
-        assert_key_of_type(claims, 'id', int)
-        return UserClaims(id=claims['id'])
+class ExpiredTokenError(Exception):
+    def __init__(self):
+        super().__init__('token was expired')
 
 
 @dataclass
-class EmailConfirmationClaim(Claim):
-    id: int
+class BaseClaims(Claim):
+    @classmethod
+    def from_user(cls, user: User, expire_in_secs: int) -> 'BaseClaims':
+        now = datetime.now(tz=UTC)
+        expires_at = now + timedelta(seconds=expire_in_secs)
+        return BaseClaims(sub=user.id, exp=expires_at, iat=now)
+
+    @classmethod
+    def from_dict(cls, claims: dict[str, Any]) -> 'BaseClaims':
+        assert_key_of_type(claims, 'sub', int)
+        assert_key_of_type(claims, 'exp', int)
+        assert_key_of_type(claims, 'iat', int)
+        exp = datetime.fromtimestamp(claims['exp'], UTC)
+        now = datetime.now(UTC)
+        if exp < now:
+            raise ExpiredTokenError()
+
+        return BaseClaims(
+            sub=claims['sub'],
+            exp=exp,
+            iat=datetime.fromtimestamp(claims['iat'], UTC),
+        )
+
+
+@dataclass
+class EmailConfirmationClaim(BaseClaims):
     email: EmailAddress
 
     @classmethod
-    def from_user(cls, user: User) -> 'EmailConfirmationClaim':
-        return EmailConfirmationClaim(id=user.id, email=user.email)
+    def from_user(cls, user: User, expire_in_secs) -> 'EmailConfirmationClaim':
+        base = BaseClaims.from_user(user, expire_in_secs)
+        return EmailConfirmationClaim(**asdict(base), email=user.email)
 
     @classmethod
     def from_dict(cls, claims: dict[str, Any]) -> 'EmailConfirmationClaim':
-        assert_key_of_type(claims, 'id', int)
+        base = BaseClaims.from_dict(claims)
         assert_key_of_type(claims, 'email', str)
 
-        return EmailConfirmationClaim(id=claims['id'], email=claims['email'])
+        return EmailConfirmationClaim(**asdict(base), email=claims['email'])
 
 
 class AuthService:
@@ -65,8 +84,8 @@ class AuthService:
         self._private_key = Ed25519PrivateKey.from_private_bytes(config.ed25519_private_key.expose_secret())
         self._public_key = self._private_key.public_key()
 
-    def mint_claim_from_user(self, claim: type[Claim], user: User) -> JWT:
-        claims = claim.from_user(user)
+    def mint_claim_from_user(self, claim: type[Claim], user: User, expires_in_secs=600) -> JWT:
+        claims = claim.from_user(user, expires_in_secs)
         return jwt.encode(asdict(claims), self._private_key, algorithm='EdDSA')
 
     def validate_token[T: Claim](self, claim: type[T], token: JWT) -> T | None:

src/services/users/data.py

@@ -32,3 +32,8 @@ class LoginError(Exception):
 
 class SignupError(Exception):
     pass
+
+
+class EmailConfirmationTokenInvalid(Exception):
+    def __init__(self):
+        super().__init__('Email confirmation token invalid')

src/services/users/service.py

@@ -1,12 +1,21 @@
-from src.services.router import Router
-from typing import TypedDict, Unpack
-from src.constants.routes import ROUTES
-from src.services.auth import AuthService, EmailConfirmationClaim, JWT
 from email.headerregistry import Address
+from typing import TypedDict, Unpack
+
 from jinja2 import Environment, PackageLoader, select_autoescape
 
+from src.constants.routes import ROUTES
+from src.services.auth import AuthService, EmailConfirmationClaim, JWT
+from src.services.router import Router
 from src.services.email import EmailService, EmailDTO
-from .data import UserDTO, User, UserProfile, SignupError, LoginError
+
+from .data import (
+    UserDTO,
+    User,
+    UserProfile,
+    SignupError,
+    LoginError,
+    EmailConfirmationTokenInvalid,
+)
 from .repo import UserRepo
 
 
@@ -47,29 +56,23 @@ class UserService:
         else:
             raise LoginError('No user found for that email or password')
 
-    def confirm_email_for_user(self, user_id: int, confirmation_token: JWT) -> bool:
+    def confirm_email_for_user(self, user_id: int, confirmation_token: JWT):
         """
         Attempt to confirm that the user at the given ID has confirmed their email by returning the JWT that was
         minted for this purpose.
-
-        Returns whether or not the confirmation was performed.
         """
         claim = self._auth_service.validate_token(EmailConfirmationClaim, confirmation_token)
+        if not claim or claim.sub != user_id:
+            raise EmailConfirmationTokenInvalid
 
-        # First check that the claim could be parsed and that it refers to the expected user
-        if claim and claim.id == user_id:
         user = self._repo.get_user_by_id(user_id)
-            # Double check that:
-            # 1. The user exists in the database
-            # 2. The claim refers to the email address on file
-            # 3. The email was not already confirmed
-            if user and user.email == claim.email and not user.email_confirmed:
+        # If there is no user or the users emails don't match or the user is already confirmed then consider
+        # the token invalid for this request.
+        if not user or user.email != claim.email or user.email_confirmed:
+            raise EmailConfirmationTokenInvalid
+
         user.email_confirmed = True
         self._repo.update_user(user)
-                return True
-
-        # In all other cases, the confirmation was not possible so return False
-        return False
 
     def _send_confirmation_email(self, user: User):
         env = Environment(loader=PackageLoader('src'), autoescape=select_autoescape())