Restrict local file operations to descendents of uploads_base_path

Fix lifecycle e2e test so it honours these requirements

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse Whitehouse

src/databricks/sql/client.py

@@ -4,6 +4,7 @@
 import pyarrow
 import requests
 import json
+import os
 
 from databricks.sql import __version__
 from databricks.sql import *
@@ -301,15 +302,23 @@ def _check_not_closed(self):
 
     def _handle_staging_operation(self, uploads_base_path: str):
         """Fetch the HTTP request instruction from a staging ingestion command
-        and call the designated handler."""
+        and call the designated handler.
+        
+        Raise an exception if localFile is specified by the server but the localFile
+        is not descended from uploads_base_path.
+        """
 
         if uploads_base_path is None:
             raise Error(
                 "You must provide an uploads_base_path when initialising a connection to perform ingestion commands"
             )
-
+        
         row = self.active_result_set.fetchone()
 
+        if getattr(row, "localFile", None):
+            if os.path.commonpath([row.localFile, uploads_base_path]) != uploads_base_path:
+                raise Error("Local file operations are restricted to paths within the configured uploads_base_path")
+
         # TODO: Experiment with DBR sending real headers.
         # The specification says headers will be in JSON format but the current null value is actually an empty list []
         handler_args = {

tests/e2e/driver_tests.py

@@ -667,29 +667,31 @@ def test_staging_ingestion_life_cycle(self):
             query = f"PUT '{temp_path}' INTO 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
             cursor.execute(query)
 
-            # GET should succeed
+        # GET should succeed
 
-            new_fh, new_temp_path = tempfile.mkstemp()
+        new_fh, new_temp_path = tempfile.mkstemp()
 
+        with self.connection(extra_params={"uploads_base_path": new_temp_path}) as conn:
             cursor = conn.cursor()
             query = f"GET 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
             cursor.execute(query)
 
-            with open(new_fh, "rb") as fp:
-                fetched_text = fp.read()
+        with open(new_fh, "rb") as fp:
+            fetched_text = fp.read()
 
-            assert fetched_text == original_text
+        assert fetched_text == original_text
 
-            # REMOVE should succeed
+        # REMOVE should succeed
 
-            remove_query = (
-                f"REMOVE 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv'"
-            )
+        remove_query = (
+            f"REMOVE 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv'"
+        )
 
+        with self.connection(extra_params={"uploads_base_path": "/"}) as conn:
             cursor = conn.cursor()
             cursor.execute(remove_query)
 
-            # GET after REMOVE should fail
+        # GET after REMOVE should fail
 
             with pytest.raises(Error):
                 cursor = conn.cursor()
@@ -718,6 +720,27 @@ def test_staging_ingestion_put_fails_without_uploadsbasepath(self):
                 query = f"PUT '{temp_path}' INTO 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
                 cursor.execute(query)
 
+    def test_staging_ingestion_put_fails_if_localFile_not_in_uploads_base_path(self):
+
+
+        fh, temp_path = tempfile.mkstemp()
+
+        original_text = "hello world!".encode("utf-8")
+
+        with open(fh, "wb") as fp:
+            fp.write(original_text)
+
+        base_path, filename = os.path.split(temp_path)
+
+        # Add junk to base_path
+        base_path = os.path.join(base_path, "temp")
+
+        with pytest.raises(Error):
+            with self.connection(extra_params={"uploads_base_path": base_path}) as conn:
+                cursor = conn.cursor()
+                query = f"PUT '{temp_path}' INTO 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
+                cursor.execute(query)
+
 
 def main(cli_args):
     global get_args_from_env