server: Rewrite the WebSocket procol handling

samples/GraphQL.Samples.SG.Subscription/GraphQL.Samples.SG.Subscription.csproj

@@ -12,8 +12,9 @@
     <ProjectReference Include="..\..\src\GraphQL\GraphQL.csproj" />
   </ItemGroup>
 
-	<ItemGroup>
-		<ProjectReference Include="..\..\src\GraphQL.Server.SourceGenerators\GraphQL.Server.SourceGenerators.csproj" OutputItemType="Analyzer" ReferenceOutputAssembly="false" />
-	</ItemGroup>
+    <ItemGroup>
+        <ProjectReference Include="..\..\src\GraphQL.Server.SourceGenerators\GraphQL.Server.SourceGenerators.csproj"
+                          OutputItemType="Analyzer" ReferenceOutputAssembly="false" />
+    </ItemGroup>
 
 </Project>

src/GraphQL.Server.SourceGenerators/GraphQL.Server.SourceGenerators.csproj

@@ -18,14 +18,14 @@
     </PropertyGroup>
 
     <ItemGroup>
-        <PackageReference Include="Polyfill" Version="2.6.5">
+        <PackageReference Include="Polyfill" Version="3.0.0">
             <PrivateAssets>all</PrivateAssets>
             <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
         </PackageReference>
     </ItemGroup>
 
     <ItemGroup>
-        <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="4.8.0" PrivateAssets="all" />
+        <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="4.9.2" PrivateAssets="all" />
         <PackageReference Include="Microsoft.CodeAnalysis.Analyzers" Version="3.3.4" PrivateAssets="all" />
         <PackageReference Include="System.Text.Json" Version="8.0.2" PrivateAssets="all" GeneratePathProperty="true" />
         <PackageReference Include="Scriban" Version="5.9.1" IncludeAssets="Build" />

src/GraphQL.Server.SourceGenerators/TypeHelper.cs

@@ -1,5 +1,4 @@
-using System;
-using System.Collections.Generic;
+using System.Collections.Generic;
 
 using Microsoft.CodeAnalysis;
 using System.Linq;

src/GraphQL.Server/GraphQL.Server.csproj

@@ -17,6 +17,7 @@
 	<ItemGroup>
 		<FrameworkReference Include="Microsoft.AspNetCore.App" />
 		<PackageReference Include="Microsoft.Extensions.Options" Version="8.0.2" />
+		<PackageReference Include="Microsoft.Extensions.Telemetry" Version="8.2.0" />
 		<PackageReference Include="System.IO.Pipelines" Version="8.0.0" />
 		<PackageReference Include="System.Net.WebSockets" Version="4.3.0" />
 	</ItemGroup>

src/GraphQL.Server/GraphQLWSTransport.cs

@@ -19,7 +19,7 @@ public class GraphQLWSTransport : IGraphQLTransport
     ///     Due to historical reasons this actually is the protocol name used
     ///     by the newer protocol.
     /// </summary>
-    public static string SubProtocol = "graphql-transport-ws";
+    public const string GraphQLTransportWSProtocol = "graphql-transport-ws";
 
     public IEndpointConventionBuilder Map(string pattern, IEndpointRouteBuilder routes,
         GraphQLRequestDelegate requestDelegate)
@@ -36,10 +36,12 @@ private async Task HandleProtocol(
         WebSocket webSocket,
         GraphQLRequestDelegate requestPipeline)
     {
-        var connection = new GraphQLWSConnection(webSocket, requestPipeline, httpContext);
-        await connection.Connect(httpContext.RequestAborted);
-    }
+        var handler = new WebSocketTransportHandler(
+            requestPipeline,
+            httpContext);
 
+        await handler.Handle(webSocket);
+    }
 
     private RequestDelegate ProcessRequest(GraphQLRequestDelegate pipeline)
     {
@@ -56,19 +58,28 @@ await httpContext.Response.WriteAsJsonAsync(new ProblemDetails
                 return;
             }
 
-            if (httpContext.WebSockets.WebSocketRequestedProtocols?.Contains(SubProtocol) == false)
+            if (httpContext.WebSockets.WebSocketRequestedProtocols?.Contains(EchoProtocol.Protocol) == true)
+            {
+                using WebSocket echoWebSocket = await httpContext.WebSockets
+                    .AcceptWebSocketAsync(EchoProtocol.Protocol);
+
+                await EchoProtocol.Run(echoWebSocket);
+                return;
+            }
+
+            if (httpContext.WebSockets.WebSocketRequestedProtocols?.Contains(GraphQLTransportWSProtocol) == false)
             {
                 httpContext.Response.StatusCode = StatusCodes.Status400BadRequest;
                 await httpContext.Response.WriteAsJsonAsync(new ProblemDetails
                 {
-                    Detail = $"Request does not contain sub-protocol '{SubProtocol}'."
+                    Detail = $"Request does not contain sub-protocol '{GraphQLTransportWSProtocol}'."
                 });
 
                 return;
             }
 
-            WebSocket webSocket = await httpContext.WebSockets
-                .AcceptWebSocketAsync(SubProtocol);
+            using WebSocket webSocket = await httpContext.WebSockets
+                .AcceptWebSocketAsync(GraphQLTransportWSProtocol);
 
             await HandleProtocol(httpContext, webSocket, pipeline);
         };

