instrumentor.rb

require 'optparse'
require 'ostruct'
require 'Mono.Cecil'

class System::Reflection::Assembly
  def get_type_from_name(name)
    method(:GetType).overload(System::String).call(name)
  end
end

class System::Type
  def get_method_from_name_and_type(name, type)
    self.method(:GetMethod).
      overload(System::String, System::Array.of(System::Type)).
      call(name, System::Array.of(System::Type).new([type.to_clr_type]))
  end
end

module Log4NetInfo 
  def log_manager_get_logger
    log_manager = assembly.get_type_from_name('log4net.LogManager')
    raise "Could not load LogManager" if log_manager.nil?
    
    log_manager.get_method_from_name_and_type('GetLogger', System::String)
  end
  
  def log_interface_debug
    log_interface = assembly.get_type_from_name('log4net.ILog')
    raise "Could not load ILog" if log_interface.nil?
    
    log_interface.get_method_from_name_and_type('Debug', System::Object)
  end
  
  private  
    def assembly
      IronRuby.require('log4net')
    end
end

class CilWorkerHelper
  def initialize(assembly, worker, insert_point)
    @assembly = assembly
    @worker = worker
    @insert_point = insert_point
    @last_instruction = nil
  end
  
  def ldstr(message)
    insert(create_ldstr(message))    
  end
  
  def call(clr_method)
    insert(create_call(clr_method))
  end
  
  def call_virt(clr_method)
    insert(create_call_virt(clr_method))
  end
  
  private 
    def insert(instruction)
      if @last_instruction.nil?
        @worker.InsertBefore(@insert_point, instruction)
      else
        @worker.InsertAfter(@last_instruction, instruction)
      end
      
      @last_instruction = instruction
    end
    
    def create_call(clr_method)
      op_code = Mono::Cecil::Cil::OpCodes.Call     
      method_reference = @assembly.MainModule.Import(clr_method)
      
      @worker.
        method(:Create).
        overload(Mono::Cecil::Cil::OpCode, Mono::Cecil::MethodReference).
        call(op_code, method_reference)
    end
    
    def create_call_virt(clr_method)
      op_code = Mono::Cecil::Cil::OpCodes.Callvirt
      method_reference = @assembly.MainModule.Import(clr_method)
      
      @worker.
        method(:Create).
        overload(Mono::Cecil::Cil::OpCode, Mono::Cecil::MethodReference).
        call(op_code, method_reference)
    end
    
    def create_ldstr(string)
      op_code = Mono::Cecil::Cil::OpCodes.Ldstr
      
      @worker.
        method(:Create).
        overload(Mono::Cecil::Cil::OpCode, System::String).
        call(op_code, string)
    end
end

