Guarding against leaking GCHandles in read/write operations

See #19

Author: loudej

src/Microsoft.AspNet.Server.Kestrel/Http/SocketOutput.cs

@@ -53,11 +53,9 @@ private static void WriteCallback(UvWriteReq req, int status, Exception error, o
 
             SocketOutput _self;
             ArraySegment<byte> _buffer;
-            Action<Exception> _drained;
             UvStreamHandle _socket;
             Action<Exception, object> _callback;
             object _state;
-            GCHandle _pin;
 
             internal void Contextualize(
                 SocketOutput socketOutput,
@@ -87,8 +85,14 @@ private void OnWrite(UvWriteReq req, int status, Exception error)
             {
                 KestrelTrace.Log.ConnectionWriteCallback(0, status);
                 //NOTE: pool this?
+
+                var callback = _callback;
+                _callback = null;
+                var state = _state;
+                _state = null;
+
                 Dispose();
-                _callback(error, _state);
+                callback(error, state);
             }
         }
 

src/Microsoft.AspNet.Server.Kestrel/Networking/UvStreamHandle.cs

@@ -37,10 +37,27 @@ protected override bool ReleaseHandle()
 
         public void Listen(int backlog, Action<UvStreamHandle, int, Exception, object> callback, object state)
         {
-            _listenCallback = callback;
-            _listenState = state;
-            _listenVitality = GCHandle.Alloc(this, GCHandleType.Normal);
-            _uv.listen(this, 10, _uv_connection_cb);
+            if (_listenVitality.IsAllocated)
+            {
+                throw new InvalidOperationException("TODO: Listen may not be called more than once");
+            }
+            try
+            {
+                _listenCallback = callback;
+                _listenState = state;
+                _listenVitality = GCHandle.Alloc(this, GCHandleType.Normal);
+                _uv.listen(this, 10, _uv_connection_cb);
+            }
+            catch
+            {
+                _listenCallback = null;
+                _listenState = null;
+                if (_listenVitality.IsAllocated)
+                {
+                    _listenVitality.Free();
+                }
+                throw;
+            }
         }
 
         public void Accept(UvStreamHandle handle)
@@ -53,15 +70,37 @@ public void ReadStart(
             Action<UvStreamHandle, int, Exception, object> readCallback,
             object state)
         {
-            _allocCallback = allocCallback;
-            _readCallback = readCallback;
-            _readState = state;
-            _readVitality = GCHandle.Alloc(this, GCHandleType.Normal);
-            _uv.read_start(this, _uv_alloc_cb, _uv_read_cb);
+            if (_readVitality.IsAllocated)
+            {
+                throw new InvalidOperationException("TODO: ReadStop must be called before ReadStart may be called again");
+            }
+            try
+            {
+                _allocCallback = allocCallback;
+                _readCallback = readCallback;
+                _readState = state;
+                _readVitality = GCHandle.Alloc(this, GCHandleType.Normal);
+                _uv.read_start(this, _uv_alloc_cb, _uv_read_cb);
+            }
+            catch
+            {
+                _allocCallback = null;
+                _readCallback = null;
+                _readState = null;
+                if (_readVitality.IsAllocated)
+                {
+                    _readVitality.Free();
+                }
+                throw;
+            }
         }
 
         public void ReadStop()
         {
+            if (!_readVitality.IsAllocated)
+            {
+                throw new InvalidOperationException("TODO: ReadStart must be called before ReadStop may be called");
+            }
             _allocCallback = null;
             _readCallback = null;
             _readState = null;

src/Microsoft.AspNet.Server.Kestrel/Networking/UvWriteRequest.cs

@@ -40,45 +40,60 @@ public unsafe void Write(
             Action<UvWriteReq, int, Exception, object> callback,
             object state)
         {
-            // add GCHandle to keeps this SafeHandle alive while request processing
-            _pins.Add(GCHandle.Alloc(this, GCHandleType.Normal));
-
-            var pBuffers = (Libuv.uv_buf_t*)_bufs;
-            var nBuffers = bufs.Count;
-            if (nBuffers > BUFFER_COUNT)
+            try
             {
-                // create and pin buffer array when it's larger than the pre-allocated one
-                var bufArray = new Libuv.uv_buf_t[nBuffers];
-                var gcHandle = GCHandle.Alloc(bufArray, GCHandleType.Pinned);
-                _pins.Add(gcHandle);
-                pBuffers = (Libuv.uv_buf_t*)gcHandle.AddrOfPinnedObject();
+                // add GCHandle to keeps this SafeHandle alive while request processing
+                _pins.Add(GCHandle.Alloc(this, GCHandleType.Normal));
+
+                var pBuffers = (Libuv.uv_buf_t*)_bufs;
+                var nBuffers = bufs.Count;
+                if (nBuffers > BUFFER_COUNT)
+                {
+                    // create and pin buffer array when it's larger than the pre-allocated one
+                    var bufArray = new Libuv.uv_buf_t[nBuffers];
+                    var gcHandle = GCHandle.Alloc(bufArray, GCHandleType.Pinned);
+                    _pins.Add(gcHandle);
+                    pBuffers = (Libuv.uv_buf_t*)gcHandle.AddrOfPinnedObject();
+                }
+
+                for (var index = 0; index != nBuffers; ++index)
+                {
+                    // create and pin each segment being written
+                    var buf = bufs.Array[bufs.Offset + index];
+
+                    var gcHandle = GCHandle.Alloc(buf.Array, GCHandleType.Pinned);
+                    _pins.Add(gcHandle);
+                    pBuffers[index] = Libuv.buf_init(
+                        gcHandle.AddrOfPinnedObject() + buf.Offset,
+                        buf.Count);
+                }
+
+                _callback = callback;
+                _state = state;
+                _uv.write(this, handle, pBuffers, nBuffers, _uv_write_cb);
             }
-
-            for (var index = 0; index != nBuffers; ++index)
+            catch
             {
-                // create and pin each segment being written
-                var buf = bufs.Array[bufs.Offset + index];
-
-                var gcHandle = GCHandle.Alloc(buf.Array, GCHandleType.Pinned);
-                _pins.Add(gcHandle);
-                pBuffers[index] = Libuv.buf_init(
-                    gcHandle.AddrOfPinnedObject() + buf.Offset,
-                    buf.Count);
+                _callback = null;
+                _state = null;
+                Unpin(this);
+                throw;
             }
-
-            _callback = callback;
-            _state = state;
-            _uv.write(this, handle, pBuffers, nBuffers, _uv_write_cb);
         }
 
-        private static void UvWriteCb(IntPtr ptr, int status)
+        private static void Unpin(UvWriteReq req)
         {
-            var req = FromIntPtr<UvWriteReq>(ptr);
             foreach (var pin in req._pins)
             {
                 pin.Free();
             }
             req._pins.Clear();
+        }
+
+        private static void UvWriteCb(IntPtr ptr, int status)
+        {
+            var req = FromIntPtr<UvWriteReq>(ptr);
+            Unpin(req);
 
             var callback = req._callback;
             req._callback = null;