Fix nested type builder ownership (#141)

Author: tinysec

Handle/BNEnumerationBuilder.cs

@@ -67,6 +67,12 @@ internal static EnumerationBuilder MustBorrowHandle(IntPtr handle)
 
 		    return new EnumerationBuilder(handle, false);
 	    }
+
+		internal void TransferOwnershipToCore()
+		{
+			this.SetHandle(IntPtr.Zero);
+			this.SetHandleAsInvalid();
+		}
 
 	    protected override bool ReleaseHandle()
 	    {
@@ -165,4 +171,4 @@ public Enumeration Build()
 	}
 
 
-}
+}

Handle/BNNamedTypeReferenceBuilder.cs

@@ -81,6 +81,12 @@ internal static NamedTypeReferenceBuilder MustBorrowHandle(IntPtr handle)
 
 		    return new NamedTypeReferenceBuilder(handle, false);
 	    }
+
+		internal void TransferOwnershipToCore()
+		{
+			this.SetHandle(IntPtr.Zero);
+			this.SetHandleAsInvalid();
+		}
 
         /// <summary>
         /// Releases the native BNNamedTypeReferenceBuilder handle when this instance is disposed or finalized.
@@ -156,4 +162,4 @@ public void SetName(QualifiedName name)
             );
         }
     }
-}
+}

Handle/BNStructureBuilder.cs

@@ -94,6 +94,12 @@ internal static StructureBuilder MustBorrowHandle(IntPtr handle)
 
 			return new StructureBuilder(handle, false);
 		}
