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 ce46173..16e8846 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 @@ -94,7 +94,11 @@ public boolean applyTo(Program obj) { // If we have a service then push the name. If not then it was explicitly not provided, i.e. the caller // is responsible for pushing the names in batch if (service != null) { - service.getApi().renameFunction(match.functionMatch().origin_function_id(), match.functionMatch().nearest_neighbor_function_name()); + service.getApi().renameFunction( + match.functionMatch().origin_function_id(), + match.functionMatch().nearest_neighbor_function_name(), + match.functionMatch().nearest_neighbor_mangled_function_name() + ); } 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 new file mode 100644 index 0000000..1a11e94 --- /dev/null +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/SimilarFunctionsWindow.java @@ -0,0 +1,502 @@ +package ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching; + +import ai.reveng.model.FunctionMatchingRequest; +import ai.reveng.model.FunctionMatchingResponse; +import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionMatch; +import ai.reveng.toolkit.ghidra.core.services.api.types.GhidraFunctionMatch; +import ai.reveng.toolkit.ghidra.core.services.api.types.GhidraFunctionMatchWithSignature; +import ai.reveng.toolkit.ghidra.binarysimilarity.cmds.ApplyMatchCmd; +import ai.reveng.toolkit.ghidra.plugins.BinarySimilarityPlugin; +import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage; +import ai.reveng.toolkit.ghidra.Utils; +import docking.ActionContext; +import docking.action.DockingAction; +import docking.action.MenuData; +import docking.action.ToggleDockingAction; +import docking.action.ToolBarData; +import generic.theme.GIcon; +import docking.widgets.table.GDynamicColumnTableModel; +import docking.widgets.table.TableColumnDescriptor; +import docking.widgets.table.TableSortState; +import ghidra.framework.plugintool.ComponentProviderAdapter; +import ghidra.framework.plugintool.PluginTool; +import ghidra.framework.plugintool.ServiceProvider; +import ghidra.program.model.listing.Function; +import ghidra.program.util.ProgramLocation; +import ghidra.util.Msg; +import ghidra.util.table.GhidraTable; +import ghidra.util.task.Task; +import ghidra.util.task.TaskBuilder; +import ghidra.util.task.TaskMonitor; +import ghidra.util.task.TaskMonitorComponent; + +import javax.swing.*; +import javax.swing.event.ListSelectionEvent; +import java.awt.*; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Window that automatically fetches and displays similar functions for the currently selected function. + * Shows a table of matches on the left and an assembly diff panel on the right. + */ +public class SimilarFunctionsWindow extends ComponentProviderAdapter { + + private final JComponent component; + private GhidraTable matchesTable; + private SimilarFunctionsTableModel tableModel; + private AssemblyDiffPanel assemblyDiffPanel; + private JSplitPane splitPane; + private JLabel statusLabel; + private TaskMonitorComponent taskMonitorComponent; + + private GhidraRevengService.FunctionWithID currentFunctionWithID; + private GhidraRevengService.AnalysedProgram analyzedProgram; + private final Map> matchCache = new java.util.concurrent.ConcurrentHashMap<>(); + + public SimilarFunctionsWindow(PluginTool tool) { + super(tool, ReaiPluginPackage.WINDOW_PREFIX + "Similar Functions", BinarySimilarityPlugin.class.getName()); + setIcon(ReaiPluginPackage.REVENG_16); + component = buildComponent(); + createToolbarActions(); + } + + private JComponent buildComponent() { + JPanel mainPanel = new JPanel(new BorderLayout()); + + // Status bar at top + JPanel statusPanel = new JPanel(new BorderLayout()); + statusLabel = new JLabel("No function selected"); + taskMonitorComponent = new TaskMonitorComponent(false, true); + taskMonitorComponent.setVisible(false); + taskMonitorComponent.setIndeterminate(true); + statusPanel.add(statusLabel, BorderLayout.CENTER); + statusPanel.add(taskMonitorComponent, BorderLayout.EAST); + mainPanel.add(statusPanel, BorderLayout.NORTH); + + // Split pane: matches table on left, assembly diff on right + splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT); + splitPane.setResizeWeight(0.4); // Give 40% to the table + + // Left side: Matches table + JPanel tablePanel = createMatchesTablePanel(); + splitPane.setLeftComponent(tablePanel); + + // Right side: Assembly diff panel (reuse existing component) + assemblyDiffPanel = new AssemblyDiffPanel(); + splitPane.setRightComponent(assemblyDiffPanel); + + mainPanel.add(splitPane, BorderLayout.CENTER); + + return mainPanel; + } + + private JPanel createMatchesTablePanel() { + JPanel panel = new JPanel(new BorderLayout()); + panel.setBorder(BorderFactory.createTitledBorder("Similar Functions")); + + tableModel = new SimilarFunctionsTableModel(tool); + matchesTable = new GhidraTable(tableModel); + matchesTable.setSelectionMode(ListSelectionModel.SINGLE_SELECTION); + matchesTable.getSelectionModel().addListSelectionListener(this::onMatchSelected); + + // Add context menu action to open matched function in portal + createTableActions(); + + JScrollPane scrollPane = new JScrollPane(matchesTable); + panel.add(scrollPane, BorderLayout.CENTER); + + return panel; + } + + /** + * Returns the currently selected match in the table, or null if no valid selection. + */ + private GhidraFunctionMatchWithSignature getSelectedMatch() { + int selectedRow = matchesTable.getSelectedRow(); + if (selectedRow < 0) { + return null; + } + int modelRow = matchesTable.convertRowIndexToModel(selectedRow); + return tableModel.getRowObject(modelRow); + } + + private void createTableActions() { + DockingAction openInPortalAction = new DockingAction("View Match in Portal", getName()) { + @Override + public void actionPerformed(ActionContext context) { + var match = getSelectedMatch(); + if (match != null) { + var service = tool.getService(GhidraRevengService.class); + service.openFunctionInPortal(match.functionMatch().nearest_neighbor_id()); + } + } + + @Override + public boolean isEnabledForContext(ActionContext context) { + return getSelectedMatch() != null; + } + }; + + openInPortalAction.setPopupMenuData(new MenuData( + new String[] { "View Match in Portal" }, + ReaiPluginPackage.REVENG_16, + ReaiPluginPackage.MENU_GROUP_NAME + )); + + addLocalAction(openInPortalAction); + + DockingAction applyMatchAction = new DockingAction("Apply Match", getName()) { + @Override + 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(); + } + } + + @Override + public boolean isEnabledForContext(ActionContext context) { + return getSelectedMatch() != null && analyzedProgram != null; + } + }; + + applyMatchAction.setPopupMenuData(new MenuData( + new String[] { "Apply Match" }, + ReaiPluginPackage.REVENG_16, + ReaiPluginPackage.MENU_GROUP_NAME + )); + + addLocalAction(applyMatchAction); + } + + private void createToolbarActions() { + ToggleDockingAction toggleDiffPanelAction = new ToggleDockingAction("Hide Assembly Diff", getName()) { + @Override + public void actionPerformed(ActionContext context) { + setAssemblyDiffVisible(isSelected()); + } + }; + toggleDiffPanelAction.setSelected(true); + toggleDiffPanelAction.setToolBarData(new ToolBarData(new GIcon("icon.plugin.programdiff.get.diffs"), null)); + toggleDiffPanelAction.setDescription("Toggle assembly diff panel visibility"); + + addLocalAction(toggleDiffPanelAction); + + DockingAction reloadAction = new DockingAction("Reload Matches", getName()) { + @Override + public void actionPerformed(ActionContext context) { + reloadMatches(); + } + + @Override + public boolean isEnabledForContext(ActionContext context) { + return currentFunctionWithID != null; + } + }; + reloadAction.setToolBarData(new ToolBarData(new GIcon("icon.refresh"), null)); + reloadAction.setDescription("Reload matches for current function (bypass cache)"); + + addLocalAction(reloadAction); + } + + private void reloadMatches() { + if (currentFunctionWithID == null) { + return; + } + // Remove from cache to force a fresh fetch + matchCache.remove(currentFunctionWithID.function()); + fetchSimilarFunctions(); + } + + private void setAssemblyDiffVisible(boolean visible) { + if (visible) { + splitPane.setRightComponent(assemblyDiffPanel); + splitPane.setDividerLocation(0.4); + } else { + splitPane.setRightComponent(null); + } + splitPane.revalidate(); + } + + @Override + public JComponent getComponent() { + return component; + } + + @Override + public void componentShown() { + // When the window becomes visible, fetch matches if we have a current function + // but haven't fetched yet (handles race condition where location changed before window was visible) + if (currentFunctionWithID != null && !matchCache.containsKey(currentFunctionWithID.function())) { + fetchSimilarFunctions(); + } + } + + /** + * Called when the program is not analyzed with RevEng.AI + */ + public void onNoAnalyzedProgram() { + clear(); + statusLabel.setText("Binary not analyzed with RevEng.AI"); + } + + /** + * Called when the cursor is not within a function + */ + public void onNoFunction() { + clear(); + statusLabel.setText("No function at current location"); + } + + /** + * Called when the function exists locally but is not in the RevEng.AI analysis + */ + public void onFunctionNotInAnalysis(Function func) { + clear(); + statusLabel.setText("Function not found in RevEng.AI analysis"); + } + + /** + * Called when the cursor moves to a function that is in the RevEng.AI analysis + */ + public void onFunctionChanged(GhidraRevengService.FunctionWithID functionWithID, GhidraRevengService.AnalysedProgram newAnalyzedProgram) { + var newFunction = functionWithID.function(); + // If we changed to a different function, fetch new matches + var currentFunction = currentFunctionWithID != null ? currentFunctionWithID.function() : null; + if (currentFunction != newFunction) { + clear(); + currentFunctionWithID = functionWithID; + analyzedProgram = newAnalyzedProgram; + fetchSimilarFunctions(); + } + } + + private void clear() { + currentFunctionWithID = null; + tableModel.setData(Collections.emptyList()); + assemblyDiffPanel.clear(); + } + + private void fetchSimilarFunctions() { + if (!isVisible()) { + return; + } + + if (currentFunctionWithID == null) { + return; + } + + var currentFunction = currentFunctionWithID.function(); + + // Check cache first + if (matchCache.containsKey(currentFunction)) { + displayMatches(matchCache.get(currentFunction)); + return; + } + + statusLabel.setText("Fetching similar functions for " + currentFunction.getName() + "..."); + taskMonitorComponent.setVisible(true); + + // Start background task to fetch matches + var task = new FetchSimilarFunctionsTask(analyzedProgram, currentFunctionWithID); + TaskBuilder.withTask(task).launchInBackground(taskMonitorComponent); + } + + private void displayMatches(List matches) { + tableModel.setData(matches); + + // Clear the assembly diff panel when showing new matches + assemblyDiffPanel.clear(); + + var functionName = currentFunctionWithID != null ? currentFunctionWithID.function().getName() : "unknown"; + statusLabel.setText("Found " + matches.size() + " similar functions for " + functionName); + taskMonitorComponent.setVisible(false); + } + + private void onMatchSelected(ListSelectionEvent e) { + if (e.getValueIsAdjusting()) { + return; + } + + var match = getSelectedMatch(); + if (match == null) { + assemblyDiffPanel.clear(); + return; + } + + // Use the AssemblyDiffPanel to show the comparison + var service = tool.getService(GhidraRevengService.class); + assemblyDiffPanel.showAssemblyFor(match, service); + } + + public void locationChanged(ProgramLocation loc) { + var service = tool.getService(GhidraRevengService.class); + var optAnalysedProgram = service.getAnalysedProgram(loc.getProgram()); + if (optAnalysedProgram.isEmpty()) { + clear(); + return; + } + var newAnalyzedProgram = optAnalysedProgram.get(); + + var functionMgr = loc.getProgram().getFunctionManager(); + var func = functionMgr.getFunctionContaining(loc.getAddress()); + if (func == null) { + onNoFunction(); + return; + } + + var functionWithID = newAnalyzedProgram.getIDForFunction(func); + if (functionWithID.isEmpty()) { + onFunctionNotInAnalysis(func); + return; + } + + onFunctionChanged(functionWithID.get(), newAnalyzedProgram); + + } + + /** + * Table model for displaying similar function matches. + * Extends GDynamicColumnTableModel for dynamic column support with hidden columns. + */ + private static class SimilarFunctionsTableModel extends GDynamicColumnTableModel { + private List data = Collections.emptyList(); + + private static final int COL_SIMILARITY = 1; + + public SimilarFunctionsTableModel(ServiceProvider serviceProvider) { + super(serviceProvider); + // Sort by similarity (descending) by default - highest matches first + setDefaultTableSortState(TableSortState.createDefaultSortState(COL_SIMILARITY, false)); + } + + @Override + public String getName() { + return "Similar Functions"; + } + + public void setData(List matches) { + this.data = matches != null ? new ArrayList<>(matches) : Collections.emptyList(); + fireTableDataChanged(); + } + + @Override + public List getModelData() { + return data; + } + + @Override + public Void getDataSource() { + return null; + } + + @Override + protected TableColumnDescriptor createTableColumnDescriptor() { + TableColumnDescriptor descriptor = new TableColumnDescriptor<>(); + + Utils.addRowToDescriptor(descriptor, "Function Name", String.class, + m -> m.functionMatch().nearest_neighbor_function_name()); + Utils.addRowToDescriptor(descriptor, "Similarity", BigDecimal.class, + m -> m.functionMatch().similarity()); + Utils.addRowToDescriptor(descriptor, "Confidence", BigDecimal.class, + m -> m.functionMatch().confidence()); + Utils.addRowToDescriptor(descriptor, "Binary", String.class, + m -> m.functionMatch().nearest_neighbor_binary_name()); + Utils.addRowToDescriptor(descriptor, "Signature", String.class, + m -> m.signature() != null ? m.signature().getPrototypeString() : ""); + Utils.addRowToDescriptor(descriptor, "Function ID", false, Long.class, + m -> m.functionMatch().nearest_neighbor_id().value()); + + return descriptor; + } + + @Override + public boolean isCellEditable(int row, int column) { + return false; + } + } + + /** + * Task to fetch similar functions from the API + */ + private class FetchSimilarFunctionsTask extends Task { + private final GhidraRevengService.AnalysedProgram program; + private final GhidraRevengService.FunctionWithID functionWithID; + + public FetchSimilarFunctionsTask(GhidraRevengService.AnalysedProgram program, GhidraRevengService.FunctionWithID functionWithID) { + super("Fetch Similar Functions", true, false, false); + this.functionWithID = functionWithID; + this.program = program; + } + + @Override + public void run(TaskMonitor monitor) { + try { + var service = tool.getService(GhidraRevengService.class); + + // Get basic analysis info for model ID + var analysisBasicInfo = service.getBasicDetailsForAnalysis(program.analysisID()); + + var request = new FunctionMatchingRequest(); + request.setMinSimilarity(BigDecimal.valueOf(70)); // Default threshold + request.setResultsPerFunction(25); + request.setModelId(analysisBasicInfo.getModelId()); + + var functionIds = new ArrayList(); + functionIds.add(functionWithID.functionID().value()); + request.setFunctionIds(functionIds); + + monitor.setMessage("Fetching matches from RevEng.AI..."); + FunctionMatchingResponse response; + while (true) { + if (monitor.isCancelled()) { + return; + } + response = service.getFunctionMatchingForFunction(request); + if (!"IN_PROGRESS".equals(response.getStatus())) { + break; + } else { + Thread.sleep(100); + } + } + + // Process results + List matches = new ArrayList<>(); + if (response.getMatches() != null) { + for (var matchResult : response.getMatches()) { + for (var matched : matchResult.getMatchedFunctions()) { + var funcId = new TypedApiInterface.FunctionID(matchResult.getFunctionId()); + FunctionMatch fm = FunctionMatch.fromMatchedFunctionAPIType(matched, funcId); + GhidraFunctionMatch match = new GhidraFunctionMatch(functionWithID.function(), fm); + matches.add(match); + } + } + } + var ghidraResultsWithSignatures = service.getSignatures(matches); + // For all matches we got, create a GhidraFunctionMatchWithSignature object (signature can be null!) + var result = matches.stream().map( + m -> m.withSignature(ghidraResultsWithSignatures.get(m)) + ).toList(); + // Cache and display + matchCache.put(functionWithID.function(), result); + + SwingUtilities.invokeLater(() -> displayMatches(result)); + + } catch (Exception e) { + Msg.error(this, "Failed to fetch similar functions: " + e.getMessage(), e); + SwingUtilities.invokeLater(() -> { + statusLabel.setText("Error fetching similar functions: " + e.getMessage()); + taskMonitorComponent.setVisible(false); + }); + } + } + } +} 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 8e3655b..005a0a9 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 @@ -733,7 +733,14 @@ public static Optional getFunctionSignature(Function // Create Data Type Manager with all dependencies var d = FunctionDependencies.fromOpenAPI(functionDataTypeMessage.getFuncDeps()); - var dtm = loadDependencyDataTypes(d); + DataTypeManager tmpDtm = null; + try { + tmpDtm = loadDependencyDataTypes(d); + } catch (EndlessTypeParsingException e) { + Msg.error("getFunctionSignature", null, e); + return Optional.empty(); + } + DataTypeManager dtm = tmpDtm; if (functionDataTypeMessage.getFuncTypes() == null){ return Optional.empty(); @@ -774,7 +781,19 @@ public static Optional getFunctionSignature(Function return Optional.of(f); } - public static DataTypeManager loadDependencyDataTypes(FunctionDependencies dependencies){ + public static class EndlessTypeParsingException extends Exception { + + public FunctionDependencies deps; + public List remaining; + private EndlessTypeParsingException(FunctionDependencies dependencies, List remainingTypes) { + super("Endless type parsing detected for function dependencies: " + dependencies); + deps = dependencies; + remaining = remainingTypes; + + } + } + + public static DataTypeManager loadDependencyDataTypes(FunctionDependencies dependencies) throws EndlessTypeParsingException{ DataTypeManager dtm = new StandAloneDataTypeManager("transient"); if (dependencies == null){ @@ -810,7 +829,9 @@ public static DataTypeManager loadDependencyDataTypes(FunctionDependencies depen int retries = 0; while (!typeDefsToAdd.isEmpty()){ if (retries > 1000){ - throw new RuntimeException("Dependency loading failed: %s".formatted(typeDefsToAdd)); + dtm.endTransaction(transactionId, false); + dtm.close(); + throw new EndlessTypeParsingException(dependencies, typeDefsToAdd.stream().toList()); } var typeDef = typeDefsToAdd.remove(); var path = TypePathAndName.fromString(typeDef.name()); 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 976e426..cc5741e 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 @@ -469,18 +469,21 @@ public AIDecompilationStatus pollAIDecompileStatus(FunctionID functionID) { /** * https://api.reveng.ai/v2/docs#tag/Functions-overview/operation/rename_function_id_v2_functions_rename__function_id__post + * * @param id * @param newName + * @param newNameMangled */ @Override - public void renameFunction(FunctionID id, String newName) { - JSONObject params = new JSONObject(); - params.put("new_name", newName); - - HttpRequest request = requestBuilderForEndpoint("functions", "rename", String.valueOf(id.value())) - .POST(HttpRequest.BodyPublishers.ofString(params.toString())) - .build(); - sendRequest(request); + public void renameFunction(FunctionID id, String newName, String newNameMangled) { + var fn = new FunctionRename(); + fn.setNewName(newName); + fn.setNewMangledName(newNameMangled); + try { + functionsRenamingHistoryApi.renameFunctionId((int) id.value(), fn); + } catch (ApiException e) { + throw new RuntimeException(e); + } } @Override diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiInterface.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiInterface.java index 8d47363..3128a41 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiInterface.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiInterface.java @@ -157,7 +157,7 @@ default AIDecompilationStatus pollAIDecompileStatus(FunctionID functionID) { throw new UnsupportedOperationException("pollAIDecompileStatus not implemented yet"); } - void renameFunction(FunctionID id, String newName); + void renameFunction(FunctionID id, String newName, String newNameMangled); default FunctionNameScore getNameScore(FunctionMatch match) { throw new UnsupportedOperationException("getNameScore not implemented yet"); diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/MockApi.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/MockApi.java index 90f1bf5..2610743 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/MockApi.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/MockApi.java @@ -52,7 +52,7 @@ public void authenticate() throws APIAuthenticationException { } @Override - public void renameFunction(FunctionID id, String newName) { + public void renameFunction(FunctionID id, String newName, String newNameMangled) { } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/UnimplementedAPI.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/UnimplementedAPI.java index 6a8c339..11b20b3 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/UnimplementedAPI.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/mocks/UnimplementedAPI.java @@ -36,7 +36,7 @@ public void authenticate() { } @Override - public void renameFunction(FunctionID id, String newName) { + public void renameFunction(FunctionID id, String newName, String newNameMangled) { } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/types/GhidraFunctionMatch.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/types/GhidraFunctionMatch.java index 8dfb3fc..c86edc3 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/types/GhidraFunctionMatch.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/types/GhidraFunctionMatch.java @@ -1,6 +1,7 @@ package ai.reveng.toolkit.ghidra.core.services.api.types; import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; +import ghidra.program.model.data.FunctionDefinitionDataType; import ghidra.program.model.listing.Function; /** @@ -17,4 +18,9 @@ public record GhidraFunctionMatch( public TypedApiInterface.FunctionID nearest_neighbor_id() { return functionMatch.nearest_neighbor_id(); } + + public GhidraFunctionMatchWithSignature withSignature(FunctionDefinitionDataType sig) { + return new GhidraFunctionMatchWithSignature(function, functionMatch, sig); + } + } 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 8ee9a1c..d5b8fe8 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java @@ -20,6 +20,7 @@ import ai.reveng.toolkit.ghidra.binarysimilarity.ui.collectiondialog.DataSetControlPanelComponent; import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching.BinaryLevelFunctionMatchingDialog; import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching.FunctionLevelFunctionMatchingDialog; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching.SimilarFunctionsWindow; import ai.reveng.toolkit.ghidra.core.RevEngAIAnalysisResultsLoaded; import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; @@ -61,6 +62,7 @@ //@formatter:on public class BinarySimilarityPlugin extends ProgramPlugin { private final AIDecompilationdWindow decompiledWindow; + private final SimilarFunctionsWindow similarFunctionsWindow; private GhidraRevengService apiService; @@ -83,6 +85,7 @@ protected void locationChanged(ProgramLocation loc) { } decompiledWindow.locationChanged(loc); + similarFunctionsWindow.locationChanged(loc); } /** @@ -99,8 +102,10 @@ public BinarySimilarityPlugin(PluginTool tool) { decompiledWindow = new AIDecompilationdWindow(tool, REVENG_AI_NAMESPACE); decompiledWindow.addToTool(); + similarFunctionsWindow = new SimilarFunctionsWindow(tool); + similarFunctionsWindow.addToTool(); - // Install Collections Control Panel + // Install Collections Control Panel DataSetControlPanelComponent collectionsComponent = new DataSetControlPanelComponent(tool, REVENG_AI_NAMESPACE); collectionsComponent.addToTool(); } diff --git a/src/test/java/ConvertBinSyncArtifactTests.java b/src/test/java/ConvertBinSyncArtifactTests.java index 511fef1..1842671 100644 --- a/src/test/java/ConvertBinSyncArtifactTests.java +++ b/src/test/java/ConvertBinSyncArtifactTests.java @@ -51,7 +51,7 @@ public void testSimpleGhidraSignatureGeneration() throws DataTypeDependencyExcep @Test - public void testDependencyToDtm() { + public void testDependencyToDtm() throws GhidraRevengService.EndlessTypeParsingException { var mockResponse = getMockResponseFromFile("confirmmatch_fdupes_77846700.json"); FunctionDataTypeStatus functionDataTypeStatus = FunctionDataTypeStatus.fromJson(mockResponse.getJsonData()); var dtm = GhidraRevengService.loadDependencyDataTypes(functionDataTypeStatus.data_types().get().func_deps()); diff --git a/src/test/java/ai/reveng/SimilarFunctionsWindowTest.java b/src/test/java/ai/reveng/SimilarFunctionsWindowTest.java new file mode 100644 index 0000000..98955d2 --- /dev/null +++ b/src/test/java/ai/reveng/SimilarFunctionsWindowTest.java @@ -0,0 +1,279 @@ +package ai.reveng; + +import ai.reveng.invoker.ApiException; +import ai.reveng.model.*; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching.SimilarFunctionsWindow; +import ai.reveng.toolkit.ghidra.core.services.api.AnalysisOptionsBuilder; +import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; +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.plugins.BinarySimilarityPlugin; +import ghidra.program.database.ProgramBuilder; +import ghidra.program.model.data.Undefined; +import ghidra.program.model.listing.Function; +import ghidra.util.task.TaskMonitor; +import org.junit.Test; + +import javax.swing.*; +import java.math.BigDecimal; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.*; + +public class SimilarFunctionsWindowTest extends RevEngMockableHeadedIntegrationTest { + + @Test + public void testSimilarFunctionsWindowBasics() throws Exception { + var tool = env.getTool(); + + var mockApi = new SimilarFunctionsMockAPI(); + var service = addMockedService(tool, mockApi); + + var binarySimilarityPlugin = env.addPlugin(BinarySimilarityPlugin.class); + + // Create a program with two functions + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + var func1 = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); + var func2 = builder.createEmptyFunction(null, "0x2000", 10, Undefined.getUndefinedDataType(4)); + + var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + + env.showTool(programWithID.program()); + + // Get the SimilarFunctionsWindow + var similarFunctionsWindow = getComponentProvider(SimilarFunctionsWindow.class); + assertNotNull("SimilarFunctionsWindow should exist", similarFunctionsWindow); + + // Get internal fields for testing + @SuppressWarnings("unchecked") + Map> matchCache = + (Map>) getInstanceField("matchCache", similarFunctionsWindow); + JTable matchesTable = (JTable) getInstanceField("matchesTable", similarFunctionsWindow); + JLabel statusLabel = (JLabel) getInstanceField("statusLabel", similarFunctionsWindow); + + // Make sure the window is hidden initially + similarFunctionsWindow.setVisible(false); + + // Navigate to function 1 while window is hidden - should not fetch + goTo(tool, programWithID.program(), func1.getEntryPoint()); + waitForTasks(); + assertFalse("Cache should be empty when window is hidden", matchCache.containsKey(func1)); + + // Now make the window visible and navigate again - should trigger fetch + similarFunctionsWindow.setVisible(true); + goTo(tool, programWithID.program(), func2.getEntryPoint()); + waitForTasks(); + waitForSwing(); + + // Check that matches were fetched and cached + assertTrue("Cache should contain func2 after navigation", matchCache.containsKey(func2)); + + // Check that the table is populated + assertEquals("Table should have 2 rows (2 matches)", 2, matchesTable.getRowCount()); + + // Check status label shows the function name + assertTrue("Status label should mention func2", + statusLabel.getText().contains("2") || statusLabel.getText().contains("similar")); + + // Navigate to func1 (window still visible) + goTo(tool, programWithID.program(), func1.getEntryPoint()); + waitForTasks(); + waitForSwing(); + + // Check that func1 is now cached + assertTrue("Cache should contain func1", matchCache.containsKey(func1)); + assertEquals("Table should still have 2 rows", 2, matchesTable.getRowCount()); + } + + @Test + public void testSimilarFunctionsWindowCaching() throws Exception { + var tool = env.getTool(); + + var mockApi = new SimilarFunctionsMockAPI(); + var service = addMockedService(tool, mockApi); + + env.addPlugin(BinarySimilarityPlugin.class); + + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + var func1 = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); + var func2 = builder.createEmptyFunction(null, "0x2000", 10, Undefined.getUndefinedDataType(4)); + + var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(programWithID.program()); + + var similarFunctionsWindow = getComponentProvider(SimilarFunctionsWindow.class); + similarFunctionsWindow.setVisible(true); + + // Navigate to func1 + goTo(tool, programWithID.program(), func1.getEntryPoint()); + waitForTasks(); + waitForSwing(); + + int callCountAfterFirst = mockApi.functionMatchingCallCount; + assertTrue("Should have called function matching API", callCountAfterFirst > 0); + + // Navigate to func2 + goTo(tool, programWithID.program(), func2.getEntryPoint()); + waitForTasks(); + waitForSwing(); + + int callCountAfterSecond = mockApi.functionMatchingCallCount; + assertTrue("Should have called function matching API again for func2", + callCountAfterSecond > callCountAfterFirst); + + // Navigate back to func1 - should use cache, not call API + goTo(tool, programWithID.program(), func1.getEntryPoint()); + waitForTasks(); + waitForSwing(); + + assertEquals("Should not call API when using cache", + callCountAfterSecond, mockApi.functionMatchingCallCount); + } + + @Test + public void testSimilarFunctionsWindowTableSelection() throws Exception { + var tool = env.getTool(); + + var mockApi = new SimilarFunctionsMockAPI(); + var service = addMockedService(tool, mockApi); + + 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 similarFunctionsWindow = getComponentProvider(SimilarFunctionsWindow.class); + similarFunctionsWindow.setVisible(true); + + // Navigate to func1 + goTo(tool, programWithID.program(), func1.getEntryPoint()); + waitForTasks(); + waitForSwing(); + + JTable matchesTable = (JTable) getInstanceField("matchesTable", similarFunctionsWindow); + + // Select the first row + runSwing(() -> matchesTable.setRowSelectionInterval(0, 0)); + waitForTasks(); + waitForSwing(); + + // Check that assembly was fetched (the API should have been called) + assertTrue("Assembly should have been fetched for selected match", + mockApi.assemblyCallCount > 0); + } + + /** + * Mock API that provides function matching and assembly responses + */ + static class SimilarFunctionsMockAPI extends UnimplementedAPI { + int functionMatchingCallCount = 0; + int assemblyCallCount = 0; + + @Override + public TypedApiInterface.AnalysisID analyse(AnalysisOptionsBuilder options) throws ApiException { + return new TypedApiInterface.AnalysisID(1); + } + + @Override + public AnalysisStatus status(TypedApiInterface.AnalysisID analysisID) { + return AnalysisStatus.Complete; + } + + @Override + public List getFunctionInfo(TypedApiInterface.AnalysisID analysisID) { + return List.of( + new FunctionInfo( + new TypedApiInterface.FunctionID(1), + "portal_func_1", + "portal_func_1_mangled", + 0x1000L, + 10), + new FunctionInfo( + new TypedApiInterface.FunctionID(2), + "portal_func_2", + "portal_func_2_mangled", + 0x2000L, + 10) + ); + } + + @Override + public Basic getAnalysisBasicInfo(TypedApiInterface.AnalysisID analysisID) throws ApiException { + var basic = new Basic(); + basic.setModelId(1); + basic.setSha256Hash("abc123"); + basic.setBinaryName("test_binary"); + return basic; + } + + @Override + public FunctionMatchingResponse functionFunctionMatching(FunctionMatchingRequest request) throws ApiException { + functionMatchingCallCount++; + + var response = new FunctionMatchingResponse(); + + // Get the function ID from the request + Long originFunctionId = request.getFunctionIds().get(0); + + var functionMatch = new ai.reveng.model.FunctionMatch(); + functionMatch.setFunctionId(originFunctionId); + + // Create two mock matched functions + var match1 = new MatchedFunction(); + match1.setFunctionId(100L); + match1.setFunctionName("similar_func_1"); + match1.setMangledName("similar_func_1_mangled"); + match1.setBinaryName("library.so"); + match1.setSha256Hash("def456"); + match1.setDebug(true); + match1.setSimilarity(BigDecimal.valueOf(0.95)); + match1.setConfidence(BigDecimal.valueOf(0.90)); + + var match2 = new MatchedFunction(); + match2.setFunctionId(101L); + match2.setFunctionName("similar_func_2"); + match2.setMangledName("similar_func_2_mangled"); + match2.setBinaryName("other_lib.so"); + match2.setSha256Hash("ghi789"); + match2.setDebug(false); + match2.setSimilarity(BigDecimal.valueOf(0.85)); + match2.setConfidence(BigDecimal.valueOf(0.80)); + + functionMatch.setMatchedFunctions(List.of(match1, match2)); + response.setMatches(List.of(functionMatch)); + + return response; + } + + @Override + public List getAssembly(TypedApiInterface.FunctionID functionID) throws ApiException { + assemblyCallCount++; + + // Return mock assembly instructions + return List.of( + "push rbp", + "mov rbp, rsp", + "sub rsp, 0x20", + "mov [rbp-0x8], rdi", + "mov eax, [rbp-0x8]", + "add eax, 1", + "leave", + "ret" + ); + } + + @Override + public FunctionDataTypesList listFunctionDataTypesForFunctions(List functionIDs) { + // Return empty list - no signatures available in mock + var result = new FunctionDataTypesList(); + result.setItems(List.of()); + return result; + } + } +}