Repository: apache/airflow
File: airflow-core/src/airflow/models/connection.py (Python)

Code before:
from typing import Any
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit

from sqlalchemy import Boolean, ForeignKey, Integer, String, Text, select
from sqlalchemy.orm import Mapped, declared_attr, mapped_column, reconstructor, synonym

from airflow._shared.module_loading import import_string
from airflow._shared.secrets_masker import mask_secret
from airflow.configuration import conf, ensure_secrets_loaded
from airflow.exceptions import AirflowException, AirflowNotFoundException
from airflow.models.base import ID_LEN, Base
from airflow.models.crypto import get_fernet
from airflow.utils.helpers import prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import NEW_SESSION, provide_session

log = logging.getLogger(__name__)
# sanitize the `conn_id` pattern by allowing alphanumeric characters plus
# the symbols #,!,-,_,.,:,\,/ and () requiring at least one match.
#
# You can try the regex here: https://regex101.com/r/69033B/1
RE_SANITIZE_CONN_ID = re.compile(r"^[\w#!()\-.:/\\]{1,}$")
# the conn ID max len should be 250
CONN_ID_MAX_LEN: int = 250
# Pattern to mask URI password in log strings
RE_SAFE_LOG_URI = re.compile(r"://(.*):(.*)@(.*?)(://.*?)?(:\d+?)?(\?.*?)?")


def sanitize_conn_id(conn_id: str | None, max_length=CONN_ID_MAX_LEN) -> str | None:
    r"""
    Sanitizes the connection id and allows only specific characters to be within.

    Namely, it allows alphanumeric characters plus the symbols #,!,-,_,.,:,\,/ and () from 1 and up to
    250 consecutive matches. If desired, the max length can be adjusted by setting `max_length`.

    You can try to play with the regex here: https://regex101.com/r/69033B/1

    The character selection is such that it prevents the injection of javascript or
    executable bits to avoid any awkward behaviour in the front-end.

    :param conn_id: The connection id to sanitize.
    :param max_length: The max length of the connection ID, by default it is 250.
    :return: the sanitized string, `None` otherwise.
    """
    # check if `conn_id` or our match group is `None` and the `conn_id` is within the specified length.
    if (not isinstance(conn_id, str) or len(conn_id) > max_length) or (
        res := re.match(RE_SANITIZE_CONN_ID, conn_id)
    ) is None:
        return None

    # if we reach here, then we matched something, return the first match

Review comment (bug):
This will work for most cases but will break under edge cases like passwords with non-standard characters like ':' or the URI is malformed (for e.g. http://username:password@abc.com@xyz.com). I think regex is too brittle for this specific use case.

Code after:
from typing import Any
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit

from sqlalchemy import Boolean, ForeignKey, Integer, String, Text, select
from sqlalchemy.orm import Mapped, declared_attr, mapped_column, reconstructor, synonym

from airflow._shared.module_loading import import_string
from airflow._shared.secrets_masker import mask_secret
from airflow.configuration import conf, ensure_secrets_loaded
from airflow.exceptions import AirflowException, AirflowNotFoundException
from airflow.models.base import ID_LEN, Base
from airflow.models.crypto import get_fernet
from airflow.utils.helpers import prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import NEW_SESSION, provide_session

log = logging.getLogger(__name__)
# sanitize the `conn_id` pattern by allowing alphanumeric characters plus
# the symbols #,!,-,_,.,:,\,/ and () requiring at least one match.
#
# You can try the regex here: https://regex101.com/r/69033B/1
RE_SANITIZE_CONN_ID = re.compile(r"^[\w#!()\-.:/\\]{1,}$")
# the conn ID max len should be 250
CONN_ID_MAX_LEN: int = 250


def sanitize_conn_id(conn_id: str | None, max_length=CONN_ID_MAX_LEN) -> str | None:
    r"""
    Sanitizes the connection id and allows only specific characters to be within.

    Namely, it allows alphanumeric characters plus the symbols #,!,-,_,.,:,\,/ and () from 1 and up to
    250 consecutive matches. If desired, the max length can be adjusted by setting `max_length`.

    You can try to play with the regex here: https://regex101.com/r/69033B/1

    The character selection is such that it prevents the injection of javascript or
    executable bits to avoid any awkward behaviour in the front-end.

    :param conn_id: The connection id to sanitize.
    :param max_length: The max length of the connection ID, by default it is 250.
    :return: the sanitized string, `None` otherwise.
    """
    # check if `conn_id` or our match group is `None` and the `conn_id` is within the specified length.
    if (not isinstance(conn_id, str) or len(conn_id) > max_length) or (
        res := re.match(RE_SANITIZE_CONN_ID, conn_id)
    ) is None:
        return None

    # if we reach here, then we matched something, return the first match
    return res.group(0)