From 77467e606aa439194f4a0863b64ab23351c6c5c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christian=20K=C3=B6llner?= <christian.koellner.80@gmail.com>
Date: Sun, 25 Aug 2019 22:06:29 +0200
Subject: [PATCH] Potential race condition when client releases capability, but
 the same capability is re-imported from server

---
 Capnp.Net.Runtime/Rpc/Impatient.cs            | 19 +++++----
 .../Rpc/RefCountingCapability.cs              | 40 +++++++++----------
 2 files changed, 29 insertions(+), 30 deletions(-)

diff --git a/Capnp.Net.Runtime/Rpc/Impatient.cs b/Capnp.Net.Runtime/Rpc/Impatient.cs
index 42155c6..bc497fa 100644
--- a/Capnp.Net.Runtime/Rpc/Impatient.cs
+++ b/Capnp.Net.Runtime/Rpc/Impatient.cs
@@ -42,18 +42,17 @@ namespace Capnp.Rpc
             }
             catch (ArgumentException)
             {
-                // Force .NET to create a new Task instance
-
-                var stask = rtask;
-
-                async Task<T> AwaitAgain()
+                if (rtask.IsCompleted)
                 {
-                    return await stask;
+                    // Force .NET to create a new Task instance
+                    var stask = rtask;
+                    rtask = new Task<T>(() => stask.Result);
+                    _taskTable.Add(rtask, promise);
+                }
+                else
+                {
+                    throw new InvalidOperationException("What the heck is wrong with Task?");
                 }
-
-                rtask = AwaitAgain();
-
-                _taskTable.Add(rtask, promise);
             }
 
             return rtask;
diff --git a/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs b/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs
index 2eb233b..31979f3 100644
--- a/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs
+++ b/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs
@@ -6,6 +6,8 @@ namespace Capnp.Rpc
 {
     abstract class RefCountingCapability: ConsumedCapability
     {
+        readonly object _reentrancyBlocker = new object();
+
         // Note on reference counting: Works in analogy to COM. AddRef() adds a reference,
         // Release() removes it. When the reference count reaches zero, the capability must be
         // released remotely, i.e. we need to tell the remote peer that it can remove this
@@ -64,36 +66,34 @@ namespace Capnp.Rpc
 
         internal sealed override void AddRef()
         {
-            if (Interlocked.Increment(ref _refCount) <= 1)
+            lock (_reentrancyBlocker)
             {
-                throw new ObjectDisposedException(nameof(ConsumedCapability));
+                if (++_refCount <= 1)
+                {
+                    throw new ObjectDisposedException(nameof(ConsumedCapability));
+                }
             }
         }
 
         internal sealed override void Release()
         {
-            while (true)
+            lock (_reentrancyBlocker)
             {
-                if ((Interlocked.CompareExchange(ref _refCount, 0, 2) == 2) ||
-                    (Interlocked.CompareExchange(ref _refCount, 0, 1) == 1))
+                switch (_refCount)
                 {
-                    Dispose(true);
-                    GC.SuppressFinalize(this);
-                    return;
-                }
+                    case 1: // initial state, actually ref. count 0
+                    case 2: // actually ref. count 1
+                        _refCount = 0;
+                        Dispose(true);
+                        GC.SuppressFinalize(this);
+                        break;
 
-                long oldCount = Interlocked.Read(ref _refCount);
+                    case var _ when _refCount > 2:
+                        --_refCount;
+                        break;
 
-                if (oldCount > 2)
-                {
-                    if (Interlocked.CompareExchange(ref _refCount, oldCount - 1, oldCount) == oldCount)
-                    {
-                        return;
-                    }
-                }
-                else if (oldCount <= 0)
-                {
-                    throw new InvalidOperationException("Capability is already disposed");
+                    default:
+                        throw new InvalidOperationException("Capability is already disposed");
                 }
             }
         }