SNOW-693548 fixing redos attack (#1327)

Author: sfc-gh-mkeller

DESCRIPTION.md

@@ -11,6 +11,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 - v2.8.2(November 18,2022)
 
   - Improved performance of OCSP response caching
+  - Improved performance of regexes used for PUT/GET SQL statement detection
 
 - v2.8.1(October 30,2022)
 

src/snowflake/connector/_sql_util.py

@@ -0,0 +1,52 @@
+#
+# Copyright (c) 2012-2021 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import re
+
+from .constants import FileTransferType
+
+COMMENT_START_SQL_RE = re.compile(
+    r"""
+                                  ^\s*(?:
+                                      /\*[\w\W]*?\*/
+                                  )""",
+    re.VERBOSE,
+)
+
+PUT_SQL_RE = re.compile(r"^\s*put", flags=re.IGNORECASE)
+GET_SQL_RE = re.compile(r"^\s*get", flags=re.IGNORECASE)
+
+
+def remove_starting_comments(sql: str) -> str:
+    """Remove all comments from the start of a SQL statement."""
+    commentless_sql = sql
+    while True:
+        start_comment = COMMENT_START_SQL_RE.match(commentless_sql)
+        if start_comment is None:
+            break
+        commentless_sql = commentless_sql[start_comment.end() :]
+    return commentless_sql
+
+
+def get_file_transfer_type(sql: str) -> FileTransferType | None:
+    """Decide whether a SQL is a file transfer and return its type.
+
+    None is returned if the SQL isn't a file transfer so that this function can be
+    used in an if-statement.
+    """
+    commentless_sql = remove_starting_comments(sql)
+    if PUT_SQL_RE.match(commentless_sql):
+        return FileTransferType.PUT
+    elif GET_SQL_RE.match(commentless_sql):
+        return FileTransferType.GET
+
+
+def is_put_statement(sql: str) -> bool:
+    return get_file_transfer_type(sql) == FileTransferType.PUT
+
+
+def is_get_statement(sql: str) -> bool:
+    return get_file_transfer_type(sql) == FileTransferType.GET

src/snowflake/connector/cursor.py

@@ -32,6 +32,7 @@
 from snowflake.connector.result_set import ResultSet
 
 from . import compat
+from ._sql_util import get_file_transfer_type
 from .bind_upload_agent import BindUploadAgent, BindUploadError
 from .constants import (
     FIELD_NAME_TO_ID,
@@ -183,8 +184,6 @@ class SnowflakeCursor:
         Calling a function is expensive in Python and most of these getters are unnecessary.
     """
 
-    PUT_SQL_RE = re.compile(r"^(?:/\*.*\*/\s*)*put\s+", flags=re.IGNORECASE)
-    GET_SQL_RE = re.compile(r"^(?:/\*.*\*/\s*)*get\s+", flags=re.IGNORECASE)
     INSERT_SQL_RE = re.compile(r"^insert\s+into", flags=re.IGNORECASE)
     COMMENT_SQL_RE = re.compile(r"/\*.*\*/")
     INSERT_SQL_VALUES_RE = re.compile(
@@ -202,11 +201,7 @@ def get_file_transfer_type(sql: str) -> FileTransferType | None:
         None is returned if the SQL isn't a file transfer so that this function can be
         used in an if-statement.
         """
-        if SnowflakeCursor.PUT_SQL_RE.match(sql):
-            return FileTransferType.PUT
-        elif SnowflakeCursor.GET_SQL_RE.match(sql):
-            return FileTransferType.GET
-        return None
+        return get_file_transfer_type(sql)
 
     def __init__(
         self,
@@ -464,9 +459,7 @@ def _execute_helper(
             self._is_file_transfer = _is_put_get
         else:
             # or detect it.
-            self._is_file_transfer = self.PUT_SQL_RE.match(
-                query
-            ) or self.GET_SQL_RE.match(query)
+            self._is_file_transfer = get_file_transfer_type(query) is not None
         logger.debug("is_file_transfer: %s", self._is_file_transfer is not None)
 
         real_timeout = (

src/snowflake/connector/gcs_storage_client.py

@@ -10,6 +10,7 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, NamedTuple
 
+from ._sql_util import is_put_statement
 from .compat import quote
 from .constants import (
     FILE_PROTOCOL,
@@ -269,9 +270,7 @@ def _get_local_file_path_from_put_command(self) -> str | None:
             The local file path.
         """
         command = self._command
-        if FILE_PROTOCOL not in self._command or not self._cursor.PUT_SQL_RE.match(
-            command
-        ):
+        if FILE_PROTOCOL not in self._command or not is_put_statement(command):
             return None
 
         file_path_begin_index = command.find(FILE_PROTOCOL)

test/unit/test_cursor.py

@@ -27,9 +27,24 @@ def __init__(self):
 @pytest.mark.parametrize(
     "sql,_type",
     (
+        ("", None),
+        ("select 1;", None),
         ("PUT file:///tmp/data/mydata.csv @my_int_stage;", FileTransferType.PUT),
         ("GET @%mytable file:///tmp/data/;", FileTransferType.GET),
-        ("select 1;", None),
+        ("/**/PUT file:///tmp/data/mydata.csv @my_int_stage;", FileTransferType.PUT),
+        ("/**/ GET @%mytable file:///tmp/data/;", FileTransferType.GET),
+        pytest.param(
+            "/**/\n"
+            + "\t/*/get\t*/\t/**/\n" * 10000
+            + "\t*/get @~/test.csv file:///tmp\n",
+            None,
+            id="long_incorrect",
+        ),
+        pytest.param(
+            "/**/\n" + "\t/*/put\t*/\t/**/\n" * 10000 + "put file:///tmp/data.csv @~",
+            FileTransferType.PUT,
+            id="long_correct",
+        ),
     ),
 )
 def test_get_filetransfer_type(sql, _type):