Work around Uri round-tripping issue in VS Code for sourcekit-lsp scheme

ahoppen · ahoppen · commit 57af01049047 · 2024-08-26T11:43:12.000-07:00
`vscode.uri` fails to round-trip URIs that have both a `=` and `%3D` (percent-encoded `=`) in the query component.

```ts
vscode.Uri.parse("scheme://host?outer=inner%3Dvalue").toString() -&gt; 'scheme://host?outer%3Dinner%3Dvalue'
vscode.Uri.parse("scheme://host?outer=inner%3Dvalue").toString(/*skipEncoding*/ true) -&gt; 'scheme://host?outer=inner=value'
```

The SourceKit-LSP scheme relies heavily on encoding options in the query parameters, eg. for Swift macro expansions and the values of those query parameters might contain percent-encoded `=` signs.

To work around the round-trip issue, use the URL type from Node.js to parse the URI and then map the URL components to the Uri components in VS Code. There’s still some mapping logic needed because the two don't agree on whether to include eg. `?` as part of the query or not, but it’s fairly straightforward.
diff --git a/src/sourcekit-lsp/LanguageClientManager.ts b/src/sourcekit-lsp/LanguageClientManager.ts
@@ -30,6 +30,7 @@ import { LSPLogger, LSPOutputChannel } from "./LSPOutputChannel";
 import { SwiftOutputChannel } from "../ui/SwiftOutputChannel";
 import { promptForDiagnostics } from "../commands/captureDiagnostics";
 import { activateGetReferenceDocument } from "./getReferenceDocument";
+import { uriConverters } from "./uriConverters";
 
 interface SourceKitLogMessageParams extends langclient.LogMessageParams {
     logName?: string;
@@ -574,6 +575,7 @@ export class LanguageClientManager {
                     };
                 })(),
             },
+            uriConverters: uriConverters,
             errorHandler,
             // Avoid attempting to reinitialize multiple times. If we fail to initialize
             // we aren't doing anything different the second time and so will fail again.
