Jitex - Parte 2 - Refatoração...

Simplificando e resolvendo alguns erros


Autor: Flavio

Publicado em: 27 de Junho de 2020

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), se tornou uma classe própria para poder facilitar o controle de chamadas envolvendo P/Invoke.

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:

  1. Obter o handle do assembly através do info.scope
  2. Obter o ponteiro do nome do assembly através do handle
  3. Fazíamos a leitura do nome do assembly
  4. Em todos os assemblies em execução, procurávamos o primeiro que correspondia ao nome do passo 3.
  5. 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.


Imagem do post