Fix URI handling when comparing encoded and unencoded URIs

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -42,7 +42,7 @@ namespace Roslyn.Test.Utilities
     [UseExportProvider]
     public abstract partial class AbstractLanguageServerProtocolTests
     {
-        private static readonly SystemTextJsonFormatter s_messageFormatter = RoslynLanguageServer.CreateJsonMessageFormatter();
+        protected static readonly JsonSerializerOptions JsonSerializerOptions = RoslynLanguageServer.CreateJsonMessageFormatter().JsonSerializerOptions;
 
         private protected readonly AbstractLspLogger TestOutputLspLogger;
         protected AbstractLanguageServerProtocolTests(ITestOutputHelper? testOutputHelper)
@@ -124,8 +124,8 @@ private protected static LSP.ClientCapabilities GetCapabilities(bool isVS)
         /// <param name="actual">the actual object to be converted to JSON.</param>
         public static void AssertJsonEquals<T1, T2>(T1 expected, T2 actual)
         {
-            var expectedStr = JsonSerializer.Serialize(expected, s_messageFormatter.JsonSerializerOptions);
-            var actualStr = JsonSerializer.Serialize(actual, s_messageFormatter.JsonSerializerOptions);
+            var expectedStr = JsonSerializer.Serialize(expected, JsonSerializerOptions);
+            var actualStr = JsonSerializer.Serialize(actual, JsonSerializerOptions);
             AssertEqualIgnoringWhitespace(expectedStr, actualStr);
         }
 
@@ -269,7 +269,7 @@ private protected static LSP.CompletionParams CreateCompletionParams(
                 SortText = sortText,
                 InsertTextFormat = LSP.InsertTextFormat.Plaintext,
                 Kind = kind,
-                Data = JsonSerializer.SerializeToElement(new CompletionResolveData(resultId, ProtocolConversions.DocumentToTextDocumentIdentifier(document)), s_messageFormatter.JsonSerializerOptions),
+                Data = JsonSerializer.SerializeToElement(new CompletionResolveData(resultId, ProtocolConversions.DocumentToTextDocumentIdentifier(document)), JsonSerializerOptions),
                 Preselect = preselect,
                 VsResolveTextEditOnCommit = vsResolveTextEditOnCommit,
                 LabelDetails = labelDetails
@@ -642,6 +642,11 @@ public Task ExecuteNotification0Async(string methodName)
                 return _clientRpc.NotifyWithParameterObjectAsync(methodName);
             }
 
+            public Task ExecutePreSerializedRequestAsync(string methodName, JsonDocument serializedRequest)
+            {
+                return _clientRpc.InvokeWithParameterObjectAsync(methodName, serializedRequest);
+            }
+
             public async Task OpenDocumentAsync(Uri documentUri, string? text = null, string languageId = "")
             {
                 if (text == null)

src/LanguageServer/Protocol/Workspaces/LspWorkspaceManager.cs

@@ -44,6 +44,47 @@ namespace Microsoft.CodeAnalysis.LanguageServer;
 /// </remarks>
 internal sealed class LspWorkspaceManager : IDocumentChangeTracker, ILspService
 {
+    private class LspUriComparer : IEqualityComparer<Uri>
+    {
+        public static readonly LspUriComparer Instance = new();
+        public bool Equals(Uri? x, Uri? y)
+        {
+            // Compare the absolute URIs to handle the case where one URI is encoded and the other is not.
+            // By default, Uri.Equals will not consider the encoded version of a URI equal to the unencoded version.
+            //
+            // The client is expected to be consistent in how it sends the URIs (either encoded or unencoded).
+            // So we normally can safely store the URIs as they send us in our map and expect subsequent requests to be encoded in the same way and match.
+            // See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#uri
+            //
+            // However when we serialize URIs to the client, we serialize the AbsoluteUri property which is always % encoded (no matter the original representation).
+            // For some requests, the client sends us exactly back what we sent (e.g. the data in a codelens/resolve request).
+            // This means that for these requests, the URI we will get from the client is the encoded version (that we sent).
+            // If the client sent us an unencoded URI originally, Uri.Equals will not consider it equal to the encoded version and we will fail to find the document
+            //
+            // So in order to resolve the encoded URI to the correct text, we can compare the AbsoluteUri properties (which are always encoded).
+            if (x is not null && y is not null && x.IsAbsoluteUri && y.IsAbsoluteUri && x.AbsoluteUri == y.AbsoluteUri)
+            {
+                return true;
+            }
+            else
+            {
+                return Uri.Equals(x, y);
+            }
+        }
+
+        public int GetHashCode(Uri obj)
+        {
+            if (obj.IsAbsoluteUri)
+            {
+                return obj.AbsoluteUri.GetHashCode();
+            }
+            else
+            {
+                return obj.GetHashCode();
+            }
+        }
+    }
+
     /// <summary>
     /// A cache from workspace to the last solution we returned for LSP.
     /// <para/> The forkedFromVersion is not null when the solution was created from a fork of the workspace with LSP
@@ -61,7 +102,7 @@ internal sealed class LspWorkspaceManager : IDocumentChangeTracker, ILspService
     /// the URI.
     /// <para/> Access to this is guaranteed to be serial by the <see cref="RequestExecutionQueue{RequestContextType}"/>
     /// </summary>
-    private ImmutableDictionary<Uri, (SourceText Text, string LanguageId)> _trackedDocuments = ImmutableDictionary<Uri, (SourceText, string)>.Empty;
+    private ImmutableDictionary<Uri, (SourceText Text, string LanguageId)> _trackedDocuments = ImmutableDictionary<Uri, (SourceText, string)>.Empty.WithComparers(LspUriComparer.Instance);
 
     private readonly ILspLogger _logger;
     private readonly LspMiscellaneousFilesWorkspace? _lspMiscellaneousFilesWorkspace;

src/LanguageServer/ProtocolUnitTests/UriTests.cs

@@ -3,9 +3,16 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Composition;
 using System.Linq;
+using System.Text.Json;
+using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.LanguageServer.Handler;
+using Microsoft.CodeAnalysis.Test.Utilities;
+using Microsoft.CommonLanguageServerProtocol.Framework;
 using Roslyn.Test.Utilities;
 using Xunit;
 using Xunit.Abstractions;
@@ -18,6 +25,8 @@ public UriTests(ITestOutputHelper? testOutputHelper) : base(testOutputHelper)
     {
     }
 
+    protected override TestComposition Composition => base.Composition.AddParts(typeof(CustomResolveHandler));
+
     [Theory, CombinatorialData]
     [WorkItem("https://github.com/dotnet/runtime/issues/89538")]
     public async Task TestMiscDocument_WithFileScheme(bool mutatingLspWorkspace)
@@ -150,4 +159,74 @@ public async Task TestWorkspaceDocument_WithFileAndGitScheme(bool mutatingLspWor
         Assert.Equal(gitDocumentUri, gitDocument.GetURI());
         Assert.Equal(gitDocumentText, gitText.ToString());
     }
+
+    [Theory, CombinatorialData]
+    public async Task TestFindsExistingDocumentWhenUriHasDifferentEncodingAsync(bool mutatingLspWorkspace)
+    {
+        await using var testLspServer = await CreateTestLspServerAsync(string.Empty, mutatingLspWorkspace, new InitializationOptions { ServerKind = WellKnownLspServerKinds.CSharpVisualBasicLspServer });
+
+        // Execute the request as JSON directly to avoid the test client serializing System.Uri using the encoded Uri to send to the server.
+        var requestJson = """
+            {
+                "textDocument": {
+                    "uri": "git:/c:/Users/dabarbet/source/repos/ConsoleApp10/ConsoleApp10/Program.cs?{{\"path\":\"c:\\\\Users\\\\dabarbet\\\\source\\\\repos\\\\ConsoleApp10\\\\ConsoleApp10\\\\Program.cs\",\"ref\":\"~\"}}",
+                    "languageId": "csharp",
+                    "text": "LSP text"
+                }
+            }
+            """;
+        var jsonDocument = JsonDocument.Parse(requestJson);
+        await testLspServer.ExecutePreSerializedRequestAsync(LSP.Methods.TextDocumentDidOpenName, jsonDocument);
+
+        // Retrieve the URI from the json - this is the unencoded (and not JSON escaped) version of the URI.
+        var unencodedUri = JsonSerializer.Deserialize<LSP.DidOpenTextDocumentParams>(jsonDocument, JsonSerializerOptions)!.TextDocument.Uri;
+
+        // Access the document using the unencoded URI to make sure we find it in the C# misc files.
+        var (workspace, _, lspDocument) = await testLspServer.GetManager().GetLspDocumentInfoAsync(new LSP.TextDocumentIdentifier { Uri = unencodedUri }, CancellationToken.None).ConfigureAwait(false);
+        AssertEx.NotNull(lspDocument);
+        Assert.Equal(WorkspaceKind.MiscellaneousFiles, workspace?.Kind);
+        Assert.Equal(LanguageNames.CSharp, lspDocument.Project.Language);
+        var originalText = await lspDocument.GetTextAsync(CancellationToken.None);
+        Assert.Equal("LSP text", originalText.ToString());
+
+        // Now make a request using the encoded document to ensure the server is able to find the document in misc C# files.
+        var encodedUriString = @"git:/c:/Users/dabarbet/source/repos/ConsoleApp10/ConsoleApp10/Program.cs?%7B%7B%22path%22:%22c:%5C%5CUsers%5C%5Cdabarbet%5C%5Csource%5C%5Crepos%5C%5CConsoleApp10%5C%5CConsoleApp10%5C%5CProgram.cs%22,%22ref%22:%22~%22%7D%7D";
+#pragma warning disable RS0030 // Do not use banned APIs
+        var encodedUri = new Uri(encodedUriString, UriKind.Absolute);
+#pragma warning restore RS0030 // Do not use banned APIs
+        var info = await testLspServer.ExecuteRequestAsync<CustomResolveParams, ResolvedDocumentInfo>(CustomResolveHandler.MethodName,
+                new CustomResolveParams(new LSP.TextDocumentIdentifier { Uri = encodedUri }), CancellationToken.None);
+        Assert.Equal(WorkspaceKind.MiscellaneousFiles, workspace?.Kind);
+        Assert.Equal(LanguageNames.CSharp, lspDocument.Project.Language);
+
+        var (encodedWorkspace, _, encodedDocument) = await testLspServer.GetManager().GetLspDocumentInfoAsync(new LSP.TextDocumentIdentifier { Uri = encodedUri }, CancellationToken.None).ConfigureAwait(false);
+        Assert.Same(workspace, encodedWorkspace);
+        AssertEx.NotNull(encodedDocument);
+        Assert.Equal(LanguageNames.CSharp, encodedDocument.Project.Language);
+        var encodedText = await encodedDocument.GetTextAsync(CancellationToken.None);
+        Assert.Equal("LSP text", encodedText.ToString());
+
+        // The text we get back should be the exact same instance that was originally saved by the unencoded request.
+        Assert.Same(originalText, encodedText);
+    }
+
+    private record class ResolvedDocumentInfo(string WorkspaceKind, string ProjectLanguage);
+    private record class CustomResolveParams([property: JsonPropertyName("textDocument")] LSP.TextDocumentIdentifier TextDocument);
+
+    [ExportCSharpVisualBasicStatelessLspService(typeof(CustomResolveHandler)), PartNotDiscoverable, Shared]
+    [LanguageServerEndpoint(MethodName, LanguageServerConstants.DefaultLanguageName)]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    private class CustomResolveHandler() : ILspServiceDocumentRequestHandler<CustomResolveParams, ResolvedDocumentInfo>
+    {
+        public const string MethodName = nameof(CustomResolveHandler);
+
+        public bool MutatesSolutionState => false;
+        public bool RequiresLSPSolution => true;
+        public LSP.TextDocumentIdentifier GetTextDocumentIdentifier(CustomResolveParams request) => request.TextDocument;
+        public Task<ResolvedDocumentInfo> HandleRequestAsync(CustomResolveParams request, RequestContext context, CancellationToken cancellationToken)
+        {
+            return Task.FromResult(new ResolvedDocumentInfo(context.Workspace!.Kind!, context.GetRequiredDocument().Project.Language));
+        }
+    }
 }