Skip to content

Commit a70d5c7

Browse files
fmaginfmagin
committed
A lot more cleanup and refactoring, to include datatypes when applying matches
1 parent 5c9ebc3 commit a70d5c7

File tree

51 files changed

+798
-832
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

51 files changed

+798
-832
lines changed

‎build.gradle‎

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

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

Lines changed: 48 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,10 @@
55
import ghidra.app.cmd.function.ApplyFunctionSignatureCmd;
66
import ghidra.app.cmd.label.RenameLabelCmd;
77
import ghidra.framework.cmd.Command;
8-
import ghidra.framework.model.DomainObject;
9-
import ghidra.program.model.data.DataTypeDependencyException;
10-
import ghidra.program.model.data.FunctionDefinitionDataType;
118
import ghidra.program.model.listing.CircularDependencyException;
129
import ghidra.program.model.listing.Program;
1310
import ghidra.program.model.symbol.Namespace;
1411
import ghidra.program.model.symbol.SourceType;
15-
import ghidra.util.Msg;
1612
import ghidra.util.exception.DuplicateNameException;
1713
import ghidra.util.exception.InvalidInputException;
1814
import org.jetbrains.annotations.NotNull;
@@ -21,31 +17,61 @@
2117

2218
import static ai.reveng.toolkit.ghidra.plugins.BinarySimilarityPlugin.REVENG_AI_NAMESPACE;
2319

