Skip to content

Commit 0d861b4

Browse files
fmaginfmagin
committed
Support filtering functions for analysis creation and attaching
Filter out functions without instructions Update SDK and use new includeParameter for function boundaries
1 parent a291e93 commit 0d861b4

23 files changed

Lines changed: 1829 additions & 391 deletions

‎build.gradle‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ dependencies {
6666
implementation 'org.json:json:20250107'
6767
implementation "com.google.guava:guava:33.2.0-jre"
6868
implementation group: 'com.fifesoft', name: 'rsyntaxtextarea', version: '3.5.2'
69-
implementation('ai.reveng:sdk:3.0.0')
69+
implementation('ai.reveng:sdk:3.45.0')
7070
testImplementation('junit:junit:4.13.1')
7171
testRuntimeOnly("org.junit.vintage:junit-vintage-engine:5.8.2")
7272

‎src/main/java/ai/reveng/toolkit/ghidra/Utils.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,43 @@ public static <ROW_TYPE, COLUMN_TYPE> void addRowToDescriptor(
5555
addRowToDescriptor(descriptor, columnName, true, columnTypeClass, rowObjectAccessor);
5656
}
5757

58+
/**
59+
* Helper method to add a column with sort ordinal specification.
60+
* @param sortOrdinal 1-based sort priority (1 = primary sort), or -1 for no default sort
61+
* @param ascending true for ascending sort, false for descending
62+
*/
63+
public static <ROW_TYPE, COLUMN_TYPE> void addRowToDescriptor(
64+
TableColumnDescriptor<ROW_TYPE> descriptor,
65+
String columnName,
66+
Class<COLUMN_TYPE> columnTypeClass,
67+
RowObjectAccessor<ROW_TYPE, COLUMN_TYPE> rowObjectAccessor,
68+
int sortOrdinal,
69+
boolean ascending) {
70+
71+
var column = new AbstractDynamicTableColumn<ROW_TYPE, COLUMN_TYPE, Object>() {
72+
@Override
73+
public String getColumnName() {
74+
return columnName;
75+
}
76+
77+
@Override
78+
public COLUMN_TYPE getValue(ROW_TYPE rowObject, Settings settings, Object data, ServiceProvider serviceProvider) throws IllegalArgumentException {
79+
return rowObjectAccessor.access(rowObject);
80+
}
81+
82+
@Override
83+
public Class<COLUMN_TYPE> getColumnClass() {
84+
return columnTypeClass;
85+
}
86+
87+
@Override
88+
public Class<ROW_TYPE> getSupportedRowType() {
89+
return null;
90+
}
91+
};
92+
descriptor.addVisibleColumn(column, sortOrdinal, ascending);
93+
}
94+
5895

5996

6097
@FunctionalInterface

‎src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,9 +48,7 @@ private boolean shouldApplyMatch() {
4848
return func != null &&
4949
// Do not override user-defined function names
5050
func.getSymbol().getSource() != SourceType.USER_DEFINED &&
51-
// Exclude thunks and external functions
52-
!func.isThunk() &&
53-
!func.isExternal() &&
51+
GhidraRevengService.isRelevantForAnalysis(func) &&
5452
// Only accept valid names (no spaces)
5553
!match.functionMatch().nearest_neighbor_mangled_function_name().contains(" ") &&
5654
!match.functionMatch().nearest_neighbor_function_name().contains(" ")

‎src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/aidecompiler/AIDecompilationdWindow.java‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package ai.reveng.toolkit.ghidra.binarysimilarity.ui.aidecompiler;
22

33
import ai.reveng.invoker.ApiException;
4+
import ai.reveng.model.AiDecompilationTaskStatus;
45
import ai.reveng.model.GetAiDecompilationTask;
56
import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService;
67
import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface;
@@ -207,7 +208,7 @@ public JComponent getComponent() {
207208

208209
public void setDisplayedValuesBasedOnStatus(Function function, GetAiDecompilationTask status) {
209210
this.function = function;
210-
if (status.getStatus().equals("success")) {
211+
if (status.getStatus() == AiDecompilationTaskStatus.SUCCESS) {
211212
setCode(status.getDecompilation());
212213
descriptionArea.setText("<html>%s</html>".formatted(status.getSummary()));
213214

@@ -219,7 +220,7 @@ public void setDisplayedValuesBasedOnStatus(Function function, GetAiDecompilatio
219220
} else {
220221
predictedNamePanel.setVisible(false);
221222
}
222-
} else if (status.getStatus().equals("error")) {
223+
} else if (status.getStatus() == AiDecompilationTaskStatus.ERROR) {
223224
setCode("");
224225
descriptionArea.setText("Decompilation failed");
225226
predictedNamePanel.setVisible(false);
@@ -286,21 +287,21 @@ void newStatusForFunction(Function function, GetAiDecompilationTask status) {
286287
setDisplayedValuesBasedOnStatus(function, status)
287288
);
288289
}
289-
if (status.getStatus().equals("success")) {
290+
if (status.getStatus() == AiDecompilationTaskStatus.SUCCESS) {
290291
var logger = tool.getService(ReaiLoggingService.class);
291292
logger.info("AI Decompilation finished for function %s: %s".formatted(function.getName(), status.getDecompilation()));
292293
if (!hasPendingDecompilations()) {
293294
taskMonitorComponent.setVisible(false);
294295
}
295-
} else if (status.getStatus().equals("error")) {
296+
} else if (status.getStatus() == AiDecompilationTaskStatus.ERROR) {
296297
if (!hasPendingDecompilations()) {
297298
taskMonitorComponent.setVisible(false);
298299
}
299300
}
300301
}
301302

302303
private boolean hasPendingDecompilations() {
303-
return cache.values().stream().anyMatch(s -> s.getStatus().equals("pending") || s.getStatus().equals("running") || s.getStatus().equals("queued"));
304+
return cache.values().stream().anyMatch(s -> s.getStatus() == AiDecompilationTaskStatus.PENDING);
304305
}
305306
class AIDecompTask extends Task {
306307

@@ -317,7 +318,7 @@ public AIDecompTask(PluginTool tool, GhidraRevengService.FunctionWithID function
317318
public void run(TaskMonitor monitor) throws CancelledException {
318319
var fID = functionWithID.functionID();
319320
// Check if there is an existing process already, because the trigger API will fail with 400 if there is
320-
if (service.getApi().pollAIDecompileStatus(fID).getStatus().equals("uninitialised")) {
321+
if (service.getApi().pollAIDecompileStatus(fID).getStatus() == AiDecompilationTaskStatus.UNINITIALISED) {
321322
// Trigger the decompilation
322323
service.getApi().triggerAIDecompilationForFunctionID(fID);
323324
}
@@ -340,21 +341,19 @@ private void waitForDecomp(TypedApiInterface.FunctionID id, TaskMonitor monitor)
340341
monitor.setMessage("Waiting for AI Decompilation for %s ... Current status: %s".formatted(functionWithID.function().getName(), lastDecompStatus.getStatus()));
341342
monitor.checkCancelled();
342343
switch (newStatus.getStatus()) {
343-
case "pending":
344-
case "uninitialised":
345-
case "queued":
346-
case "running":
344+
case PENDING:
345+
case UNINITIALISED:
347346
try {
348347
// Wait a second before polling again. We don't want to spam the API with requests too often
349348
Thread.sleep(1000);
350349
} catch (InterruptedException e) {
351350
throw new RuntimeException(e);
352351
}
353352
break;
354-
case "success":
353+
case SUCCESS:
355354
monitor.setProgress(monitor.getMaximum());
356355
return;
357-
case "error":
356+
case ERROR:
358357
logger.error("Decompilation failed: %s".formatted(newStatus.getDecompilation()));
359358
return;
360359
default:

‎src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/analysiscreation/RevEngAIAnalysisOptionsDialog.java‎

Lines changed: 40 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22

33
import ai.reveng.model.ConfigResponse;
44
import ai.reveng.toolkit.ghidra.binarysimilarity.ui.dialog.RevEngDialogComponentProvider;
5+
import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection.FunctionSelectionPanel;
56
import ai.reveng.toolkit.ghidra.core.services.api.AnalysisOptionsBuilder;
67
import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService;
78
import ai.reveng.toolkit.ghidra.core.services.api.types.AnalysisScope;
89
import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage;
10+
import ghidra.framework.plugintool.PluginTool;
911
import ghidra.program.model.listing.Program;
1012
import ghidra.util.Msg;
1113
import ghidra.util.Swing;
@@ -25,6 +27,7 @@ public class RevEngAIAnalysisOptionsDialog extends RevEngDialogComponentProvider
2527
private JCheckBox dynamicExecutionCheckBox;
2628
private final Program program;
2729
private final GhidraRevengService service;
30+
private final PluginTool tool;
2831
private JRadioButton privateScope;
2932
private JRadioButton publicScope;
3033
private JTextField tagsTextBox;
@@ -33,22 +36,25 @@ public class RevEngAIAnalysisOptionsDialog extends RevEngDialogComponentProvider
3336
private JCheckBox identifyCVECheckBox;
3437
private JCheckBox generateSBOMCheckBox;
3538
private JComboBox<String> architectureComboBox;
39+
private FunctionSelectionPanel functionSelectionPanel;
3640
private boolean okPressed = false;
41+
private boolean configCheckPassed = false;
3742

3843
private JLabel fileSizeWarningLabel;
3944
private JLabel loadingLabel;
4045

41-
public static RevEngAIAnalysisOptionsDialog withModelsFromServer(Program program, GhidraRevengService reService) {
42-
return new RevEngAIAnalysisOptionsDialog(program, reService);
46+
public static RevEngAIAnalysisOptionsDialog withModelsFromServer(Program program, GhidraRevengService reService, PluginTool tool) {
47+
return new RevEngAIAnalysisOptionsDialog(program, tool, reService);
4348
}
4449

45-
public RevEngAIAnalysisOptionsDialog(Program program, GhidraRevengService service) {
50+
public RevEngAIAnalysisOptionsDialog(Program program, PluginTool tool, GhidraRevengService service) {
4651
super(ReaiPluginPackage.WINDOW_PREFIX + "Configure Analysis for %s".formatted(program.getName()), true);
4752
this.program = program;
4853
this.service = service;
54+
this.tool = tool;
4955

5056
buildInterface();
51-
setPreferredSize(320, 420);
57+
setPreferredSize(600, 550);
5258

5359
fetchConfigAsync();
5460
}
@@ -176,18 +182,26 @@ private void buildInterface() {
176182
loadingLabel.setBorder(BorderFactory.createEmptyBorder(5, 5, 5, 5));
177183
workPanel.add(loadingLabel);
178184

185+
workPanel.add(new JSeparator(SwingConstants.HORIZONTAL));
186+
187+
// Add function selection panel
188+
functionSelectionPanel = new FunctionSelectionPanel(tool);
189+
functionSelectionPanel.initForProgram(program);
190+
functionSelectionPanel.getTableModel().addTableModelListener(e -> updateStartButtonState());
191+
workPanel.add(functionSelectionPanel);
192+
193+
179194
addCancelButton();
180195
addOKButton();
181196

182197
okButton.setText("Start Analysis");
183-
okButton.setEnabled(false); // Disabled until config check completes
198+
okButton.setEnabled(false); // Disabled until config check completes and functions are selected
184199
}
185200

186-
public @Nullable AnalysisOptionsBuilder getOptionsFromUI() {
187-
if (!okPressed) {
188-
return null;
189-
}
190-
var options = AnalysisOptionsBuilder.forProgram(program);
201+
public AnalysisOptionsBuilder getOptionsFromUI() {
202+
// Use the selected functions from the function selection panel
203+
var selectedFunctions = functionSelectionPanel.getSelectedFunctions();
204+
var options = AnalysisOptionsBuilder.forProgramWithFunctions(program, selectedFunctions);
191205

192206
options.skipScraping(!scrapeExternalTagsBox.isSelected());
193207
options.skipCapabilities(!identifyCapabilitiesCheckBox.isSelected());
@@ -216,6 +230,10 @@ protected void okCallback() {
216230
close();
217231
}
218232

233+
public boolean isOkPressed() {
234+
return okPressed;
235+
}
236+
219237
@Override
220238
public JComponent getComponent() {
221239
return super.getComponent();
@@ -239,7 +257,8 @@ private void handleConfigResponse(@Nullable ConfigResponse config) {
239257

240258
if (config == null) {
241259
// Config fetch failed, allow upload attempt (server will reject if too large)
242-
okButton.setEnabled(true);
260+
configCheckPassed = true;
261+
updateStartButtonState();
243262
return;
244263
}
245264

@@ -251,7 +270,8 @@ private void validateFileSize(long maxFileSizeBytes) {
251270
long fileSize = getProgramFileSize();
252271
if (fileSize < 0) {
253272
// Could not determine file size, allow upload attempt
254-
okButton.setEnabled(true);
273+
configCheckPassed = true;
274+
updateStartButtonState();
255275
return;
256276
}
257277

@@ -262,13 +282,19 @@ private void validateFileSize(long maxFileSizeBytes) {
262282
"<html><center>File size (%s) exceeds<br>server limit (%s)</center></html>"
263283
.formatted(fileSizeStr, maxSizeStr));
264284
fileSizeWarningLabel.setVisible(true);
265-
okButton.setEnabled(false);
285+
configCheckPassed = false;
286+
updateStartButtonState();
266287
} else {
267288
fileSizeWarningLabel.setVisible(false);
268-
okButton.setEnabled(true);
289+
configCheckPassed = true;
290+
updateStartButtonState();
269291
}
270292
}
271293

294+
private void updateStartButtonState() {
295+
okButton.setEnabled(configCheckPassed && functionSelectionPanel.getSelectedCount() > 0);
296+
}
297+
272298
private long getProgramFileSize() {
273299
try {
274300
Path filePath;

‎src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionRowObject.java‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
package ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection;
2+
3+
import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionInfo;
4+
import ghidra.program.model.address.Address;
5+
import ghidra.program.model.listing.Function;
6+
7+
import javax.annotation.Nullable;
8+
9+
/**
10+
* Wrapper around a Ghidra {@link Function} with a mutable selection flag.
11+
* Used to display functions in a table where users can select which functions
12+
* to include in analysis.
13+
*/
14+
public class FunctionRowObject {
15+
private final Function function;
16+
private boolean selected;
17+
private boolean enabled = true;
18+
@Nullable
19+
private FunctionInfo remoteFunctionInfo;
20+
21+
public FunctionRowObject(Function function, boolean selected) {
22+
this.function = function;
23+
this.selected = selected;
24+
}
25+
26+
public Function getFunction() {
27+
return function;
28+
}
29+
30+
public String getName() {
31+
return function.getName();
32+
}
33+
34+
public Address getAddress() {
35+
return function.getEntryPoint();
36+
}
37+
38+
/**
39+
* Returns the size of the function based on address count.
40+
*/
41+
public long getSize() {
42+
return function.getBody().getNumAddresses();
43+
}
44+
45+
public boolean isExternal() {
46+
return function.isExternal();
47+
}
48+
49+
public boolean isThunk() {
50+
return function.isThunk();
51+
}
52+
53+
/**
54+
* Returns a human-readable type string for the function.
55+
*/
56+
public String getType() {
57+
if (isExternal()) {
58+
return "External";
59+
} else if (isThunk()) {
60+
return "Thunk";
61+
} else {
62+
return "Normal";
63+
}
64+
}
65+
66+
public boolean isSelected() {
67+
return selected;
68+
}
69+
70+
public void setSelected(boolean selected) {
71+
if (!enabled) {
72+
return;
73+
}
74+
this.selected = selected;
75+
}
76+
77+
public boolean isEnabled() {
78+
return enabled;
79+
}
80+
81+
public void setEnabled(boolean enabled) {
82+
this.enabled = enabled;
83+
if (!enabled) {
84+
this.selected = false;
85+
}
86+
}
87+
88+
@Nullable
89+
public FunctionInfo getRemoteFunctionInfo() {
90+
return remoteFunctionInfo;
91+
}
92+
93+
public void setRemoteFunctionInfo(@Nullable FunctionInfo remoteFunctionInfo) {
94+
this.remoteFunctionInfo = remoteFunctionInfo;
95+
}
96+
97+
@Nullable
98+
public String getRemoteFunctionName() {
99+
return remoteFunctionInfo != null ? remoteFunctionInfo.functionName() : null;
100+
}
101+
102+
@Nullable
103+
public String getRemoteMangledName() {
104+
return remoteFunctionInfo != null ? remoteFunctionInfo.functionMangledName() : null;
105+
}
106+
107+
@Nullable
108+
public Long getRemoteFunctionID() {
109+
return remoteFunctionInfo != null ? remoteFunctionInfo.functionID().value() : null;
110+
}
111+
}

0 commit comments

Comments
 (0)