adding injection order checks to EventInjectionTransformer to prevent listeners …

…being appended after event injection

adding injection order checks to EventInjectionTransformer to prevent listeners …
…being appended after event injection
Mumfrey
1 parent bbee82bb
Showing 4 changed files with 155 additions and 51 deletions
java/common/com/mumfrey/liteloader/transformers/event/Event.java
java/common/com/mumfrey/liteloader/transformers/event/EventAlreadyInjectedException.java
java/common/com/mumfrey/liteloader/transformers/event/EventInjectionTransformer.java
java/common/com/mumfrey/liteloader/transformers/event/MethodInfo.java
@@ -93,6 +93,8 @@ public class Event implements Comparable&lt;Event&gt;
  
 	protected String eventInfoClass;
  
+	protected Set<MethodInfo> pendingInjections;
+	
 	Event(String name, boolean cancellable, int priority)
 	{
 		this.name = name.toLowerCase();
@@ -226,6 +228,38 @@ public class Event implements Comparable&lt;Event&gt;
 		this.method = null;
 	}
  
+	void addPendingInjection(MethodInfo targetMethod)
+	{
+		if (this.pendingInjections == null)
+		{
+			this.pendingInjections = new HashSet<MethodInfo>();
+		}
+		
+		this.pendingInjections.add(targetMethod);
+	}
+	
+	void notifyInjected(String method, String desc, String className)
+	{
+		MethodInfo thisInjection = null;
+		
+		if (this.pendingInjections != null)
+		{
+			for (MethodInfo pendingInjection : this.pendingInjections)
+			{
+				if (pendingInjection.matches(method, desc, className))
+				{
+					thisInjection = pendingInjection;
+					break;
+				}
+			}
+		}
+		
+		if (thisInjection != null)
+		{
+			this.pendingInjections.remove(thisInjection);
+		}
+	}
+	
 	/**
 	 * Pre-flight check
 	 * 
@@ -371,6 +405,11 @@ public class Event implements Comparable&lt;Event&gt;
 			throw new IllegalArgumentException("Descriptor is not allowed for listener methods");
 		}
  
+		if (this.pendingInjections != null && this.pendingInjections.size() == 0)
+		{
+			throw new EventAlreadyInjectedException("The event " + this.name + " was already injected and has 0 pending injections, addListener() is not allowed at this point");
+		}
+		
 		this.listeners.add(listener);
  
 		return this;
+package com.mumfrey.liteloader.transformers.event;
+
+public class EventAlreadyInjectedException extends RuntimeException
+{
+	private static final long serialVersionUID = 1L;
+
+	public EventAlreadyInjectedException(String message)
+	{
+		super(message);
+	}
+	
+	public EventAlreadyInjectedException(Throwable cause)
+	{
+		super(cause);
+	}
+	
+	public EventAlreadyInjectedException(String message, Throwable cause)
+	{
+		super(message, cause);
+	}
+}
@@ -77,7 +77,14 @@ public abstract class EventInjectionTransformer extends ClassTransformer
 			EventInjectionTransformer.master = this;
 		}
  
-		this.addEvents();
+		try
+		{
+			this.addEvents();
+		}
+		catch (Exception ex)
+		{
+			ex.printStackTrace();
+		}
 	}
  
 	/**
@@ -120,6 +127,8 @@ public abstract class EventInjectionTransformer extends ClassTransformer
 		this.addEvent(event, targetMethod.owner, targetMethod.sigSrg, injectionPoint);
 		this.addEvent(event, targetMethod.ownerObf, targetMethod.sigObf, injectionPoint);
  
+		event.addPendingInjection(targetMethod);
+		
 		return event;
 	}
  
@@ -215,65 +224,100 @@ public abstract class EventInjectionTransformer extends ClassTransformer
 			Map<Event, InjectionPoint> methodInjections = mappings.get(signature);
 			if (methodInjections != null)
 			{
-				ReadOnlyInsnList insns = new ReadOnlyInsnList(method.instructions);
-				Map<AbstractInsnNode, Set<Event>> injectionPoints = new LinkedHashMap<AbstractInsnNode, Set<Event>>();
-				Collection<AbstractInsnNode> nodes = new ArrayList<AbstractInsnNode>(32);
-				for (Entry<Event, InjectionPoint> eventEntry : methodInjections.entrySet())
-				{
-					Event event = eventEntry.getKey();
-					event.attach(method);
-					InjectionPoint injectionPoint = eventEntry.getValue();
-					nodes.clear();
-					if (injectionPoint.find(method.desc, insns, nodes, event))
-					{
-						for (AbstractInsnNode node : nodes)
-						{
-							Set<Event> nodeEvents = injectionPoints.get(node);
-							if (nodeEvents == null)
-							{
-								nodeEvents = new TreeSet<Event>();
-								injectionPoints.put(node, nodeEvents);
-							}
-							
-							nodeEvents.add(event);
-						}
-					}
-				}
-				
-				for (Entry<AbstractInsnNode, Set<Event>> injectionPoint : injectionPoints.entrySet())
-				{
-					AbstractInsnNode insn = injectionPoint.getKey();
-					Set<Event> events = injectionPoint.getValue();
+				this.injectIntoMethod(classNode, signature, method, methodInjections);
+			}
+		}
+		
+		if (true || this.runValidator)
+		{
+			ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
+			classNode.accept(new CheckClassAdapter(writer));
+		}
+		
+		return this.writeClass(classNode);
+	}
  
-					// Injection is cancellable if ANY of the events on this insn are cancellable
-					boolean cancellable = false;
-					for (Event event : events)
-						cancellable |= event.isCancellable();
-					
-					Event head = events.iterator().next();
-					MethodNode handler = head.inject(insn, cancellable, this.globalEventID);
+	/**
+	 * @param classNode
+	 * @param signature
+	 * @param method
+	 * @param methodInjections
+	 */
+	void injectIntoMethod(ClassNode classNode, String signature, MethodNode method, Map<Event, InjectionPoint> methodInjections)
+	{
+		Map<AbstractInsnNode, Set<Event>> injectionPoints = this.findInjectionPoints(classNode, method, methodInjections);
+		
+		for (Entry<AbstractInsnNode, Set<Event>> injectionPoint : injectionPoints.entrySet())
+		{
+			this.injectEventsAt(classNode, method, injectionPoint.getKey(), injectionPoint.getValue());
+		}
  
-					LiteLoaderLogger.info("Injecting event %s with %d handlers in method %s in class %s", head.getName(), events.size(), method.name, classNode.name.replace('/', '.'));
-					
-					for (Event event : events)
-						event.addToHandler(handler);
-					
-					this.globalEventID++;
-				}
+		for (Event event : methodInjections.keySet())
+		{
+			event.notifyInjected(method.name, method.desc, classNode.name);
+			event.detach();
+		}
+	}
  
