diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java index 16e88465..08a3c2b3 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java @@ -108,8 +108,12 @@ public boolean applyTo(Program obj) { public void applyWithTransaction() { var program = this.analyzedProgram.program(); var tID = program.startTransaction("RevEng.AI: Apply Match"); - var status = applyTo(program); - program.endTransaction(tID, status); + boolean status = false; + try { + status = applyTo(program); + } finally { + program.endTransaction(tID, status); + } } private Namespace getRevEngAINameSpace() { diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/aidecompiler/AIDecompilationdWindow.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/aidecompiler/AIDecompilationdWindow.java index 17020a60..75bc5042 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/aidecompiler/AIDecompilationdWindow.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/aidecompiler/AIDecompilationdWindow.java @@ -57,23 +57,7 @@ public AIDecompilationdWindow(PluginTool tool, String owner) { @Override public void actionPerformed(ActionContext context) { - if (function != null) { - var service = tool.getService(GhidraRevengService.class); - var analyzedProgram = service.getAnalysedProgram(function.getProgram()); - if (analyzedProgram.isEmpty()) { - Msg.error(this, "Failed to send positive feedback: Program is not known to RevEng.AI"); - return; - } - var fID = analyzedProgram.get().getIDForFunction(function); - fID.ifPresent(id -> { - try { - service.getApi().aiDecompRating(id.functionID(), "POSITIVE", ""); - } catch (ApiException e) { - // Fail silently because this is not a critical feature - Msg.error(this, "Failed to send positive feedback for function %s: %s".formatted(function.getName(), e.getMessage())); - } - }); - } + sendFeedbackInBackground(function, "POSITIVE", ""); } @Override @@ -93,28 +77,14 @@ public boolean isEnabled() { } @Override public void actionPerformed(ActionContext context) { - // Spawn textbox to enter reason for negative feedback - + final Function target = function; + if (target == null) { + return; + } var dialog = new InputDialog("Negative Feedback", "Please provide details about what was wrong with the decompilation:", ""); tool.showDialog(dialog); if (!dialog.isCanceled()) { - if (function != null) { - var service = tool.getService(GhidraRevengService.class); - var programWithID = service.getAnalysedProgram(function.getProgram()); - if (programWithID.isEmpty()) { - Msg.error(this, "Failed to send negative feedback: Program is not known to RevEng.AI"); - return; - } - var fID = programWithID.get().getIDForFunction(function); - fID.ifPresent(id -> { - try { - service.getApi().aiDecompRating(id.functionID(), "NEGATIVE", dialog.getValue()); - } catch (ApiException e) { - // Fail silently because this is not a critical feature - Msg.error(this, "Failed to send negative feedback for function %s: %s".formatted(function.getName(), e.getMessage())); - } - }); - } + sendFeedbackInBackground(target, "NEGATIVE", dialog.getValue()); } } @@ -126,6 +96,37 @@ public boolean isEnabled() { }); } + private void sendFeedbackInBackground(Function target, String rating, String reason) { + if (target == null) { + return; + } + final String detail = reason == null ? "" : reason; + tool.execute(new Task("Send AI Decompilation Feedback", false, false, false) { + @Override + public void run(TaskMonitor monitor) { + var service = tool.getService(GhidraRevengService.class); + var analyzedProgram = service.getAnalysedProgram(target.getProgram()); + if (analyzedProgram.isEmpty()) { + Msg.error(AIDecompilationdWindow.this, + "Failed to send %s feedback: Program is not known to RevEng.AI".formatted(rating)); + return; + } + var fID = analyzedProgram.get().getIDForFunction(target); + if (fID.isEmpty()) { + Msg.error(AIDecompilationdWindow.this, + "Failed to send %s feedback: function %s not known to RevEng.AI".formatted(rating, target.getName())); + return; + } + try { + service.getApi().aiDecompRating(fID.get().functionID(), rating, detail); + } catch (ApiException e) { + Msg.error(AIDecompilationdWindow.this, + "Failed to send %s feedback for function %s: %s".formatted(rating, target.getName(), e.getMessage())); + } + } + }, 0); + } + private JComponent buildComponent() { @@ -317,7 +318,7 @@ public void locationChanged(ProgramLocation loc) { var newFuncLocation = functionMgr.getFunctionContaining(loc.getAddress()); // If we changed to a different function, we want to clear the output of the old function - if (function != null && newFuncLocation != function) { + if (function != null && !isSameFunction(newFuncLocation, function)) { clear(); } @@ -329,7 +330,7 @@ public void locationChanged(ProgramLocation loc) { void newStatusForFunction(Function function, AIDecompilationStatus status) { cache.put(function, status); - if (function == this.function) { + if (isCurrentFunction(function)) { SwingUtilities.invokeLater(() -> setDisplayedValuesBasedOnStatus(function, status) ); @@ -352,6 +353,28 @@ private boolean hasPendingDecompilations() { s.status() == DecompilationData.StatusEnum.PENDING || s.status() == DecompilationData.StatusEnum.RUNNING); } + + private static boolean isSameFunction(Function a, Function b) { + return a != null && b != null && a.getEntryPoint().equals(b.getEntryPoint()); + } + + private boolean isCurrentFunction(Function candidate) { + return isSameFunction(candidate, this.function); + } + + private void reportDecompFailure(Function failedFunction, Exception e) { + var logger = tool.getService(ReaiLoggingService.class); + logger.error("AI Decompilation failed for function %s: %s".formatted(failedFunction.getName(), e.getMessage())); + Msg.error(this, "AI Decompilation failed for function %s".formatted(failedFunction.getName()), e); + SwingUtilities.invokeLater(() -> { + if (isCurrentFunction(failedFunction)) { + descriptionArea.setText("AI Decompilation failed: " + e.getMessage()); + if (!hasPendingDecompilations()) { + taskMonitorComponent.setVisible(false); + } + } + }); + } class AIDecompTask extends Task { private final GhidraRevengService service; @@ -365,14 +388,20 @@ public AIDecompTask(PluginTool tool, GhidraRevengService.FunctionWithID function @Override public void run(TaskMonitor monitor) throws CancelledException { - var fID = functionWithID.functionID(); - // Check if there is an existing process already, because the trigger API will fail with 400 if there is - if (service.getApi().pollAIDecompileStatus(fID).status() == DecompilationData.StatusEnum.UNINITIALISED) { - // Trigger the decompilation - service.getApi().triggerAIDecompilationForFunctionID(fID); + try { + var fID = functionWithID.functionID(); + // Check if there is an existing process already, because the trigger API will fail with 400 if there is + if (service.getApi().pollAIDecompileStatus(fID).status() == DecompilationData.StatusEnum.UNINITIALISED) { + // Trigger the decompilation + service.getApi().triggerAIDecompilationForFunctionID(fID); + } + waitForDecomp(fID, monitor); + // TODO: Inform the component that something is finished + } catch (CancelledException e) { + throw e; + } catch (Exception e) { + reportDecompFailure(functionWithID.function(), e); } - waitForDecomp(fID, monitor); - // TODO: Inform the component that something is finished } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AbstractFunctionMatchingDialog.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AbstractFunctionMatchingDialog.java index ccde8fd5..0e2d8d01 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AbstractFunctionMatchingDialog.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AbstractFunctionMatchingDialog.java @@ -13,6 +13,9 @@ import ai.reveng.toolkit.ghidra.core.services.api.types.GhidraFunctionMatchWithSignature; import com.google.common.collect.BiMap; import ghidra.program.model.listing.Function; +import ghidra.util.task.Task; +import ghidra.util.task.TaskBuilder; +import ghidra.util.task.TaskMonitor; import ghidra.util.task.TaskMonitorComponent; import ghidra.util.Msg; import resources.ResourceManager; @@ -893,8 +896,7 @@ protected void onMatchButtonClicked() { protected void onRenameAllButtonClicked() { - batchRenameFunctions(functionMatchResults); - importFunctionNames(functionMatchResults); + renameInBackground(functionMatchResults); } protected void onRenameSelectedButtonClicked() { @@ -916,8 +918,18 @@ protected void onRenameSelectedButtonClicked() { } } - batchRenameFunctions(selectedMatches); - importFunctionNames(selectedMatches); + renameInBackground(selectedMatches); + } + + private void renameInBackground(List matches) { + var task = new Task("Applying Renames", true, false, false) { + @Override + public void run(TaskMonitor monitor) { + batchRenameFunctions(matches); + importFunctionNames(matches); + } + }; + TaskBuilder.withTask(task).launchInBackground(taskMonitorComponent); } protected void batchRenameFunctions(List functionMatches) { @@ -925,7 +937,7 @@ protected void batchRenameFunctions(List funct revengService.batchRenamingGhidraMatchesWithSignatures(functionMatches); } catch (Exception e) { Msg.error(this, "Failed to rename functions: " + e.getMessage(), e); - showError("Failed to rename functions: " + e.getMessage()); + SwingUtilities.invokeLater(() -> showError("Failed to rename functions: " + e.getMessage())); } } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/SimilarFunctionsWindow.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/SimilarFunctionsWindow.java index e212d84b..ef27eb6d 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/SimilarFunctionsWindow.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/SimilarFunctionsWindow.java @@ -134,7 +134,17 @@ public void actionPerformed(ActionContext context) { var match = getSelectedMatch(); if (match != null) { var service = tool.getService(GhidraRevengService.class); - service.openFunctionInPortal(match.functionMatch().nearest_neighbor_id()); + var neighborId = match.functionMatch().nearest_neighbor_id(); + tool.execute(new Task("View Match in Portal", false, false, false) { + @Override + public void run(TaskMonitor monitor) { + try { + service.openFunctionInPortal(neighborId); + } catch (Exception e) { + Msg.error(SimilarFunctionsWindow.this, "Failed to open match in portal: " + e.getMessage(), e); + } + } + }, 0); } } @@ -158,8 +168,14 @@ public void actionPerformed(ActionContext context) { var match = getSelectedMatch(); if (match != null && analyzedProgram != null) { var service = tool.getService(GhidraRevengService.class); - var cmd = new ApplyMatchCmd(service, analyzedProgram, match, true); - cmd.applyWithTransaction(); + var program = analyzedProgram; + tool.execute(new Task("Apply Match", true, false, false) { + @Override + public void run(TaskMonitor monitor) { + var cmd = new ApplyMatchCmd(service, program, match, true); + cmd.applyWithTransaction(); + } + }, 0); } } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java index 51f1d487..21b65b46 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java @@ -8,7 +8,10 @@ import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage; import ghidra.framework.plugintool.PluginTool; import ghidra.program.model.listing.Program; +import ghidra.util.Msg; import ghidra.util.table.GhidraFilterTable; +import ghidra.util.task.Task; +import ghidra.util.task.TaskMonitor; import javax.swing.*; import java.awt.*; @@ -64,8 +67,18 @@ public void mouseClicked(MouseEvent e) { if ("Analysis ID".equals(columnName)) { LegacyAnalysisResult result = recentAnalysesTable.getModel().getRowObject(row); if (result != null) { - var analysisID = ghidraRevengService.getApi().getAnalysisIDfromBinaryID(result.binary_id()); - ghidraRevengService.openPortalFor(analysisID); + var binaryID = result.binary_id(); + tool.execute(new Task("Open analysis in portal", false, false, false) { + @Override + public void run(TaskMonitor monitor) { + try { + var analysisID = ghidraRevengService.getApi().getAnalysisIDfromBinaryID(binaryID); + ghidraRevengService.openPortalFor(analysisID); + } catch (Exception ex) { + Msg.error(RecentAnalysisDialog.this, "Failed to open analysis in portal: " + ex.getMessage(), ex); + } + } + }, 0); } } } @@ -97,15 +110,26 @@ public void mouseClicked(MouseEvent e) { private void pickAnalysis(LegacyAnalysisResult result) { var service = tool.getService(GhidraRevengService.class); - var analysisID = service.getApi().getAnalysisIDfromBinaryID(result.binary_id()); - var programWithID = service.registerAnalysisForProgram(program, analysisID); - tool.firePluginEvent( - new RevEngAIAnalysisStatusChangedEvent( - "Recent Analysis Dialog", - programWithID, - result.status() - ) - ); - close(); + tool.execute(new Task("Attach to analysis", true, false, false) { + @Override + public void run(TaskMonitor monitor) { + try { + var analysisID = service.getApi().getAnalysisIDfromBinaryID(result.binary_id()); + var programWithID = service.registerAnalysisForProgram(program, analysisID); + SwingUtilities.invokeLater(() -> { + tool.firePluginEvent( + new RevEngAIAnalysisStatusChangedEvent( + "Recent Analysis Dialog", + programWithID, + result.status() + ) + ); + close(); + }); + } catch (Exception ex) { + Msg.error(RecentAnalysisDialog.this, "Failed to attach to analysis: " + ex.getMessage(), ex); + } + } + }, 0); } } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java index f8e42150..3cc1212a 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java @@ -206,23 +206,7 @@ private Optional getAnalysisIDFromOptions( if (bid == ReaiPluginPackage.INVALID_ANALYSIS_ID) { return Optional.empty(); } - var analysisID = new TypedApiInterface.AnalysisID((int) bid); - if (!statusCache.containsKey(analysisID)) { - // Check that it's really valid in the context of the currently configured API - try { - var status = api.status(analysisID); - statusCache.put(analysisID, status); - } catch (APIAuthenticationException | ApiException e) { - Msg.showError(this, null, "Invalid Analysis ID", - ("The Analysis ID %s stored in the program options is invalid for the currently configured RevEng.AI server %s. " - + "This could be an intermittent error, or you switched the servers") - .formatted(analysisID, this.apiInfo.hostURI()), e); - return Optional.empty(); - } - // Now it's certain that it is a valid binary ID - } - - return Optional.of(analysisID); + return Optional.of(new TypedApiInterface.AnalysisID((int) bid)); } @Deprecated @@ -241,7 +225,7 @@ private Optional getBinaryIDfromOptions( try { status = api.status(binID); } catch (APIAuthenticationException | ApiException e) { - Msg.showError(this, null, "Invalid Binary ID", + Msg.error(this, ("The Binary ID %s stored in the program options is invalid for the currently configured RevEng.AI server %s. " + "This could be an intermittent error, or you switched servers") .formatted(binID, this.apiInfo.hostURI()), e); diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiImplementation.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiImplementation.java index 644e0f11..b7f73514 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiImplementation.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiImplementation.java @@ -95,6 +95,10 @@ public TypedApiImplementation(String baseUrl, String apiKey) { // Set withResponseBody to true if debugging issues with the API to see the full response body apiClient.setHttpClient(apiClient.getHttpClient().newBuilder().addInterceptor(ghidraLogger(false)).build()); + apiClient.setConnectTimeout(5000); + apiClient.setReadTimeout(15000); + apiClient.setWriteTimeout(15000); + ApiKeyAuth APIKey = (ApiKeyAuth) apiClient.getAuthentication("APIKey"); APIKey.setApiKey(apiKey); @@ -532,6 +536,11 @@ public AIDecompilationStatus pollAIDecompileStatus(FunctionID functionID) { inlineCommentsStatus, inlineComments); } catch (ApiException e) { + if (e.getCode() == 404) { + return new AIDecompilationStatus( + DecompilationData.StatusEnum.UNINITIALISED, + null, null, null, null, null, List.of()); + } throw new RuntimeException("Failed to poll AI decompilation status", e); } } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/ui/wizard/panels/UserCredentialsPanel.java b/src/main/java/ai/reveng/toolkit/ghidra/core/ui/wizard/panels/UserCredentialsPanel.java index ff3355e2..1bf4bc06 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/ui/wizard/panels/UserCredentialsPanel.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/ui/wizard/panels/UserCredentialsPanel.java @@ -12,6 +12,8 @@ import docking.wizard.WizardPanelDisplayability; import docking.wizard.WizardState; import ghidra.framework.plugintool.PluginTool; +import ghidra.util.task.Task; +import ghidra.util.task.TaskMonitor; import java.awt.BorderLayout; import java.awt.GridBagLayout; import java.awt.GridBagConstraints; @@ -30,10 +32,12 @@ public class UserCredentialsPanel extends AbstractMageJPanel { var apiInfo = new ApiInfo(tfApiHostname.getText(), tfPortalHostname.getText(), tfApiKey.getText()); - try { - credentialValidator.validate(apiInfo); - credentialsValidated = true; - notifyListenersOfValidityChanged(); - } catch (InvalidAPIInfoException ex) { - credentialsValidated = false; - notifyListenersOfStatusMessage("Problem with user info:\n" + ex.getMessage()); - } + runTestsButton.setEnabled(false); + tool.execute(new Task("Validate Credentials", false, false, false) { + @Override + public void run(TaskMonitor monitor) { + try { + credentialValidator.validate(apiInfo); + SwingUtilities.invokeLater(() -> { + credentialsValidated = true; + notifyListenersOfValidityChanged(); + runTestsButton.setEnabled(true); + }); + } catch (InvalidAPIInfoException ex) { + SwingUtilities.invokeLater(() -> { + credentialsValidated = false; + notifyListenersOfStatusMessage("Problem with user info:\n" + ex.getMessage()); + runTestsButton.setEnabled(true); + }); + } + } + }, 0); }); userDetailsPanel.add(runTestsButton, gbc); } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java b/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java index cda6bea0..4b40280c 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java @@ -39,8 +39,10 @@ import ghidra.program.model.listing.Program; import ghidra.program.util.GhidraProgramUtilities; import ghidra.util.Msg; +import ghidra.util.task.Task; import ghidra.util.task.TaskBuilder; import ghidra.util.task.TaskMonitor; +import javax.swing.SwingUtilities; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -237,12 +239,23 @@ private void setupActions() { .onAction(context -> { var currentProgram = tool.getService(ProgramManager.class).getCurrentProgram(); var knownProgram = revengService.getKnownProgram(currentProgram).orElseThrow(); - var logs = revengService.getAnalysisLog(knownProgram.analysisID()); - analysisLogComponent.setLogs(logs); - AnalysisStatus status = revengService.status(knownProgram); - tool.getService(ReaiLoggingService.class).info("Check Status: " + status); - Msg.showInfo(this, null, ReaiPluginPackage.WINDOW_PREFIX + "Check status", - "Status of analysis " + knownProgram + ": " + status); + tool.execute(new Task("Check status", true, false, false) { + @Override + public void run(TaskMonitor monitor) { + try { + var logs = revengService.getAnalysisLog(knownProgram.analysisID()); + AnalysisStatus status = revengService.status(knownProgram); + tool.getService(ReaiLoggingService.class).info("Check Status: " + status); + SwingUtilities.invokeLater(() -> { + analysisLogComponent.setLogs(logs); + Msg.showInfo(AnalysisManagementPlugin.this, null, ReaiPluginPackage.WINDOW_PREFIX + "Check status", + "Status of analysis " + knownProgram + ": " + status); + }); + } catch (Exception e) { + Msg.error(AnalysisManagementPlugin.this, "Failed to check analysis status: " + e.getMessage(), e); + } + } + }, 0); }) .menuPath(new String[] { ReaiPluginPackage.MENU_GROUP_NAME, "Analysis", "Check status" }) .menuGroup(REAI_ANALYSIS_MANAGEMENT_MENU_GROUP, "400") diff --git a/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java b/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java index 362b905e..bcaa0f03 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java @@ -38,6 +38,8 @@ import ghidra.framework.plugintool.util.PluginStatus; import ghidra.program.util.ProgramLocation; import ghidra.util.Msg; +import ghidra.util.task.Task; +import ghidra.util.task.TaskMonitor; import java.util.Arrays; @@ -231,7 +233,17 @@ private void setupActions() { return; } - apiService.openFunctionInPortal(functionWithID.get().functionID()); + var functionID = functionWithID.get().functionID(); + tool.execute(new Task("View function in portal", false, false, false) { + @Override + public void run(TaskMonitor monitor) { + try { + apiService.openFunctionInPortal(functionID); + } catch (Exception e) { + Msg.error(BinarySimilarityPlugin.this, "Failed to open function in portal: " + e.getMessage(), e); + } + } + }, 0); }) .popupMenuPath(new String[] { "View function in portal" }) .popupMenuIcon(ReaiPluginPackage.REVENG_16) diff --git a/src/test/java/ai/reveng/AIDecompilerComponentTest.java b/src/test/java/ai/reveng/AIDecompilerComponentTest.java index 9ee8bc7a..db4f9671 100644 --- a/src/test/java/ai/reveng/AIDecompilerComponentTest.java +++ b/src/test/java/ai/reveng/AIDecompilerComponentTest.java @@ -186,11 +186,95 @@ public void testAIDecompFeedbackMechanism() throws Exception { dialog.setValue(reason); dialog.close(); waitForSwing(); + waitForTasks(); assertEquals("NEGATIVE", ratingsAPI.lastFeedback); assertEquals(reason, ratingsAPI.lastReason); } + @Test + public void testFeedbackDoesNotBlockSwingThread() throws Exception { + var tool = env.getTool(); + + var ratingsAPI = new BlockingRatingsAPI(); + var service = addMockedService(tool, ratingsAPI); + + env.addPlugin(BinarySimilarityPlugin.class); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + var func1 = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction(null, "0x2000", 10, Undefined.getUndefinedDataType(4)); + + var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(programWithID.program()); + waitForSwing(); + + var aiDecompComponent = getComponentProvider(AIDecompilationdWindow.class); + aiDecompComponent.setVisible(true); + setInstanceField("function", aiDecompComponent, func1); + + var positiveFeedbackAction = getLocalAction(aiDecompComponent, "Positive Feedback Action"); + performAction(positiveFeedbackAction); + + assertTrue("Feedback send should run off the Swing thread", + ratingsAPI.ratingStarted.await(5, java.util.concurrent.TimeUnit.SECONDS)); + assertNull("Rating must not be recorded while the network call is blocked", ratingsAPI.lastFeedback); + + ratingsAPI.blockRating.countDown(); + waitForTasks(); + assertEquals("POSITIVE", ratingsAPI.lastFeedback); + } + + @Test + public void testDecompilationFailureIsShownInWindow() throws Exception { + var tool = env.getTool(); + + var service = addMockedService(tool, new UnimplementedAPI() { + @Override + public AnalysisStatus status(AnalysisID analysisID) { + return AnalysisStatus.Complete; + } + + @Override + public AnalysisID analyse(AnalysisOptionsBuilder options) throws ApiException { + return new AnalysisID(1); + } + + @Override + public List getFunctionInfo(AnalysisID analysisID) { + return List.of(new FunctionInfo(new FunctionID(1), "portal_func_1", "portal_func_1_mangled", 0x1000L, 10)); + } + + @Override + public AIDecompilationStatus pollAIDecompileStatus(FunctionID functionID) { + throw new RuntimeException("decompilation backend unavailable"); + } + }); + + env.addPlugin(BinarySimilarityPlugin.class); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + var func1 = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); + var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(programWithID.program()); + + var aiDecompComponent = getComponentProvider(AIDecompilationdWindow.class); + aiDecompComponent.setVisible(true); + + var action = getAction(tool, "AI Decompilation"); + var context = new ProgramLocationActionContext( + null, + programWithID.program(), + new ProgramLocation(programWithID.program(), func1.getEntryPoint()), + null, null); + performAction(action, context, true); + waitForTasks(); + waitForSwing(); + + javax.swing.JEditorPane descriptionArea = + (javax.swing.JEditorPane) getInstanceField("descriptionArea", aiDecompComponent); + assertTrue("decompilation failure should be surfaced in the window, was: " + descriptionArea.getText(), + descriptionArea.getText().contains("AI Decompilation failed")); + } + static class RatingsAPI extends UnimplementedAPI { String lastFeedback; String lastReason; @@ -249,4 +333,20 @@ public void aiDecompRating(FunctionID functionID, String rating, String reason) lastReason = reason; } } + + static class BlockingRatingsAPI extends RatingsAPI { + final java.util.concurrent.CountDownLatch ratingStarted = new java.util.concurrent.CountDownLatch(1); + final java.util.concurrent.CountDownLatch blockRating = new java.util.concurrent.CountDownLatch(1); + + @Override + public void aiDecompRating(FunctionID functionID, String rating, String reason) { + ratingStarted.countDown(); + try { + blockRating.await(10, java.util.concurrent.TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + super.aiDecompRating(functionID, rating, reason); + } + } } diff --git a/src/test/java/ai/reveng/ApplyMatchCmdTest.java b/src/test/java/ai/reveng/ApplyMatchCmdTest.java new file mode 100644 index 00000000..10eb4653 --- /dev/null +++ b/src/test/java/ai/reveng/ApplyMatchCmdTest.java @@ -0,0 +1,83 @@ +package ai.reveng; + +import ai.reveng.invoker.ApiException; +import ai.reveng.toolkit.ghidra.binarysimilarity.cmds.ApplyMatchCmd; +import ai.reveng.toolkit.ghidra.core.services.api.AnalysisOptionsBuilder; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface.AnalysisID; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface.BinaryHash; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface.FunctionID; +import ai.reveng.toolkit.ghidra.core.services.api.mocks.UnimplementedAPI; +import ai.reveng.toolkit.ghidra.core.services.api.types.AnalysisStatus; +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionInfo; +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionMatch; +import ai.reveng.toolkit.ghidra.core.services.api.types.GhidraFunctionMatchWithSignature; +import ghidra.program.database.ProgramBuilder; +import ghidra.program.model.data.Undefined; +import ghidra.util.task.TaskMonitor; +import org.junit.Test; + +import java.math.BigDecimal; +import java.util.List; + +import static org.junit.Assert.assertNull; +import static org.junit.Assert.fail; + +public class ApplyMatchCmdTest extends RevEngMockableHeadedIntegrationTest { + + @Test + public void testFailedServerRenameDoesNotLeakTransaction() throws Exception { + var tool = env.getTool(); + var service = addMockedService(tool, new RenameFailsAPI()); + + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + var func1 = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); + var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + var program = programWithID.program(); + var analysedProgram = service.getAnalysedProgram(program).orElseThrow(); + + var match = new FunctionMatch( + new FunctionID(999), + new FunctionID(2), + "renamed_func", + "renamed_func_mangled", + "libfoo.so", + new BinaryHash("abc123"), + false, + BigDecimal.valueOf(0.9), + BigDecimal.valueOf(0.9)); + var ghidraMatch = new GhidraFunctionMatchWithSignature(func1, match, null); + + var cmd = new ApplyMatchCmd(service, analysedProgram, ghidraMatch, true); + try { + cmd.applyWithTransaction(); + fail("expected the server rename to throw while offline"); + } catch (RuntimeException expected) { + // the whole point: it throws, but the transaction must still be closed + } + + assertNull("transaction must be closed even when the server rename fails", + program.getCurrentTransactionInfo()); + } + + static class RenameFailsAPI extends UnimplementedAPI { + @Override + public AnalysisID analyse(AnalysisOptionsBuilder options) throws ApiException { + return new AnalysisID(1); + } + + @Override + public AnalysisStatus status(AnalysisID analysisID) { + return AnalysisStatus.Complete; + } + + @Override + public List getFunctionInfo(AnalysisID analysisID) { + return List.of(new FunctionInfo(new FunctionID(1), "f1", "f1_mangled", 0x1000L, 10)); + } + + @Override + public void renameFunction(FunctionID id, String newName, String newNameMangled) { + throw new RuntimeException("simulated offline rename failure"); + } + } +} diff --git a/src/test/java/ai/reveng/GhidraRevengServiceTest.java b/src/test/java/ai/reveng/GhidraRevengServiceTest.java new file mode 100644 index 00000000..12fb1599 --- /dev/null +++ b/src/test/java/ai/reveng/GhidraRevengServiceTest.java @@ -0,0 +1,69 @@ +package ai.reveng; + +import ai.reveng.invoker.ApiException; +import ai.reveng.toolkit.ghidra.core.services.api.AnalysisOptionsBuilder; +import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface.AnalysisID; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface.FunctionID; +import ai.reveng.toolkit.ghidra.core.services.api.mocks.UnimplementedAPI; +import ai.reveng.toolkit.ghidra.core.services.api.types.AnalysisStatus; +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionInfo; +import ghidra.program.database.ProgramBuilder; +import ghidra.program.model.data.Undefined; +import ghidra.util.task.TaskMonitor; +import org.junit.Test; + +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class GhidraRevengServiceTest extends RevEngMockableHeadedIntegrationTest { + + @Test + public void testKnownProgramLookupIsNetworkFree() throws Exception { + var tool = env.getTool(); + var mock = new OfflineAfterSetupAPI(); + var service = addMockedService(tool, mock); + + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); + var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + var program = programWithID.program(); + + mock.offline = true; + var freshService = new GhidraRevengService(mock); + int statusCallsBefore = mock.statusCalls; + + assertTrue("known program must resolve from local state without the network", + freshService.getKnownProgram(program).isPresent()); + assertTrue("analysed program must resolve from local state without the network", + freshService.getAnalysedProgram(program).isPresent()); + assertEquals("navigation/enablement lookups must not call status()", + statusCallsBefore, mock.statusCalls); + } + + static class OfflineAfterSetupAPI extends UnimplementedAPI { + volatile boolean offline = false; + int statusCalls = 0; + + @Override + public AnalysisID analyse(AnalysisOptionsBuilder options) throws ApiException { + return new AnalysisID(1); + } + + @Override + public AnalysisStatus status(AnalysisID analysisID) throws ApiException { + statusCalls++; + if (offline) { + throw new ApiException("simulated offline"); + } + return AnalysisStatus.Complete; + } + + @Override + public List getFunctionInfo(AnalysisID analysisID) { + return List.of(new FunctionInfo(new FunctionID(1), "f1", "f1_mangled", 0x1000L, 10)); + } + } +} diff --git a/src/test/java/ai/reveng/SetupWizardTest.java b/src/test/java/ai/reveng/SetupWizardTest.java index a8612908..f0be50a5 100644 --- a/src/test/java/ai/reveng/SetupWizardTest.java +++ b/src/test/java/ai/reveng/SetupWizardTest.java @@ -61,6 +61,7 @@ public void testWizardCompletesSuccessfully() throws Exception { JButton validateButton = findButtonByText(dialog.getComponent(), "Validate Credentials"); assertNotNull("Validate button should exist", validateButton); pressButton(validateButton); + waitForTasks(); waitForSwing(); // Finish button should now be enabled @@ -112,6 +113,7 @@ public void testValidationFailurePreventsFinish() throws Exception { // Click Validate — should fail JButton validateButton = findButtonByText(dialog.getComponent(), "Validate Credentials"); pressButton(validateButton); + waitForTasks(); waitForSwing(); // Finish button should remain disabled @@ -151,6 +153,7 @@ public void testEditingAfterValidationResetsState() throws Exception { // Validate successfully JButton validateButton = findButtonByText(dialog.getComponent(), "Validate Credentials"); pressButton(validateButton); + waitForTasks(); waitForSwing(); JButton finishButton = findButtonByText(dialog.getComponent(), "Finish"); @@ -164,6 +167,7 @@ public void testEditingAfterValidationResetsState() throws Exception { // Re-validate and finish pressButton(validateButton); + waitForTasks(); waitForSwing(); assertTrue("Finish should be re-enabled after re-validation", finishButton.isEnabled());