Durante o período de desenvolvimento do Jitex, algumas mudanças foram feitas na estrutura para poder facilitar a implementação. Infelizmente, quando eu resolvi fazer este post mostrando as modificações nas estruturas, a biblioteca já estava uns 3 posts adiantada. Caso tenha alguma dúvida, recomendo explorar o repositório nesta versão: Jitex 1.3.0. Tenha em mente que nesta versão já está implementada a Parte 3 - Injetando variáveis locais e a Parte 4 - Resolvendo Tokens.
CompileContext
Para facilitar o controle dos métodos que serão modificados, a classe ReplaceInfo se tornou CompileContext. Ela será responsável por todo contexto de injeção/modificação dos métodos.
Resolvers
Foi implementada uma nova forma expor os métodos que estão sendo compilados.
Antes:
foreach (Module module in assemblyFound.Modules) {
try {
MethodBase methodFound = module.ResolveMethod (methodToken);
replaceInfo = OnPreCompile (methodFound);
} catch {
// ignored
}
}
Depois:
//ManagedJit.cs
private ResolveCompileHandle _resolversCompile;
public delegate void ResolveCompileHandle (CompileContext context);
public void AddCompileResolver (ResolveCompileHandle compileResolver) {
_resolversCompile += compileResolver;
}
public void RemoveCompileResolver (ResolveCompileHandle compileResolver) {
_resolversCompile -= compileResolver;
}
private CorJitResult CompileMethod (IntPtr thisPtr, IntPtr comp, ref CORINFO_METHOD_INFO info, uint flags, out IntPtr nativeEntry, out int nativeSizeOfCode) {
//...
foreach (ResolveCompileHandle resolver in _resolversCompile.GetInvocationList ()) {
resolver (compileContext);
if (compileContext.IsResolved)
break;
}
//...
}
Nesta nova implementação, damos suporte a adição de mais de um resolver. Agora para fazer o hook do compile, basta apenas informamos o método através do AddCompileResolver:
ManagedJit jit = ManagedJit.GetInstance();
jit.AddCompileResolver(CompileResolver);
CEEInfo
O que antes era um método dentro do ManagedJit (ExecuteCEEInfo
Antes:
private static TOut ExecuteCEEInfo<TDelegate, TOut, TValue> (TValue value, int offset) {
IntPtr delegatePtr = Marshal.ReadIntPtr (_corJitInfoPtr, IntPtr.Size * offset);
Delegate delegateMethod = Marshal.GetDelegateForFunctionPointer (delegatePtr, typeof (TDelegate));
return (TOut) delegateMethod.DynamicInvoke (_corJitInfoPtr, value);
}
Depois:
internal class CEEInfo {
private readonly IntPtr _corJitInfo;
private readonly GetMethodDefFromMethodDelegate _getMethodDefFromMethod;
[UnmanagedFunctionPointer (default)]
public delegate uint GetMethodDefFromMethodDelegate (IntPtr thisHandle, IntPtr hMethod);
public CEEInfo (IntPtr corJitInfo) {
_corJitInfo = corJitInfo;
Version version = new Version (3, 1, 1);
IntPtr getMethodDefFromMethodIndex = IntPtr.Zero;
if (Environment.Version >= version) {
getMethodDefFromMethodIndex = _corJitInfo + IntPtr.Size * 0x74;
}
IntPtr getMethodDefFromMethodPtr = Marshal.ReadIntPtr (getMethodDefFromMethodIndex);
_getMethodDefFromMethod = Marshal.GetDelegateForFunctionPointer<GetMethodDefFromMethodDelegate> (getMethodDefFromMethodPtr);
}
public uint GetMethodDefFromMethod (IntPtr hMethod) {
return _getMethodDefFromMethod (_corJitInfo, hMethod);
}
}
Com isso, a classe VTable foi excluída.
Isso nos permite isolar e ter um maior controle com as chamadas diretas pro JIT.
Agora o lock contém somente a instância do CEEInfo:
//ManagedJit.cs
//CompileMethod()
if (compileEntry.EnterCount == 1 && _resolversCompile != null) {
lock (JitLock) {
if (_corJitInfoPtr == IntPtr.Zero) {
_corJitInfoPtr = Marshal.ReadIntPtr (comp);
_ceeInfo = new CEEInfo (_corJitInfoPtr);
}
}
//...
}
InjectHook
Esta classe foi criada para facilitar o controle de hooks. Ela somente faz o gerenciamento de adicionar e remover um hook em determinado endereço:
internal sealed class HookManager {
private readonly IList<VTableHook> _hooks = new List<VTableHook> ();
public void InjectHook (IntPtr pointerAddress, Delegate delToInject) {
IntPtr originalAddress = Marshal.ReadIntPtr (pointerAddress);
IntPtr hookAddress = Marshal.GetFunctionPointerForDelegate (delToInject);
VTableHook hook = new VTableHook (delToInject, originalAddress, pointerAddress);
WritePointer (pointerAddress, hookAddress);
_hooks.Add (hook);
}
public bool RemoveHook (Delegate del) {
VTableHook hookFound = _hooks.FirstOrDefault (h => h.Delegate.Method.Equals (del.Method));
if (hookFound == null)
return false;
return RemoveHook (hookFound);
}
private bool RemoveHook (VTableHook hook) {
WritePointer (hook.Address, hook.OriginalAddress);
_hooks.Remove (hook);
return true;
}
private void WritePointer (IntPtr address, IntPtr pointer) {
VirtualProtect (address, new IntPtr (IntPtr.Size), MemoryProtection.ReadWrite, out MemoryProtection oldFlags);
Marshal.WriteIntPtr (address, pointer);
VirtualProtect (address, new IntPtr (IntPtr.Size), oldFlags, out _);
}
}
AppModule
Classe criada para gerenciar os módulos em execução da aplicação:
internal static class AppModules {
private static readonly IDictionary<IntPtr, Module> MapScopeToHandle = new Dictionary<IntPtr, Module> (IntPtrEqualityComparer.Instance);
private static readonly FieldInfo m_pData;
static AppModules () {
m_pData = Type.GetType ("System.Reflection.RuntimeModule").GetField ("m_pData", BindingFlags.NonPublic | BindingFlags.Instance);
foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies ()) {
AddAssembly (assembly);
}
AppDomain.CurrentDomain.AssemblyLoad += CurrentDomainOnAssemblyLoad;
}
private static void AddAssembly (Assembly assembly) {
Module module = assembly.Modules.First ();
IntPtr scope = GetPointerFromModule (module);
MapScopeToHandle.Add (scope, module);
}
private static void CurrentDomainOnAssemblyLoad (object? sender, AssemblyLoadEventArgs args) {
AddAssembly (args.LoadedAssembly);
}
public static Module GetModuleByPointer (IntPtr scope) {
return MapScopeToHandle.TryGetValue (scope, out Module module) ? module : null;
}
public static IntPtr GetPointerFromModule (Module module) {
return (IntPtr) m_pData.GetValue (module);
}
}
Durante o desenvolvimento, usaremos bastante ponteiro para módulo e vice-versa. Esta classe facilitará nosso trabalho.
Melhoria de como obter o método
Na primeira implementação do compileMethod, para identificar o método que estava sendo compilado, fazíamos os seguintes passos:
- Obter o handle do assembly através do info.scope
- Obter o ponteiro do nome do assembly através do handle
- Fazíamos a leitura do nome do assembly
- Em todos os assemblies em execução, procurávamos o primeiro que correspondia ao nome do passo 3.
- Obter o token do método
Este processo além de ser lento, era extremamente suscetível a falhas. Agora, uma forma melhorada de identificar o método foi implementada.
O info.scope é o handle do módulo que o método pertence. Junto com a classe AppModuleS, é possível obter o módulo somente com o ponteiro:
Module module = AppModules.GetModuleByPointer(info.scope);
Desta forma é possível obter o método com maior precisão e menos esforço.
Antes:
IntPtr assemblyHandle = ExecuteCEEInfo<GetModuleAssemblyDelegate, IntPtr, IntPtr> (info.scope, VTable.GetModuleAssembly);
if (!MapHandleToAssembly.TryGetValue (assemblyHandle, out Assembly assemblyFound)) {
IntPtr assemblyNamePtr = ExecuteCEEInfo<GetAssemblyName, IntPtr, IntPtr> (assemblyHandle, VTable.GetAssemblyName);
string assemblyName = Marshal.PtrToStringAnsi (assemblyNamePtr);
assemblyFound = AppDomain.CurrentDomain.GetAssemblies ().FirstOrDefault (assembly => assembly.GetName ().Name == assemblyName);
MapHandleToAssembly.Add (assemblyHandle, assemblyFound);
}
if (assemblyFound != null) {
int methodToken = ExecuteCEEInfo<GetMethodDefFromMethodDelegate, int, IntPtr> (info.ftn, VTable.GetMethodDefFromMethod);
foreach (Module module in assemblyFound.Modules) {
try {
MethodBase methodFound = module.ResolveMethod (methodToken);
//...
} catch {
// ignored
}
}
}
Depois:
Module module = AppModules.GetModuleByPointer (info.scope);
if (module != null) {
uint methodToken = _ceeInfo.GetMethodDefFromMethod (info.ftn);
MethodBase methodFound = module.ResolveMethod ((int) methodToken);
//...
}
Com isso, o único P/Invoke que continuou implementado é o GetMethodDefFromMethod. Os outros foram removidos por não serem mais utilizados.