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

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,
             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.

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();
+    },
+};

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("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");
+            });
+        });
+    });
+});