fix: Support half-close on native transport hijacked streams

Directory.Build.props

@@ -8,9 +8,7 @@
     <DebugType>embedded</DebugType>
     <SignAssembly>true</SignAssembly>
     <AssemblyOriginatorKeyFile>$(MSBuildThisFileDirectory)\Key.snk</AssemblyOriginatorKeyFile>
+    <!-- The public key of Key.snk. The SDK uses it for InternalsVisibleTo items that do not specify a key. -->
+    <PublicKey>002400000480000094000000060200000024000052534131000400000100010013a6d952388480a1ce272e8c8ac11d710668c8723e696a421190445a1e6198288112f5e04eb99a626f8bb1454cdf30ebfb0a09cb7fc7b299cb03aa6fea1ae9a58f05f9fb92a85ce82ad4490bb2f0074822d8b0a786684f26a6eb1765f9026dae4857925b4e077d04b6311bec7dacf8e8a031dcc9f7e0384bca914256abee25d9</PublicKey>
   </PropertyGroup>
-  <ItemGroup>
-    <!-- Defines the StrongNamePublicKeys class with constant which contains the repo identity public key formatted for use in InternalsVisibleTo attributes. -->
-    <Compile Include="$(MSBuildThisFileDirectory)\StrongNamePublicKeys.cs" Visible="false" />
-  </ItemGroup>
 </Project>

src/Docker.DotNet.Handler.Abstractions/Docker.DotNet.Handler.Abstractions.csproj

@@ -7,11 +7,17 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" />
   </ItemGroup>
+  <ItemGroup>
+    <InternalsVisibleTo Include="Docker.DotNet" />
+    <InternalsVisibleTo Include="Docker.DotNet.NativeHttp" />
+  </ItemGroup>
   <ItemGroup>
     <Using Include="System" />
     <Using Include="System.Collections.Generic" />
     <Using Include="System.IO" />
     <Using Include="System.Net.Http" />
+    <Using Include="System.Net.Sockets" />
+    <Using Include="System.Threading" />
     <Using Include="System.Threading.Tasks" />
     <Using Include="Microsoft.Extensions.Logging" />
   </ItemGroup>

src/Docker.DotNet.Handler.Abstractions/HijackScope.cs

@@ -0,0 +1,61 @@
+namespace Docker.DotNet.Handler.Abstractions;
+
+/// <summary>
+/// Tracks the socket of the connection that carries a hijacked stream.
+/// </summary>
+/// <remarks>
+/// The client begins a scope before it sends a request whose connection it
+/// hijacks. A transport that owns its sockets tracks the socket the request is
+/// written to, so it can close only the write side of the hijacked stream.
+/// </remarks>
+internal sealed class HijackScope : IDisposable
+{
+    private static readonly AsyncLocal<HijackScope?> CurrentScope = new();
+
+    private HijackScope()
+    {
+    }
+
+    /// <summary>
+    /// Gets the scope of the current asynchronous control flow, or
+    /// <see langword="null"/> if none was begun.
+    /// </summary>
+    public static HijackScope? Current
+        => CurrentScope.Value;
+
+    /// <summary>
+    /// Gets the socket the request was last written to, or
+    /// <see langword="null"/> if none was tracked.
+    /// </summary>
+    public Socket? Socket { get; private set; }
+
+    /// <summary>
+    /// Begins a scope for the current asynchronous control flow.
+    /// </summary>
+    /// <returns>The scope, which ends when it is disposed.</returns>
+    public static HijackScope Begin()
+    {
+        var scope = new HijackScope();
+        CurrentScope.Value = scope;
+        return scope;
+    }
+
+    /// <summary>
+    /// Tracks the socket the request is written to.
+    /// </summary>
+    /// <param name="socket">The socket.</param>
+    /// <remarks>
+    /// The last tracked socket wins, because a request that fails
+    /// on a pooled connection is retried on a new one.
+    /// </remarks>
+    public void TrackSocket(Socket socket)
+    {
+        Socket = socket;
+    }
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        CurrentScope.Value = null;
+    }
+}

