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: