diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e67f884..e6380f9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -52,5 +52,11 @@ jobs: - name: Build run: ./gradlew buildExtension + - name: Check SLA timestamps + run: ls -l $GHIDRA_INSTALL_DIR/Ghidra/Processors/x86/data/languages/* + + - name: Precompile slaspec + run: $GHIDRA_INSTALL_DIR/support/sleigh $GHIDRA_INSTALL_DIR/Ghidra/Processors/x86/data/languages/x86-64.slaspec + - name: Run tests run: xvfb-run ./gradlew test --info diff --git a/build.gradle b/build.gradle index 5f92d5d..3903754 100644 --- a/build.gradle +++ b/build.gradle @@ -62,6 +62,7 @@ repositories { dependencies { // Any external dependencies added here will automatically be copied to the lib/ directory when // this extension is built. + implementation 'io.github.java-diff-utils:java-diff-utils:4.12' implementation 'org.json:json:20250107' implementation "com.google.guava:guava:33.2.0-jre" implementation group: 'com.fifesoft', name: 'rsyntaxtextarea', version: '3.5.2' diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/components/ItemSelectionPanel.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/components/ItemSelectionPanel.java index 44933f5..f09647b 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/components/ItemSelectionPanel.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/components/ItemSelectionPanel.java @@ -53,7 +53,7 @@ private void initializeUI(int minimumQueryLength) { selectedItemsPanel.setLayout(new WrapLayout(FlowLayout.LEFT)); JScrollPane scrollPane = new JScrollPane(selectedItemsPanel); - scrollPane.setPreferredSize(new Dimension(0, 100)); + scrollPane.setPreferredSize(new Dimension(0, 60)); scrollPane.setVerticalScrollBarPolicy(JScrollPane.VERTICAL_SCROLLBAR_AS_NEEDED); scrollPane.setHorizontalScrollBarPolicy(JScrollPane.HORIZONTAL_SCROLLBAR_NEVER); scrollPane.setBorder(null); // Remove any border from the scroll pane 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 43383aa..9b3ad24 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,9 +13,6 @@ import ai.reveng.toolkit.ghidra.core.services.api.types.GhidraFunctionMatchWithSignature; import com.google.common.collect.BiMap; import ghidra.program.model.listing.Function; -import ghidra.program.model.symbol.SourceType; -import ghidra.util.exception.DuplicateNameException; -import ghidra.util.exception.InvalidInputException; import ghidra.util.task.TaskMonitorComponent; import ghidra.util.Msg; import resources.ResourceManager; @@ -28,8 +25,6 @@ import java.util.Set; import java.util.stream.Collectors; -import static ai.reveng.toolkit.ghidra.plugins.BinarySimilarityPlugin.REVENG_AI_NAMESPACE; - public abstract class AbstractFunctionMatchingDialog extends RevEngDialogComponentProvider { protected final GhidraRevengService revengService; protected final GhidraRevengService.AnalysedProgram analyzedProgram; @@ -52,6 +47,9 @@ public abstract class AbstractFunctionMatchingDialog extends RevEngDialogCompone protected Timer pollTimer; protected JPanel renameButtonsPanel; + // Assembly comparison panel + protected AssemblyDiffPanel assemblyDiffPanel; + // Data protected Basic analysisBasicInfo; protected FunctionMatchingResponse functionMatchingResponse; @@ -61,26 +59,6 @@ public abstract class AbstractFunctionMatchingDialog extends RevEngDialogCompone // Polling configuration protected static final int POLL_INTERVAL_MS = 2000; // Poll every 2 seconds -// /// Inner class to hold function match results -// /// @deprecated {@link ai.reveng.toolkit.ghidra.core.services.api.types.FunctionMatch} -// public record FunctionMatchResult( -// String virtualAddress, -// String functionName, -// String bestMatchName, -// String bestMatchMangledName, -// String similarity, -// String confidence, -// String matchedHash, -// String binary, -// Long matcherFunctionId -// ) { -// // Constructor for function-level dialog (without virtual address and function name) -// public FunctionMatchResult(String bestMatchName, String bestMatchMangledName, String similarity, -// String confidence, String matchedHash, String binary, Long matcherFunctionId) { -// this("", "", bestMatchName, bestMatchMangledName, similarity, confidence, matchedHash, binary, matcherFunctionId); -// } -// } - protected AbstractFunctionMatchingDialog(String title, Boolean isModal, GhidraRevengService revengService, GhidraRevengService.AnalysedProgram analyzedProgram) { super(title, isModal); @@ -106,7 +84,7 @@ protected AbstractFunctionMatchingDialog(String title, Boolean isModal, GhidraRe addWorkPanel(buildMainPanel()); // Set dialog size to be wider - setPreferredSize(1000, 800); + setPreferredSize(1200, 1000); // Don't start function matching automatically - wait for user to click Match button statusLabel.setText("Ready - adjust filters and click 'Match Functions' to begin search"); @@ -521,7 +499,25 @@ protected JPanel createResultsContainer() { resultsTable.setAutoCreateRowSorter(true); resultsScrollPane = new JScrollPane(resultsTable); resultsScrollPane.setBorder(BorderFactory.createTitledBorder("Function Matching Results")); - resultsContainer.add(resultsScrollPane, BorderLayout.CENTER); + + // Add selection listener for assembly comparison + resultsTable.getSelectionModel().addListSelectionListener(e -> { + if (!e.getValueIsAdjusting()) { + onTableSelectionChanged(); + } + }); + + // Create assembly comparison panel + assemblyDiffPanel = new AssemblyDiffPanel(); + + // Create vertical split pane with results table on top and assembly comparison below + JSplitPane mainSplitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT); + mainSplitPane.setTopComponent(resultsScrollPane); + mainSplitPane.setBottomComponent(assemblyDiffPanel); + mainSplitPane.setResizeWeight(0.4); + mainSplitPane.setDividerLocation(250); + + resultsContainer.add(mainSplitPane, BorderLayout.CENTER); // Rename buttons panel renameButtonsPanel = createRenameButtonsPanel(); @@ -531,6 +527,29 @@ protected JPanel createResultsContainer() { return resultsContainer; } + protected void onTableSelectionChanged() { + int selectedRow = resultsTable.getSelectedRow(); + if (selectedRow < 0) { + assemblyDiffPanel.clear(); + return; + } + + // Convert view index to model index (in case table is sorted) + int modelRow = resultsTable.convertRowIndexToModel(selectedRow); + + // Get the appropriate results list + String filterText = functionFilterField != null ? functionFilterField.getText().trim() : ""; + List resultsToShow = filterText.isEmpty() ? + functionMatchResults : filteredFunctionMatchResults; + + if (modelRow >= resultsToShow.size()) { + return; + } + + GhidraFunctionMatchWithSignature selectedMatch = resultsToShow.get(modelRow); + assemblyDiffPanel.showAssemblyFor(selectedMatch, revengService); + } + protected JPanel createRenameButtonsPanel() { JPanel panel = new JPanel(new FlowLayout(FlowLayout.CENTER)); diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AssemblyDiffPanel.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AssemblyDiffPanel.java new file mode 100644 index 0000000..b51d410 --- /dev/null +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionmatching/AssemblyDiffPanel.java @@ -0,0 +1,324 @@ +package ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching; + +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.GhidraFunctionMatchWithSignature; +import com.github.difflib.DiffUtils; +import com.github.difflib.patch.AbstractDelta; +import com.github.difflib.patch.DeltaType; +import com.github.difflib.patch.Patch; +import ghidra.util.Msg; +import org.fife.ui.rsyntaxtextarea.RSyntaxTextArea; +import org.fife.ui.rsyntaxtextarea.SyntaxConstants; +import org.fife.ui.rtextarea.RTextScrollPane; + +import javax.swing.*; +import java.awt.*; +import java.util.ArrayList; +import java.util.List; + +/// A panel that displays side-by-side assembly comparison with diff highlighting. +/// Supports synchronized scrolling between the two panes. +/// The panel requires a [GhidraFunctionMatchWithSignature] so it can be easily extended into a general diffing panel +public class AssemblyDiffPanel extends JPanel { + private static final String PLACEHOLDER_TEXT = "Select a row to view assembly comparison"; + + // Diff highlighting colors + private static final Color DIFF_COLOR_CHANGED = new Color(255, 255, 180); // Light yellow + private static final Color DIFF_COLOR_DELETED = new Color(255, 200, 200); // Light red + private static final Color DIFF_COLOR_INSERTED = new Color(200, 255, 200); // Light green + private static final Color DIFF_COLOR_PLACEHOLDER = new Color(240, 240, 240); // Light gray + + private final RSyntaxTextArea localAssemblyTextArea; + private final RSyntaxTextArea matchedAssemblyTextArea; + private final RTextScrollPane localAssemblyScrollPane; + private final RTextScrollPane matchedAssemblyScrollPane; + + private boolean isSyncingScroll = false; + + public AssemblyDiffPanel() { + super(new BorderLayout()); + setBorder(BorderFactory.createTitledBorder("Assembly Comparison")); + + // Create text areas + localAssemblyTextArea = createAssemblyTextArea(); + matchedAssemblyTextArea = createAssemblyTextArea(); + + // Create scroll panes + localAssemblyScrollPane = new RTextScrollPane(localAssemblyTextArea); + localAssemblyScrollPane.setBorder(BorderFactory.createTitledBorder("Local Function Assembly")); + + matchedAssemblyScrollPane = new RTextScrollPane(matchedAssemblyTextArea); + matchedAssemblyScrollPane.setBorder(BorderFactory.createTitledBorder("Matched Function Assembly")); + + // Set up synchronized scrolling + setupSynchronizedScrolling(); + + // Create horizontal split pane for side-by-side comparison + JSplitPane splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT); + splitPane.setLeftComponent(localAssemblyScrollPane); + splitPane.setRightComponent(matchedAssemblyScrollPane); + splitPane.setResizeWeight(0.5); + splitPane.setDividerLocation(0.5); + + add(splitPane, BorderLayout.CENTER); + + // Set size to show approximately 10 lines of assembly by default + setPreferredSize(new Dimension(0, 300)); + setMinimumSize(new Dimension(0, 200)); + } + + private RSyntaxTextArea createAssemblyTextArea() { + RSyntaxTextArea textArea = new RSyntaxTextArea(); + textArea.setSyntaxEditingStyle(SyntaxConstants.SYNTAX_STYLE_ASSEMBLER_X86); + textArea.setCodeFoldingEnabled(false); + textArea.setEditable(false); + textArea.setLineWrap(false); + textArea.setHighlightCurrentLine(false); + textArea.setText(PLACEHOLDER_TEXT); + return textArea; + } + + private void setupSynchronizedScrolling() { + syncScrollBars(localAssemblyScrollPane.getVerticalScrollBar(), + matchedAssemblyScrollPane.getVerticalScrollBar()); + syncScrollBars(matchedAssemblyScrollPane.getVerticalScrollBar(), + localAssemblyScrollPane.getVerticalScrollBar()); + syncScrollBars(localAssemblyScrollPane.getHorizontalScrollBar(), + matchedAssemblyScrollPane.getHorizontalScrollBar()); + syncScrollBars(matchedAssemblyScrollPane.getHorizontalScrollBar(), + localAssemblyScrollPane.getHorizontalScrollBar()); + } + + private void syncScrollBars(JScrollBar source, JScrollBar target) { + source.addAdjustmentListener(e -> { + if (!isSyncingScroll) { + isSyncingScroll = true; + target.setValue(e.getValue()); + isSyncingScroll = false; + } + }); + } + + /** + * Clear the assembly panels and show placeholder text. + */ + public void clear() { + localAssemblyTextArea.removeAllLineHighlights(); + matchedAssemblyTextArea.removeAllLineHighlights(); + localAssemblyTextArea.setText(PLACEHOLDER_TEXT); + matchedAssemblyTextArea.setText(PLACEHOLDER_TEXT); + } + + /** + * Fetch and display assembly comparison for the given function match. + * Fetching is done in a background thread. + */ + public void showAssemblyFor(GhidraFunctionMatchWithSignature match, GhidraRevengService revengService) { + TypedApiInterface.FunctionID localFunctionId = match.functionMatch().origin_function_id(); + TypedApiInterface.FunctionID matchedFunctionId = match.functionMatch().nearest_neighbor_id(); + + // Clear existing highlights + localAssemblyTextArea.removeAllLineHighlights(); + matchedAssemblyTextArea.removeAllLineHighlights(); + + // Show loading state + localAssemblyTextArea.setText("Loading assembly for " + match.function().getName() + "..."); + matchedAssemblyTextArea.setText("Loading assembly for " + match.functionMatch().nearest_neighbor_function_name() + "..."); + + // Fetch assembly in background thread + new AssemblyFetchWorker(localFunctionId, matchedFunctionId, revengService).execute(); + } + + private void displayAssemblyWithDiff(List localAssembly, List matchedAssembly, Patch patch) { + // Build aligned content with placeholders for proper side-by-side comparison + List alignedLocal = new ArrayList<>(); + List alignedMatched = new ArrayList<>(); + List localLineTypes = new ArrayList<>(); + List matchedLineTypes = new ArrayList<>(); + + int localIndex = 0; + int matchedIndex = 0; + + // Process each delta in the patch + for (AbstractDelta delta : patch.getDeltas()) { + int sourceStart = delta.getSource().getPosition(); + int targetStart = delta.getTarget().getPosition(); + + // Add unchanged lines before this delta + while (localIndex < sourceStart && matchedIndex < targetStart) { + alignedLocal.add(localAssembly.get(localIndex++)); + alignedMatched.add(matchedAssembly.get(matchedIndex++)); + localLineTypes.add(DiffLineType.EQUAL); + matchedLineTypes.add(DiffLineType.EQUAL); + } + + // Handle the delta based on its type + List sourceLines = delta.getSource().getLines(); + List targetLines = delta.getTarget().getLines(); + + switch (delta.getType()) { + case DELETE -> { + for (String line : sourceLines) { + alignedLocal.add(line); + alignedMatched.add(""); + localLineTypes.add(DiffLineType.DELETED); + matchedLineTypes.add(DiffLineType.PLACEHOLDER); + } + localIndex += sourceLines.size(); + } + case INSERT -> { + for (String line : targetLines) { + alignedLocal.add(""); + alignedMatched.add(line); + localLineTypes.add(DiffLineType.PLACEHOLDER); + matchedLineTypes.add(DiffLineType.INSERTED); + } + matchedIndex += targetLines.size(); + } + case CHANGE -> { + int maxLines = Math.max(sourceLines.size(), targetLines.size()); + for (int i = 0; i < maxLines; i++) { + alignedLocal.add(i < sourceLines.size() ? sourceLines.get(i) : ""); + localLineTypes.add(i < sourceLines.size() ? DiffLineType.CHANGED : DiffLineType.PLACEHOLDER); + alignedMatched.add(i < targetLines.size() ? targetLines.get(i) : ""); + matchedLineTypes.add(i < targetLines.size() ? DiffLineType.CHANGED : DiffLineType.PLACEHOLDER); + } + localIndex += sourceLines.size(); + matchedIndex += targetLines.size(); + } + default -> {} + } + } + + // Add remaining unchanged lines after the last delta + while (localIndex < localAssembly.size() && matchedIndex < matchedAssembly.size()) { + alignedLocal.add(localAssembly.get(localIndex++)); + alignedMatched.add(matchedAssembly.get(matchedIndex++)); + localLineTypes.add(DiffLineType.EQUAL); + matchedLineTypes.add(DiffLineType.EQUAL); + } + + // Handle any trailing lines + while (localIndex < localAssembly.size()) { + alignedLocal.add(localAssembly.get(localIndex++)); + alignedMatched.add(""); + localLineTypes.add(DiffLineType.DELETED); + matchedLineTypes.add(DiffLineType.PLACEHOLDER); + } + while (matchedIndex < matchedAssembly.size()) { + alignedLocal.add(""); + alignedMatched.add(matchedAssembly.get(matchedIndex++)); + localLineTypes.add(DiffLineType.PLACEHOLDER); + matchedLineTypes.add(DiffLineType.INSERTED); + } + + // Set text content + localAssemblyTextArea.setText(String.join("\n", alignedLocal)); + matchedAssemblyTextArea.setText(String.join("\n", alignedMatched)); + localAssemblyTextArea.setCaretPosition(0); + matchedAssemblyTextArea.setCaretPosition(0); + + // Apply line highlights + applyDiffHighlights(localAssemblyTextArea, localLineTypes); + applyDiffHighlights(matchedAssemblyTextArea, matchedLineTypes); + } + + private void applyDiffHighlights(RSyntaxTextArea textArea, List lineTypes) { + textArea.removeAllLineHighlights(); + + for (int i = 0; i < lineTypes.size(); i++) { + Color highlightColor = getColorForLineType(lineTypes.get(i)); + if (highlightColor != null) { + try { + textArea.addLineHighlight(i, highlightColor); + } catch (Exception e) { + // Line index out of bounds - ignore + } + } + } + } + + private Color getColorForLineType(DiffLineType lineType) { + return switch (lineType) { + case DELETED -> DIFF_COLOR_DELETED; + case INSERTED -> DIFF_COLOR_INSERTED; + case CHANGED -> DIFF_COLOR_CHANGED; + case PLACEHOLDER -> DIFF_COLOR_PLACEHOLDER; + case EQUAL -> null; + }; + } + + private enum DiffLineType { + EQUAL, + DELETED, + INSERTED, + CHANGED, + PLACEHOLDER + } + + /** + * SwingWorker to fetch assembly in the background and update the UI when done. + */ + private class AssemblyFetchWorker extends SwingWorker { + private final TypedApiInterface.FunctionID localFunctionId; + private final TypedApiInterface.FunctionID matchedFunctionId; + private final GhidraRevengService revengService; + + private List localAssembly; + private List matchedAssembly; + private String localError; + private String matchedError; + private Patch patch; + + AssemblyFetchWorker(TypedApiInterface.FunctionID localFunctionId, + TypedApiInterface.FunctionID matchedFunctionId, + GhidraRevengService revengService) { + this.localFunctionId = localFunctionId; + this.matchedFunctionId = matchedFunctionId; + this.revengService = revengService; + } + + @Override + protected Void doInBackground() { + // Fetch local function assembly + try { + localAssembly = revengService.getApi().getAssembly(localFunctionId); + } catch (Exception e) { + localError = "Failed to fetch assembly: " + e.getMessage(); + Msg.error(this, "Failed to fetch local function assembly", e); + } + + // Fetch matched function assembly + try { + matchedAssembly = revengService.getApi().getAssembly(matchedFunctionId); + } catch (Exception e) { + matchedError = "Failed to fetch assembly: " + e.getMessage(); + Msg.error(this, "Failed to fetch matched function assembly", e); + } + + // Compute diff if both assemblies were fetched successfully + if (localAssembly != null && matchedAssembly != null) { + patch = DiffUtils.diff(localAssembly, matchedAssembly); + } + + return null; + } + + @Override + protected void done() { + if (localError != null) { + localAssemblyTextArea.setText(localError); + return; + } + if (matchedError != null) { + matchedAssemblyTextArea.setText(matchedError); + return; + } + + if (localAssembly != null && matchedAssembly != null) { + displayAssemblyWithDiff(localAssembly, matchedAssembly, patch); + } + } + } +} \ No newline at end of file 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 b8fc0ae..976e426 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 @@ -639,5 +639,25 @@ public FunctionMatchingResponse functionFunctionMatching(FunctionMatchingRequest public void batchRenameFunctions(FunctionsListRename functionsList) throws ApiException { this.functionsRenamingHistoryApi.batchRenameFunction(functionsList); } + + @Override + public List getAssembly(FunctionID id) { + + FunctionBlocksResponse blocks; + List result = new ArrayList<>(); + try { + blocks = this.functionsCoreApi.getFunctionBlocks(id.asInteger()).getData(); + } catch (ApiException e) { + throw new RuntimeException(e); + } + blocks.getBlocks().stream() + .sorted( (b1, b2) -> b1.getMinAddr().compareTo(b2.getMinAddr()) ) + .forEach(block -> { + result.addAll(block.getAsm()); + }); + + return result; + } + } 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 e7b1250..8d47363 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 @@ -29,7 +29,11 @@ */ public interface TypedApiInterface { /// Data type to represent the RevEng.AI API concept of a function ID - record FunctionID(long value){} + record FunctionID(long value){ + public Integer asInteger() { + return Math.toIntExact(value); + } + } /// This is a special box type for an analysis ID /// It enforces that the integer is specifically an analysis ID, @@ -216,5 +220,10 @@ default FunctionMatchingResponse functionFunctionMatching(FunctionMatchingReques default void batchRenameFunctions(FunctionsListRename functionsList) throws ApiException { throw new UnsupportedOperationException("batchRenameFunctions not implemented yet"); } + + default List getAssembly(FunctionID functionID) throws ApiException { + throw new UnsupportedOperationException("getAssembly not implemented yet"); + } + } diff --git a/src/test/java/ai/reveng/AIDecompilerComponentTest.java b/src/test/java/ai/reveng/AIDecompilerComponentTest.java index 7f7ec1f..fa29aa7 100644 --- a/src/test/java/ai/reveng/AIDecompilerComponentTest.java +++ b/src/test/java/ai/reveng/AIDecompilerComponentTest.java @@ -99,7 +99,7 @@ public boolean triggerAIDecompilationForFunctionID(FunctionID functionID) { var binarySimilarityPlugin = env.addPlugin(BinarySimilarityPlugin.class); - var builder = new ProgramBuilder("mock", ProgramBuilder._8051, this); + 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)); @@ -157,7 +157,7 @@ public void testAIDecompFeedbackMechanism() throws Exception { var service = addMockedService(tool, ratingsAPI); var binarySimilarityPlugin = env.addPlugin(BinarySimilarityPlugin.class); - var builder = new ProgramBuilder("mock", ProgramBuilder._8051, this); + 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)); diff --git a/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java b/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java index 8693fec..f5e2d6b 100644 --- a/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java +++ b/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java @@ -44,7 +44,7 @@ public AnalysisOptionsDialogTest() { public void testWithMockModels() throws Exception { var reService = new GhidraRevengService( new MockApi() {}); - var builder = new ProgramBuilder("mock", ProgramBuilder._8051, this); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); var program = builder.getProgram(); var dialog = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService); diff --git a/src/test/java/ai/reveng/FunctionLevelFunctionMatchingDialogTest.java b/src/test/java/ai/reveng/FunctionLevelFunctionMatchingDialogTest.java new file mode 100644 index 0000000..a46cccf --- /dev/null +++ b/src/test/java/ai/reveng/FunctionLevelFunctionMatchingDialogTest.java @@ -0,0 +1,356 @@ +package ai.reveng; + +import ai.reveng.model.*; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionmatching.FunctionLevelFunctionMatchingDialog; +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.plugins.BinarySimilarityPlugin; +import docking.DockingWindowManager; +import ghidra.program.database.ProgramBuilder; +import ghidra.program.model.data.Undefined; +import ghidra.util.task.TaskMonitor; +import org.junit.Test; + +import javax.swing.*; +import java.math.BigDecimal; +import java.util.List; + +import static org.junit.Assert.*; + +/** + * Integration tests for the FunctionLevelFunctionMatchingDialog. + * Tests dialog opening, display, and interaction with mocked API responses. + */ +public class FunctionLevelFunctionMatchingDialogTest extends RevEngMockableHeadedIntegrationTest { + + @Test + public void testDialogOpensWithMockedService() throws Exception { + var tool = env.getTool(); + + // Create a mock API that returns the necessary data for the dialog + var mockApi = new FunctionMatchingMockApi(); + var service = addMockedService(tool, mockApi); + + // Add the BinarySimilarity plugin which provides the dialog actions + env.addPlugin(BinarySimilarityPlugin.class); + + // Create a test program with a function + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + var testFunction = builder.createEmptyFunction("test_function", "0x1000", 50, Undefined.getUndefinedDataType(4)); + + // Register the program as analyzed (this triggers associateFunctionInfo internally) + var analysedProgram = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + + // Show the tool with the program + env.showTool(analysedProgram.program()); + waitForSwing(); + + // Create and show the dialog + FunctionLevelFunctionMatchingDialog dialog = runSwing(() -> + new FunctionLevelFunctionMatchingDialog(tool, analysedProgram, testFunction) + ); + + // Show dialog without blocking, then wait for it to appear + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(FunctionLevelFunctionMatchingDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Verify the dialog is visible + assertTrue("Dialog should be visible", foundDialog.isVisible()); + + // Verify the dialog has the correct title + assertTrue("Dialog title should contain 'Function Matching'", + foundDialog.getTitle().contains("Function Matching")); + + // Get internal components for verification + JTable resultsTable = (JTable) getInstanceField("resultsTable", foundDialog); + assertNotNull("Results table should exist", resultsTable); + + JSlider thresholdSlider = (JSlider) getInstanceField("thresholdSlider", foundDialog); + assertNotNull("Threshold slider should exist", thresholdSlider); + assertEquals("Default threshold should be 70", 70, thresholdSlider.getValue()); + + // Close the dialog + close(foundDialog); + waitForSwing(); + } + + @Test + public void testDialogHasResultsTableConfigured() throws Exception { + var tool = env.getTool(); + + // Create a mock API with function matching results + var mockApi = new FunctionMatchingMockApi(); + var service = addMockedService(tool, mockApi); + + env.addPlugin(BinarySimilarityPlugin.class); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + var testFunction = builder.createEmptyFunction("test_function", "0x1000", 50, Undefined.getUndefinedDataType(4)); + + var analysedProgram = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(analysedProgram.program()); + waitForSwing(); + + // Create the dialog + FunctionLevelFunctionMatchingDialog dialog = runSwing(() -> + new FunctionLevelFunctionMatchingDialog(tool, analysedProgram, testFunction) + ); + + // Show dialog without blocking, then wait for it to appear + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(FunctionLevelFunctionMatchingDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Verify the table exists and is configured for sorting + JTable resultsTable = (JTable) getInstanceField("resultsTable", foundDialog); + assertNotNull("Results table should exist", resultsTable); + assertTrue("Results table should have auto row sorter enabled", resultsTable.getAutoCreateRowSorter()); + + // Close the dialog + close(foundDialog); + waitForSwing(); + } + + @Test + public void testAssemblyComparisonPanelExists() throws Exception { + var tool = env.getTool(); + + var mockApi = new FunctionMatchingMockApi(); + var service = addMockedService(tool, mockApi); + + env.addPlugin(BinarySimilarityPlugin.class); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + var testFunction = builder.createEmptyFunction("test_function", "0x1000", 50, Undefined.getUndefinedDataType(4)); + + var analysedProgram = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(analysedProgram.program()); + waitForSwing(); + + FunctionLevelFunctionMatchingDialog dialog = runSwing(() -> + new FunctionLevelFunctionMatchingDialog(tool, analysedProgram, testFunction) + ); + + // Show dialog without blocking, then wait for it to appear + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(FunctionLevelFunctionMatchingDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Verify assembly diff panel exists + var assemblyDiffPanel = getInstanceField("assemblyDiffPanel", foundDialog); + assertNotNull("Assembly diff panel should exist", assemblyDiffPanel); + + // Close the dialog + close(foundDialog); + waitForSwing(); + } + + @Test + public void testFunctionMatchingTriggersAPICall() throws Exception { + var tool = env.getTool(); + + var mockApi = new FunctionMatchingMockApi(); + var service = addMockedService(tool, mockApi); + + env.addPlugin(BinarySimilarityPlugin.class); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + var testFunction = builder.createEmptyFunction("test_function", "0x1000", 50, Undefined.getUndefinedDataType(4)); + + var analysedProgram = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(analysedProgram.program()); + waitForSwing(); + + FunctionLevelFunctionMatchingDialog dialog = runSwing(() -> + new FunctionLevelFunctionMatchingDialog(tool, analysedProgram, testFunction) + ); + + // Show dialog without blocking, then wait for it to appear + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(FunctionLevelFunctionMatchingDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Verify the initial status message + JLabel statusLabel = (JLabel) getInstanceField("statusLabel", foundDialog); + assertNotNull("Status label should exist", statusLabel); + assertTrue("Status should show ready message", + statusLabel.getText().contains("Ready")); + + // Close the dialog + close(foundDialog); + waitForSwing(); + } + + @Test + public void testClickMatchButtonPopulatesResultsTable() throws Exception { + var tool = env.getTool(); + + var mockApi = new FunctionMatchingMockApi(); + var service = addMockedService(tool, mockApi); + + env.addPlugin(BinarySimilarityPlugin.class); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + var testFunction = builder.createEmptyFunction("test_function", "0x1000", 50, Undefined.getUndefinedDataType(4)); + + var analysedProgram = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); + env.showTool(analysedProgram.program()); + waitForSwing(); + + FunctionLevelFunctionMatchingDialog dialog = runSwing(() -> + new FunctionLevelFunctionMatchingDialog(tool, analysedProgram, testFunction) + ); + + // Show dialog without blocking, then wait for it to appear + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(FunctionLevelFunctionMatchingDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Get the results table before clicking + JTable resultsTable = (JTable) getInstanceField("resultsTable", foundDialog); + assertNotNull("Results table should exist", resultsTable); + assertEquals("Results table should be empty initially", 0, resultsTable.getRowCount()); + + // Find and click the "Match Functions" button + JButton matchButton = findButtonByText(foundDialog.getComponent(), "Match Functions"); + assertNotNull("Match Functions button should exist", matchButton); + + pressButton(matchButton); + + // Wait for the API call to complete and results to be processed + // The mock API returns COMPLETED status immediately, so we just need to wait for Swing + waitForTasks(); + waitForSwing(); + + // Wait for the table to be populated (poll with timeout) + waitForCondition(() -> resultsTable.getRowCount() > 0, + "Results table should have rows after matching"); + + // Verify the table has the expected data from our mock + assertTrue("Results table should have at least one row", resultsTable.getRowCount() > 0); + + // Verify the API was actually called + assertTrue("Function matching API should have been called", mockApi.functionMatchingCalled); + + // Close the dialog + close(foundDialog); + waitForSwing(); + } + + /** + * Mock API implementation for function matching dialog tests. + * Provides necessary responses for the dialog to function without a real server. + */ + static class FunctionMatchingMockApi extends UnimplementedAPI { + private final AnalysisStatus currentStatus = AnalysisStatus.Complete; + boolean functionMatchingCalled = false; + + @Override + public TypedApiInterface.AnalysisID analyse(AnalysisOptionsBuilder options) { + return new TypedApiInterface.AnalysisID(12345); + } + + @Override + public AnalysisStatus status(TypedApiInterface.AnalysisID analysisID) { + return currentStatus; + } + + @Override + public List getFunctionInfo(TypedApiInterface.AnalysisID analysisID) { + // Return function info matching the test function we create + return List.of( + new FunctionInfo( + new TypedApiInterface.FunctionID(100), + "test_function", + "test_function", + 0x1000L, + 50 + ) + ); + } + + @Override + public Basic getAnalysisBasicInfo(TypedApiInterface.AnalysisID analysisID) { + // Create a Basic object with required fields + var basic = new Basic(); + basic.setModelId(1); + basic.setModelName("test-model"); + basic.setBinaryName("test_binary"); + basic.setSha256Hash("0".repeat(64)); + return basic; + } + + @Override + public FunctionMatchingResponse functionFunctionMatching(FunctionMatchingRequest request) { + functionMatchingCalled = true; + + // Create a response with completed status and mock matches + var response = new FunctionMatchingResponse(); + response.setStatus("COMPLETED"); + response.setProgress(100); + + // Create a match result using the SDK model types + var functionMatch = new ai.reveng.model.FunctionMatch(); + functionMatch.setFunctionId(100L); + + // Create a matched function + var matchedFunc = new MatchedFunction(); + matchedFunc.setFunctionId(200L); + matchedFunc.setFunctionName("similar_function"); + matchedFunc.setMangledName("similar_function"); + matchedFunc.setSha256Hash("1".repeat(64)); + matchedFunc.setBinaryName("libc.so"); + matchedFunc.setBinaryId(1); + matchedFunc.setFunctionVaddr(0x2000L); + matchedFunc.setAnalysisId(12345); + matchedFunc.setDebug(false); + matchedFunc.setSimilarity(BigDecimal.valueOf(0.95)); + matchedFunc.setConfidence(BigDecimal.valueOf(0.87)); + + functionMatch.setMatchedFunctions(List.of(matchedFunc)); + response.setMatches(List.of(functionMatch)); + + return response; + } + + @Override + public List getAssembly(TypedApiInterface.FunctionID functionID) { + // Return some mock assembly for the diff display + if (functionID.value() == 100) { + // Local function assembly + return List.of( + "push rbp", + "mov rbp, rsp", + "sub rsp, 0x20", + "mov [rbp-0x8], rdi", + "call 0x1234", + "leave", + "ret" + ); + } else if (functionID.value() == 200) { + // Matched function assembly (slightly different) + return List.of( + "push rbp", + "mov rbp, rsp", + "sub rsp, 0x30", + "mov [rbp-0x8], rdi", + "mov [rbp-0x10], rsi", + "call 0x5678", + "leave", + "ret" + ); + } + return List.of(); + } + + @Override + public FunctionDataTypesList listFunctionDataTypesForFunctions(List functionIDs) { + // Return empty list - no type info available + return new FunctionDataTypesList(); + } + } +} \ No newline at end of file diff --git a/src/test/java/ai/reveng/PortalAnalysisIntegrationTest.java b/src/test/java/ai/reveng/PortalAnalysisIntegrationTest.java index 6e22f6d..6b918b7 100644 --- a/src/test/java/ai/reveng/PortalAnalysisIntegrationTest.java +++ b/src/test/java/ai/reveng/PortalAnalysisIntegrationTest.java @@ -135,7 +135,7 @@ public AnalysisID analyse(AnalysisOptionsBuilder options) throws ApiException { return new AnalysisID(1); } }); - var builder = new ProgramBuilder("mock", ProgramBuilder._8051, this); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); // Add an example function var exampleFunc = builder.createEmptyFunction(null, "0x4000", 0x100, Undefined.getUndefinedDataType(8)); /// Tell Ghidra that the function signature source is just default, diff --git a/src/test/java/ai/reveng/TestAnalysisLogComponent.java b/src/test/java/ai/reveng/TestAnalysisLogComponent.java index 51eb196..0e65ee6 100644 --- a/src/test/java/ai/reveng/TestAnalysisLogComponent.java +++ b/src/test/java/ai/reveng/TestAnalysisLogComponent.java @@ -22,7 +22,7 @@ public class TestAnalysisLogComponent extends RevEngMockableHeadedIntegrationTest { private GhidraRevengService.ProgramWithID getPlaceHolderID() throws Exception{ - var builder = new ghidra.program.database.ProgramBuilder("mock", ProgramBuilder._8051, this); + var builder = new ghidra.program.database.ProgramBuilder("mock", ProgramBuilder._X64, this); // Add an example function var program = builder.getProgram(); return new GhidraRevengService.ProgramWithID( diff --git a/src/test/java/ai/reveng/TestUpgradeFromBinaryID.java b/src/test/java/ai/reveng/TestUpgradeFromBinaryID.java index 7294458..705ae2a 100644 --- a/src/test/java/ai/reveng/TestUpgradeFromBinaryID.java +++ b/src/test/java/ai/reveng/TestUpgradeFromBinaryID.java @@ -18,7 +18,7 @@ public class TestUpgradeFromBinaryID extends RevEngMockableHeadedIntegrationTest /// Tests the logic for handling a program that has only a binary ID stored in its properties @Test public void test() throws Exception { - var builder = new ProgramBuilder("upgrade-test", ProgramBuilder._8051, this); + var builder = new ProgramBuilder("upgrade-test", ProgramBuilder._X64, this); var tId = builder.getProgram().startTransaction("Set Binary ID"); builder.getProgram() .getOptions(ReaiPluginPackage.REAI_OPTIONS_CATEGORY) diff --git a/src/test/java/ai/reveng/UnstripTest.java b/src/test/java/ai/reveng/UnstripTest.java index b48f8ec..020a389 100644 --- a/src/test/java/ai/reveng/UnstripTest.java +++ b/src/test/java/ai/reveng/UnstripTest.java @@ -71,7 +71,7 @@ public List getFunctionInfo(AnalysisID analysisID) { }); addPlugin(tool, BinarySimilarityPlugin.class); - var builder = new ProgramBuilder("mock", ProgramBuilder._8051, this); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); var func = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); builder.createMemory("test", "01000", 0x100); var programWithID = service.analyse(builder.getProgram(), null, TaskMonitor.DUMMY); @@ -160,7 +160,7 @@ public FunctionDataTypesList listFunctionDataTypesForAnalysis(AnalysisID analysi }); addPlugin(tool, BinarySimilarityPlugin.class); - var builder = new ProgramBuilder("mock", ProgramBuilder._8051, this); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); // We provide no function name, so ghidra will assign the default "FUN_1000" name var func = builder.createEmptyFunction(null, "0x1000", 10, Undefined.getUndefinedDataType(4)); builder.createMemory("test", "01000", 0x100);