Update SharpFuzz to use the new OnBranch tracing

Author: Metalnem

README.md

@@ -238,7 +238,7 @@ dotnet add package Sigil --version 4.7.0
 the following command:
 
 ```shell
-dotnet add package SharpFuzz --version 1.4.3
+dotnet add package SharpFuzz --version 1.5.0
 ```
 
 **6.** Now it's time to write some code. The **Main**

build/Common.props

@@ -11,7 +11,7 @@
     <PublishRepositoryUrl>true</PublishRepositoryUrl>
     <PackageLicenseExpression>MIT</PackageLicenseExpression>
     <PackageRequireLicenseAcceptance>false</PackageRequireLicenseAcceptance>
-    <Copyright>Copyright © 2018, Nemanja Mijailovic</Copyright>
-    <PackageTags>fuzzing testing fuzz-testing fuzzer</PackageTags>
+    <Copyright>Copyright © 2018-2019, Nemanja Mijailovic</Copyright>
+    <PackageTags>fuzzing testing fuzz-testing fuzzer afl afl-fuzz libfuzzer</PackageTags>
   </PropertyGroup>
 </Project>

src/SharpFuzz/Fuzzer.cs

@@ -28,8 +28,15 @@ public static partial class Fuzzer
 		/// A function that accepts the full name of the class and returns
 		/// true if the class should be instrumented, false otherwise.
 		/// </param>
+		/// <param name="enableOnBranchCallback">
+		/// True if <see cref="SharpFuzz.Common.Trace.OnBranch"/> callback
+		/// should be called each time a branch is hit, false otherwise.
+		/// </param>
 		/// <returns>An ordered collection of instrumented types.</returns>