src/Docker.DotNet.NativeHttp/Docker.DotNet.NativeHttp.csproj

@@ -15,6 +15,7 @@
     <Using Include="System.Net.Http" />
     <Using Include="System.Net.Http.Headers" />
     <Using Include="System.Net.Sockets" />
+    <Using Include="System.Runtime.CompilerServices" />
     <Using Include="System.Runtime.InteropServices" />
     <Using Include="System.Security" />
     <Using Include="System.Text" />

src/Docker.DotNet.NativeHttp/DockerHandlerFactory.cs

@@ -38,6 +38,7 @@ public ResolvedTransport CreateHandler(NativeHttpTransportOptions transportOptio
             MaxConnectionsPerServer = MaxConnectionsPerServer,
             PooledConnectionLifetime = PooledConnectionLifetime,
             PooledConnectionIdleTimeout = PooledConnectionIdleTimeout,
+            PlaintextStreamFilter = TrackingNetworkStream.EnableTrackingAsync,
         };
 
         if (IsUnixSocket(uri))
@@ -48,11 +49,12 @@ public ResolvedTransport CreateHandler(NativeHttpTransportOptions transportOptio
 
             // Proxies do not apply to Unix domain sockets.
             handler.UseProxy = false;
-            handler.ConnectCallback = (_, cancellationToken) => ConnectUnixSocketAsync(socketPath, cancellationToken);
+            handler.ConnectCallback = (context, cancellationToken) => ConnectUnixSocketAsync(context, socketPath, cancellationToken);
         }
         else
         {
             uri = CreateHttpUri(uri, clientOptions);
+            handler.ConnectCallback = ConnectTcpSocketAsync;
         }
 
         transportOptions.ConfigureHandler(handler);
@@ -70,7 +72,7 @@ public async Task<WriteClosableStream> HijackStreamAsync(HttpContent content)
         var stream = await content.ReadAsStreamAsync()
             .ConfigureAwait(false);
 
-        return new WriteClosableStreamWrapper(stream);
+        return new WriteClosableStreamWrapper(stream, HijackScope.Current?.Socket);
     }
 
     private static void Validate(Uri endpoint)
@@ -102,7 +104,7 @@ private static bool IsUnixSocket(Uri endpoint)
         return string.Equals(endpoint.Scheme, "unix", StringComparison.OrdinalIgnoreCase);
     }
 
