Enable setting session cookie attributes from ui.run() (zauberzeug#5213)

### Motivation

When using `app.storage.user` it creates a session cookie with
permissions that disallow cross-origin usage. While this is a good
secure default, there's currently no way to override it when needed,
e.g. if you want your NiceGUI site to be embeddable as an iframe from
another domain. See discussion
zauberzeug#4252.

### Implementation

A new argument is added to `ui.run()` which allows passing keyword
arguments that will be forwarded to the
[SessionMiddleware](https://www.starlette.dev/middleware/#sessionmiddleware)
that is used to create the session cookie.

### Progress

- [x] I chose a meaningful title that completes the sentence: "If
applied, this PR will..."
- [x] The implementation is complete.
- [x] Pytests have been added (or are not necessary).
- [x] Documentation has been added (or is not necessary).

---------

Co-authored-by: Falko Schindler <[email protected]>

Author: samuller

nicegui/server.py

@@ -2,6 +2,7 @@
 
 import multiprocessing
 import socket
+from typing import Any
 
 import uvicorn
 
@@ -13,6 +14,7 @@ class CustomServerConfig(uvicorn.Config):
     storage_secret: str | None = None
     method_queue: multiprocessing.Queue | None = None
     response_queue: multiprocessing.Queue | None = None
+    session_middleware_kwargs: dict[str, Any] | None = None
 
 
 class Server(uvicorn.Server):
@@ -31,5 +33,5 @@ def run(self, sockets: list[socket.socket] | None = None) -> None:
             native.method_queue = self.config.method_queue
             native.response_queue = self.config.response_queue
 
-        storage.set_storage_secret(self.config.storage_secret)
+        storage.set_storage_secret(self.config.storage_secret, self.config.session_middleware_kwargs)
         super().run(sockets=sockets)

nicegui/storage.py

@@ -3,7 +3,7 @@
 import uuid
 from datetime import timedelta
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 from starlette.middleware import Middleware
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
@@ -35,14 +35,15 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
         return response
 
 
-def set_storage_secret(storage_secret: Optional[str] = None) -> None:
+def set_storage_secret(storage_secret: Optional[str] = None,
+                       session_middleware_kwargs: Optional[dict[str, Any]] = None) -> None:
     """Set storage_secret and add request tracking middleware."""
     if any(m.cls == SessionMiddleware for m in core.app.user_middleware):
         # NOTE not using "add_middleware" because it would be the wrong order
         core.app.user_middleware.append(Middleware(RequestTrackingMiddleware))
     elif storage_secret is not None:
         core.app.add_middleware(RequestTrackingMiddleware)
-        core.app.add_middleware(SessionMiddleware, secret_key=storage_secret)
+        core.app.add_middleware(SessionMiddleware, secret_key=storage_secret, **(session_middleware_kwargs or {}))
     Storage.secret = storage_secret
 
 

nicegui/ui_run.py

@@ -75,6 +75,7 @@ def run(root: Optional[Callable] = None, *,
         prod_js: bool = True,
         endpoint_documentation: Literal['none', 'internal', 'page', 'all'] = 'none',
         storage_secret: Optional[str] = None,
+        session_middleware_kwargs: Optional[dict[str, Any]] = None,
         show_welcome_message: bool = True,
         **kwargs: Any,
         ) -> None:
@@ -111,6 +112,7 @@ def run(root: Optional[Callable] = None, *,
     :param prod_js: whether to use the production version of Vue and Quasar dependencies (default: `True`)
     :param endpoint_documentation: control what endpoints appear in the autogenerated OpenAPI docs (default: 'none', options: 'none', 'internal', 'page', 'all')
     :param storage_secret: secret key for browser-based storage (default: `None`, a value is required to enable ui.storage.individual and ui.storage.browser)
+    :param session_middleware_kwargs: additional keyword arguments passed to SessionMiddleware that creates the session cookies used for browser-based storage
     :param show_welcome_message: whether to show the welcome message (default: `True`)
     :param kwargs: additional keyword arguments are passed to `uvicorn.run`
     """
@@ -181,7 +183,7 @@ def run_script() -> None:
         core.app.setup()
 
     if helpers.is_user_simulation():
-        set_storage_secret(storage_secret)
+        set_storage_secret(storage_secret, session_middleware_kwargs)
         return
 
     if on_air:
@@ -256,6 +258,7 @@ def split_args(args: str) -> list[str]:
     config.storage_secret = storage_secret
     config.method_queue = native_module.native.method_queue if native else None
     config.response_queue = native_module.native.response_queue if native else None
+    config.session_middleware_kwargs = session_middleware_kwargs
     Server.create_singleton(config)
 
     if (reload or config.workers > 1) and not isinstance(config.app, str):

nicegui/ui_run_with.py

@@ -1,6 +1,6 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Callable, Literal, Optional, Union
+from typing import Any, Callable, Literal, Optional, Union
 
 from fastapi import FastAPI
 from fastapi.middleware.gzip import GZipMiddleware
@@ -29,6 +29,7 @@ def run_with(
     tailwind: bool = True,
     prod_js: bool = True,
     storage_secret: Optional[str] = None,
+    session_middleware_kwargs: Optional[dict[str, Any]] = None,
     show_welcome_message: bool = True,
 ) -> None:
     """Run NiceGUI with FastAPI.
@@ -49,6 +50,7 @@ def run_with(
     :param tailwind: whether to use Tailwind CSS (experimental, default: `True`)
     :param prod_js: whether to use the production version of Vue and Quasar dependencies (default: `True`)
     :param storage_secret: secret key for browser-based storage (default: `None`, a value is required to enable ui.storage.individual and ui.storage.browser)
+    :param session_middleware_kwargs: additional keyword arguments passed to SessionMiddleware that creates the session cookies used for browser-based storage
     :param show_welcome_message: whether to show the welcome message (default: `True`)
     """
     core.app.config.add_run_config(
@@ -67,7 +69,7 @@ def run_with(
         cache_control_directives=cache_control_directives,
     )
     core.root = root
-    storage.set_storage_secret(storage_secret)
+    storage.set_storage_secret(storage_secret, session_middleware_kwargs)
     core.app.add_middleware(GZipMiddleware)
     core.app.add_middleware(RedirectWithPrefixMiddleware)
     core.app.add_middleware(SetCacheControlMiddleware)

tests/test_storage.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import httpx
+import pytest
 
 from nicegui import Client, app, background_tasks, context, core, nicegui, ui
 from nicegui.persistence.file_persistent_dict import FilePersistentDict
@@ -385,3 +386,23 @@ def page():
     await background_tasks.teardown()
     assert path.exists(), 'backup should be written during teardown'
     assert path.read_text(encoding='utf-8') == '{"key":"value"}'
+
+
+@pytest.mark.parametrize('custom_cookie_headers', [False, True])
+def test_storage_cookie_headers(screen: Screen, custom_cookie_headers: bool):
+    @ui.page('/')
+    def page():
+        ui.label('Hello, world!')
+
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
+    if custom_cookie_headers:
+        screen.ui_run_kwargs['session_middleware_kwargs'] = {'same_site': 'none', 'https_only': True}
+    screen.open('/')
+    with httpx.Client() as http_client:
+        response = http_client.get(f'http://localhost:{Screen.PORT}/')
+        assert response.status_code == 200
+        cookie_settings = str(response.headers.get('set-cookie')).lower()
+        if custom_cookie_headers:
+            assert cookie_settings.endswith('httponly; samesite=none; secure')
+        else:
+            assert cookie_settings.endswith('httponly; samesite=lax')