From 592fd021356d36b46a3b93506156eee622ab0514 Mon Sep 17 00:00:00 2001
From: Matt Usifer <mattusifer@gmail.com>
Date: Sat, 6 May 2023 08:34:06 -0400
Subject: [PATCH 1/3] Validate connection_ids

---
 airflow/www/forms.py                     |  3 +-
 airflow/www/validators.py                | 22 +++++++++++++++
 tests/www/test_validators.py             | 35 ++++++++++++++++++++++++
 tests/www/views/test_views_connection.py | 12 ++++++++
 4 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/airflow/www/forms.py b/airflow/www/forms.py
index c7d8a0ce786bb..d925498f574b3 100644
--- a/airflow/www/forms.py
+++ b/airflow/www/forms.py
@@ -41,6 +41,7 @@
 from airflow.providers_manager import ProvidersManager
 from airflow.utils import timezone
 from airflow.utils.types import DagRunType
+from airflow.www.validators import ValidKey
 from airflow.www.widgets import (
     AirflowDateTimePickerROWidget,
     AirflowDateTimePickerWidget,
@@ -205,7 +206,7 @@ def _iter_connection_types() -> Iterator[tuple[str, str]]:
     class ConnectionForm(DynamicForm):
         conn_id = StringField(
             lazy_gettext("Connection Id"),
-            validators=[InputRequired()],
+            validators=[InputRequired(), ValidKey()],
             widget=BS3TextFieldWidget(),
         )
         conn_type = SelectField(
diff --git a/airflow/www/validators.py b/airflow/www/validators.py
index 0c08697f57a4d..71b5a82a002a8 100644
--- a/airflow/www/validators.py
+++ b/airflow/www/validators.py
@@ -22,6 +22,9 @@
 
 from wtforms.validators import EqualTo, ValidationError
 
+from airflow import AirflowException
+from airflow.utils import helpers
+
 
 class GreaterEqualThan(EqualTo):
     """Compares the values of two fields.
@@ -76,3 +79,22 @@ def __call__(self, form, field):
             except JSONDecodeError as ex:
                 message = self.message or f"JSON Validation Error: {ex}"
                 raise ValidationError(message=field.gettext(message.format(field.data)))
+
+
+class ValidKey:
+    """
+    Validates values that will be used as keys
+
+    :param max_length:
+        The maximum length of the given key
+    """
+
+    def __init__(self, max_length=200):
+        self.max_length = max_length
+
+    def __call__(self, form, field):
+        if field.data:
+            try:
+                helpers.validate_key(field.data, self.max_length)
+            except AirflowException as e:
+                raise ValidationError(e)
diff --git a/tests/www/test_validators.py b/tests/www/test_validators.py
index 0f6e9f5f31c0d..718091214ce58 100644
--- a/tests/www/test_validators.py
+++ b/tests/www/test_validators.py
@@ -120,3 +120,38 @@ def test_validation_raises_custom_message(self):
             self._validate(
                 message="Invalid JSON: {}",
             )
+
+
+class TestValidKey:
+    def setup_method(self):
+        self.form_field_mock = mock.MagicMock(data="valid_key")
+        self.form_field_mock.gettext.side_effect = lambda msg: msg
+        self.form_mock = mock.MagicMock(spec_set=dict)
+
+    def _validate(self):
+        validator = validators.ValidKey()
+
+        return validator(self.form_mock, self.form_field_mock)
+
+    def test_form_field_is_none(self):
+        self.form_field_mock.data = None
+
+        assert self._validate() is None
+
+    def test_validation_pass(self):
+        assert self._validate() is None
+
+    def test_validation_fails_with_trailing_whitespace(self):
+        self.form_field_mock.data = "invalid key  "
+
+        with pytest.raises(validators.ValidationError):
+            self._validate()
+
+    def test_validation_fails_with_too_many_characters(self):
+        self.form_field_mock.data = "".join("x" for _ in range(1000))
+
+        with pytest.raises(
+            validators.ValidationError,
+            match=r"The key has to be less than [0-9]+ characters",
+        ):
+            self._validate()
diff --git a/tests/www/views/test_views_connection.py b/tests/www/views/test_views_connection.py
index befe9d028139b..37e631ea8b87d 100644
--- a/tests/www/views/test_views_connection.py
+++ b/tests/www/views/test_views_connection.py
@@ -57,6 +57,18 @@ def test_create_connection(admin_client, session):
     _check_last_log(session, dag_id=None, event="connection.create", execution_date=None)
 
 
+def test_invalid_connection_id_trailing_blanks(admin_client, session):
+    invalid_conn_id = "conn_id_with_trailing_blanks   "
+    invalid_connection = {**CONNECTION, "conn_id": invalid_conn_id}
+    resp = admin_client.post(
+        "/connection/add", data=invalid_connection, follow_redirects=True
+    )
+    check_content_in_response(
+        f"The key '{invalid_conn_id}' has to be made of alphanumeric characters, dashes, dots and underscores exclusively",
+        resp,
+    )
+
+
 def test_action_logging_connection_masked_secrets(session, admin_client):
     admin_client.post("/connection/add", data=conn_with_extra(), follow_redirects=True)
     _check_last_log_masked_connection(session, dag_id=None, event="connection.create", execution_date=None)

From b52ae46c75181a68b0c0cb9a7eeb02b2af3cbeeb Mon Sep 17 00:00:00 2001
From: Matt Usifer <mattusifer@gmail.com>
Date: Tue, 9 May 2023 08:50:52 -0400
Subject: [PATCH 2/3] Catch all exceptions

---
 airflow/www/validators.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/airflow/www/validators.py b/airflow/www/validators.py
index 71b5a82a002a8..1a460d2f0a067 100644
--- a/airflow/www/validators.py
+++ b/airflow/www/validators.py
@@ -22,7 +22,6 @@
 
 from wtforms.validators import EqualTo, ValidationError
 
-from airflow import AirflowException
 from airflow.utils import helpers
 
 
@@ -96,5 +95,5 @@ def __call__(self, form, field):
         if field.data:
             try:
                 helpers.validate_key(field.data, self.max_length)
-            except AirflowException as e:
+            except Exception as e:
                 raise ValidationError(e)

From 1c5411eb3ce9fe2b7ade57c638c469d97428012b Mon Sep 17 00:00:00 2001
From: Matt Usifer <mattusifer@gmail.com>
Date: Wed, 10 May 2023 10:19:07 -0400
Subject: [PATCH 3/3] PR feedback and black formatting

---
 airflow/www/validators.py                | 2 +-
 tests/www/views/test_views_connection.py | 7 +++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/airflow/www/validators.py b/airflow/www/validators.py
index 1a460d2f0a067..d3e75b3888724 100644
--- a/airflow/www/validators.py
+++ b/airflow/www/validators.py
@@ -96,4 +96,4 @@ def __call__(self, form, field):
             try:
                 helpers.validate_key(field.data, self.max_length)
             except Exception as e:
-                raise ValidationError(e)
+                raise ValidationError(str(e))
diff --git a/tests/www/views/test_views_connection.py b/tests/www/views/test_views_connection.py
index 37e631ea8b87d..fb6566d93616d 100644
--- a/tests/www/views/test_views_connection.py
+++ b/tests/www/views/test_views_connection.py
@@ -60,11 +60,10 @@ def test_create_connection(admin_client, session):
 def test_invalid_connection_id_trailing_blanks(admin_client, session):
     invalid_conn_id = "conn_id_with_trailing_blanks   "
     invalid_connection = {**CONNECTION, "conn_id": invalid_conn_id}
-    resp = admin_client.post(
-        "/connection/add", data=invalid_connection, follow_redirects=True
-    )
+    resp = admin_client.post("/connection/add", data=invalid_connection, follow_redirects=True)
     check_content_in_response(
-        f"The key '{invalid_conn_id}' has to be made of alphanumeric characters, dashes, dots and underscores exclusively",
+        f"The key '{invalid_conn_id}' has to be made of alphanumeric characters, "
+        + "dashes, dots and underscores exclusively",
         resp,
     )
 