src/GraphQL.Server/WebSockets/EchoProtocol.cs

@@ -0,0 +1,28 @@
+using System.Net.WebSockets;
+using System.Text.Json;
+
+namespace Tanka.GraphQL.Server.WebSockets;
+
+public static class EchoProtocol
+{
+    public const string Protocol = "echo-ws";
+
+    public static async Task Run(WebSocket webSocket)
+    {
+        var channel = new WebSocketChannel(webSocket, new JsonSerializerOptions(JsonSerializerDefaults.Web));
+        var echo = Echo(channel);
+
+        await Task.WhenAll(channel.Run(), echo);
+    }
+
+    private static async Task Echo(WebSocketChannel channel)
+    {
+        while (await channel.Reader.WaitToReadAsync())
+        {
+            if (channel.Reader.TryRead(out var message))
+                await channel.Writer.WriteAsync(message);
+        }
+
+        channel.Complete();
+    }
+}

src/GraphQL.Server/WebSockets/GraphQLTransportWSProtocol.cs

@@ -0,0 +1,36 @@
+using Microsoft.Extensions.Logging;
+
+using Tanka.GraphQL.Server.WebSockets.Results;
+
+namespace Tanka.GraphQL.Server.WebSockets;
+
+public class GraphQLTransportWSProtocol(
+    SubscriptionManager subscriptions,
+    ILoggerFactory loggerFactory)
+{
+    public bool ConnectionInitReceived = false;
+
+    public IMessageResult Accept(MessageBase message)
+    {
+        if (!ConnectionInitReceived)
+            return new ConnectionAckResult(
+                this,
+                loggerFactory.CreateLogger<ConnectionAckResult>()
+            );
+
+        return message.Type switch
+        {
+            MessageTypes.ConnectionInit => new WebSocketCloseResult(
+                CloseCode.TooManyInitialisationRequests,
+                loggerFactory.CreateLogger<WebSocketCloseResult>()),
+            MessageTypes.Ping => new PongResult(loggerFactory.CreateLogger<PongResult>()),
+            MessageTypes.Subscribe => new Results.SubscribeResult(
+                subscriptions, 
+                loggerFactory.CreateLogger<Results.SubscribeResult>()),
+            MessageTypes.Complete => new Results.CompleteSubscriptionResult(
+                subscriptions,
+                loggerFactory.CreateLogger<Results.CompleteSubscriptionResult>()),
+            _ => new UnknownMessageResult(loggerFactory.CreateLogger<UnknownMessageResult>())
+        };
+    }
+}

src/GraphQL.Server/WebSockets/IMessageContext.cs

@@ -0,0 +1,10 @@
+namespace Tanka.GraphQL.Server.WebSockets;
+
+public interface IMessageContext
+{
+    Task Write<T>(T message) where T: MessageBase;
+
+    Task Close(Exception? error = default);
+
+    MessageBase Message { get; }
+}

src/GraphQL.Server/WebSockets/IMessageResult.cs

@@ -0,0 +1,6 @@
+namespace Tanka.GraphQL.Server.WebSockets;
+
+public interface IMessageResult
+{
+    Task Execute(IMessageContext context);
+}

src/GraphQL.Server/WebSockets/MessageContext.cs

@@ -0,0 +1,23 @@
+namespace Tanka.GraphQL.Server.WebSockets;
+
+public class MessageContext(
+    WebSocketChannel channel,
+    MessageBase contextMessage,
+    GraphQLRequestDelegate requestPipeline) : IMessageContext
+{
+    public async Task Write<T>(T message) where T: MessageBase
+    {
+        await channel.Writer.WriteAsync(message);
+    }
+
+    public Task Close(Exception? error = default)
+    {
+        channel.Complete(error);
+        return Task.CompletedTask;
+    }
+
+    public MessageBase Message => contextMessage;
+
+    public GraphQLRequestDelegate RequestPipeline => requestPipeline;
+
+}