public void Rewrite(MethodDefinition methodDefinition, ILProcessor il, ObserverRewriterContext context)
{
// Start from 2 - we added the 2 load variable instructions in ObserverRewriter
for (int i = 2; i < methodDefinition.Body.Instructions.Count; i++)
{
// Keep current number of instructions to detect changes.
int instructionCount = methodDefinition.Body.Instructions.Count;
Instruction instruction = methodDefinition.Body.Instructions[i];
if (instruction.OpCode.Code == Code.Newarr)
{
this.CheckArrayCreationSize(instruction, il, context);
}
if (instruction.Operand is MethodReference called)
{
this.PossiblyRewriteCalledMethod(called, instruction, il, context);
}
// If we rewrote, need to increase our iterator by the number of instructions we inserted.
int instructionsAdded = methodDefinition.Body.Instructions.Count - instructionCount;
i += instructionsAdded;
}
}
/// <summary>
/// Insert instructions to check on the size of an array before it is created.
/// Assumes the item on top of the stack at the moment is an int for the length of the array we're creating.
/// </summary>
private void CheckArrayCreationSize(Instruction instruction, ILProcessor il, ObserverRewriterContext context)
{
il.Body.SimplifyMacros();
il.InsertBefore(instruction, il.CreateLdlocBest(context.ObserverVariable));
il.InsertBefore(instruction, il.Create(OpCodes.Call, context.Observer.FlowThroughMemoryInt32Method));
il.Body.OptimizeMacros();
}
/// <inheritdoc />
public void Rewrite(MethodDefinition methodDefinition, ILProcessor il, ObserverRewriterContext context)
{
List <Instruction> branches = GetBranchingOps(methodDefinition).ToList();
List <Instruction> branchTos = branches.Select(x => (Instruction)x.Operand).ToList();
// Start from 2 because we setup Observer in first 2 instructions
int position = 2;
List <CodeSegment> segments = new List <CodeSegment>();
var codeSegment = new CodeSegment(methodDefinition);
while (position < methodDefinition.Body.Instructions.Count)
{
Instruction instruction = methodDefinition.Body.Instructions[position];
// End of a segment. Add this as the last instruction and move onwards with a new segment
if (branches.Contains(instruction))
{
codeSegment.Instructions.Add(instruction);
segments.Add(codeSegment);
codeSegment = new CodeSegment(methodDefinition);
}
// Start of a new segment. End last segment and start new one with this as the first instruction
else if (branchTos.Contains(instruction))
{
if (codeSegment.Instructions.Any())
{
segments.Add(codeSegment);
}
codeSegment = new CodeSegment(methodDefinition);
codeSegment.Instructions.Add(instruction);
}
// Just an in-between instruction. Add to current segment.
else
{
codeSegment.Instructions.Add(instruction);
}
position++;
}
// Got to end of the method. Add the last one if necessary
if (!segments.Contains(codeSegment) && codeSegment.Instructions.Any())
{
segments.Add(codeSegment);
}
foreach (CodeSegment segment in segments)
{
AddSpendGasMethodBeforeInstruction(il, context.Observer, context.ObserverVariable, segment);
}
// All of the branches now need to point to the place 3 instructions earlier!
foreach (Instruction branch in branches)
{
Instruction currentlyPointingTo = (Instruction)branch.Operand;
branch.Operand = currentlyPointingTo.Previous.Previous.Previous;
}
}
/// <summary>
/// Makes the <see cref="Observer"/> available to the given method as a variable and then
/// applies all of the individual rewriters to the method.
/// </summary>
private void RewriteMethod(MethodDefinition methodDefinition, ObserverReferences observer)
{
if (methodDefinition.DeclaringType == observer.InstanceField.DeclaringType)
{
return; // don't inject on our injected type.
}
if (!methodDefinition.HasBody || methodDefinition.Body.Instructions.Count == 0)
{
return; // don't inject on method without a Body
}
// Inject observer instance to method.
ILProcessor il = methodDefinition.Body.GetILProcessor();
var observerVariable = new VariableDefinition(observer.InstanceField.FieldType);
il.Body.Variables.Add(observerVariable);
Instruction start = methodDefinition.Body.Instructions[0];
il.InsertBefore(start, il.Create(OpCodes.Ldsfld, observer.InstanceField));
il.InsertBefore(start, il.CreateStlocBest(observerVariable));
var context = new ObserverRewriterContext(observer, observerVariable);
foreach (IObserverMethodRewriter rewriter in methodRewriters)
{
rewriter.Rewrite(methodDefinition, il, context);
}
}
/// <summary>
/// If is the specific string constructor that takes in a 'count' parameter, we need to check its size,
/// </summary>
private void CheckStringConstructor(Instruction instruction, ILProcessor il, ObserverRewriterContext context)
{
MethodDefinition method = ((MethodReference)instruction.Operand).Resolve();
// Ensure is the constructor with a count param (not all string constructors have a count param)
if (method.Parameters.Any(x => x.Name == "count"))
{
this.CheckArrayCreationSize(instruction, il, context);
}
}
/// <summary>
/// Insert instructions to check on the size of an array that has just been pushed onto the top of the stack.
/// </summary>
private void CheckArrayReturnSize(Instruction instruction, ILProcessor il, ObserverRewriterContext context)
{
// TODO: We could do away with the pop on the end and not return anything from the observer method but we would need to cast to long correctly.
il.Body.SimplifyMacros();
il.InsertAfter(instruction,
il.Create(OpCodes.Dup),
il.Create(OpCodes.Ldlen),
il.CreateLdlocBest(context.ObserverVariable),
il.Create(OpCodes.Call, context.Observer.FlowThroughMemoryInt32Method),
il.Create(OpCodes.Pop)
);
il.Body.OptimizeMacros();
}
/// <summary>
/// Checks if it is one of a few method calls we need to check on.
/// If so, does the necessary IL rewrite.
/// </summary>
private void PossiblyRewriteCalledMethod(MethodReference called, Instruction instruction, ILProcessor il, ObserverRewriterContext context)
{
if (called.DeclaringType.FullName == typeof(Array).FullName && called.Name == nameof(Array.Resize))
{
this.CheckArrayCreationSize(instruction, il, context);
return;
}
if (called.DeclaringType.FullName == typeof(string).FullName)
{
if (called.Name == nameof(string.ToCharArray))
{
this.CheckArrayReturnSize(instruction, il, context);
}
else if (called.Name == nameof(string.Split))
{
this.CheckArrayReturnSize(instruction, il, context);
}
else if (called.Name == nameof(string.Concat))
{
this.CheckArrayReturnSize(instruction, il, context);
}
else if (called.Name == nameof(string.Join))
{
this.CheckArrayReturnSize(instruction, il, context);
}
else if (called.Name == ".ctor")
{
this.CheckStringConstructor(instruction, il, context);
}
}
}