diff --git a/src/sourcekit-lsp/uriConverters.ts b/src/sourcekit-lsp/uriConverters.ts
@@ -0,0 +1,125 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the VS Code Swift open source project
+//
+// Copyright (c) 2024 the VS Code Swift project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of VS Code Swift project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import * as vscode from "vscode";
+
+export const uriConverters = {
+    protocol2Code: (value: string): vscode.Uri => {
+        if (!value.startsWith("sourcekit-lsp:")) {
+            // Use the default implementation for all schemes other than sourcekit-lsp, as defined here:
+            // https://github.com/microsoft/vscode-languageserver-node/blob/14ddabfc22187b698e83ecde072247aa40727308/client/src/common/protocolConverter.ts#L286
+            return vscode.Uri.parse(value);
+        }
+
+        // vscode.uri fails to round-trip URIs that have both a `=` and `%3D` (percent-encoded `=`) in the query component.
+        // ```ts
+        // vscode.Uri.parse("scheme://host?outer=inner%3Dvalue").toString() -> 'scheme://host?outer%3Dinner%3Dvalue'
+        // vscode.Uri.parse("scheme://host?outer=inner%3Dvalue").toString(/*skipEncoding*/ true) -> 'scheme://host?outer=inner=value'
+        // ```
+        // The SourceKit-LSP scheme relies heavily on encoding options in the query parameters, eg. for Swift macro
+        // expansions and the values of those query parameters might contain percent-encoded `=` signs.
+        //
+        // To work around the round-trip issue, use the URL type from Node.js to parse the URI and then map the URL
+        // components to the Uri components in VS Code.
+        const url = new URL(value);
+        let scheme = url.protocol;
+        if (scheme.endsWith(":")) {
+            // URL considers ':' part of the protocol, `vscode.URI` does not consider it part of the scheme.
+            scheme = scheme.substring(0, scheme.length - 1);
+        }
+
+        let auth = url.username;
+        if (url.password) {
+            auth += ":" + url.password;
+        }
+        let host = url.host;
+        if (auth) {
+            host = auth + "@" + host;
+        }
+
+        let query = url.search;
+        if (query.startsWith("?")) {
+            // URL considers '?' not part of the search, `vscode.URI` does consider '?' part of the query.
+            query = query.substring(1);
+        }
+
+        let fragment = url.hash;
+        if (fragment.startsWith("#")) {
+            // URL considers '#' not part of the hash, `vscode.URI` does consider '#' part of the fragment.
+            fragment = fragment.substring(1);
+        }
+
+        return vscode.Uri.from({
+            scheme: scheme,
+            authority: host,
+            path: url.pathname,
+            query: query,
+            fragment: fragment,
+        });
+    },
+    code2Protocol: (value: vscode.Uri): string => {
+        if (value.scheme !== "sourcekit-lsp") {
+            // Use the default implementation for all schemes other than sourcekit-lsp, as defined here:
+            // https://github.com/microsoft/vscode-languageserver-node/blob/14ddabfc22187b698e83ecde072247aa40727308/client/src/common/codeConverter.ts#L155
+            return value.toString();
+        }
+        // Create a dummy URL. We set all the components below.
+        const url = new URL(value.scheme + "://");
+
+        // Uri encodes username and password in `authority`. Url has its custom fields for those.
+        let host: string;
+        let username: string;
+        let password: string;
+        const atInAuthority = value.authority.indexOf("@");
+        if (atInAuthority != -1) {
+            host = value.authority.substring(atInAuthority + 1);
+            const auth = value.authority.substring(0, atInAuthority);
+            const colonInAuth = auth.indexOf(":");
+            if (colonInAuth == -1) {
+                username = auth;
+                password = "";
+            } else {
+                username = auth.substring(0, colonInAuth);
+                password = auth.substring(colonInAuth + 1);
+            }
+        } else {
+            host = value.authority;
+            username = "";
+            password = "";
+        }
+
+        // Need to set host before username and password because otherwise setting username + password is a no-op (probably
+        // because a URL can't have a username without a host).
+        url.host = host;
+        url.username = username;
+        url.password = password;
+        url.pathname = value.path;
+
+        let search = value.query;
+        if (search) {
+            // URL considers '?' not part of the search, vscode.URI does '?' part of the query.
+            search = "?" + search;
+        }
+        url.search = search;
+
+        let hash = value.fragment;
+        if (hash) {
+            // URL considers '#' not part of the hash, vscode.URI does '#' part of the fragment.
+            hash = "#" + hash;
+        }
+        url.hash = hash;
+
+        return url.toString();
+    },
+};
diff --git a/test/unit-tests/sourcekit-lsp/uriConverters.test.ts b/test/unit-tests/sourcekit-lsp/uriConverters.test.ts
@@ -0,0 +1,112 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the VS Code Swift open source project
+//
+// Copyright (c) 2024 the VS Code Swift project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of VS Code Swift project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import * as assert from "assert";
+import * as vscode from "vscode";
+import { uriConverters } from "../../../src/sourcekit-lsp/uriConverters";
+
+/// Check that decoding the given URI string and re-encoding it results in the original string and that the decoded Uri
+/// does not cause any assertion failures in `verifyUri`.
+function checkUri(input: string, verifyUri: (uri: vscode.Uri) => void) {
+    const uri = uriConverters.protocol2Code(input);
+    verifyUri(uri);
+    assert.equal(uriConverters.code2Protocol(uri), input);
+}
+
+suite.only("uriConverters Suite", () => {
+    suite("Default Coding", () => {
+        test("Space in host", () => {
+            checkUri("file://host%20with%20space/", uri => {
+                assert.equal(uri.authority, "host with space");
+            });
+        });
+
+        test("Space in path", () => {
+            checkUri("file://host/with%20space", uri => {
+                assert.equal(uri.path, "/with space");
+            });
+        });
+
+        test("Query does not round-trip", () => {
+            // If this test starts passing, the underlying VS Code issue that requires us to have custom URI coding
+            // has been fixed and we should be able to remove our custom uri converter.
+            const uri = uriConverters.protocol2Code("scheme://host?outer=inner%3Dvalue");
+            assert.equal(
+                uri.toString(/*skipEncoding*/ false),
+                "scheme://host?outer%3Dinner%3Dvalue"
+            );
+            assert.equal(uri.toString(/*skipEncoding*/ true), "scheme://host?outer=inner=value");
+        });
+    });
+
+    suite("Custom Coding", () => {
+        test("Basic", () => {
+            checkUri("sourcekit-lsp://host?outer=inner%3Dvalue", uri => {
+                assert.equal(uri.query, "outer=inner%3Dvalue");
+            });
+        });
+
+        test("Percent-encoded hash in query", () => {
+            checkUri("sourcekit-lsp://host?outer=with%23hash", uri => {
+                assert.equal(uri.query, "outer=with%23hash");
+            });
+        });
+
+        test("Query and fragment", () => {
+            checkUri("sourcekit-lsp://host?outer=with%23hash#fragment", uri => {
+                assert.equal(uri.query, "outer=with%23hash");
+                assert.equal(uri.fragment, "fragment");
+            });
+        });
+
+        test("Percent encoding in host", () => {
+            // Technically, it would be nice to percent-decode the authority and path here but then we get into
+            // ambiguities around username in the authority (see the `Encoded '@' in host` test).
+            // For now, rely on SourceKit-LSP not using any characters that need percent-encoding here.
+            checkUri("sourcekit-lsp://host%20with%20space", uri => {
+                assert.equal(uri.authority, "host%20with%20space");
+            });
+        });
+
+        test("Encoded '@' in host", () => {
+            checkUri("sourcekit-lsp://user%40with-at@host%40with-at", uri => {
+                assert.equal(uri.authority, "user%40with-at@host%40with-at");
+            });
+        });
+
+        test("Percent encoding in path", () => {
+            checkUri("sourcekit-lsp://host/with%20space", uri => {
+                assert.equal(uri.path, "/with%20space");
+            });
+        });
+
+        test("No query", () => {
+            checkUri("sourcekit-lsp://host/with/path", uri => {
+                assert.equal(uri.query, "");
+            });
+        });
+
+        test("With username", () => {
+            checkUri("sourcekit-lsp://user@host", uri => {
+                assert.equal(uri.authority, "user@host");
+            });
+        });
+
+        test("With username and password", () => {
+            checkUri("sourcekit-lsp://user:pass@host", uri => {
+                assert.equal(uri.authority, "user:pass@host");
+            });
+        });
+    });
+});