-    private static async ValueTask<Stream> ConnectUnixSocketAsync(string socketPath, CancellationToken cancellationToken)
+    private static async ValueTask<Stream> ConnectUnixSocketAsync(SocketsHttpConnectionContext context, string socketPath, CancellationToken cancellationToken)
     {
         var socket = new Socket(AddressFamily.Unix, SocketType.Stream, ProtocolType.Unspecified);
 
@@ -111,7 +113,25 @@ private static async ValueTask<Stream> ConnectUnixSocketAsync(string socketPath,
             await socket.ConnectAsync(new UnixDomainSocketEndPoint(socketPath), cancellationToken)
                 .ConfigureAwait(false);
 
-            return new NetworkStream(socket, ownsSocket: true);
+            return TrackingNetworkStream.Create(socket, context);
+        }
+        catch
+        {
+            socket.Dispose();
+            throw;
+        }
+    }
+
+    private static async ValueTask<Stream> ConnectTcpSocketAsync(SocketsHttpConnectionContext context, CancellationToken cancellationToken)
+    {
+        var socket = new Socket(SocketType.Stream, ProtocolType.Tcp) { NoDelay = true };
+
+        try
+        {
+            await socket.ConnectAsync(context.DnsEndPoint, cancellationToken)
+                .ConfigureAwait(false);
+
+            return TrackingNetworkStream.Create(socket, context);
         }
         catch
         {

src/Docker.DotNet.NativeHttp/TrackingNetworkStream.cs

@@ -0,0 +1,76 @@
+#if NET
+namespace Docker.DotNet.NativeHttp;
+
+// SocketsHttpHandler does not expose the socket of an upgraded connection, so its
+// hijacked stream cannot close only its write side. As a workaround, this stream
+// records its socket in the current HijackScope when a hijack request is written to
+// it, which relies on SocketsHttpHandler writing a request in the asynchronous
+// control flow of the caller. A cleaner solution would send hijack requests through
+// a minimal HTTP/1.1 client of its own, like the legacy transport, but that client
+// would need its own TLS, proxy, and authentication handling.
+internal sealed class TrackingNetworkStream : NetworkStream
+{
+    private static readonly ConditionalWeakTable<HttpRequestMessage, TrackingNetworkStream> PendingStreams = new();
+
+    private volatile bool _isTrackingEnabled;
+
+    private TrackingNetworkStream(Socket socket)
+        : base(socket, ownsSocket: true)
+    {
+    }
+
+    public static Stream Create(Socket socket, SocketsHttpConnectionContext context)
+    {
+        var stream = new TrackingNetworkStream(socket);
+        PendingStreams.AddOrUpdate(context.InitialRequestMessage, stream);
+        return stream;
+    }
+
+    public static ValueTask<Stream> EnableTrackingAsync(SocketsHttpPlaintextStreamFilterContext context, CancellationToken cancellationToken)
+    {
+        // The plaintext stream filter runs once the connection is set up,
+        // after the TLS handshake.
+        if (PendingStreams.TryGetValue(context.InitialRequestMessage, out var stream))
+        {
+            PendingStreams.Remove(context.InitialRequestMessage);
+            stream._isTrackingEnabled = true;
+        }
+
+        return new ValueTask<Stream>(context.PlaintextStream);
+    }
+
+    public override void Write(byte[] buffer, int offset, int count)
+    {
+        TrackSocket();
+        base.Write(buffer, offset, count);
+    }
+
+    public override void Write(ReadOnlySpan<byte> buffer)
+    {
+        TrackSocket();
+        base.Write(buffer);
+    }
+
+    public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+    {
+        TrackSocket();
+        return base.WriteAsync(buffer, offset, count, cancellationToken);
+    }
+
+    public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+    {
+        TrackSocket();
+        return base.WriteAsync(buffer, cancellationToken);
+    }
+
+    private void TrackSocket()
+    {
+        // The connection setup runs in the scope of the request that started it, even
+        // if that request is sent over another connection, so setup writes do not count.
+        if (_isTrackingEnabled)
+        {
+            HijackScope.Current?.TrackSocket(Socket);
+        }
+    }
+}
+#endif

src/Docker.DotNet.NativeHttp/WriteClosableStreamWrapper.cs

@@ -1,9 +1,11 @@
 namespace Docker.DotNet.NativeHttp;
 
-internal sealed class WriteClosableStreamWrapper(Stream stream) : WriteClosableStream
+internal sealed class WriteClosableStreamWrapper(Stream stream, Socket? socket) : WriteClosableStream
 {
     private readonly Stream _stream = stream ?? throw new ArgumentNullException(nameof(stream));
 
+    private readonly Socket? _socket = socket;
+
     public override bool CanRead
         => _stream.CanRead;
 
@@ -41,7 +43,17 @@ public override void Write(byte[] buffer, int offset, int count)
         => _stream.Write(buffer, offset, count);
 
     public override void CloseWrite()
-        => _stream.Close();
+    {
+        if (_socket != null)
+        {
+            // Send buffered data before closing the write side of the socket.
+            _stream.Flush();
+            _socket.Shutdown(SocketShutdown.Send);
+            return;
+        }
+
+        _stream.Close();
+    }
 
     protected override void Dispose(bool disposing)
     {

src/Docker.DotNet/Docker.DotNet.csproj

@@ -23,6 +23,9 @@
     <ProjectReference Include="..\Docker.DotNet.Unix\Docker.DotNet.Unix.csproj" />
     <ProjectReference Include="..\..\tools\Docker.DotNet.SourceAnalyzers\Docker.DotNet.SourceAnalyzers.csproj" OutputItemType="Analyzer" ReferenceOutputAssembly="false" />
   </ItemGroup>
+  <ItemGroup>
+    <InternalsVisibleTo Include="Docker.DotNet.Tests" />
+  </ItemGroup>
   <ItemGroup>
     <Using Include="System" />
     <Using Include="System.Buffers" />

src/Docker.DotNet/DockerClient.cs

@@ -337,6 +337,10 @@ internal async Task<HijackedStreamResponse> MakeRequestForHijackedStreamAsync(
         headers.Add("Upgrade", "tcp");
         headers.Add("Connection", "upgrade");
 
+        // Lets the transport track the socket of the hijacked connection,
+        // so the stream can close only its write side.
+        using var hijackScope = HijackScope.Begin();
+
         var response = await PrivateMakeRequestAsync(timeout, HttpCompletionOption.ResponseHeadersRead, method, path, queryString, headers, body, cancellationToken)
             .ConfigureAwait(false);
 

test/Docker.DotNet.Tests/IContainerOperationsTests.cs

@@ -778,4 +778,42 @@ public async Task WriteAsync_OnMultiplexedStream_ForwardsInputToExecStdin_Comple
         Assert.Equal(0, containerExecInspectResponse.ExitCode);
         Assert.Equal("Done\n", stdout);
     }
+
+    [Fact]
+    public async Task CloseWrite_OnMultiplexedStream_ClosesExecStdin_CompletesExecProcess()
+    {
+        // Given
+        var inputBytes = Encoding.ASCII.GetBytes("Done\n");
+
+        var createContainerParameters = new CreateContainerParameters();
+        createContainerParameters.Image = _testFixture.Image.ID;
+        createContainerParameters.Entrypoint = CommonCommands.SleepInfinity;
+
+        var containerExecCreateParameters = new ContainerExecCreateParameters();
+        containerExecCreateParameters.AttachStdout = true;
+        containerExecCreateParameters.AttachStderr = true;
+        containerExecCreateParameters.AttachStdin = true;
+        containerExecCreateParameters.Cmd = new[] { "cat" };
+
+        var containerExecStartParameters = new ContainerExecStartParameters();
+
+        var createContainerResponse = await _testFixture.DockerClient.Containers.CreateContainerAsync(createContainerParameters, TestContext.Current.CancellationToken);
+        _ = await _testFixture.DockerClient.Containers.StartContainerAsync(createContainerResponse.ID, new ContainerStartParameters(), TestContext.Current.CancellationToken);
+
+        // When
+        var containerExecCreateResponse = await _testFixture.DockerClient.Exec.CreateContainerExecAsync(createContainerResponse.ID, containerExecCreateParameters, TestContext.Current.CancellationToken);
+        using var stream = await _testFixture.DockerClient.Exec.StartContainerExecAsync(containerExecCreateResponse.ID, containerExecStartParameters, TestContext.Current.CancellationToken);
+
+        await stream.WriteAsync(inputBytes, 0, inputBytes.Length, _testFixture.Cts.Token);
+        stream.CloseWrite();
+
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        var (stdout, _) = await stream.ReadOutputToEndAsync(cts.Token);
+
+        var containerExecInspectResponse = await _testFixture.DockerClient.Exec.InspectContainerExecAsync(containerExecCreateResponse.ID, _testFixture.Cts.Token);
+
+        // Then
+        Assert.Equal(0, containerExecInspectResponse.ExitCode);
+        Assert.Equal("Done\n", stdout);
+    }
 }