-		public static IEnumerable<string> Instrument(string source, Func<string, bool> matcher)
+		public static IEnumerable<string> Instrument(
+			string source,
+			Func<string, bool> matcher,
+			bool enableOnBranchCallback)
 		{
 			ThrowIfNull(source, nameof(source));
 			ThrowIfNull(matcher, nameof(matcher));
@@ -62,14 +69,14 @@ public static IEnumerable<string> Instrument(string source, Func<string, bool> m
 					{
 						var traceType = GenerateTraceType(src);
 						src.Types.Add(traceType);
-						types = Instrument(src, dst, matcher, traceType);
+						types = Instrument(src, dst, matcher, enableOnBranchCallback, traceType);
 					}
 					else
 					{
 						using (var commonMod = ModuleDefMD.Load(common.Location))
 						{
 							var traceType = commonMod.Types.Single(t => t.FullName == typeof(Common.Trace).FullName);
-							types = Instrument(src, dst, matcher, traceType);
+							types = Instrument(src, dst, matcher, enableOnBranchCallback, traceType);
 						}
 					}
 				}
@@ -85,13 +92,20 @@ public static IEnumerable<string> Instrument(string source, Func<string, bool> m
 			return types;
 		}
 
-		private static SortedSet<string> Instrument(ModuleDefMD src, Stream dst, Func<string, bool> matcher, TypeDef traceType)
+		private static SortedSet<string> Instrument(
+			ModuleDefMD src,
+			Stream dst,
+			Func<string, bool> matcher,
+			bool enableOnBranchCallback,
+			TypeDef traceType)
 		{
 			var sharedMemDef = traceType.Fields.Single(f => f.Name == nameof(Common.Trace.SharedMem));
 			var prevLocationDef = traceType.Fields.Single(f => f.Name == nameof(Common.Trace.PrevLocation));
+			var onBranchDef = traceType.Fields.Single(f => f.Name == nameof(Common.Trace.OnBranch));
 
 			var sharedMemRef = src.Import(sharedMemDef);
 			var prevLocationRef = src.Import(prevLocationDef);
+			var onBranchRef = src.Import(onBranchDef);
 
 			var types = new SortedSet<string>();
 
@@ -105,7 +119,7 @@ private static SortedSet<string> Instrument(ModuleDefMD src, Stream dst, Func<st
 					{
 						if (method.HasBody)
 						{
-							Method.Instrument(sharedMemRef, prevLocationRef, method);
+							Method.Instrument(sharedMemRef, prevLocationRef, onBranchRef, enableOnBranchCallback, method);
 
 							if (!instrumented)
 							{

src/SharpFuzz/Method.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using dnlib.DotNet;
@@ -12,15 +13,30 @@ internal sealed class Method
 	{
 		private readonly MemberRef sharedMem;
 		private readonly MemberRef prevLocation;
+		private readonly MemberRef onBranch;
+		private readonly bool enableOnBranchCallback;
 
+		private readonly IMethod invoke;
 		private readonly CilBody body;
 		private readonly List<Instruction> instructions;
 		private readonly Dictionary<Instruction, Instruction> instrumented;
 
-		private Method(MemberRef sharedMem, MemberRef prevLocation, MethodDef method)
+		private Method(
+			MemberRef sharedMem,
+			MemberRef prevLocation,
+			MemberRef onBranch,
+			bool enableOnBranchCallback,
+			MethodDef method)
 		{
 			this.sharedMem = sharedMem;
 			this.prevLocation = prevLocation;
+			this.onBranch = onBranch;
+			this.enableOnBranchCallback = enableOnBranchCallback;
+
+			if (enableOnBranchCallback)
+			{
+				invoke = method.Module.Import(typeof(Action<int, string>).GetMethod(nameof(Action.Invoke)));
+			}
 
 			body = method.Body;
 			instructions = body.Instructions.ToList();
@@ -30,16 +46,21 @@ private Method(MemberRef sharedMem, MemberRef prevLocation, MethodDef method)
 			body.Instructions.Clear();
 
 			FindInstrumentationTargets();
-			Instrument();
+			Instrument(method.FullName);
 			UpdateBranchTargets();
 			UpdateExceptionHandlers();
 
 			body.OptimizeBranches();
 		}
 
-		public static void Instrument(MemberRef sharedMem, MemberRef prevLocation, MethodDef method)
+		public static void Instrument(
+			MemberRef sharedMem,
+			MemberRef prevLocation,
+			MemberRef onBranch,
+			bool enableOnBranchCallback,
+			MethodDef method)
 		{
-			new Method(sharedMem, prevLocation, method);
+			new Method(sharedMem, prevLocation, onBranch, enableOnBranchCallback, method);
 		}
 
 		// Find all the locations that we want to instrument. These are:
@@ -83,13 +104,13 @@ private void FindInstrumentationTargets()
 		// Regenerate the IL for the method. If some instruction was
 		// previously marked as an instrumentation target, generate
 		// the instrumentation code and put it before the instruction.
-		private void Instrument()
+		private void Instrument(string methodName)
 		{
 			foreach (var ins in instructions)
 			{
 				if (instrumented.ContainsKey(ins))
 				{
-					using (var it = GenerateInstrumentationInstructions().GetEnumerator())
+					using (var it = GenerateInstrumentationInstructions(methodName).GetEnumerator())
 					{
 						it.MoveNext();
 						instrumented[ins] = it.Current;
@@ -111,7 +132,7 @@ private void Instrument()
 		// var id = IdGenerator.Next();
 		// SharpFuzz.Common.Trace.SharedMem[id ^ SharpFuzz.Common.Trace.PrevLocation]++;
 		// SharpFuzz.Common.Trace.PrevLocation = id >> 1;
-		private IEnumerable<Instruction> GenerateInstrumentationInstructions()
+		private IEnumerable<Instruction> GenerateInstrumentationInstructions(string methodName)
 		{
 			int id = IdGenerator.Next();
 
@@ -128,6 +149,14 @@ private IEnumerable<Instruction> GenerateInstrumentationInstructions()
 			yield return Instruction.Create(OpCodes.Stind_I1);
 			yield return Instruction.Create(OpCodes.Ldc_I4, id >> 1);
 			yield return Instruction.Create(OpCodes.Stsfld, prevLocation);
+
+			if (enableOnBranchCallback)
+			{
+				yield return Instruction.Create(OpCodes.Ldsfld, onBranch);
+				yield return Instruction.Create(OpCodes.Ldc_I4, id);
+				yield return Instruction.Create(OpCodes.Ldstr, methodName);
+				yield return Instruction.Create(OpCodes.Callvirt, invoke);
+			}
 		}
 
 		// Change all branch destinations to point to the first instruction

src/SharpFuzz/SharpFuzz.csproj

@@ -9,8 +9,8 @@
     <GenerateDocumentationFile>true</GenerateDocumentationFile>
     <PackageId>SharpFuzz</PackageId>
     <Title>SharpFuzz</Title>
-    <PackageVersion>1.4.3</PackageVersion>
-    <AssemblyVersion>1.4.3.0</AssemblyVersion>
+    <PackageVersion>1.5.0</PackageVersion>
+    <AssemblyVersion>1.5.0.0</AssemblyVersion>
     <Description>AFL-based fuzz testing for .NET</Description>
     <EmbedUntrackedSources>true</EmbedUntrackedSources>
     <AllowedOutputExtensionsInPackageBuildOutputFolder>$(AllowedOutputExtensionsInPackageBuildOutputFolder);.pdb</AllowedOutputExtensionsInPackageBuildOutputFolder>