From 9051bbf322a226c0ff9bcf2bfa7009d715b28fd4 Mon Sep 17 00:00:00 2001
From: Cesar Blum Silveira <cesars@microsoft.com>
Date: Wed, 25 Jan 2017 14:11:21 -0800
Subject: [PATCH] Fix deadlock in SocketOutput (#1304).

---
 .../Internal/Http/SocketOutput.cs             |  2 +-
 .../CancellationTokenExtensions.cs            | 76 +++++++++++++++++++
 .../Infrastructure/DisposableAction.cs        | 40 ++++++++++
 3 files changed, 117 insertions(+), 1 deletion(-)
 create mode 100644 src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/CancellationTokenExtensions.cs
 create mode 100644 src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/DisposableAction.cs

diff --git a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/SocketOutput.cs b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/SocketOutput.cs
index 4a6cac500..14c7393f2 100644
--- a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/SocketOutput.cs
+++ b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/SocketOutput.cs
@@ -169,7 +169,7 @@ public Task WriteAsync(
                             _tasksPending.Enqueue(new WaitingTask()
                             {
                                 CancellationToken = cancellationToken,
-                                CancellationRegistration = cancellationToken.Register(_connectionCancellation, this),
+                                CancellationRegistration = cancellationToken.SafeRegister(_connectionCancellation, this),
                                 BytesToWrite = buffer.Count,
                                 CompletionSource = tcs
                             });
diff --git a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/CancellationTokenExtensions.cs b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/CancellationTokenExtensions.cs
new file mode 100644
index 000000000..2f9b6c1af
--- /dev/null
+++ b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/CancellationTokenExtensions.cs
@@ -0,0 +1,76 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Threading;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
+{
+    internal static class CancellationTokenExtensions
+    {
+        public static IDisposable SafeRegister(this CancellationToken cancellationToken, Action<object> callback, object state)
+        {
+            var callbackWrapper = new CancellationCallbackWrapper(callback, state);
+            var registration = cancellationToken.Register(s => InvokeCallback(s), callbackWrapper);
+            var disposeCancellationState = new DisposeCancellationState(callbackWrapper, registration);
+
+            return new DisposableAction(s => Dispose(s), disposeCancellationState);
+        }
+
+        private static void InvokeCallback(object state)
+        {
+            ((CancellationCallbackWrapper)state).TryInvoke();
+        }
+
+        private static void Dispose(object state)
+        {
+            ((DisposeCancellationState)state).TryDispose();
+        }
+
+        private class DisposeCancellationState
+        {
+            private readonly CancellationCallbackWrapper _callbackWrapper;
+            private readonly CancellationTokenRegistration _registration;
+
+            public DisposeCancellationState(CancellationCallbackWrapper callbackWrapper, CancellationTokenRegistration registration)
+            {
+                _callbackWrapper = callbackWrapper;
+                _registration = registration;
+            }
+
+            public void TryDispose()
+            {
+                if (_callbackWrapper.TrySetInvoked())
+                {
+                    _registration.Dispose();
+                }
+            }
+        }
+
+        private class CancellationCallbackWrapper
+        {
+            private readonly Action<object> _callback;
+            private readonly object _state;
+            private int _callbackInvoked;
+
+            public CancellationCallbackWrapper(Action<object> callback, object state)
+            {
+                _callback = callback;
+                _state = state;
+            }
+
+            public bool TrySetInvoked()
+            {
+                return Interlocked.Exchange(ref _callbackInvoked, 1) == 0;
+            }
+
+            public void TryInvoke()
+            {
+                if (TrySetInvoked())
+                {
+                    _callback(_state);
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/DisposableAction.cs b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/DisposableAction.cs
new file mode 100644
index 000000000..466e66b4a
--- /dev/null
+++ b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/DisposableAction.cs
@@ -0,0 +1,40 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Threading;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
+{
+    internal class DisposableAction : IDisposable
+    {
+        public static readonly DisposableAction Empty = new DisposableAction(() => { });
+
+        private Action<object> _action;
+        private readonly object _state;
+
+        public DisposableAction(Action action)
+            : this(state => ((Action)state).Invoke(), state: action)
+        {
+        }
+
+        public DisposableAction(Action<object> action, object state)
+        {
+            _action = action;
+            _state = state;
+        }
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                Interlocked.Exchange(ref _action, (state) => { }).Invoke(_state);
+            }
+        }
+
+        public void Dispose()
+        {
+            Dispose(true);
+        }
+    }
+}
\ No newline at end of file