-				for (Event event : methodInjections.keySet())
+	/**
+	 * @param classNode
+	 * @param method
+	 * @param methodInjections
+	 * @return
+	 */
+	private Map<AbstractInsnNode, Set<Event>> findInjectionPoints(ClassNode classNode, MethodNode method, Map<Event, InjectionPoint> methodInjections)
+	{
+		ReadOnlyInsnList insns = new ReadOnlyInsnList(method.instructions);
+		Collection<AbstractInsnNode> nodes = new ArrayList<AbstractInsnNode>(32);
+		Map<AbstractInsnNode, Set<Event>> injectionPoints = new LinkedHashMap<AbstractInsnNode, Set<Event>>();
+		for (Entry<Event, InjectionPoint> eventEntry : methodInjections.entrySet())
+		{
+			Event event = eventEntry.getKey();
+			event.attach(method);
+			InjectionPoint injectionPoint = eventEntry.getValue();
+			nodes.clear();
+			if (injectionPoint.find(method.desc, insns, nodes, event))
+			{
+				for (AbstractInsnNode node : nodes)
 				{
-					event.detach();
+					Set<Event> nodeEvents = injectionPoints.get(node);
+					if (nodeEvents == null)
+					{
+						nodeEvents = new TreeSet<Event>();
+						injectionPoints.put(node, nodeEvents);
+					}
+					
+					nodeEvents.add(event);
 				}
 			}
 		}
  
-		if (true || this.runValidator)
+		return injectionPoints;
+	}
+
+	/**
+	 * @param classNode
+	 * @param method
+	 * @param injectionPoint
+	 * @param events
+	 */
+	private void injectEventsAt(ClassNode classNode, MethodNode method, AbstractInsnNode injectionPoint, Set<Event> events)
+	{
+		// Injection is cancellable if ANY of the events on this insn are cancellable
+		boolean cancellable = false;
+		for (Event event : events)
+			cancellable |= event.isCancellable();
+		
+		Event head = events.iterator().next();
+		MethodNode handler = head.inject(injectionPoint, cancellable, this.globalEventID);
+
+		LiteLoaderLogger.info("Injecting event %s with %d handlers in method %s in class %s", head.getName(), events.size(), method.name, classNode.name.replace('/', '.'));
+		
+		for (Event event : events)
 		{
-			ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
-			classNode.accept(new CheckClassAdapter(writer));
+			event.addToHandler(handler);
 		}
  
-		return this.writeClass(classNode);
+		this.globalEventID++;
 	}
 }
 \ No newline at end of file
@@ -290,7 +290,7 @@ public class MethodInfo
  
 	public boolean matches(String method, String desc, String className)
 	{
-		if ((className == null || this.owner.equals(className)) && (this.name.equals(method) || this.nameSrg.equals(method)))
+		if ((className == null || this.ownerRef.equals(className)) && (this.name.equals(method) || this.nameSrg.equals(method)))
 		{
 			return this.desc == null || this.desc.equals(desc);
 		}