Merge pull request #575 from koic/remove_dot_segments_in_linear_time

Resolve dot segments in linear time when canonicalizing a URL

Author: koic

lib/mcp/client/oauth/discovery.rb

@@ -579,47 +579,46 @@ def normalize_query(query)
               end.join("&")
           end
 
-          # Implements RFC 3986 Section 5.2.4 `remove_dot_segments`. Walks the input
-          # buffer one segment at a time, popping the previous output segment
-          # whenever a `..` is encountered, so that `/api/../mcp` collapses to
-          # `/mcp` and `/foo/./bar` collapses to `/foo/bar`.
+          # Implements RFC 3986 Section 5.2.4 `remove_dot_segments` over the path's segments, so that `/api/../mcp` collapses to
+          # `/mcp` and `/foo/./bar` collapses to `/foo/bar`. Each segment is visited once: the RFC's buffer rewriting,
+          # applied literally, copies the remaining input for every dot segment, and the path is the server's to choose.
+          #
+          # The output matches the RFC's algorithm for a relative path as well, including its quirk that a `..` popping
+          # the first segment leaves the result absolute (`a/../b` becomes `/b`), although `URI#path` never hands over a relative path.
+          # The rule letters below are the RFC's own: steps A through E of its loop.
           # https://www.rfc-editor.org/rfc/rfc3986#section-5.2.4
           def remove_dot_segments(path)
             return path if path.nil? || path.empty?
 
-            input = path.dup
-            output = +""
-            until input.empty?
-              if input.start_with?("../")
-                input = input[3..]
-              elsif input.start_with?("./")
-                input = input[2..]
-              elsif input.start_with?("/./")
-                input = "/#{input[3..]}"
-              elsif input == "/."
-                input = "/"
-              elsif input.start_with?("/../")
-                input = "/#{input[4..]}"
-                output = remove_last_segment(output)
-              elsif input == "/.."
-                input = "/"
-                output = remove_last_segment(output)
-              elsif input == "." || input == ".."
-                input = ""
+            absolute = path.start_with?("/")
+            segments = path.split("/", -1)
+            segments.shift if absolute
+
+            output = []
+
+            # True until a segment other than `.` or `..` is kept: a relative path's leading `./` and `../` are dropped together with
+            # the slash after them (Rule A), so the segment that follows is still the slash-less first one.
+            leading = !absolute
+            segments.each_with_index do |segment, index|
+              last = index == segments.length - 1
+
+              # Rule A, and Rule D for a path that is nothing but `.` or `..`.
+              next if leading && (segment == "." || segment == "..")
+
+              if segment == "."
+                # Rule B: `/./` disappears; a final `/.` leaves its slash behind.
+                output << "/" if last
+              elsif segment == ".."
+                # Rule C: `/../` removes the previous segment; a final `/..` leaves its slash behind.
+                output.pop
+                output << "/" if last
               else
-                segment = input.match(%r{\A/?[^/]*})[0]
-                output << segment
-                input = input[segment.length..]
+                # Rule E: the first segment of a relative path carries no slash.
+                output << (leading ? segment : "/#{segment}")
               end
+              leading = false
             end
-            output
-          end
-
-          def remove_last_segment(output)
-            idx = output.rindex("/")
-            return +"" if idx.nil?
-
-            output[0...idx]
+            output.join
           end
         end
       end

test/mcp/client/oauth/discovery_test.rb

@@ -291,6 +291,29 @@ def test_canonicalize_url_resolves_percent_encoded_dot_segments
           )
         end
 
+        def test_canonicalize_url_keeps_the_slash_a_final_dot_segment_leaves_behind
+          # RFC 3986 Section 5.2.4 turns a final `/.` or `/..` into `/`, so the path ends in a slash rather than
+          # losing it with the segment; an empty segment is a segment too.
+          assert_equal("https://srv.example.com/a/", Discovery.canonicalize_url("https://srv.example.com/a/."))
+          assert_equal("https://srv.example.com/a/", Discovery.canonicalize_url("https://srv.example.com/a/b/.."))
+          assert_equal("https://srv.example.com/a//b", Discovery.canonicalize_url("https://srv.example.com/a//b"))
+          assert_equal("https://srv.example.com", Discovery.canonicalize_url("https://srv.example.com/.."))
+        end
+
+        def test_canonicalize_url_resolves_a_path_with_many_dot_segments_in_linear_time
+          # The path is the server's to choose (a PRM `resource`, an endpoint URL). Rewriting the input buffer
+          # for every dot segment copied the remainder each time, so 300,000 of them took tens of seconds;
+          # the bound below is loose enough for a slow CI machine and far below that.
+          url = "https://srv.example.com/#{"a/../" * 300_000}mcp"
+
+          started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          canonical = Discovery.canonicalize_url(url)
+          elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started
+
+          assert_equal("https://srv.example.com/mcp", canonical)
+          assert_operator(elapsed, :<, 5)
+        end
+
         def test_canonicalize_url_drops_userinfo
           # The canonicalized URL is sent on the wire as the RFC 8707 `resource`
           # claim and is surfaced in error messages, so credentials in