From 3ee755128fc4c4c5befa810b9e6793a6f83b978e Mon Sep 17 00:00:00 2001
From: antiduh <antiduh@e114e3f1-24f7-e311-a78f-001bfcf5e160>
Date: Wed, 25 Jun 2014 01:26:59 +0000
Subject: [PATCH] More CER goodness.

---
 Contexts/Context.cs              | 87 ++++++++++++++++++++++++--------
 Contexts/ContextNativeMethods.cs |  4 +-
 2 files changed, 68 insertions(+), 23 deletions(-)

diff --git a/Contexts/Context.cs b/Contexts/Context.cs
index 1741838..244efff 100644
--- a/Contexts/Context.cs
+++ b/Contexts/Context.cs
@@ -239,14 +239,36 @@ namespace NSspi
         internal SecPkgContext_Sizes QueryBufferSizes()
         {
             SecPkgContext_Sizes sizes = new SecPkgContext_Sizes();
-            SecurityStatus status;
+            SecurityStatus status = SecurityStatus.InternalError;
+            bool gotRef = false;
 
-            // TODO SAFE_CER
-            status = ContextNativeMethods.QueryContextAttributes_Sizes(
-                ref this.ContextHandle.rawHandle,
-                ContextQueryAttrib.Sizes,
-                ref sizes 
-            );
+            RuntimeHelpers.PrepareConstrainedRegions();
+            try
+            {
+                this.ContextHandle.DangerousAddRef( ref gotRef );
+            }
+            catch ( Exception )
+            {
+                if ( gotRef )
+                {
+                    this.ContextHandle.DangerousRelease();
+                    gotRef = false;
+                }
+
+                throw;
+            }
+            finally
+            {
+                if ( gotRef )
+                {
+                    status = ContextNativeMethods.QueryContextAttributes_Sizes(
+                        ref this.ContextHandle.rawHandle,
+                        ContextQueryAttrib.Sizes,
+                        ref sizes
+                    );
+                    this.ContextHandle.DangerousRelease();
+                }
+            }
 
             if( status != SecurityStatus.OK )
             {
@@ -259,8 +281,9 @@ namespace NSspi
         internal string QueryContextString(ContextQueryAttrib attrib)
         {
             SecPkgContext_String stringAttrib;
-            SecurityStatus status;
-            string result;
+            SecurityStatus status = SecurityStatus.InternalError;
+            string result = null;
+            bool gotRef = false;
 
             if( attrib != ContextQueryAttrib.Names && attrib != ContextQueryAttrib.Authority )
             {
@@ -269,25 +292,45 @@ namespace NSspi
 
             stringAttrib = new SecPkgContext_String();
 
-            // TODO SAFE_CER
-            status = ContextNativeMethods.QueryContextAttributes_String(
-                ref this.ContextHandle.rawHandle,
-                attrib,
-                ref stringAttrib
-            );
+            RuntimeHelpers.PrepareConstrainedRegions();
+            try
+            {
+                this.ContextHandle.DangerousAddRef( ref gotRef );
+            }
+            catch ( Exception )
+            {
+                if ( gotRef )
+                {
+                    this.ContextHandle.DangerousRelease();
+                    gotRef = false;
+                }
+                throw;
+            }
+            finally
+            {
+                if ( gotRef )
+                {
+                    status = ContextNativeMethods.QueryContextAttributes_String(
+                        ref this.ContextHandle.rawHandle,
+                        attrib,
+                        ref stringAttrib
+                    );
 
+                    this.ContextHandle.DangerousRelease();
+
+                    if ( status == SecurityStatus.OK )
+                    {
+                        result = Marshal.PtrToStringUni( stringAttrib.StringResult );
+                        ContextNativeMethods.FreeContextBuffer( stringAttrib.StringResult );
+                    }
+                }
+            }
             
             if( status == SecurityStatus.Unsupported )
             {
                 return null;
             }
-            else if( status == SecurityStatus.OK )
-            {
-                // TODO handle this safely.
-                result = Marshal.PtrToStringUni( stringAttrib.StringResult );
-                ContextNativeMethods.FreeContextBuffer( stringAttrib.StringResult );
-            }
-            else
+            else if( status != SecurityStatus.OK )
             {
                 throw new SSPIException( "Failed to query the context's associated user name", status );
             }
diff --git a/Contexts/ContextNativeMethods.cs b/Contexts/ContextNativeMethods.cs
index 91f1225..0483ff7 100644
--- a/Contexts/ContextNativeMethods.cs
+++ b/Contexts/ContextNativeMethods.cs
@@ -125,6 +125,7 @@ namespace NSspi
             int qualityOfProtection
         );
 
+        [ReliabilityContract( Consistency.WillNotCorruptState, Cer.Success )]
         [DllImport( "Secur32.dll", EntryPoint = "QueryContextAttributes", CharSet = CharSet.Unicode )]
         public static extern SecurityStatus QueryContextAttributes_Sizes(
             ref RawSspiHandle contextHandle,
@@ -132,6 +133,7 @@ namespace NSspi
             ref SecPkgContext_Sizes sizes
         );
 
+        [ReliabilityContract( Consistency.WillNotCorruptState, Cer.Success)]
         [DllImport( "Secur32.dll", EntryPoint = "QueryContextAttributes", CharSet = CharSet.Unicode )]
         public static extern SecurityStatus QueryContextAttributes_String(
             ref RawSspiHandle contextHandle,
@@ -139,7 +141,7 @@ namespace NSspi
             ref SecPkgContext_String names
         );
 
-
+        [ReliabilityContract( Consistency.WillNotCorruptState, Cer.Success )]
         [DllImport( "Secur32.dll", EntryPoint = "FreeContextBuffer", CharSet = CharSet.Unicode )]
         public static extern SecurityStatus FreeContextBuffer( IntPtr handle );