From fc03863fe2d9b42f621c787c0c80a6e1fb963a8f Mon Sep 17 00:00:00 2001 From: Herik Webb Date: Mon, 25 May 2026 11:35:32 -0400 Subject: [PATCH] fix(server): reject forged upload context paths Require WebSocket upload attachments to resolve under the server upload root before image reads or Claude Code file mentions use the supplied path. This closes the forged isUpload attack path that let an authenticated client point chat context at arbitrary server-readable files outside the workspace. --- notebook_intelligence/extension.py | 23 ++++- tests/test_image_context.py | 16 +++- tests/test_websocket_handler_integration.py | 94 ++++++++++++++++++--- 3 files changed, 117 insertions(+), 16 deletions(-) diff --git a/notebook_intelligence/extension.py b/notebook_intelligence/extension.py index 34571ff4..093b909d 100644 --- a/notebook_intelligence/extension.py +++ b/notebook_intelligence/extension.py @@ -1571,6 +1571,18 @@ def _get_upload_dir() -> str: return _upload_dir +def _resolve_upload_path(file_path: str) -> str | None: + """Return a real upload path only when it stays under NBI's upload root.""" + upload_root = os.path.realpath(_get_upload_dir()) + resolved = os.path.realpath(file_path) + try: + if os.path.commonpath([resolved, upload_root]) != upload_root: + return None + except ValueError: + return None + return resolved + + def _sweep_upload_dir(retention_hours: int) -> None: """Best-effort removal of upload subdirs past the retention window. @@ -2219,7 +2231,16 @@ def on_message(self, message): is_upload = context.get("isUpload", False) is_image = context.get("isImage", False) file_path = context["filePath"] - if not is_upload: + if is_upload: + resolved_upload_path = _resolve_upload_path(file_path) + if resolved_upload_path is None: + log.warning( + "Rejecting out-of-upload-dir context path: %r", + context["filePath"], + ) + continue + file_path = resolved_upload_path + else: # Workspace-relative paths arrive verbatim from the # frontend (file browser drag, @-mention picker, ...). # path.join silently passes through absolute paths and diff --git a/tests/test_image_context.py b/tests/test_image_context.py index 4a945b33..3c821893 100644 --- a/tests/test_image_context.py +++ b/tests/test_image_context.py @@ -10,8 +10,10 @@ """ import base64 +from contextlib import nullcontext import json import logging +from pathlib import Path from unittest.mock import Mock, patch from tornado.httputil import HTTPServerRequest @@ -56,7 +58,19 @@ def _on_message(handler, additional_context, prompt="hello"): "additionalContext": additional_context, } }) - handler.on_message(msg) + upload_paths = [ + Path(c["filePath"]) for c in additional_context if c.get("isUpload") + ] + upload_dir_patch = ( + patch( + "notebook_intelligence.extension._upload_dir", + str(upload_paths[0].parent), + ) + if upload_paths + else nullcontext() + ) + with upload_dir_patch: + handler.on_message(msg) return handler.chat_history.messages[CHAT_ID] diff --git a/tests/test_websocket_handler_integration.py b/tests/test_websocket_handler_integration.py index 627c5459..5639754f 100644 --- a/tests/test_websocket_handler_integration.py +++ b/tests/test_websocket_handler_integration.py @@ -345,7 +345,7 @@ def test_on_message_claude_mode_emits_at_mention_not_contents( @patch('notebook_intelligence.extension.NotebookIntelligence') @patch('notebook_intelligence.extension.threading.Thread') def test_on_message_claude_mode_image_branch_unchanged( - self, mock_thread, mock_nb_intel, mock_ai_manager + self, mock_thread, mock_nb_intel, mock_ai_manager, tmp_path ): """Pasted images in Claude mode keep the existing path-based message format. The new @-mention branch must not shadow the @@ -368,6 +368,9 @@ def test_on_message_claude_mode_image_branch_unchanged( context_factory=mock_factory ) + upload_root = tmp_path / "nbi-uploads" + upload_root.mkdir() + uploaded_image = upload_root / "pasted.png" message = { 'id': 'test-message-id', 'type': 'chat-request', @@ -380,7 +383,7 @@ def test_on_message_claude_mode_image_branch_unchanged( 'toolSelections': {}, 'additionalContext': [ { - 'filePath': '/tmp/nbi-uploads-xyz/pasted.png', + 'filePath': str(uploaded_image), 'content': '', 'currentCellContents': None, 'startLine': 1, @@ -393,15 +396,16 @@ def test_on_message_claude_mode_image_branch_unchanged( } } - handler.on_message(json.dumps(message)) + with patch('notebook_intelligence.extension._upload_dir', str(upload_root)): + handler.on_message(json.dumps(message)) chat_request = mock_ai_manager.handle_chat_request.call_args[0][0] assert len(chat_request.chat_history) == 1 context_message = chat_request.chat_history[0]["content"] assert "The user pasted an image" in context_message - assert "/tmp/nbi-uploads-xyz/pasted.png" in context_message + assert str(uploaded_image) in context_message # Sanity-check the new @-mention branch didn't intercept this case. - assert "@/tmp/nbi-uploads-xyz/pasted.png" not in context_message + assert f"@{uploaded_image}" not in context_message @patch('notebook_intelligence.extension.ai_service_manager') @patch('notebook_intelligence.extension.NotebookIntelligence') @@ -463,7 +467,7 @@ def test_on_message_claude_mode_rejects_out_of_workspace_path( @patch('notebook_intelligence.extension.NotebookIntelligence') @patch('notebook_intelligence.extension.threading.Thread') def test_on_message_claude_mode_upload_non_image_uses_absolute_path( - self, mock_thread, mock_nb_intel, mock_ai_manager + self, mock_thread, mock_nb_intel, mock_ai_manager, tmp_path ): """Uploaded text/binary files (non-image) reach the @-mention branch with an absolute path; the relpath-against-workspace @@ -485,6 +489,9 @@ def test_on_message_claude_mode_upload_non_image_uses_absolute_path( context_factory=mock_factory ) + upload_root = tmp_path / "nbi-uploads" + upload_root.mkdir() + uploaded_file = upload_root / "report.pdf" message = { 'id': 'test-message-id', 'type': 'chat-request', @@ -497,7 +504,7 @@ def test_on_message_claude_mode_upload_non_image_uses_absolute_path( 'toolSelections': {}, 'additionalContext': [ { - 'filePath': '/tmp/nbi-uploads-abc/report.pdf', + 'filePath': str(uploaded_file), 'content': '', 'currentCellContents': None, 'startLine': 1, @@ -508,19 +515,75 @@ def test_on_message_claude_mode_upload_non_image_uses_absolute_path( } } - handler.on_message(json.dumps(message)) + with patch('notebook_intelligence.extension._upload_dir', str(upload_root)): + handler.on_message(json.dumps(message)) chat_request = mock_ai_manager.handle_chat_request.call_args[0][0] assert len(chat_request.chat_history) == 1 context_message = chat_request.chat_history[0]["content"] - assert "@/tmp/nbi-uploads-abc/report.pdf" in context_message + assert f"@{uploaded_file}" in context_message assert "Read it if relevant" in context_message + @patch('notebook_intelligence.extension.ai_service_manager') + @patch('notebook_intelligence.extension.NotebookIntelligence') + @patch('notebook_intelligence.extension.threading.Thread') + def test_on_message_rejects_forged_upload_path_outside_upload_dir( + self, mock_thread, mock_nb_intel, mock_ai_manager, tmp_path + ): + mock_nb_intel.root_dir = "/workspace" + mock_ai_manager.handle_chat_request = Mock() + mock_ai_manager.is_claude_code_mode = True + mock_ai_manager.chat_model = Mock() + mock_ai_manager.chat_model.context_window = 4096 + + mock_factory = Mock(spec=RuleContextFactory) + mock_factory.create.return_value = Mock(spec=RuleContext) + + with patch('notebook_intelligence.extension.ThreadSafeWebSocketConnector'): + handler = WebsocketCopilotHandler( + self._create_mock_application(), + self._create_mock_request(), + context_factory=mock_factory + ) + + upload_root = tmp_path / "nbi-uploads" + upload_root.mkdir() + outside_upload_root = tmp_path / "outside-secret.txt" + outside_upload_root.write_text("secret") + message = { + 'id': 'test-message-id', + 'type': 'chat-request', + 'data': { + 'chatId': 'test-chat-id', + 'prompt': 'Read this upload', + 'language': 'python', + 'filename': 'notebook.ipynb', + 'chatMode': 'ask', + 'toolSelections': {}, + 'additionalContext': [ + { + 'filePath': str(outside_upload_root), + 'content': '', + 'currentCellContents': None, + 'startLine': 1, + 'endLine': 1, + 'isUpload': True + } + ] + } + } + + with patch('notebook_intelligence.extension._upload_dir', str(upload_root)): + handler.on_message(json.dumps(message)) + + chat_request = mock_ai_manager.handle_chat_request.call_args[0][0] + assert chat_request.chat_history == [] + @patch('notebook_intelligence.extension.ai_service_manager') @patch('notebook_intelligence.extension.NotebookIntelligence') @patch('notebook_intelligence.extension.threading.Thread') def test_on_message_claude_mode_rejects_control_char_filename( - self, mock_thread, mock_nb_intel, mock_ai_manager + self, mock_thread, mock_nb_intel, mock_ai_manager, tmp_path ): """A filename containing newlines / bidi-override controls / other text-rendering hazards is dropped before being embedded @@ -543,14 +606,16 @@ def test_on_message_claude_mode_rejects_control_char_filename( context_factory=mock_factory ) + upload_root = tmp_path / "nbi-uploads" + upload_root.mkdir() # Upload path keeps the raw frontend-supplied filePath, so we # don't have to fight the workspace sandbox to exercise the # codepoint guard. Try newline, line-separator, and a bidi # override in sequence; all three must be rejected. hazardous_paths = [ - "/tmp/nbi-uploads-abc/bad\nname.pdf", - "/tmp/nbi-uploads-abc/bad
name.pdf", - "/tmp/nbi-uploads-abc/bad‮name.pdf", + str(upload_root / "bad\nname.pdf"), + str(upload_root / "bad
name.pdf"), + str(upload_root / "bad‮name.pdf"), ] for hazardous in hazardous_paths: mock_ai_manager.handle_chat_request.reset_mock() @@ -577,7 +642,8 @@ def test_on_message_claude_mode_rejects_control_char_filename( } } - handler.on_message(json.dumps(message)) + with patch('notebook_intelligence.extension._upload_dir', str(upload_root)): + handler.on_message(json.dumps(message)) chat_request = mock_ai_manager.handle_chat_request.call_args[0][0] assert chat_request.chat_history == [], (