Browse Source

only enable session storage if secret is provided

Rodja Trappe 2 years ago
parent
commit
1e425d66ff
3 changed files with 21 additions and 5 deletions
  1. 7 3
      nicegui/run.py
  2. 6 2
      nicegui/storage.py
  3. 8 0
      tests/test_storage.py

+ 7 - 3
nicegui/run.py

@@ -18,8 +18,9 @@ from .storage import RequestTrackingMiddleware
 class Server(uvicorn.Server):
 
     def run(self, sockets: List[Any] = None) -> None:
-        globals.app.add_middleware(RequestTrackingMiddleware)
-        globals.app.add_middleware(SessionMiddleware, secret_key='some_random_string')  # TODO real random string
+        if self.config.storage_secret is not None:
+            globals.app.add_middleware(RequestTrackingMiddleware)
+            globals.app.add_middleware(SessionMiddleware, secret_key=self.config.storage_secret)
         super().run(sockets=sockets)
 
 
@@ -43,6 +44,7 @@ def run(*,
         uvicorn_reload_excludes: str = '.*, .py[cod], .sw.*, ~*',
         exclude: str = '',
         tailwind: bool = True,
+        storage_secret: Optional[str] = None,
         **kwargs: Any,
         ) -> None:
     '''ui.run
@@ -69,7 +71,8 @@ def run(*,
     :param exclude: comma-separated string to exclude elements (with corresponding JavaScript libraries) to save bandwidth
       (possible entries: aggrid, audio, chart, colors, interactive_image, joystick, keyboard, log, markdown, mermaid, plotly, scene, video)
     :param tailwind: whether to use Tailwind (experimental, default: `True`)
-    :param kwargs: additional keyword arguments are passed to `uvicorn.run`
+    :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 kwargs: additional keyword arguments are passed to `uvicorn.run`    
     '''
     globals.ui_run_has_been_called = True
     globals.reload = reload
@@ -125,6 +128,7 @@ def run(*,
         log_level=uvicorn_logging_level,
         **kwargs,
     )
+    config.storage_secret = storage_secret
     globals.server = Server(config=config)
 
     if (reload or config.workers > 1) and not isinstance(config.app, str):

+ 6 - 2
nicegui/storage.py

@@ -96,6 +96,8 @@ class Storage:
         The data is shared between all browser tab and can only be modified before the initial request has been submitted.
         Normally it is better to use `app.storage.individual` instead to reduce payload, improved security and larger storage capacity)."""
         request: Request = request_contextvar.get()
+        if request is None:
+            raise RuntimeError('storage.browser needs a storage_secret passed in ui.run()')
         if request.state.responded:
             return ReadOnlyDict(
                 request.session,
@@ -105,19 +107,21 @@ class Storage:
 
     @property
     def individual(self) -> Dict:
-        """Individual user storage that is persisted on the server.
+        """Individual user storage that is persisted on the server (where NiceGUI is executed).
 
         The data is stored in a file on the server.
         It is shared between all browser tabs by identifying the user via session cookie id.
         """
         request: Request = request_contextvar.get()
+        if request is None:
+            raise RuntimeError('storage.individual needs a storage_secret passed in ui.run()')
         if request.session['id'] not in self._individuals:
             self._individuals[request.session['id']] = {}
         return self._individuals[request.session['id']]
 
     @property
     def general(self) -> Dict:
-        """General storage shared between all users that is persisted on the server."""
+        """General storage shared between all users that is persisted on the server (where NiceGUI is executed)."""
         return self._general
 
     async def backup(self):

+ 8 - 0
tests/test_storage.py

@@ -15,6 +15,7 @@ def test_browser_data_is_stored_in_the_browser(screen: Screen):
     def count():
         return 'count = ' + str(app.storage.browser['count'])
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.should_contain('1')
     screen.open('/')
@@ -32,6 +33,7 @@ def test_browser_storage_supports_asyncio(screen: Screen):
         await asyncio.sleep(0.5)
         ui.label(app.storage.browser['count'])
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.switch_to(1)
     screen.open('/')
@@ -50,6 +52,7 @@ def test_browser_storage_modifications_after_page_load_are_forbidden(screen: Scr
         except TypeError as e:
             ui.label(str(e))
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.should_contain('response to the browser has already been build')
 
@@ -62,6 +65,7 @@ def test_individual_storage_modifications(screen: Screen):
         app.storage.individual['count'] = app.storage.individual.get('count', 0) + 1
         ui.label().bind_text_from(app.storage.individual, 'count')
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.should_contain('1')
     screen.open('/?delayed=True')
@@ -77,6 +81,7 @@ async def test_access_individual_storage_on_interaction(screen: Screen):
             app.storage.individual['test_switch'] = False
         ui.switch('switch').bind_value(app.storage.individual, 'test_switch')
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.click('switch')
     screen.wait(1)
@@ -93,6 +98,7 @@ def test_access_individual_storage_from_button_click_handler(screen: Screen):
 
         ui.button('test', on_click=inner)
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.click('test')
     screen.wait(1)
@@ -108,6 +114,7 @@ async def test_access_individual_storage_from_background_task(screen: Screen):
             await app.storage.backup()
         background_tasks.create(subtask())
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     assert '{"subtask": "works"}' in app.storage._individuals.filename.read_text()
 
@@ -121,6 +128,7 @@ def test_individual_and_general_storage_is_persisted(screen: Screen):
         ui.label(f'general: {app.storage.general["count"]}')
         ui.button('backup', on_click=app.storage.backup)
 
+    screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.open('/')
     screen.open('/')