Patched code sample

import re

# The following code is a self-contained example demonstrating the validation logic
# introduced in BlackSheep version 2.4.6 to fix CVE-2024-22779.
# The core of the fix is the addition of functions that validate header names and
# values to prevent CRLF (Carriage Return Line Feed) injection.

# Regular expressions to find invalid characters in header names and values.
# This mirrors the implementation in the patch.
_INVALID_HEADER_NAME_RE = re.compile(b"[\r\n:]")
_INVALID_HEADER_VALUE_RE = re.compile(b"[\r\n]")


def _validate_header_name(name: bytes):
    """
    Validates an HTTP header name, raising ValueError if it contains
    invalid characters like CR, LF, or colon.
    """
    if _INVALID_HEADER_NAME_RE.search(name):
        raise ValueError(f"Invalid header name: {name.decode('latin-1')}")


def _validate_header_value(value: bytes):
    """
    Validates an HTTP header value, raising ValueError if it contains
    invalid characters like CR or LF.
    """
    if _INVALID_HEADER_VALUE_RE.search(value):
        raise ValueError(f"Invalid header value: {value.decode('latin-1')}")


def set_header_safely(headers: dict, name: str, value: str):
    """
    A function that simulates how the fixed BlackSheep client would
    set a header, by encoding and validating its parts.
    """
    name_bytes = name.encode("latin-1")
    value_bytes = value.encode("latin-1")

    # This is the fix: validating inputs before using them.
    _validate_header_name(name_bytes)
    _validate_header_value(value_bytes)

    headers[name_bytes] = value_bytes
    print(f"Successfully set header: '{name}: {value}'")


def demonstrate_crlf_fix():
    """
    Demonstrates the fix by attempting to set both a malicious and a
    valid header.
    """
    print("--- Demonstrating fix for CVE-2024-22779 ---")

    # Example 1: Malicious header value with CRLF injection attempt.
    # The attacker tries to inject a new header ("Injected-Header: InjectedValue").
    malicious_value = "NormalValue\r\nInjected-Header: InjectedValue"
    headers = {}

    print(f"\nAttempting to set a header with a malicious value: '{malicious_value}'")
    try:
        set_header_safely(headers, "X-Normal-Header", malicious_value)
    except ValueError as e:
        print(f"FIX APPLIED: Blocked malicious header. Reason: {e}")

    # Example 2: Malicious header name with CRLF injection attempt.
    malicious_name = "X-Normal-Header\r\nX-Injected-Header: Injected"
    headers = {}
    print(f"\nAttempting to set a malicious header name: '{malicious_name}'")
    try:
        set_header_safely(headers, malicious_name, "some-value")
    except ValueError as e:
        print(f"FIX APPLIED: Blocked malicious header name. Reason: {e}")

    # Example 3: A valid, non-malicious header.
    valid_value = "This is a legitimate value"
    headers = {}
    print(f"\nAttempting to set a valid header with value: '{valid_value}'")
    try:
        set_header_safely(headers, "X-Legit-Header", valid_value)
        print("Validation passed for legitimate header.")
        print("Final headers:", headers)
    except ValueError as e:
        print(f"ERROR: A valid header was unexpectedly blocked. Reason: {e}")


if __name__ == "__main__":
    demonstrate_crlf_fix()