Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment