From ca9032dc2051a4329ab9878528ccc745fa5d17ec Mon Sep 17 00:00:00 2001
From: Benjamin Carr <becarr@microsoft.com>
Date: Tue, 21 Jun 2016 20:16:21 -0700
Subject: [PATCH] Apply mitigation to CimInstance disposal race condition

---
 .../CimInstance.cs                            |  5 +++-
 .../Native/MI_Instance.cs                     | 17 +-------------
 .../Native/MI_NativeObject.cs                 | 23 ++++++++++++++++---
 .../Native/Structures/MI_Value.cs             |  4 +++-
 4 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/src/Microsoft.Management.Infrastructure/CimInstance.cs b/src/Microsoft.Management.Infrastructure/CimInstance.cs
index 05cfdea..3a3cc4c 100644
--- a/src/Microsoft.Management.Infrastructure/CimInstance.cs
+++ b/src/Microsoft.Management.Infrastructure/CimInstance.cs
@@ -384,7 +384,10 @@ private void Dispose(bool disposing)
             {
                 if (this.InstanceHandle != null && !this.InstanceHandle.IsNull)
                 {
-                    this.InstanceHandle.Delete();
+                    using (this.InstanceHandle)
+                    {
+                        this.InstanceHandle.Delete();
+                    }
                 }
             }
         }
diff --git a/src/Microsoft.Management.Infrastructure/Native/MI_Instance.cs b/src/Microsoft.Management.Infrastructure/Native/MI_Instance.cs
index b270478..96de591 100644
--- a/src/Microsoft.Management.Infrastructure/Native/MI_Instance.cs
+++ b/src/Microsoft.Management.Infrastructure/Native/MI_Instance.cs
@@ -3,7 +3,7 @@
 
 namespace Microsoft.Management.Infrastructure.Native
 {
-    internal class MI_Instance : MI_NativeObjectWithFT<MI_Instance.MI_InstanceFT>
+    internal class MI_Instance : MI_NativeObjectWithFT<MI_Instance.MI_InstanceFT>, IDisposable
     {
         internal MI_Result GetElement(
             string name,
@@ -124,21 +124,6 @@ internal MI_Result Delete()
             return localResult;
         }
 
-        private void ZeroPtr()
-        {
-            if (this.isDirect)
-            {
-                this.allocatedData = IntPtr.Zero;
-            }
-            else if (this.allocatedData != IntPtr.Zero)
-            {
-                unsafe
-                {
-                    *(IntPtr*)this.allocatedData = IntPtr.Zero;
-                }
-            }
-        }
-
         internal MI_Result IsA(
             MI_ClassDecl classDecl,
             out bool flag
diff --git a/src/Microsoft.Management.Infrastructure/Native/MI_NativeObject.cs b/src/Microsoft.Management.Infrastructure/Native/MI_NativeObject.cs
index d9156df..1d759dd 100644
--- a/src/Microsoft.Management.Infrastructure/Native/MI_NativeObject.cs
+++ b/src/Microsoft.Management.Infrastructure/Native/MI_NativeObject.cs
@@ -3,7 +3,7 @@
 
 namespace Microsoft.Management.Infrastructure.Native
 {
-    internal abstract class MI_NativeObject
+    internal abstract class MI_NativeObject : IDisposable
     {
         [StructLayout(LayoutKind.Sequential, CharSet = MI_PlatformSpecific.AppropriateCharSet)]
         protected struct MI_NativeObjectNormalMembersLayout
@@ -60,11 +60,13 @@ internal struct ArrayPtr
         protected IntPtr allocatedData;
         protected bool isDirect;
 
-        ~MI_NativeObject()
+        public virtual void Dispose()
         {
             if (this.allocatedData != IntPtr.Zero)
             {
-                Marshal.FreeHGlobal(this.allocatedData);
+                IntPtr tmp = this.allocatedData;
+                this.ZeroPtr();
+                Marshal.FreeHGlobal(tmp);
             }
         }
 
@@ -126,5 +128,20 @@ internal IntPtr Ptr
                 return structurePtr;
             }
         }
+
+        protected void ZeroPtr()
+        {
+            if (this.isDirect)
+            {
+                this.allocatedData = IntPtr.Zero;
+            }
+            else if (this.allocatedData != IntPtr.Zero)
+            {
+                unsafe
+                {
+                    *(IntPtr*)this.allocatedData = IntPtr.Zero;
+                }
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/src/Microsoft.Management.Infrastructure/Native/Structures/MI_Value.cs b/src/Microsoft.Management.Infrastructure/Native/Structures/MI_Value.cs
index 60064f5..8010acd 100644
--- a/src/Microsoft.Management.Infrastructure/Native/Structures/MI_Value.cs
+++ b/src/Microsoft.Management.Infrastructure/Native/Structures/MI_Value.cs
@@ -77,7 +77,7 @@ internal MI_Value() : base(true)
             this.Dispose(false);
         }
 
-        internal void Dispose()
+        public override void Dispose()
         {
             this.Dispose(true);
             GC.SuppressFinalize(this);
@@ -105,6 +105,8 @@ private void Dispose(bool disposing)
                     this.allocatedData = IntPtr.Zero;
                 }
             }
+
+            base.Dispose();
         }
 
         internal string String