[airflow] branch main updated: Fix an issue that crashes Airflow Webserver when passed invalid private key path to Snowflake (#32016)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new c6f151fa8f Fix an issue that crashes Airflow Webserver when passed invalid private key path to Snowflake (#32016)
c6f151fa8f is described below

commit c6f151fa8f569687f3d889bce04bc270f114d208
Author: Bartosz Jankiewicz <bj...@users.noreply.github.com>
AuthorDate: Sun Jul 9 00:32:13 2023 +0200

    Fix an issue that crashes Airflow Webserver when passed invalid private key path to Snowflake (#32016)
    
    * Validate private key file path and size in Snowflake hook
    
    * Test update for snowflake hook
    
    * Code review fixes
    
    * Fix import
    
    Co-authored-by: xrmr <17...@users.noreply.github.com>
    
    * Fixes
    
    * Fix error message.
    
    ---------
    
    Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
    Co-authored-by: xrmr <17...@users.noreply.github.com>
    Co-authored-by: bjankiewicz <bj...@google.com>
---
 airflow/providers/snowflake/hooks/snowflake.py    |  7 ++++++-
 tests/providers/snowflake/hooks/test_snowflake.py | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py
index 2a202441a7..59199cf8cd 100644
--- a/airflow/providers/snowflake/hooks/snowflake.py
+++ b/airflow/providers/snowflake/hooks/snowflake.py
@@ -248,7 +248,12 @@ class SnowflakeHook(DbApiHook):
                 "Please remove one."
             )
         elif private_key_file:
-            private_key_pem = Path(private_key_file).read_bytes()
+            private_key_file_path = Path(private_key_file)
+            if not private_key_file_path.is_file() or private_key_file_path.stat().st_size == 0:
+                raise ValueError("The private_key_file path points to an empty or invalid file.")
+            if private_key_file_path.stat().st_size > 4096:
+                raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
+            private_key_pem = Path(private_key_file_path).read_bytes()
         elif private_key_content:
             private_key_pem = private_key_content.encode()
 
diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py
index e1105fbb59..6a738952d9 100644
--- a/tests/providers/snowflake/hooks/test_snowflake.py
+++ b/tests/providers/snowflake/hooks/test_snowflake.py
@@ -393,6 +393,24 @@ class TestPytestSnowflakeHook:
         ), pytest.raises(TypeError, match="Password was given but private key is not encrypted."):
             SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
 
+    def test_get_conn_params_should_fail_on_invalid_key(self):
+        connection_kwargs = {
+            **BASE_CONNECTION_KWARGS,
+            "password": None,
+            "extra": {
+                "database": "db",
+                "account": "airflow",
+                "warehouse": "af_wh",
+                "region": "af_region",
+                "role": "af_role",
+                "private_key_file": "/dev/urandom",
+            },
+        }
+        with mock.patch.dict(
+            "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()
+        ), pytest.raises(ValueError, match="The private_key_file path points to an empty or invalid file."):
+            SnowflakeHook(snowflake_conn_id="test_conn").get_conn()
+
     def test_should_add_partner_info(self):
         with mock.patch.dict(
             "os.environ",