+
+		internal void TransferOwnershipToCore()
+		{
+			this.SetHandle(IntPtr.Zero);
+			this.SetHandleAsInvalid();
+		}
 
 		protected override bool ReleaseHandle()
 		{

Handle/BNTypeBuilder.cs

@@ -5,7 +5,7 @@
 
 namespace BinaryNinja
 {
-	public class TypeBuilder : AbstractSafeHandle<TypeBuilder>
+	public partial class TypeBuilder : AbstractSafeHandle<TypeBuilder>
 	{
 		internal TypeBuilder(IntPtr handle , bool owner)
 			: base(handle , owner)
@@ -81,6 +81,7 @@ public static TypeBuilder CreateStructureType(Structure structure)
         /// </summary>
         /// <param name="builder">The structure builder to create a type builder from.</param>
         /// <returns>A new owned TypeBuilder for the structure type.</returns>
+        /// <remarks>The core takes ownership and invalidates <paramref name="builder"/>.</remarks>
         public static TypeBuilder CreateStructureTypeWithBuilder(StructureBuilder builder)
         {
             // 1. Validate the required builder parameter.
@@ -90,9 +91,13 @@ public static TypeBuilder CreateStructureTypeWithBuilder(StructureBuilder builde
             }
 
             // 2. Create a structure type builder from the builder; the returned handle is owned.
-            return TypeBuilder.MustTakeHandle(
-                NativeMethods.BNCreateStructureTypeBuilderWithBuilder(builder.DangerousGetHandle())
-            );
+			IntPtr handle = NativeMethods.BNCreateStructureTypeBuilderWithBuilder(
+				builder.DangerousGetHandle()
+			);
+			TypeBuilder result = TypeBuilder.MustTakeHandle(handle);
+			builder.TransferOwnershipToCore();
+
+			return result;
         }
 
         /// <summary>
@@ -713,6 +718,7 @@ public void SetIntegerTypeDisplayType(IntegerDisplayType displayType)
 		/// <param name="width">The width of the enumeration type in bytes.</param>
 		/// <param name="isSigned">The signedness with confidence.</param>
 		/// <returns>A new owned TypeBuilder for the enumeration type.</returns>
+		/// <remarks>The core takes ownership and invalidates <paramref name="builder"/>.</remarks>
 		public static unsafe TypeBuilder CreateEnumerationTypeWithBuilder(
 			Architecture arch ,
 			EnumerationBuilder builder ,
@@ -735,14 +741,16 @@ BoolWithConfidence isSigned
 			BNBoolWithConfidence nativeSigned = isSigned.ToNative();
 
 			// 3. Create an enumeration type builder from the builder; the returned handle is owned.
-			return TypeBuilder.MustTakeHandle(
-				NativeMethods.BNCreateEnumerationTypeBuilderWithBuilder(
+			IntPtr handle = NativeMethods.BNCreateEnumerationTypeBuilderWithBuilder(
 					arch.DangerousGetHandle() ,
 					builder.DangerousGetHandle() ,
 					width ,
 					(IntPtr)(&nativeSigned)
-				)
-			);
+				);
+			TypeBuilder result = TypeBuilder.MustTakeHandle(handle);
+			builder.TransferOwnershipToCore();
+
+			return result;
 		}
 
 		/// <summary>
@@ -815,6 +823,7 @@ BinaryNinja.Type type
 		/// <param name="cnst">The const qualifier with confidence.</param>
 		/// <param name="vltl">The volatile qualifier with confidence.</param>
 		/// <returns>A new owned TypeBuilder for the named type reference.</returns>
+		/// <remarks>The core takes ownership and invalidates <paramref name="builder"/>.</remarks>
 		public static unsafe TypeBuilder CreateNamedTypeReferenceBuilderWithBuilder(
 			NamedTypeReferenceBuilder builder ,
 			ulong width ,
@@ -834,15 +843,18 @@ BoolWithConfidence vltl
 			BNBoolWithConfidence nativeVltl = vltl.ToNative();
 
 			// 3. Create the builder; the returned handle is owned.
-			return TypeBuilder.MustTakeHandle(
+			IntPtr handle =
 				NativeMethods.BNCreateNamedTypeReferenceBuilderWithBuilder(
 					builder.DangerousGetHandle() ,
 					width ,
 					align ,
 					(IntPtr)(&nativeCnst) ,
 					(IntPtr)(&nativeVltl)
-				)
-			);
+				);
+			TypeBuilder result = TypeBuilder.MustTakeHandle(handle);
+			builder.TransferOwnershipToCore();
+
+			return result;
 		}
 
 		/// <summary>

Handle/BNTypeBuilderComposition.cs

@@ -0,0 +1,71 @@
+using System;
+
+namespace BinaryNinja
+{
+    public partial class TypeBuilder
+    {
+        /// <summary>Finalizes this generic builder into an immutable type.</summary>
+        public BinaryNinja.Type FinalizeType()
+        {
+            return BinaryNinja.Type.MustTakeHandle(
+                NativeMethods.BNFinalizeTypeBuilder(this.handle)
+            );
+        }
+
+        /// <summary>Replaces the mutable structure owned by this type builder.</summary>
+        /// <remarks>
+        /// The core takes ownership and invalidates <paramref name="builder"/>.
+        /// </remarks>
+        public void SetStructureBuilder(StructureBuilder builder)
+        {
+            if (null == builder)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+
+            NativeMethods.BNSetStructureBuilder(
+                this.handle,
+                builder.DangerousGetHandle()
+            );
+            builder.TransferOwnershipToCore();
+        }
+
+        /// <summary>Replaces the mutable enumeration owned by this type builder.</summary>
+        /// <remarks>
+        /// The core takes ownership and invalidates <paramref name="builder"/>.
+        /// </remarks>
+        public void SetEnumerationBuilder(EnumerationBuilder builder)
+        {
+            if (null == builder)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+
+            NativeMethods.BNSetEnumerationBuilder(
+                this.handle,
+                builder.DangerousGetHandle()
+            );
+            builder.TransferOwnershipToCore();
+        }
+
+        /// <summary>Replaces the mutable named reference owned by this type builder.</summary>
+        /// <remarks>
+        /// The core takes ownership and invalidates <paramref name="builder"/>.
+        /// </remarks>
+        public void SetNamedTypeReferenceBuilder(
+            NamedTypeReferenceBuilder builder
+        )
+        {
+            if (null == builder)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+
+            NativeMethods.BNSetNamedTypeReferenceBuilder(
+                this.handle,
+                builder.DangerousGetHandle()
+            );
+            builder.TransferOwnershipToCore();
+        }
+    }
+}