From 3e3b144a800fef3ea0c0323bd576cfac9ac76889 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christian=20K=C3=B6llner?= <christian.koellner.80@gmail.com>
Date: Mon, 26 Aug 2019 20:03:56 +0200
Subject: [PATCH] Troubleshooting / robustness

---
 Capnp.Net.Runtime/Rpc/IRpcEndpoint.cs         |  1 -
 Capnp.Net.Runtime/Rpc/ImportedCapability.cs   |  1 -
 Capnp.Net.Runtime/Rpc/PromisedCapability.cs   |  2 +-
 .../Rpc/RefCountingCapability.cs              | 11 +++++++++++
 Capnp.Net.Runtime/Rpc/RpcEngine.cs            | 19 ++++++++-----------
 5 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/Capnp.Net.Runtime/Rpc/IRpcEndpoint.cs b/Capnp.Net.Runtime/Rpc/IRpcEndpoint.cs
index 12a4aca..e72e140 100644
--- a/Capnp.Net.Runtime/Rpc/IRpcEndpoint.cs
+++ b/Capnp.Net.Runtime/Rpc/IRpcEndpoint.cs
@@ -13,7 +13,6 @@ namespace Capnp.Rpc
         void RequestPostAction(Action postAction);
         void Finish(uint questionId);
         void ReleaseImport(uint importId);
-        void RemoveImport(uint importId);
         void Resolve(uint preliminaryId, Skeleton preliminaryCap, Func<ConsumedCapability> resolvedCapGetter);
 
         Task RequestSenderLoopback(Action<MessageTarget.WRITER> writer);
diff --git a/Capnp.Net.Runtime/Rpc/ImportedCapability.cs b/Capnp.Net.Runtime/Rpc/ImportedCapability.cs
index 9773e34..c821634 100644
--- a/Capnp.Net.Runtime/Rpc/ImportedCapability.cs
+++ b/Capnp.Net.Runtime/Rpc/ImportedCapability.cs
@@ -17,7 +17,6 @@ namespace Capnp.Rpc
         protected override void ReleaseRemotely()
         {
             _ep.ReleaseImport(_remoteId);
-            _ep.RemoveImport(_remoteId);
         }
 
         protected override Call.WRITER SetupMessage(DynamicSerializerState args, ulong interfaceId, ushort methodId)
diff --git a/Capnp.Net.Runtime/Rpc/PromisedCapability.cs b/Capnp.Net.Runtime/Rpc/PromisedCapability.cs
index 7eefd0a..87718a7 100644
--- a/Capnp.Net.Runtime/Rpc/PromisedCapability.cs
+++ b/Capnp.Net.Runtime/Rpc/PromisedCapability.cs
@@ -242,7 +242,7 @@ namespace Capnp.Rpc
                 _ep.ReleaseImport(_remoteId);
             }
 
-            _ep.RemoveImport(_remoteId);
+            _ep.ReleaseImport(_remoteId);
 
             this.DisposeWhenResolved();
         }
diff --git a/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs b/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs
index 31979f3..aa05db0 100644
--- a/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs
+++ b/Capnp.Net.Runtime/Rpc/RefCountingCapability.cs
@@ -97,5 +97,16 @@ namespace Capnp.Rpc
                 }
             }
         }
+
+        internal void Validate()
+        {
+            lock (_reentrancyBlocker)
+            {
+                if (_refCount <= 0)
+                {
+                    throw new ObjectDisposedException(nameof(ConsumedCapability), "Validation failed, capability is already disposed");
+                }
+            }
+        }
     }
 }
diff --git a/Capnp.Net.Runtime/Rpc/RpcEngine.cs b/Capnp.Net.Runtime/Rpc/RpcEngine.cs
index c88d0f4..d9e2f0a 100644
--- a/Capnp.Net.Runtime/Rpc/RpcEngine.cs
+++ b/Capnp.Net.Runtime/Rpc/RpcEngine.cs
@@ -1182,6 +1182,7 @@ namespace Capnp.Rpc
                             {
                                 if (rcw.Cap.TryGetTarget(out var impCap))
                                 {
+                                    impCap.Validate();
                                     rcw.AddRef();
                                     return impCap;
                                 }
@@ -1209,6 +1210,7 @@ namespace Capnp.Rpc
                             {
                                 if (rcwp.Cap.TryGetTarget(out var impCap))
                                 {
+                                    impCap.Validate();
                                     rcwp.AddRef();
                                     return impCap;
                                 }
@@ -1281,6 +1283,7 @@ namespace Capnp.Rpc
                                 if (rcv.Cap.TryGetTarget(out var impCap))
                                 {
                                     rcv.AddRef();
+                                    impCap.Validate();
                                     return impCap;
                                 }
                                 else
@@ -1430,6 +1433,11 @@ namespace Capnp.Rpc
                         count = rc.RefCount;
                         rc.ReleaseAll();
                     }
+
+                    if (exists)
+                    {
+                        _importTable.Remove(importId);
+                    }
                 }
 
                 if (exists && count > 0)
@@ -1451,17 +1459,6 @@ namespace Capnp.Rpc
                 }
             }
 
-            void IRpcEndpoint.RemoveImport(uint importId)
-            {
-                lock (_reentrancyBlocker)
-                {
-                    if (!_importTable.Remove(importId))
-                    {
-                        throw new ArgumentException("Given ID does not exist in import table");
-                    }
-                }
-            }
-
             Task IRpcEndpoint.RequestSenderLoopback(Action<MessageTarget.WRITER> describe)
             {
                 (var tcs, uint id) = AllocateDisembargo();