Fix agent server remote vm resource not delegating listener calls

Col-E · Aug 12, 2023 · da896c6 · da896c6
1 parent c8a81cb
commit da896c6
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 23 deletions.
diff --git a/.../main/java/software/coley/recaf/workspace/model/resource/AgentServerRemoteVmResource.java b/.../main/java/software/coley/recaf/workspace/model/resource/AgentServerRemoteVmResource.java
@@ -25,10 +25,7 @@
 import software.coley.recaf.workspace.model.bundle.JvmClassBundle;
 
 import java.io.IOException;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.stream.Stream;
 
@@ -41,6 +38,7 @@ public class AgentServerRemoteVmResource extends BasicWorkspaceResource implemen
 	private static final DebuggingLogger logger = Logging.get(AgentServerRemoteVmResource.class);
 	private final Map<Integer, RemoteJvmClassBundle> remoteBundleMap = new HashMap<>();
 	private final Map<Integer, ClassLoaderInfo> remoteLoaders = new HashMap<>();
+	private final Map<Integer, Set<ClassData>> queuedClasses = new HashMap<>();
 	private final Set<String> queuedRedefines = new ConcurrentSkipListSet<>();
 	private final VirtualMachine virtualMachine;
 	private final Client client;
@@ -114,7 +112,17 @@ public void connect() throws IOException {
 					// New loader reported
 					BroadcastClassloaderMessage loaderMessage = (BroadcastClassloaderMessage) message;
 					ClassLoaderInfo loaderInfo = loaderMessage.getClassLoader();
-					remoteLoaders.put(loaderInfo.getId(), loaderInfo);
+					int loaderId = loaderInfo.getId();
+					remoteLoaders.put(loaderId, loaderInfo);
+
+					// If the loader was one that we saw classes for earlier, now we can
+					// link those to this loader.
+					Set<ClassData> pendingClasses = queuedClasses.remove(loaderId);
+					if (pendingClasses != null) {
+						for (ClassData pendingClass : pendingClasses) {
+							handleReceiveClassData(pendingClass, null);
+						}
+					}
 					break;
 				case MessageConstants.ID_BROADCAST_CLASS:
 					// New class, or update to existing class reported
@@ -147,7 +155,7 @@ public void connect() throws IOException {
 				// Get/create bundle for loader
 				ClassLoaderInfo loaderInfo = remoteLoaders.get(loaderId);
 				RemoteJvmClassBundle bundle = remoteBundleMap
-						.computeIfAbsent(loaderId, id -> new RemoteJvmClassBundle(loaderInfo));
+						.computeIfAbsent(loaderId, id -> createRemoteBundle(loaderInfo));
 
 				// Request all classes from classloader
 				logger.info("Sending initial request for class names in classloader {}...", loader.getName());
@@ -158,7 +166,7 @@ public void connect() throws IOException {
 					for (String className : classes) {
 						// If class does not exist in bundle, then request it from remote server
 						if (bundle.get(className) == null) {
-							client.sendBlocking(new RequestClassMessage(loaderId, className), reply -> {
+							client.sendAsync(new RequestClassMessage(loaderId, className), reply -> {
 								if (reply.hasData()) {
 									ClassData data = reply.getData();
 									handleReceiveClassData(data, bundle);
@@ -190,23 +198,45 @@ private void handleReceiveClassData(@Nonnull ClassData data, @Nullable RemoteJvm
 			// Get the bundle for the remote classloader if not specified by parameter
 			if (bundle == null) {
 				ClassLoaderInfo loaderInfo = remoteLoaders.get(loaderId);
-				bundle = remoteBundleMap
-						.computeIfAbsent(loaderId, id -> new RemoteJvmClassBundle(loaderInfo));
+				if (loaderInfo == null) {
+					queuedClasses.computeIfAbsent(loaderId, i -> new HashSet<>()).add(data);
+					return;
+				}
+				bundle = remoteBundleMap.computeIfAbsent(loaderId, id -> createRemoteBundle(loaderInfo));
 			}
 
 			// Add the class
 			JvmClassInfo classInfo = new JvmClassInfoBuilder(new ClassReader(data.getCode())).build();
 			RemoteClassloaderProperty.set(classInfo, loaderId);
-			bundle.initialPut(classInfo);
+			bundle.put(classInfo);
 		}
 	}
 
+	/**
+	 * Creates a new {@link RemoteJvmClassBundle} and sets up listener delegation.
+	 *
+	 * @param loaderInfo
+	 * 		Loader to associate with some classes in a bundle.
+	 *
+	 * @return New bundle for classes within that loader.
+	 */
+	@Nonnull
+	private RemoteJvmClassBundle createRemoteBundle(@Nonnull ClassLoaderInfo loaderInfo) {
+		RemoteJvmClassBundle bundle = new RemoteJvmClassBundle(loaderInfo);
+		delegateJvmClassBundle(this, bundle);
+		return bundle;
+	}
+
 	/**
 	 * JVM bundle extension adding a listener to handle syncing local changes with the remote server.
 	 */
-	private class RemoteJvmClassBundle extends BasicJvmClassBundle {
-		private RemoteJvmClassBundle(ClassLoaderInfo loaderInfo) {
-			this.addBundleListener(new BundleListener<>() {
+	public class RemoteJvmClassBundle extends BasicJvmClassBundle {
+		private final ClassLoaderInfo loaderInfo;
+
+		private RemoteJvmClassBundle(@Nonnull ClassLoaderInfo loaderInfo) {
+			this.loaderInfo = loaderInfo;
+
+			addBundleListener(new BundleListener<>() {
 				@Override
 				public void onNewItem(String key, JvmClassInfo value) {
 					// Should occur when we get data from the client.
@@ -239,5 +269,13 @@ public void onRemoveItem(String key, JvmClassInfo value) {
 				}
 			});
 		}
+
+		/**
+		 * @return Loader information for this bundle.
+		 */
+		@Nonnull
+		public ClassLoaderInfo getLoaderInfo() {
+			return loaderInfo;
+		}
 	}
 }
diff --git a/...i/src/main/java/software/coley/recaf/workspace/model/resource/BasicWorkspaceResource.java b/...i/src/main/java/software/coley/recaf/workspace/model/resource/BasicWorkspaceResource.java
@@ -89,7 +89,21 @@ protected void setup() {
 	 */
 	private void setupListenerDelegation() {
 		WorkspaceResource resource = this;
-		jvmClassBundleStream().forEach(bundle -> bundle.addBundleListener(new BundleListener<>() {
+		jvmClassBundleStream().forEach(bundle -> delegateJvmClassBundle(resource, bundle));
+		androidClassBundleStream().forEach(bundle -> delegateAndroidClassBundle(resource, bundle));
+		fileBundleStream().forEach(bundle -> delegateFileBundle(resource, bundle));
+	}
+
+	/**
+	 * Adds listeners to the given bundle to forward to the appropriate resource-level listener types.
+	 *
+	 * @param resource
+	 * 		Resource to delegate to.
+	 * @param bundle
+	 * 		Bundle to delegate from.
+	 */
+	protected void delegateJvmClassBundle(@Nonnull WorkspaceResource resource, @Nonnull JvmClassBundle bundle) {
+		bundle.addBundleListener(new BundleListener<>() {
 			@Override
 			public void onNewItem(String key, JvmClassInfo cls) {
 				for (ResourceJvmClassListener listener : jvmClassListeners) {
@@ -122,8 +136,19 @@ public void onRemoveItem(String key, JvmClassInfo cls) {
 					}
 				}
 			}
-		}));
-		androidClassBundleStream().forEach(bundle -> bundle.addBundleListener(new BundleListener<>() {
+		});
+	}
+
+	/**
+	 * Adds listeners to the given bundle to forward to the appropriate resource-level listener types.
+	 *
+	 * @param resource
+	 * 		Resource to delegate to.
+	 * @param bundle
+	 * 		Bundle to delegate from.
+	 */
+	protected void delegateAndroidClassBundle(@Nonnull WorkspaceResource resource, @Nonnull AndroidClassBundle bundle) {
+		bundle.addBundleListener(new BundleListener<>() {
 			@Override
 			public void onNewItem(String key, AndroidClassInfo cls) {
 				for (ResourceAndroidClassListener listener : androidClassListeners) {
@@ -156,8 +181,19 @@ public void onRemoveItem(String key, AndroidClassInfo cls) {
 					}
 				}
 			}
-		}));
-		fileBundleStream().forEach(bundle -> bundle.addBundleListener(new BundleListener<>() {
+		});
+	}
+
+	/**
+	 * Adds listeners to the given bundle to forward to the appropriate resource-level listener types.
+	 *
+	 * @param resource
+	 * 		Resource to delegate to.
+	 * @param bundle
+	 * 		Bundle to delegate from.
+	 */
+	protected void delegateFileBundle(@Nonnull WorkspaceResource resource, @Nonnull FileBundle bundle) {
+		bundle.addBundleListener(new BundleListener<>() {
 			@Override
 			public void onNewItem(String key, FileInfo file) {
 				for (ResourceFileListener listener : fileListeners) {
@@ -190,7 +226,7 @@ public void onRemoveItem(String key, FileInfo file) {
 					}
 				}
 			}
-		}));
+		});
 	}
 
 	/**

diff --git a/recaf-ui/src/main/java/software/coley/recaf/services/cell/TextProviderService.java b/recaf-ui/src/main/java/software/coley/recaf/services/cell/TextProviderService.java
@@ -15,10 +15,7 @@
 import software.coley.recaf.util.StringUtil;
 import software.coley.recaf.workspace.model.Workspace;
 import software.coley.recaf.workspace.model.bundle.*;
-import software.coley.recaf.workspace.model.resource.WorkspaceDirectoryResource;
-import software.coley.recaf.workspace.model.resource.WorkspaceFileResource;
-import software.coley.recaf.workspace.model.resource.WorkspaceRemoteVmResource;
-import software.coley.recaf.workspace.model.resource.WorkspaceResource;
+import software.coley.recaf.workspace.model.resource.*;
 
 import java.util.Map;
 
@@ -279,6 +276,8 @@ public TextProvider getBundleTextProvider(@Nonnull Workspace workspace,
 						.orElse(null);
 				if (dexName != null)
 					return dexName;
+			} else if (bundle instanceof AgentServerRemoteVmResource.RemoteJvmClassBundle remoteBundle) {
+				return formatConfig.filter(remoteBundle.getLoaderInfo().getName());
 			}
 
 			if (bundle instanceof ClassBundle)