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

src/config/__init__.py

@@ -6,12 +6,12 @@ from .logging import Logging
 from .email import Email
 from .auth import Auth
 from .database import Database
-from .parse import assert_key_of_type, parse_nested_config
+from .parse import parse_nested_config
 
 
 @dataclass
 class Config:
-    host: str | None
+    host: str
     port: int | None
     logging: Logging
     email: Email
@@ -20,17 +20,14 @@ class Config:
 
     @classmethod
     def from_dict(cls, config: dict[str, Any]) -> 'Config':
-        assert_key_of_type(config, 'host', str)
-        assert_key_of_type(config, 'port', int)
-
         log_config = parse_nested_config(config, 'logging', Logging.from_dict)
         email_config = parse_nested_config(config, 'email', Email.from_dict)
         db_config = parse_nested_config(config, 'database', Database.from_dict)
         auth_config = parse_nested_config(config, 'auth', Auth.from_dict)
 
         return Config(
-            host=config['host'],
-            port=config['port'],
+            host=config.get('host', '0.0.0.0'),
+            port=config.get('port'),
             email=email_config,
             logging=log_config,
             database=db_config,

src/constants/routes.py

@@ -0,0 +1,47 @@
+import abc
+from typing import Any, TypedDict
+from dataclasses import dataclass, asdict
+
+
+class ParamCodec[Params]:
+    @abc.abstractmethod
+    def parse(self, path) -> Params:
+        pass
+
+
+class EmptyParams(ParamCodec[None]):
+    def parse(self, path) -> None:
+        pass
+
+
+@dataclass
+class Route[Params: Any]:
+    path: str
+    codec: ParamCodec[Params]
+
+    def interpolate(self, params: Params) -> str:
+        interpolated = self.path
+        if params:
+            for k, v in asdict(params).items():
+                interpolated.replace(f':{k}', str(v))
+
+        return interpolated
+
+
+def static(path: str) -> Route[None]:
+    return Route(path, EmptyParams())
+
+
+class Routes(TypedDict):
+    index: Route[None]
+    confirm_email: Route[None]
+    signup: Route[None]
+    dashboard: Route[None]
+
+
+ROUTES: Routes = {
+    'index': static('/'),
+    'confirm_email': static('/confirm-email'),
+    'signup': static('/register'),
+    'dashboard': static('/dashboard'),
+}

src/main.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+from src.services.router import Router
 from src.services.notifications import NotificationService, get_notification_service
 from src.services.email import EmailService, get_email_service
 from src.services.auth import AuthService
@@ -23,6 +24,7 @@ class AppServices:
     auth: AuthService
     email: EmailService
     notifications: NotificationService
+    router: Router
 
 
 class MyFlask(Flask):
@@ -31,19 +33,24 @@ class MyFlask(Flask):
 
 def create_app(name: str, config: Config) -> Flask:
     logging.basicConfig(level=config.logging.level, format=('%(asctime)s %(levelname)s [%(name)s] %(message)s'))
-
     app = MyFlask(name)
+
+    # Configure Services
     database = get_database(config.database)
     user_repo = UserRepoImpl(database)
     email_service = get_email_service(config.email)
+    auth_service = AuthService(config.auth)
+    router = Router(config.host, config.port)
+
     app.services = AppServices(
-        users=UserService(user_repo, email_service),
-        auth=AuthService(config.auth),
+        users=UserService(repo=user_repo, email_service=email_service, auth_service=auth_service, router=router),
+        auth=auth_service,
         email=email_service,
         notifications=get_notification_service(email_service),
+        router=router,
     )
 
-    if config.host:
+    if config.host != '0.0.0.0':
         app.config['SERVER_NAME'] = config.host
 
     return app
@@ -57,7 +64,7 @@ def main():
     config = parse_config(args.config)
     app = create_app('Cereal', config)
     logger.info(str(config))
-    serve(app, host='0.0.0.0' if config.host is None else config.host, port=config.port)
+    serve(app, host=config.host, port=config.port)
 
 
 if __name__ == '__main__':

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/router.py

@@ -0,0 +1,19 @@
+from src.constants.routes import Route
+
+
+class Router:
+    def __init__(self, host: str, port: int | None):
+        self._host = host
+        self._port = port
+
+    def interpolate_url[Params](self, route: Route[Params], params: Params, **query_params) -> str:
+        query_strs = []
+        for key, value in query_params.items():
+            query_strs.append(f'{key}={value}')
+
+        query_str = '&'.join(query_strs)
+        query_str = f'?{query_str}' if len(query_str) > 0 else query_str
+        return f'{self._host}{"" if self._port is None else f":{self._port}"}{route.interpolate(params)}{query_str}'
+
+    def get_url(self, route: Route[None], **query_params):
+        return self.interpolate_url(route, None, **query_params)

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,9 +1,21 @@
-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
 
 
@@ -23,11 +35,19 @@ def is_valid_password(password: str) -> bool:
     return long_enough and short_enough and has_symbol
 
 
+class ServiceDeps(TypedDict):
+    repo: UserRepo
+    email_service: EmailService
+    auth_service: AuthService
+    router: Router
+
+
 class UserService:
-    def __init__(self, repo: UserRepo, email_service: EmailService, auth_service: AuthService):
-        self._repo = repo
-        self._email_service = email_service
-        self._auth_service = auth_service
+    def __init__(self, **kwargs: Unpack[ServiceDeps]):
+        self._repo = kwargs['repo']
+        self._email_service = kwargs['email_service']
+        self._auth_service = kwargs['auth_service']
+        self._router = kwargs['router']
 
     def login(self, user: UserDTO) -> UserProfile:
         full_user = self._repo.auth_as_user(user)
@@ -36,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:
-                user.email_confirmed = True
-                self._repo.update_user(user)
-                return True
+        user = self._repo.get_user_by_id(user_id)
+        # 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
 
-        # In all other cases, the confirmation was not possible so return False
-        return False
+        user.email_confirmed = True
+        self._repo.update_user(user)
 
     def _send_confirmation_email(self, user: User):
         env = Environment(loader=PackageLoader('src'), autoescape=select_autoescape())
@@ -66,8 +80,7 @@ class UserService:
         html_template = env.get_template('mail/confirmation_email.html')
         token = self._auth_service.mint_claim_from_user(EmailConfirmationClaim, user)
 
-        # TODO: Parameterize this with configuration that also drives the API
-        url = f'/confirm?token={token}'
+        url = self._router.get_url(ROUTES['confirm_email'], token=token)
         text_content = text_template.render(confirmation_link=url)
         html_content = html_template.render(confirmation_link=url)