Add an extra percent encoding layer when encoding DocumentURIs to LSP requests

The URI standard RFC 3986 is ambiguous about whether percent encoding and their represented characters are considered equivalent. VS Code considers them equivalent and treats them the same:

```js
vscode.Uri.parse("x://a?b=xxxx%3Dyyyy").toString() -> 'x://a?b%3Dxxxx%3Dyyyy'
vscode.Uri.parse("x://a?b=xxxx%3Dyyyy").toString(/*skipEncoding=*/true) -> 'x://a?b=xxxx=yyyy'
```

This causes issues because SourceKit-LSP's macro expansion URLs encoded by URLComponents use `=` do denote the separation of a key and a value in the outer query. The value of the `parent` key may itself contain query items, which use the escaped form '%3D'. Simplified, such a URL may look like `scheme://host?parent=scheme://host?line%3D2`.

But after running this through VS Code's URI type `=` and `%3D` get canonicalized and are indistinguishable.

To avoid this ambiguity, always percent escape the characters we use to distinguish URL query parameters, producing the following URL: `scheme://host?parent%3Dscheme://host%3Fline%253D2`.

Author: ahoppen

Sources/LanguageServerProtocol/SupportTypes/DocumentURI.swift

@@ -84,7 +84,19 @@ public struct DocumentURI: Codable, Hashable, Sendable {
   }
 
   public init(from decoder: Decoder) throws {
-    try self.init(string: decoder.singleValueContainer().decode(String.self))
+    let string = try decoder.singleValueContainer().decode(String.self)
+    guard let url = URL(string: string) else {
+      throw FailedToConstructDocumentURIFromStringError(string: string)
+    }
+    if url.query() != nil, var urlComponents = URLComponents(string: url.absoluteString) {
+      // See comment in `encode(to:)`
+      urlComponents.percentEncodedQuery = urlComponents.percentEncodedQuery!.removingPercentEncoding
+      if let rewrittenQuery = urlComponents.url {
+        self.init(rewrittenQuery)
+        return
+      }
+    }
+    self.init(url)
   }
 
   /// Equality check to handle escape sequences in file URLs.
@@ -97,7 +109,36 @@ public struct DocumentURI: Codable, Hashable, Sendable {
     hasher.combine(self.pseudoPath)
   }
 
+  private static let additionalQueryEncodingCharacterSet = CharacterSet(charactersIn: "?=&%").inverted
+
   public func encode(to encoder: Encoder) throws {
-    try storage.absoluteString.encode(to: encoder)
+    let urlToEncode: URL
+    if let query = storage.query(percentEncoded: true), var components = URLComponents(string: storage.absoluteString) {
+      // The URI standard RFC 3986 is ambiguous about whether percent encoding and their represented characters are
+      // considered equivalent. VS Code considers them equivalent and treats them the same:
+      //
+      // vscode.Uri.parse("x://a?b=xxxx%3Dyyyy").toString() -> 'x://a?b%3Dxxxx%3Dyyyy'
+      // vscode.Uri.parse("x://a?b=xxxx%3Dyyyy").toString(/*skipEncoding=*/true) -> 'x://a?b=xxxx=yyyy'
+      //
+      // This causes issues because SourceKit-LSP's macro expansion URLs encoded by URLComponents use `=` do denote the
+      // separation of a key and a value in the outer query. The value of the `parent` key may itself contain query
+      // items, which use the escaped form '%3D'. Simplified, such a URL may look like
+      // scheme://host?parent=scheme://host?line%3D2
+      // But after running this through VS Code's URI type `=` and `%3D` get canonicalized and are indistinguishable.
+      // To avoid this ambiguity, always percent escape the characters we use to distinguish URL query parameters,
+      // producing the following URL.
+      // scheme://host?parent%3Dscheme://host%3Fline%253D2
+      components.percentEncodedQuery =
+        query
+        .addingPercentEncoding(withAllowedCharacters: Self.additionalQueryEncodingCharacterSet)
+      if let componentsUrl = components.url {
+        urlToEncode = componentsUrl
+      } else {
+        urlToEncode = self.storage
+      }
+    } else {
+      urlToEncode = self.storage
+    }
+    try urlToEncode.absoluteString.encode(to: encoder)
   }
 }

Sources/SKTestSupport/CheckCoding.swift

@@ -43,9 +43,12 @@ package func checkCoding<T: Codable & Equatable>(
   XCTAssertEqual(json, str, file: file, line: line)
 
   let decoder = JSONDecoder()
-  let decodedValue = try! decoder.decode(WrapFragment<T>.self, from: data).value
-
-  XCTAssertEqual(value, decodedValue, file: file, line: line)
+  XCTAssertNoThrow(
+    try {
+      let decodedValue = try decoder.decode(WrapFragment<T>.self, from: data).value
+      XCTAssertEqual(value, decodedValue, file: file, line: line)
+    }()
+  )
 }
 
 /// JSONEncoder requires the top-level value to be encoded as a JSON container (array or object). Give it one.

Tests/LanguageServerProtocolTests/CodingTests.swift

@@ -1343,6 +1343,29 @@ final class CodingTests: XCTestCase {
         """
     )
   }
+
+  func testDocumentUriQueryParameterCoding() throws {
+    checkCoding(
+      try DocumentURI(string: "scheme://host?parent=scheme://host?line%3D2"),
+      json: #"""
+        "scheme:\/\/host?parent%3Dscheme:\/\/host%3Fline%253D2"
+        """#
+    )
+
+    checkCoding(
+      try DocumentURI(string: "scheme://host?parent=with space"),
+      json: #"""
+        "scheme:\/\/host?parent%3Dwith%2520space"
+        """#
+    )
+
+    checkCoding(
+      try DocumentURI(string: "scheme://host?parent=with%23hash"),
+      json: #"""
+        "scheme:\/\/host?parent%3Dwith%2523hash"
+        """#
+    )
+  }
 }
 
 func with<T>(_ value: T, mutate: (inout T) -> Void) -> T {