24-
// We can't use Command<Program> because that breaks compatibility with Ghidra 11.0
25-
public class ApplyMatchCmd implements Command {
20+
/// The central command to apply a function match to a {@link Program}
21+
/// It centralizes product design decisions about how to apply a match, like moving it to a namespace,
22+
/// renaming it, applying the signature, etc.
23+
/// It will apply the function signature if available, otherwise it will just rename the function
24+
/// There are various considerations when ap
25+
public class ApplyMatchCmd implements Command<Program> {
2626

27-
private final Program program;
27+
private final GhidraRevengService.AnalysedProgram analyzedProgram;
2828
private final GhidraFunctionMatchWithSignature match;
2929
@Nullable private final GhidraRevengService service;
30+
private final Boolean includeBinaryNameInNameSpace;
3031

3132
public ApplyMatchCmd(
3233
@Nullable GhidraRevengService service,
33-
@NotNull GhidraFunctionMatchWithSignature match) {
34+
@NotNull GhidraRevengService.AnalysedProgram program,
35+
@NotNull GhidraFunctionMatchWithSignature match,
36+
Boolean includeBinaryNameInNameSpace
37+
38+
) {
3439
super();
35-
this.program = match.function().getProgram();
40+
this.analyzedProgram = program;
3641
this.match = match;
3742
this.service = service;
43+
this.includeBinaryNameInNameSpace = includeBinaryNameInNameSpace;
44+
}
45+
46+
private boolean shouldApplyMatch() {
47+
var func = match.function();
48+
return func != null &&
49+
// Do not override user-defined function names
50+
func.getSymbol().getSource() != SourceType.USER_DEFINED &&
51+
// Exclude thunks and external functions
52+
!func.isThunk() &&
53+
!func.isExternal() &&
54+
// Only accept valid names (no spaces)
55+
!match.functionMatch().nearest_neighbor_mangled_function_name().contains(" ") &&
56+
!match.functionMatch().nearest_neighbor_function_name().contains(" ")
57+
// Only rename if the function ID is known (boundaries matched)
58+
&& analyzedProgram.getIDForFunction(func).map(id -> id.functionID() != match.functionMatch().origin_function_id()).orElse(false);
3859
}
60+
3961
@Override
40-
public boolean applyTo(DomainObject obj) {
62+
public boolean applyTo(Program obj) {
4163
// Check that this is the same program
42-
if (obj != this.program) {
64+
if (obj != this.analyzedProgram.program()) {
4365
throw new IllegalArgumentException("This command can only be applied to the same program as the one provided in the constructor");
4466
}
45-
var libraryNamespace = getLibraryNameSpaceForName(match.functionMatch().nearest_neighbor_binary_name());
67+
if (!shouldApplyMatch()) {
68+
return false;
69+
}
70+
71+
var nameSpace = includeBinaryNameInNameSpace ? getLibraryNameSpaceForName(match.functionMatch().nearest_neighbor_binary_name()): getRevEngAINameSpace();
4672
var function = match.function();
4773
try {
48-
function.setParentNamespace(libraryNamespace);
74+
function.setParentNamespace(nameSpace);
4975
} catch (DuplicateNameException e) {
5076
throw new RuntimeException(e);
5177
} catch (InvalidInputException e) {
@@ -54,43 +80,36 @@ public boolean applyTo(DomainObject obj) {
5480
throw new RuntimeException(e);
5581
}
5682

57-
FunctionDefinitionDataType signature = null;
58-
if (match.signature().isPresent()) {
59-
try {
60-
signature = GhidraRevengService.getFunctionSignature(match.signature().get());
61-
} catch (DataTypeDependencyException e) {
62-
Msg.showError(this, null,"Failed to create function signature",
63-
"Failed to create signature for match function with type %s"
64-
.formatted(match.signature().get().func_types().getSignature()),
65-
e);
66-
}
67-
}
83+
this.analyzedProgram.setMangledNameForFunction(function, match.functionMatch().nearest_neighbor_mangled_function_name());
6884

85+
var signature = match.signature();
6986
if (signature != null) {
7087
var cmd = new ApplyFunctionSignatureCmd(function.getEntryPoint(), signature, SourceType.USER_DEFINED);
71-
cmd.applyTo(program);
88+
cmd.applyTo(analyzedProgram.program());
7289
}
7390
else {
7491
var renameCmd = new RenameLabelCmd(match.function().getSymbol(), match.functionMatch().name(), SourceType.USER_DEFINED);
75-
renameCmd.applyTo(program);
92+
renameCmd.applyTo(analyzedProgram.program());
7693
}
7794
// If we have a service then push the name. If not then it was explicitly not provided, i.e. the caller
7895
// is responsible for pushing the names in batch
7996
if (service != null) {
80-
service.getApi().renameFunction(match.functionMatch().origin_function_id(), match.functionMatch().name());
97+
service.getApi().renameFunction(match.functionMatch().origin_function_id(), match.functionMatch().nearest_neighbor_function_name());
8198
}
8299

83100

84-
return false;
101+
return true;
85102
}
86103

87104
public void applyWithTransaction() {
105+
var program = this.analyzedProgram.program();
88106
var tID = program.startTransaction("RevEng.AI: Apply Match");
89107
var status = applyTo(program);
90108
program.endTransaction(tID, status);
91109
}
92110

93111
private Namespace getRevEngAINameSpace() {
112+
var program = this.analyzedProgram.program();
94113
Namespace revengMatchNamespace = null;
95114
try {
96115
revengMatchNamespace = program.getSymbolTable().getOrCreateNameSpace(
@@ -105,6 +124,7 @@ private Namespace getRevEngAINameSpace() {
105124
}
106125

107126
private Namespace getLibraryNameSpaceForName(String name) {
127+
var program = this.analyzedProgram.program();
108128
Namespace libraryNamespace = null;
109129
try {
110130
libraryNamespace = program.getSymbolTable().getOrCreateNameSpace(

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

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

33
import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService;
4+
import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface;
45
import ai.reveng.toolkit.ghidra.core.services.api.types.*;
56
import ghidra.util.exception.CancelledException;
67
import ghidra.util.task.Task;
@@ -10,9 +11,7 @@
1011
import java.time.Duration;
1112
import java.util.HashSet;
1213
import java.util.List;
13-
import java.util.Map;
1414
import java.util.Set;
15-
import java.util.stream.Collectors;
1615

1716
import static java.util.stream.Collectors.groupingBy;
1817

@@ -23,12 +22,12 @@
2322
*
2423
*/
2524
public class ComputeTypeInfoTask extends Task {
26-
private final List<FunctionID> functions;
25+
private final List<TypedApiInterface.FunctionID> functions;
2726
private final GhidraRevengService service;
2827
private final DataTypeAvailableCallback callback;
2928

3029
public ComputeTypeInfoTask(GhidraRevengService service,
31-
List<FunctionID> functions,
30+
List<TypedApiInterface.FunctionID> functions,
3231
@Nullable DataTypeAvailableCallback callback) {
3332
super("Computing Type Info", true, true, false);
3433
this.service = service;
@@ -49,7 +48,7 @@ public void run(TaskMonitor monitor) throws CancelledException {
4948
service.getApi().generateFunctionDataTypes(analysisID, functions.stream().map(FunctionDetails::functionId).toList());
5049
});
5150

52-
Set<FunctionID> missing = new HashSet<>(functions);
51+
Set<TypedApiInterface.FunctionID> missing = new HashSet<>(functions);
5352

5453
while (!missing.isEmpty()) {
5554
try {
@@ -76,6 +75,6 @@ public void run(TaskMonitor monitor) throws CancelledException {
7675
}
7776

7877
public interface DataTypeAvailableCallback {
79-
void dataTypeAvailable(FunctionID functionID, FunctionDataTypeStatus dataTypeStatus);
78+
void dataTypeAvailable(TypedApiInterface.FunctionID functionID, FunctionDataTypeStatus dataTypeStatus);
8079
}
8180
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import ai.reveng.invoker.ApiException;
44
import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService;
5-
import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionID;
5+
import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface;
66
import ai.reveng.toolkit.ghidra.core.services.logging.ReaiLoggingService;
77
import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage;
88
import ai.reveng.toolkit.ghidra.core.services.api.types.AIDecompilationStatus;
@@ -262,7 +262,7 @@ public void run(TaskMonitor monitor) throws CancelledException {
262262
}
263263

264264

265-
private void waitForDecomp(FunctionID id, TaskMonitor monitor) throws CancelledException {
265+
private void waitForDecomp(TypedApiInterface.FunctionID id, TaskMonitor monitor) throws CancelledException {
266266
var logger = tool.getService(ReaiLoggingService.class);
267267
var api = service.getApi();
268268
AIDecompilationStatus lastDecompStatus = null;

0 commit comments

Comments
 (0)