Introducing AuthMiddleware in the authentication example

examples/authentication/main.py

@@ -1,33 +1,60 @@
 #!/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 fastapi import Request
 from fastapi.responses import RedirectResponse
+from starlette.middleware.base import BaseHTTPMiddleware
 
+import nicegui.globals
 from nicegui import app, ui
 
 # in reality users passwords would obviously need to be hashed
 passwords = {'user1': 'pass1', 'user2': 'pass2'}
 
+unrestricted_page_routes = {'/login'}
+
+
+class AuthMiddleware(BaseHTTPMiddleware):
+    """This middleware restricts access to all NiceGUI pages.
+
+    It redirects the user to the login page if they are not authenticated.
+    """
+
+    async def dispatch(self, request: Request, call_next):
+        if not app.storage.user.get('authenticated', False):
+            if request.url.path in nicegui.globals.page_routes.values() and request.url.path not in unrestricted_page_routes:
+                app.storage.user['referrer_path'] = request.url.path  # remember where the user wanted to go
+                return RedirectResponse('/login')
+        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 sub page.')
+
+
 @ui.page('/login')
-def login() -> None:
+def login() -> 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(app.storage.user.get('referrer_path', '/'))  # go back to where the user wanted to go
         else:
             ui.notify('Wrong username or password', color='negative')