Introducing AuthMiddleware in the authentication example

examples/authentication/main.py

@@ -1,33 +1,53 @@
 #!/usr/bin/env python3
-"""This is just a very simple authentication example.
+"""This is just a simple authentication example.
 
 Please see the `OAuth2 example at FastAPI <https://fastapi.tiangolo.com/tutorial/security/simple-oauth2/>`_  or
 use the great `Authlib package <https://docs.authlib.org/en/v0.13/client/starlette.html#using-fastapi>`_ to implement a classing real authentication system.
 Here we just demonstrate the NiceGUI integration.
 """
+from typing import Optional
+from urllib.parse import quote
+
+from fastapi import Request
 from fastapi.responses import RedirectResponse
+from starlette.middleware.base import BaseHTTPMiddleware
 
 from nicegui import app, ui
 
 # in reality users passwords would obviously need to be hashed
 passwords = {'user1': 'pass1', 'user2': 'pass2'}
 
 
+class AuthMiddleware(BaseHTTPMiddleware):
+    """This middleware redirects the user to the login page if they are not authenticated."""
+
+    async def dispatch(self, request: Request, call_next):
+        if request.url.path not in ['/login'] and not app.storage.user.get('authenticated', False):
+            return RedirectResponse(f'/login?referrer_path={quote(request.url.path)}')
+        return await call_next(request)
+
+
+app.add_middleware(AuthMiddleware)
+
+
 @ui.page('/')
 def main_page() -> None:
-    if not app.storage.user.get('authenticated', False):
-        return RedirectResponse('/login')
     with ui.column().classes('absolute-center items-center'):
         ui.label(f'Hello {app.storage.user["username"]}!').classes('text-2xl')
         ui.button(on_click=lambda: (app.storage.user.clear(), ui.open('/login')), icon='logout').props('outline round')
 
 
+@ui.page('/subpage')
+def test_page() -> None:
+    ui.label('This is a subpage page.')
+
+
 @ui.page('/login')
-def login() -> None:
+def login(referrer_path: str = '') -> Optional[RedirectResponse]:
     def try_login() -> None:  # local function to avoid passing username and password as arguments
         if passwords.get(username.value) == password.value:
             app.storage.user.update({'username': username.value, 'authenticated': True})
-            ui.open('/')
+            ui.open(referrer_path or '/')
         else:
             ui.notify('Wrong username or password', color='negative')