fix quadratic parsing of WWW-Authenticate challenge in _parse_challenge

sdk/containerregistry/azure-containerregistry/CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- Fixed quadratic-time parsing of the `WWW-Authenticate` challenge header in `_parse_challenge`, which let a registry send an oversized header that consumed excessive CPU during authentication.
 
 ### Other Changes
 - Added support for Python 3.12.

sdk/containerregistry/azure-containerregistry/azure/containerregistry/_helpers.py

@@ -8,15 +8,15 @@
 import re
 import time
 import json
-from typing import Any, List, Dict, MutableMapping, IO, Optional, Union
+from typing import Any, Dict, MutableMapping, IO, Optional, Union
 from urllib.parse import urlparse
 from azure.core.exceptions import ServiceRequestError
 from azure.core.pipeline import PipelineRequest
 
 JSON = MutableMapping[str, Any]
 
 BEARER = "Bearer"
-AUTHENTICATION_CHALLENGE_PARAMS_PATTERN = re.compile('(?:(\\w+)="([^""]*)")+')
+AUTHENTICATION_CHALLENGE_PARAMS_PATTERN = re.compile('(?:^|[\\s,])(\\w+)="([^"]*)"')
 SUPPORTED_API_VERSIONS = ["2019-08-15-preview", "2021-07-01"]
 DEFAULT_CHUNK_SIZE = 4 * 1024 * 1024  # 4MB
 MAX_MANIFEST_SIZE = 4 * 1024 * 1024
@@ -45,25 +45,6 @@ def _is_tag(tag_or_digest: str) -> bool:
     return not (len(tag) == 2 and tag[0].startswith("sha"))
 
 
-def _clean(matches: List[str]) -> None:
-    """This method removes empty strings and commas from the regex matching of the Challenge header.
-
-    :param list[str] matches: The regex list to clean.
-    :return: None
-    """
-    while True:
-        try:
-            matches.remove("")
-        except ValueError:
-            break
-
-    while True:
-        try:
-            matches.remove(",")
-        except ValueError:
-            return
-
-
 def _parse_challenge(header: str) -> Dict[str, str]:
     """Parse challenge header into service and scope
 
@@ -74,11 +55,8 @@ def _parse_challenge(header: str) -> Dict[str, str]:
     ret: Dict[str, str] = {}
     if header.startswith(BEARER):
         challenge_params = header[len(BEARER) + 1 :]
-
-        matches = re.split(AUTHENTICATION_CHALLENGE_PARAMS_PATTERN, challenge_params)
-        _clean(matches)
-        for i in range(0, len(matches), 2):
-            ret[matches[i]] = matches[i + 1]
+        for match in AUTHENTICATION_CHALLENGE_PARAMS_PATTERN.finditer(challenge_params):
+            ret[match.group(1)] = match.group(2)
 
     return ret
 

sdk/containerregistry/azure-containerregistry/tests/test_container_registry_client.py

@@ -8,6 +8,7 @@
 import pytest
 import hashlib
 import json
+import time
 from datetime import datetime
 from io import BytesIO
 from unittest.mock import MagicMock
@@ -27,6 +28,7 @@
     DOCKER_MANIFEST,
     OCI_IMAGE_MANIFEST,
     DEFAULT_CHUNK_SIZE,
+    _parse_challenge,
 )
 from azure.core.exceptions import (
     ResourceNotFoundError,
@@ -960,3 +962,26 @@ def send(request: PipelineRequest, **kwargs) -> MagicMock:
                     assert manifest.architecture == "unknown"
                     assert isinstance(manifest.operating_system, str)
                     assert manifest.operating_system == "unknown"
+
+    def test_parse_challenge(self):
+        header = (
+            'Bearer realm="https://fake_url.azurecr.io/oauth2/token",'
+            'service="fake_url.azurecr.io",scope="registry:catalog:*"'
+        )
+        parsed = _parse_challenge(header)
+        assert parsed == {
+            "realm": "https://fake_url.azurecr.io/oauth2/token",
+            "service": "fake_url.azurecr.io",
+            "scope": "registry:catalog:*",
+        }
+
+    def test_parse_challenge_no_catastrophic_backtracking(self):
+        # The challenge header comes from the registry's 401 response, before the
+        # client is authenticated. A long run of word characters with no key="value"
+        # pair used to drive quadratic backtracking in the parser.
+        header = "Bearer " + "a" * 100000
+        start = time.perf_counter()
+        parsed = _parse_challenge(header)
+        elapsed = time.perf_counter() - start
+        assert parsed == {}
+        assert elapsed < 1.0