class Instrumentor
  include Log4NetInfo

  def initialize(assembly_name)
    @assembly_name = assembly_name
    @assembly = Mono::Cecil::AssemblyFactory.method(:GetAssembly).overload(System::String).call(assembly_name)
  end

  def apply
    @assembly.MainModule.Types.each do |type|
      type.Constructors.each do |constructor|
        name = "#{type.Namespace}.#{type.Name}.ctor"
        puts "Instrumenting #{name}"
        instrument(constructor, "Entering #{name}", "Exiting #{name}")
      end
      
      type.Methods.each do |method|
        name = "#{type.Namespace}.#{type.Name}.#{method.Name}"
        puts "Instrumenting #{name}"
        instrument(method, "Entering #{name}", "Exiting #{name}")
      end      
    end
    
    puts "Saving instrumented assembly..."
    Mono::Cecil::AssemblyFactory.SaveAssembly(@assembly, @assembly_name)
    puts "done."
  end
  
  def remove
    modified = false
    
    @assembly.MainModule.Types.each do |type|    
      type.Constructors.each do |constructor|
        if was_instrumented?(constructor)
          name = "#{type.Namespace}.#{type.Name}.ctor"
          puts "De-Instrumenting #{name}"
          de_instrument(constructor)
          modified = true
        end
      end
      
      type.Methods.each do |method|
        if was_instrumented?(method)
          name = "#{type.Namespace}.#{type.Name}.#{method.Name}"
          puts "De-Instrumenting #{name}"
          de_instrument(method)
          modified = true
        end
      end      
    end
    
    if modified
      puts "Saving de-instrumented assembly..."
      Mono::Cecil::AssemblyFactory.SaveAssembly(@assembly, @assembly_name)
      puts "done."
    else
      puts "Assembly does not contain any instrumentation."
    end
  end
  
  private
  
    def log_message(worker, insert_point, message)
      cil = CilWorkerHelper.new(@assembly, worker, insert_point)
      # ldstr "TracingLog"
      cil.ldstr("TracingLog")
      
      # call class [log4net]log4net.ILog [log4net]log4net.LogManager::GetLogger(string)
      cil.call(log_manager_get_logger)

      # ldstr message
      cil.ldstr(message)

      # callvirt instance void [log4net]log4net.ILog::Debug(object)
      cil.call_virt(log_interface_debug)
    end

    def instrument(method, enter_message, exit_message)
      if !method.nil?
        worker = method.Body.CilWorker
        
        first_instruction = method.Body.Instructions[0]
        log_message(worker, first_instruction, enter_message)
        
        last_instruction = method.Body.Instructions[method.Body.Instructions.Count - 1]
        log_message(worker, last_instruction, exit_message)
      end
    end
    
    def was_instrumented?(method)
      # make sure that the first instruction is one that we added
      first_instruction = method.Body.Instructions[0]
      return (first_instruction.OpCode == Mono::Cecil::Cil::OpCodes.Ldstr and
        first_instruction.Operand == "TracingLog")
    end
    
    def de_instrument(method)
      if !method.nil?
        if was_instrumented?(method)
          # remove the first 4 instructions
          4.times do 
            method.Body.Instructions.RemoveAt(0)
          end
          
          # remove the last 4 instructions (before the return)
          4.times do
            method.Body.Instructions.RemoveAt(method.Body.Instructions.Count - 2)
          end
        end
      end
    end
end

class InstrumentorOptions
  def self.parse(args)
    options = OpenStruct.new
    options.assembly_name = nil
    options.operation = nil
    
    optionParser = OptionParser.new do |optionParser|
      optionParser.banner = "Usage: instrumentor.rb [options] -a assembly_name"
      optionParser.separator ""
      optionParser.separator "Specific options:"
      
      optionParser.on("-a assembly_name", "--assembly assembly_name", "The assembly to modify") do |assembly_name|
        if !options.assembly_name.nil?
          puts "Error: assembly may only be specified once"
          puts optionParser
          exit
        end
        
        options.assembly_name = assembly_name
      end
      
      optionParser.on("--apply", "Instrument the specified assembly. This is the default if no options are specified.") do
        if !options.operation.nil?
          if options.operation == :apply
            puts "Error: --apply may only be supplied once"
          elsif options.operation == :remove
            puts "Error: --apply cannot be combined with --remove"
          else
            raise "Unexpected operation value"
          end
          
          puts optionParser
          exit
        end
        
        options.operation = :apply
      end
      
      optionParser.on("--remove", "De-instrument the specified assembly") do
        if !options.operation.nil?
          if options.operation == :apply
            puts "Error: --remove cannot be combined with --apply"
          elsif options.operation == :remove
            puts "Error: --remove may only be supplied once"
          else
            raise "Unexpected operation value"
          end
          
          puts optionParser
          exit
        end

        options.operation = :remove
      end
      
      optionParser.on_tail("-h", "--help", "Show this message") do
        puts optionsParser
        exit
      end
    end
    
    optionParser.parse!(args)
    
    if options.assembly_name.nil?
      puts "Error: assembly name must be specified."
      puts optionParser
      exit
    end
    
    if options.assembly_name == true
      puts "Error: assembly name must be specified."
      puts optionParser
      exit
    end
    
    if options.operation.nil?
      options.operation = :apply
    end
    
    options
  end
end

options = InstrumentorOptions.parse(ARGV)

instrumentor = Instrumentor.new(options.assembly_name)
if options.operation == :apply
  instrumentor.apply
elsif options.operation == :remove
  instrumentor.remove
else
  raise "Unexpected operation value."
end