From b6a4a2b7d73ff9a76f18eae1d46b2b233f9515d1 Mon Sep 17 00:00:00 2001
From: Dan <46821332+nsadeveloper789@users.noreply.github.com>
Date: Mon, 9 Sep 2024 14:48:20 -0400
Subject: [PATCH] GP-4905: Fix target withdrawal on connection close.

---
 .../service/tracermi/TraceRmiHandler.java     |  3 +++
 ...TraceRmiConnectionManagerProviderTest.java | 25 +++++++++++++------
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/Ghidra/Debug/Debugger-rmi-trace/src/main/java/ghidra/app/plugin/core/debug/service/tracermi/TraceRmiHandler.java b/Ghidra/Debug/Debugger-rmi-trace/src/main/java/ghidra/app/plugin/core/debug/service/tracermi/TraceRmiHandler.java
index d05eab6408..5bb0ab1133 100644
--- a/Ghidra/Debug/Debugger-rmi-trace/src/main/java/ghidra/app/plugin/core/debug/service/tracermi/TraceRmiHandler.java
+++ b/Ghidra/Debug/Debugger-rmi-trace/src/main/java/ghidra/app/plugin/core/debug/service/tracermi/TraceRmiHandler.java
@@ -197,6 +197,9 @@ public class TraceRmiHandler implements TraceRmiConnection {
 			List<OpenTrace> all = List.copyOf(byId.values());
 			byId.clear();
 			byTrace.clear();
+			for (OpenTrace open : all) {
+				plugin.withdrawTarget(open.target);
+			}
 			return all;
 		}
 
diff --git a/Ghidra/Debug/Debugger-rmi-trace/src/test/java/ghidra/app/plugin/core/debug/gui/tracermi/connection/TraceRmiConnectionManagerProviderTest.java b/Ghidra/Debug/Debugger-rmi-trace/src/test/java/ghidra/app/plugin/core/debug/gui/tracermi/connection/TraceRmiConnectionManagerProviderTest.java
index 529e0c8e97..145fde579c 100644
--- a/Ghidra/Debug/Debugger-rmi-trace/src/test/java/ghidra/app/plugin/core/debug/gui/tracermi/connection/TraceRmiConnectionManagerProviderTest.java
+++ b/Ghidra/Debug/Debugger-rmi-trace/src/test/java/ghidra/app/plugin/core/debug/gui/tracermi/connection/TraceRmiConnectionManagerProviderTest.java
@@ -48,6 +48,11 @@ import ghidra.debug.api.tracermi.TraceRmiConnection;
 import ghidra.util.exception.CancelledException;
 
 public class TraceRmiConnectionManagerProviderTest extends AbstractGhidraHeadedDebuggerTest {
+	private static final String SCHEMA_XML = """
+			<context>
+			  <schema name="Root" elementResync="NEVER" attributeResync="NEVER" />
+			</context>
+			""";
 	TraceRmiConnectionManagerProvider provider;
 	TraceRmiService traceRmiService;
 	DebuggerControlService controlService;
@@ -178,6 +183,17 @@ public class TraceRmiConnectionManagerProviderTest extends AbstractGhidraHeadedD
 		}
 	}
 
+	@Test
+	public void testCloseWithdrawsTargets() throws Exception {
+		Target target;
+		try (Cx cx = Cx.connect(traceRmiService, "Test client")) {
+			cx.client.createTrace(0, "bash");
+			target = waitForPass(() -> Unique.assertOne(targetService.getPublishedTargets()));
+		}
+		// Outside the try, the connection has been closed
+		waitForPass(() -> assertFalse(targetService.getPublishedTargets().contains(target)));
+	}
+
 	@Test
 	public void testServerNode() throws Exception {
 		TraceRmiServerNode node = TraceRmiConnectionTreeHelper.getServerNode(provider.rootNode);
@@ -213,8 +229,7 @@ public class TraceRmiConnectionManagerProviderTest extends AbstractGhidraHeadedD
 	}
 
 	public record Cx(SocketChannel channel, TestTraceRmiClient client,
-			TraceRmiConnection connection)
-			implements AutoCloseable {
+			TraceRmiConnection connection) implements AutoCloseable {
 		public static Cx complete(TraceRmiAcceptor acceptor, String description)
 				throws IOException, CancelledException {
 			SocketChannel channel = null;
@@ -378,11 +393,7 @@ public class TraceRmiConnectionManagerProviderTest extends AbstractGhidraHeadedD
 
 	@Test
 	public void testActivateTargetNode() throws Exception {
-		SchemaContext ctx = XmlSchemaContext.deserialize("""
-				<context>
-				  <schema name="Root" elementResync="NEVER" attributeResync="NEVER" />
-				</context>
-				""");
+		SchemaContext ctx = XmlSchemaContext.deserialize(SCHEMA_XML);
 		try (Cx cx = Cx.connect(traceRmiService, "Test client")) {
 			cx.client.createTrace(1, "bash");
 			try (Tx tx = cx.client.new Tx(1, 1, "Create snapshots")) {