Enable setting session cookie attributes from ui.run()

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,16 @@ 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."""
+    session_middleware_kwargs = session_middleware_kwargs or {}
     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)
     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:
@@ -255,6 +257,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

@@ -385,3 +385,38 @@ def page():
     await background_tasks.teardown()
     assert path.exists(), 'backup should be written during teardown'
     assert path.read_text(encoding='utf-8') == '{"key":"value"}'
+
+
+def test_storage_default_cookie_headers(screen: Screen):
+    @ui.page('/')
+    def page():
+        ui.label('Hello, world!')
+
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
+    screen.open('/')
+    with httpx.Client() as http_client:
+        response = http_client.get(f'http://localhost:{Screen.PORT}/')
+        assert response.status_code == 200
+        assert 'set-cookie' in response.headers
+        cookie_settings = response.headers['set-cookie'].split('; ')
+        assert 'httponly' in cookie_settings
+        assert 'samesite=lax' in cookie_settings
+        assert 'secure' not in cookie_settings
+
+
+def test_storage_custom_cookie_headers(screen: Screen):
+    @ui.page('/')
+    def page():
+        ui.label('Hello, world!')
+
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
+    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
+        assert 'set-cookie' in response.headers
+        cookie_settings = response.headers['set-cookie'].split('; ')
+        assert 'httponly' in cookie_settings
+        assert 'samesite=none' in cookie_settings
+        assert 'secure' in cookie_settings