Update stack emulation

- Allow System.Object string decrypter return type - Fix for StackFrame.GetMethod() - Caller method is updated in stack frame
shine429 · Jan 5, 2012 · e1a81d0 · e1a81d0
1 parent 93d8019
commit e1a81d0
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 26 deletions.
diff --git a/AssemblyData/AssemblyService.cs b/AssemblyData/AssemblyService.cs
@@ -76,16 +76,17 @@ public int defineStringDecrypter(int methodToken) {
 			var methodInfo = findMethod(methodToken);
 			if (methodInfo == null)
 				throw new ApplicationException(string.Format("Could not find method {0:X8}", methodToken));
-			if (methodInfo.ReturnType != typeof(string))
-				throw new ApplicationException(string.Format("Method return type must be string: {0}", methodInfo));
+			if (methodInfo.ReturnType != typeof(string) && methodInfo.ReturnType != typeof(object))
+				throw new ApplicationException(string.Format("Method return type must be string or object: {0}", methodInfo));
 			return stringDecrypter.defineStringDecrypter(methodInfo);
 		}
 
-		public object[] decryptStrings(int stringDecrypterMethod, object[] args) {
+		public object[] decryptStrings(int stringDecrypterMethod, object[] args, int callerToken) {
 			checkStringDecrypter();
+			var caller = assembly.GetModules()[0].ResolveMethod(callerToken);
 			foreach (var arg in args)
 				SimpleData.unpack((object[])arg);
-			return SimpleData.pack(stringDecrypter.decryptStrings(stringDecrypterMethod, args));
+			return SimpleData.pack(stringDecrypter.decryptStrings(stringDecrypterMethod, args, caller));
 		}
 
 		public void exit() {

diff --git a/AssemblyData/DelegateStringDecrypter.cs b/AssemblyData/DelegateStringDecrypter.cs
@@ -32,7 +32,7 @@ public int defineStringDecrypter(MethodInfo method) {
 			return stringDecryptMethods.Count - 1;
 		}
 
-		public object[] decryptStrings(int stringDecrypterMethod, object[] args) {
+		public object[] decryptStrings(int stringDecrypterMethod, object[] args, MethodBase caller) {
 			if (stringDecrypterMethod > stringDecryptMethods.Count)
 				throw new ApplicationException("Invalid string decrypter method");
 

diff --git a/AssemblyData/EmuStringDecrypter.cs b/AssemblyData/EmuStringDecrypter.cs
@@ -41,11 +41,12 @@ public int defineStringDecrypter(MethodInfo method) {
 			return decryptInfos.Count - 1;
 		}
 
-		public object[] decryptStrings(int stringDecrypterMethod, object[] args) {
+		public object[] decryptStrings(int stringDecrypterMethod, object[] args, MethodBase caller) {
 			var decryptInfo = decryptInfos[stringDecrypterMethod];
 			if (decryptInfo.decryptString == null)
 				decryptInfo.decryptString = createDecryptString(decryptInfo.method);
 
+			methodsRewriter.setCaller(decryptInfo.decryptString, caller);
 			var result = new object[args.Length];
 			for (int i = 0; i < args.Length; i++)
 				result[i] = decryptInfo.decryptString((object[])args[i]);

diff --git a/AssemblyData/IAssemblyService.cs b/AssemblyData/IAssemblyService.cs
@@ -28,7 +28,7 @@ public interface IAssemblyService {
 		void loadAssembly(string filename);
 		void setStringDecrypterType(StringDecrypterType type);
 		int defineStringDecrypter(int methodToken);
-		object[] decryptStrings(int stringDecrypterMethod, object[] args);
+		object[] decryptStrings(int stringDecrypterMethod, object[] args, int callerToken);
 		void exit();
 	}
 }
diff --git a/AssemblyData/IStringDecrypter.cs b/AssemblyData/IStringDecrypter.cs
@@ -22,6 +22,6 @@ You should have received a copy of the GNU General Public License
 namespace AssemblyData {
 	interface IStringDecrypter {
 		int defineStringDecrypter(MethodInfo method);
-		object[] decryptStrings(int stringDecrypterMethod, object[] args);
+		object[] decryptStrings(int stringDecrypterMethod, object[] args, MethodBase caller);
 	}
 }
diff --git a/AssemblyData/methodsrewriter/MethodsRewriter.cs b/AssemblyData/methodsrewriter/MethodsRewriter.cs
@@ -78,6 +78,7 @@ public MethodBase getMethod(Module module) {
 	class MethodsRewriter : IMethodsRewriter {
 		MethodsFinder methodsFinder = new MethodsFinder();
 		Dictionary<MethodBase, NewMethodInfo> realMethodToNewMethod = new Dictionary<MethodBase, NewMethodInfo>();
+		Dictionary<NewMethodInfo, MethodBase> newStackMethodDict = new Dictionary<NewMethodInfo, MethodBase>();
 		List<NewMethodInfo> newMethodInfos = new List<NewMethodInfo>();
 
 		// There's no documented way to get a dynamic method's MethodInfo. If we name the
@@ -163,6 +164,11 @@ public RewrittenMethod createDelegate(MethodBase realMethod) {
 			return newMethodInfo.rewrittenMethod;
 		}
 
+		public void setCaller(RewrittenMethod rewrittenMethod, MethodBase caller) {
+			var newMethodInfo = getNewMethodInfo(rewrittenMethod.Method.Name);
+			newStackMethodDict[newMethodInfo] = caller;
+		}
+
 		string getDelegateMethodName(MethodBase method) {
 			string name = null;
 			do {
@@ -229,6 +235,12 @@ void update(Block block, NewMethodInfo currentMethodInfo) {
 						i += 2;
 						continue;
 					}
+					else if (MemberReferenceHelper.verifyType(ctor.DeclaringType, "mscorlib", "System.Diagnostics.StackFrame")) {
+						insertLoadThis(block, i + 1);
+						insertCallOurMethod(block, i + 2, "static_rtFixStackFrame");
+						i += 2;
+						continue;
+					}
 				}
 
 				if (instr.OpCode == OpCodes.Call || instr.OpCode == OpCodes.Callvirt) {
@@ -359,30 +371,45 @@ StackTrace rtFixStackTrace(StackTrace stackTrace) {
 
 			var newFrames = new List<StackFrame>(frames.Length);
 			foreach (var frame in frames) {
-				var method = frame.GetMethod();
-				var info = getNewMethodInfo(method.Name);
-				if (info == null) {
-					newFrames.Add(frame);
-				}
-				else if (info.isRewrittenMethod(method.Name)) {
-					// Write random method from the same module
-					writeMethodBase(frame, methodsFinder.getMethod(info.oldMethod.Module));
-					newFrames.Add(frame);
-				}
-				else if (info.isDelegateMethod(method.Name)) {
-					// Write original method
-					writeMethodBase(frame, info.oldMethod);
-					newFrames.Add(frame);
-				}
-				else {
-					throw new ApplicationException("BUG: Shouldn't be here");
-				}
+				fixStackFrame(frame);
+				newFrames.Add(frame);
 			}
 
 			framesField.SetValue(stackTrace, newFrames.ToArray());
 			return stackTrace;
 		}
 
+		static StackFrame static_rtFixStackFrame(StackFrame stackFrame, MethodsRewriter self) {
+			return self.rtFixStackFrame(stackFrame);
+		}
+
+		StackFrame rtFixStackFrame(StackFrame frame) {
+			fixStackFrame(frame);
+			return frame;
+		}
+
+		void fixStackFrame(StackFrame frame) {
+			var method = frame.GetMethod();
+			var info = getNewMethodInfo(method.Name);
+			if (info == null)
+				return;
+
+			MethodBase stackMethod;
+			if (newStackMethodDict.TryGetValue(info, out stackMethod)) {
+				writeMethodBase(frame, stackMethod);
+			}
+			else if (info.isRewrittenMethod(method.Name)) {
+				// Write random method from the same module
+				writeMethodBase(frame, methodsFinder.getMethod(info.oldMethod.Module));
+			}
+			else if (info.isDelegateMethod(method.Name)) {
+				// Write original method
+				writeMethodBase(frame, info.oldMethod);
+			}
+			else
+				throw new ApplicationException("BUG: Shouldn't be here");
+		}
+
 		// Called when the code calls GetCallingAssembly(), GetEntryAssembly(), or GetExecutingAssembly()
 		Assembly rtGetAssembly(int delegateIndex) {
 			return newMethodInfos[delegateIndex].oldMethod.Module.Assembly;