Fix channel-open cancellation deadlock in ConnectionService

src/cs/Ssh/Services/ConnectionService.cs

@@ -49,6 +49,14 @@ public PendingChannel(
 	private long channelCounter = -1;
 	private Exception? closedException = null;
 
+	/// <summary>
+	/// Test-only hook invoked while holding <see cref="lockObject"/> in the channel-open
+	/// confirmation handler, right before the pending channel's cancellation registration is
+	/// captured. Tests use this to deterministically reproduce the channel-open cancellation
+	/// race. Always null (a no-op) in production.
+	/// </summary>
+	internal Action? TestHook_ChannelOpenResponseLocked { get; set; }
+
 	public ConnectionService(SshSession session) : base(session)
 	{
 		this.channels = new Dictionary<uint, SshChannel>();
@@ -492,14 +500,24 @@ private async Task HandleMessageAsync(
 
 		TaskCompletionSource<SshChannel>? completionSource = null;
 		ChannelOpenMessage openMessage;
+		CancellationTokenRegistration cancellationRegistration = default;
 		lock (this.lockObject)
 		{
 			if (this.pendingChannels.TryGetValue(
 				message.RecipientChannel, out var pendingChannel))
 			{
 				openMessage = pendingChannel.OpenMessage;
 				completionSource = pendingChannel.CompletionSource;
-				pendingChannel.CancellationRegistration.Dispose();
+
+				this.TestHook_ChannelOpenResponseLocked?.Invoke();
+
+				// Capture the registration and dispose it AFTER releasing the lock.
+				// Disposing here would deadlock: the callback registered in
+				// OpenChannelAsync also acquires this.lockObject, and
+				// CancellationTokenRegistration.Dispose() blocks until an in-flight
+				// callback completes. Removing the pending channel under the lock
+				// makes the callback a no-op, so deferring the dispose is safe.
+				cancellationRegistration = pendingChannel.CancellationRegistration;
 				this.pendingChannels.Remove(message.RecipientChannel);
 			}
 			else if (this.channels.ContainsKey(message.RecipientChannel))
@@ -534,6 +552,8 @@ private async Task HandleMessageAsync(
 			this.channels.Add(channel.ChannelId, channel);
 		}
 
+		cancellationRegistration.Dispose();
+
 		var args = new SshChannelOpeningEventArgs(openMessage, channel, isRemoteRequest: false);
 		await Session.OnChannelOpeningAsync(args, cancellation).ConfigureAwait(false);
 
@@ -565,17 +585,24 @@ private Task HandleMessageAsync(
 		cancellation.ThrowIfCancellationRequested();
 
 		TaskCompletionSource<SshChannel>? completionSource = null;
+		CancellationTokenRegistration cancellationRegistration = default;
 		lock (this.lockObject)
 		{
 			if (this.pendingChannels.TryGetValue(
 				message.RecipientChannel, out var pendingChannel))
 			{
 				completionSource = pendingChannel.CompletionSource;
-				pendingChannel.CancellationRegistration.Dispose();
+
+				// Capture the registration and dispose it after releasing the lock;
+				// disposing under the lock can deadlock with the cancellation callback
+				// (which also acquires this.lockObject). See HandleMessageAsync above.
+				cancellationRegistration = pendingChannel.CancellationRegistration;
 				this.pendingChannels.Remove(message.RecipientChannel);
 			}
 		}
 
+		cancellationRegistration.Dispose();
+
 		if (completionSource != null)
 		{
 			completionSource.TrySetException(new SshChannelException(

test/cs/Ssh.Test/ChannelTests.cs

@@ -4,6 +4,7 @@
 using System.Globalization;
 using System.IO;
 using System.Linq;
+using System.Reflection;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -176,6 +177,146 @@ public async Task OpenChannelCancelByAcceptor()
 		}
 	}
 
+	/// <summary>
+	/// Installs a test hook on the client's internal ConnectionService that deterministically
+	/// reproduces the channel-open cancellation race. The internal hook is reached via
+	/// reflection (as <see cref="ReconnectTests"/> does for KeyRotationThreshold) so the test
+	/// otherwise uses only the public session API. The hook runs on the receive pump thread
+	/// while it holds the ConnectionService lock during open-confirmation processing: it cancels
+	/// the open from another thread, using a dual-registration handshake to deterministically
+	/// confirm the internal callback is contending for the held lock before returning.
+	/// A monitoring callback is registered both BEFORE and AFTER the internal one (via
+	/// the returned action and inside the hook). CancellationToken callbacks fire in LIFO
+	/// order on .NET Core and FIFO on .NET Framework — so one of the two monitoring
+	/// callbacks always fires before the internal one. Since Cancel() executes callbacks
+	/// sequentially on the calling thread, once the monitoring callback completes, the
+	/// internal callback starts immediately on the same thread and blocks on lockObject.
+	/// Before the fix, the pump then disposed the cancellation registration while still
+	/// holding the lock, which blocks on the in-flight callback and deadlocks the session;
+	/// after the fix the registration is disposed after the lock is released.
+	/// </summary>
+	/// <returns>An action that must be invoked after the hook is installed but before
+	/// OpenChannelAsync, to register the "early" monitoring callback.</returns>
+	private static Action InstallOpenCancelRaceHook(
+		SshSession session,
+		CancellationTokenSource cancellationSource,
+		TaskCompletionSource<bool> callbackStarted)
+	{
+		var connectionServiceType = typeof(SshSession).Assembly.GetType(
+			"Microsoft.DevTunnels.Ssh.Services.ConnectionService");
+		var activateMethod = typeof(SshSession).GetMethods()
+			.Single(m => m.Name == nameof(SshSession.ActivateService) &&
+				m.IsGenericMethodDefinition &&
+				m.GetParameters().Length == 0)
+			.MakeGenericMethod(connectionServiceType);
+		var connectionService = activateMethod.Invoke(session, null);
+
+		var hookProperty = connectionServiceType.GetProperty(
+			"TestHook_ChannelOpenResponseLocked",
+			BindingFlags.NonPublic | BindingFlags.Instance);
+
+		Action hook = () =>
+		{
+			// Fire only once, on the first (raced) open confirmation.
+			hookProperty.SetValue(connectionService, null);
+
+			// Register a "late" monitoring callback — AFTER the internal callback
+			// (which was registered inside OpenChannelAsync). On LIFO runtimes
+			// (.NET Core) this fires first; on FIFO runtimes (.NET Framework) the
+			// "early" callback registered below fires first. Either way, one of them
+			// signals before the internal callback executes.
+			cancellationSource.Token.Register(
+				() => callbackStarted.TrySetResult(true));
+
+			// Cancel on a background thread. Cancel() invokes all registered
+			// callbacks sequentially on its thread before returning.
+			Task.Run(() => cancellationSource.Cancel());
+
+			// Wait for a monitoring callback to fire. Since Cancel() executes
+			// callbacks sequentially, the internal callback starts immediately
+			// after the monitoring one completes — on the same thread — and blocks
+			// on lockObject (which the caller holds). So when Wait() returns, the
+			// internal callback is guaranteed to be contending for the lock.
+			callbackStarted.Task.Wait();
+		};
+
+		hookProperty.SetValue(connectionService, hook);
+
+		// Return an action the caller must invoke BEFORE OpenChannelAsync to register
+		// the "early" monitoring callback (covers FIFO runtimes like .NET Framework).
+		return () => cancellationSource.Token.Register(
+			() => callbackStarted.TrySetResult(true));
+	}
+
+	[Fact]
+	public async Task OpenChannelCancelDuringConfirmationDoesNotDeadlock()
+	{
+		await this.sessionPair.ConnectAsync().WithTimeout(Timeout);
+
+		using var cancellationSource = new CancellationTokenSource();
+		var callbackStarted = new TaskCompletionSource<bool>(
+			TaskCreationOptions.RunContinuationsAsynchronously);
+		var registerEarlyCallback = InstallOpenCancelRaceHook(
+			this.clientSession, cancellationSource, callbackStarted);
+
+		// Register "early" monitoring callback BEFORE OpenChannelAsync registers
+		// the internal one. On FIFO (.NET Framework) this fires first.
+		registerEarlyCallback();
+
+		var openTask = this.clientSession.OpenChannelAsync(
+			(string)null, cancellationSource.Token);
+
+		// Before the fix this times out (the pump thread deadlocks); after the fix the open
+		// completes, whether it ends up cancelled or opened.
+		try
+		{
+			await openTask.WithTimeout(Timeout);
+		}
+		catch (OperationCanceledException)
+		{
+		}
+		catch (SshChannelException)
+		{
+		}
+	}
+
+	[Fact]
+	public async Task SessionRemainsUsableAfterOpenCancelRace()
+	{
+		await this.sessionPair.ConnectAsync().WithTimeout(Timeout);
+
+		using var cancellationSource = new CancellationTokenSource();
+		var callbackStarted = new TaskCompletionSource<bool>(
+			TaskCreationOptions.RunContinuationsAsynchronously);
+		var registerEarlyCallback = InstallOpenCancelRaceHook(
+			this.clientSession, cancellationSource, callbackStarted);
+
+		registerEarlyCallback();
+
+		var racedOpenTask = this.clientSession.OpenChannelAsync(
+			(string)null, cancellationSource.Token);
+
+		try
+		{
+			await racedOpenTask.WithTimeout(Timeout);
+		}
+		catch (OperationCanceledException)
+		{
+		}
+		catch (SshChannelException)
+		{
+		}
+
+		// If the pump thread had deadlocked, the receive loop could no longer process
+		// messages, so this subsequent open would hang instead of completing end-to-end.
+		var serverChannelTask = this.serverSession.AcceptChannelAsync();
+		var clientChannel = await this.clientSession.OpenChannelAsync().WithTimeout(Timeout);
+		var serverChannel = await serverChannelTask.WithTimeout(Timeout);
+
+		Assert.NotNull(clientChannel);
+		Assert.NotNull(serverChannel);
+	}
+
 	[Theory]
 	[InlineData(false, false)]
 	[InlineData(false, true)]