Merge pull request #571 from koic/parse_www_authenticate_in_linear_time

Parse the WWW-Authenticate challenge in linear time

Author: koic

lib/mcp/client/oauth/discovery.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "ipaddr"
+require "strscan"
 require "uri"
 
 module MCP
@@ -43,6 +44,12 @@ module Discovery
         # or a bare token, per RFC 7235.
         WWW_AUTH_PARAM_PATTERN = /\A([A-Za-z0-9_-]+)\s*=\s*(?:"((?:[^"\\]|\\.)*)"|([^\s,]+))/.freeze
 
+        # The whitespace and optional comma between two `key=value` pairs, or before the first one.
+        WWW_AUTH_PARAM_SEPARATOR_PATTERN = /\s*,?\s*/.freeze
+
+        # The `Bearer` challenge: at the start of the header or after a comma.
+        WWW_AUTH_BEARER_PATTERN = /(?:\A|,)\s*Bearer(?:\s+|\z)/i.freeze
+
         class << self
           # Parses a `WWW-Authenticate` header and returns the parameters of
           # the `Bearer` challenge as a hash with lower-cased keys (e.g. `resource_metadata`,
@@ -55,25 +62,21 @@ class << self
           def parse_www_authenticate(header)
             return {} unless header
 
-            # Locate the Bearer challenge: at the start of the header or after a comma.
-            bearer = header.match(/(?:\A|,)\s*Bearer(?:\s+|\z)/i)
-            return {} unless bearer
-
             # Walk key=value pairs starting where Bearer's parameters begin.
-            # The loop stops at the first token that is not a key=value pair,
-            # which marks the next challenge (e.g. `, DPoP algs="..."`).
-            cursor = bearer.end(0)
-            params = {}
-            while cursor < header.length
-              prefix = header[cursor..]
-              prefix = prefix.sub(/\A\s*,?\s*/, "")
-              break if prefix.empty?
+            # The loop stops at the first token that is not a key=value pair, which marks the next challenge (e.g. `, DPoP algs="..."`).
+            # The scanner keeps the walk linear in the header's length: slicing off the consumed prefix instead copies the remainder
+            # for every pair, and the server chooses how many pairs it sends. The header is also the server's to fill: a byte sequence
+            # that is not valid in the string's encoding would make the patterns raise `ArgumentError`, so such bytes are replaced first.
+            scanner = StringScanner.new(header.scrub)
+            return {} unless scanner.skip_until(WWW_AUTH_BEARER_PATTERN)
 
-              match = prefix.match(WWW_AUTH_PARAM_PATTERN)
-              break unless match
+            params = {}
+            until scanner.eos?
+              scanner.skip(WWW_AUTH_PARAM_SEPARATOR_PATTERN)
+              break if scanner.eos?
+              break unless scanner.scan(WWW_AUTH_PARAM_PATTERN)
 
-              params[match[1].downcase] = match[2] ? unescape_quoted_pair(match[2]) : match[3]
-              cursor = header.length - prefix.length + match.end(0)
+              params[scanner[1].downcase] = scanner[2] ? unescape_quoted_pair(scanner[2]) : scanner[3]
             end
             params
           end

test/mcp/client/oauth/discovery_test.rb

@@ -79,6 +79,49 @@ def test_parse_www_authenticate_unescapes_quoted_pair
           assert_equal('value with "quoted" word and a back\\slash', params["error_description"])
         end
 
+        def test_parse_www_authenticate_finds_bearer_after_a_challenge_with_multibyte_text
+          # The Bearer challenge is located by byte position, so text before it that is wider than
+          # one byte per character must not shift where its parameters are read from.
+          params = Discovery.parse_www_authenticate(%(Basic realm="café", Bearer scope="s"))
+
+          assert_equal({ "scope" => "s" }, params)
+        end
+
+        def test_parse_www_authenticate_tolerates_bytes_that_are_invalid_in_the_header_encoding
+          # A value the server fills with bytes that are not valid UTF-8 must not turn the `401` into
+          # an `ArgumentError`; the bytes are replaced and the other parameters still come through.
+          header = %(Bearer error="invalid_token", scope="s\xff", realm="r").dup.force_encoding(Encoding::UTF_8)
+
+          params = Discovery.parse_www_authenticate(header)
+
+          assert_equal("invalid_token", params["error"])
+          assert_equal("s�", params["scope"])
+          assert_equal("r", params["realm"])
+        end
+
+        def test_parse_www_authenticate_reads_a_binary_header
+          # Net::HTTP hands header values over as ASCII-8BIT; high bytes are kept as they are.
+          params = Discovery.parse_www_authenticate(%(Bearer scope="s\xff", realm="r").b)
+
+          assert_equal("s\xff".b, params["scope"])
+          assert_equal("r", params["realm"])
+        end
+
+        def test_parse_www_authenticate_walks_a_header_with_many_parameters_in_linear_time
+          # The header is the server's to choose. Slicing off the consumed prefix for every pair copied
+          # the remainder each time, so 200,000 pairs took tens of seconds; the bound below is loose enough
+          # for a slow CI machine and far below that.
+          header = "Bearer " + (1..200_000).map { |i| %(k#{i}="v#{i}") }.join(", ")
+
+          started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          params = Discovery.parse_www_authenticate(header)
+          elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started
+
+          assert_equal(200_000, params.size)
+          assert_equal("v200000", params["k200000"])
+          assert_operator(elapsed, :<, 5)
+        end
+
         def test_protected_resource_metadata_urls_uses_explicit_url_first
           urls = Discovery.protected_resource_metadata_urls(
             server_url: "https://api.example.com/mcp",