Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -108,8 +108,12 @@ public boolean applyTo(Program obj) {
public void applyWithTransaction() {
var program = this.analyzedProgram.program();
var tID = program.startTransaction("RevEng.AI: Apply Match");
var status = applyTo(program);
program.endTransaction(tID, status);
boolean status = false;
try {
status = applyTo(program);
} finally {
program.endTransaction(tID, status);
}
}

private Namespace getRevEngAINameSpace() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,23 +57,7 @@ public AIDecompilationdWindow(PluginTool tool, String owner) {

@Override
public void actionPerformed(ActionContext context) {
if (function != null) {
var service = tool.getService(GhidraRevengService.class);
var analyzedProgram = service.getAnalysedProgram(function.getProgram());
if (analyzedProgram.isEmpty()) {
Msg.error(this, "Failed to send positive feedback: Program is not known to RevEng.AI");
return;
}
var fID = analyzedProgram.get().getIDForFunction(function);
fID.ifPresent(id -> {
try {
service.getApi().aiDecompRating(id.functionID(), "POSITIVE", "");
} catch (ApiException e) {
// Fail silently because this is not a critical feature
Msg.error(this, "Failed to send positive feedback for function %s: %s".formatted(function.getName(), e.getMessage()));
}
});
}
sendFeedbackInBackground(function, "POSITIVE", "");
}

@Override
Expand All @@ -93,28 +77,14 @@ public boolean isEnabled() {
}
@Override
public void actionPerformed(ActionContext context) {
// Spawn textbox to enter reason for negative feedback

final Function target = function;
if (target == null) {
return;
}
var dialog = new InputDialog("Negative Feedback", "Please provide details about what was wrong with the decompilation:", "");
tool.showDialog(dialog);
if (!dialog.isCanceled()) {
if (function != null) {
var service = tool.getService(GhidraRevengService.class);
var programWithID = service.getAnalysedProgram(function.getProgram());
if (programWithID.isEmpty()) {
Msg.error(this, "Failed to send negative feedback: Program is not known to RevEng.AI");
return;
}
var fID = programWithID.get().getIDForFunction(function);
fID.ifPresent(id -> {
try {
service.getApi().aiDecompRating(id.functionID(), "NEGATIVE", dialog.getValue());
} catch (ApiException e) {
// Fail silently because this is not a critical feature
Msg.error(this, "Failed to send negative feedback for function %s: %s".formatted(function.getName(), e.getMessage()));
}
});
}
sendFeedbackInBackground(target, "NEGATIVE", dialog.getValue());
}
}

Expand All @@ -126,6 +96,37 @@ public boolean isEnabled() {
});
}

private void sendFeedbackInBackground(Function target, String rating, String reason) {
if (target == null) {
return;
}
final String detail = reason == null ? "" : reason;
tool.execute(new Task("Send AI Decompilation Feedback", false, false, false) {
@Override
public void run(TaskMonitor monitor) {
var service = tool.getService(GhidraRevengService.class);
var analyzedProgram = service.getAnalysedProgram(target.getProgram());
if (analyzedProgram.isEmpty()) {
Msg.error(AIDecompilationdWindow.this,
"Failed to send %s feedback: Program is not known to RevEng.AI".formatted(rating));
return;
}
var fID = analyzedProgram.get().getIDForFunction(target);
if (fID.isEmpty()) {
Msg.error(AIDecompilationdWindow.this,
"Failed to send %s feedback: function %s not known to RevEng.AI".formatted(rating, target.getName()));
return;
}
try {
service.getApi().aiDecompRating(fID.get().functionID(), rating, detail);
} catch (ApiException e) {
Msg.error(AIDecompilationdWindow.this,
"Failed to send %s feedback for function %s: %s".formatted(rating, target.getName(), e.getMessage()));
}
}
}, 0);
}


private JComponent buildComponent() {

Expand Down Expand Up @@ -317,7 +318,7 @@ public void locationChanged(ProgramLocation loc) {
var newFuncLocation = functionMgr.getFunctionContaining(loc.getAddress());

// If we changed to a different function, we want to clear the output of the old function
if (function != null && newFuncLocation != function) {
if (function != null && !isSameFunction(newFuncLocation, function)) {
clear();
}

Expand All @@ -329,7 +330,7 @@ public void locationChanged(ProgramLocation loc) {

void newStatusForFunction(Function function, AIDecompilationStatus status) {
cache.put(function, status);
if (function == this.function) {
if (isCurrentFunction(function)) {
SwingUtilities.invokeLater(() ->
setDisplayedValuesBasedOnStatus(function, status)
);
Expand All @@ -352,6 +353,28 @@ private boolean hasPendingDecompilations() {
s.status() == DecompilationData.StatusEnum.PENDING
|| s.status() == DecompilationData.StatusEnum.RUNNING);
}

private static boolean isSameFunction(Function a, Function b) {
return a != null && b != null && a.getEntryPoint().equals(b.getEntryPoint());
}

private boolean isCurrentFunction(Function candidate) {
return isSameFunction(candidate, this.function);
}

private void reportDecompFailure(Function failedFunction, Exception e) {
var logger = tool.getService(ReaiLoggingService.class);
logger.error("AI Decompilation failed for function %s: %s".formatted(failedFunction.getName(), e.getMessage()));
Msg.error(this, "AI Decompilation failed for function %s".formatted(failedFunction.getName()), e);
SwingUtilities.invokeLater(() -> {
if (isCurrentFunction(failedFunction)) {
descriptionArea.setText("AI Decompilation failed: " + e.getMessage());
if (!hasPendingDecompilations()) {
taskMonitorComponent.setVisible(false);
}
}
});
}
class AIDecompTask extends Task {

private final GhidraRevengService service;
Expand All @@ -365,14 +388,20 @@ public AIDecompTask(PluginTool tool, GhidraRevengService.FunctionWithID function

@Override
public void run(TaskMonitor monitor) throws CancelledException {
var fID = functionWithID.functionID();
// Check if there is an existing process already, because the trigger API will fail with 400 if there is
if (service.getApi().pollAIDecompileStatus(fID).status() == DecompilationData.StatusEnum.UNINITIALISED) {
// Trigger the decompilation
service.getApi().triggerAIDecompilationForFunctionID(fID);
try {
var fID = functionWithID.functionID();
// Check if there is an existing process already, because the trigger API will fail with 400 if there is
if (service.getApi().pollAIDecompileStatus(fID).status() == DecompilationData.StatusEnum.UNINITIALISED) {
// Trigger the decompilation
service.getApi().triggerAIDecompilationForFunctionID(fID);
}
waitForDecomp(fID, monitor);
// TODO: Inform the component that something is finished
} catch (CancelledException e) {
throw e;
} catch (Exception e) {
reportDecompFailure(functionWithID.function(), e);
}
waitForDecomp(fID, monitor);
// TODO: Inform the component that something is finished
}


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@
import ai.reveng.toolkit.ghidra.core.services.api.types.GhidraFunctionMatchWithSignature;
import com.google.common.collect.BiMap;
import ghidra.program.model.listing.Function;
import ghidra.util.task.Task;
import ghidra.util.task.TaskBuilder;
import ghidra.util.task.TaskMonitor;
import ghidra.util.task.TaskMonitorComponent;
import ghidra.util.Msg;
import resources.ResourceManager;
Expand Down Expand Up @@ -893,8 +896,7 @@ protected void onMatchButtonClicked() {


protected void onRenameAllButtonClicked() {
batchRenameFunctions(functionMatchResults);
importFunctionNames(functionMatchResults);
renameInBackground(functionMatchResults);
}

protected void onRenameSelectedButtonClicked() {
Expand All @@ -916,16 +918,26 @@ protected void onRenameSelectedButtonClicked() {
}
}

batchRenameFunctions(selectedMatches);
importFunctionNames(selectedMatches);
renameInBackground(selectedMatches);
}

private void renameInBackground(List<GhidraFunctionMatchWithSignature> matches) {
var task = new Task("Applying Renames", true, false, false) {
@Override
public void run(TaskMonitor monitor) {
batchRenameFunctions(matches);
importFunctionNames(matches);
}
};
TaskBuilder.withTask(task).launchInBackground(taskMonitorComponent);
}

protected void batchRenameFunctions(List<GhidraFunctionMatchWithSignature> functionMatches) {
try {
revengService.batchRenamingGhidraMatchesWithSignatures(functionMatches);
} catch (Exception e) {
Msg.error(this, "Failed to rename functions: " + e.getMessage(), e);
showError("Failed to rename functions: " + e.getMessage());
SwingUtilities.invokeLater(() -> showError("Failed to rename functions: " + e.getMessage()));
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,17 @@ public void actionPerformed(ActionContext context) {
var match = getSelectedMatch();
if (match != null) {
var service = tool.getService(GhidraRevengService.class);
service.openFunctionInPortal(match.functionMatch().nearest_neighbor_id());
var neighborId = match.functionMatch().nearest_neighbor_id();
tool.execute(new Task("View Match in Portal", false, false, false) {
@Override
public void run(TaskMonitor monitor) {
try {
service.openFunctionInPortal(neighborId);
} catch (Exception e) {
Msg.error(SimilarFunctionsWindow.this, "Failed to open match in portal: " + e.getMessage(), e);
}
}
}, 0);
}
}

Expand All @@ -158,8 +168,14 @@ public void actionPerformed(ActionContext context) {
var match = getSelectedMatch();
if (match != null && analyzedProgram != null) {
var service = tool.getService(GhidraRevengService.class);
var cmd = new ApplyMatchCmd(service, analyzedProgram, match, true);
cmd.applyWithTransaction();
var program = analyzedProgram;
tool.execute(new Task("Apply Match", true, false, false) {
@Override
public void run(TaskMonitor monitor) {
var cmd = new ApplyMatchCmd(service, program, match, true);
cmd.applyWithTransaction();
}
}, 0);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,10 @@
import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage;
import ghidra.framework.plugintool.PluginTool;
import ghidra.program.model.listing.Program;
import ghidra.util.Msg;
import ghidra.util.table.GhidraFilterTable;
import ghidra.util.task.Task;
import ghidra.util.task.TaskMonitor;

import javax.swing.*;
import java.awt.*;
Expand Down Expand Up @@ -64,8 +67,18 @@ public void mouseClicked(MouseEvent e) {
if ("Analysis ID".equals(columnName)) {
LegacyAnalysisResult result = recentAnalysesTable.getModel().getRowObject(row);
if (result != null) {
var analysisID = ghidraRevengService.getApi().getAnalysisIDfromBinaryID(result.binary_id());
ghidraRevengService.openPortalFor(analysisID);
var binaryID = result.binary_id();
tool.execute(new Task("Open analysis in portal", false, false, false) {
@Override
public void run(TaskMonitor monitor) {
try {
var analysisID = ghidraRevengService.getApi().getAnalysisIDfromBinaryID(binaryID);
ghidraRevengService.openPortalFor(analysisID);
} catch (Exception ex) {
Msg.error(RecentAnalysisDialog.this, "Failed to open analysis in portal: " + ex.getMessage(), ex);
}
}
}, 0);
}
}
}
Expand Down Expand Up @@ -97,15 +110,26 @@ public void mouseClicked(MouseEvent e) {

private void pickAnalysis(LegacyAnalysisResult result) {
var service = tool.getService(GhidraRevengService.class);
var analysisID = service.getApi().getAnalysisIDfromBinaryID(result.binary_id());
var programWithID = service.registerAnalysisForProgram(program, analysisID);
tool.firePluginEvent(
new RevEngAIAnalysisStatusChangedEvent(
"Recent Analysis Dialog",
programWithID,
result.status()
)
);
close();
tool.execute(new Task("Attach to analysis", true, false, false) {
@Override
public void run(TaskMonitor monitor) {
try {
var analysisID = service.getApi().getAnalysisIDfromBinaryID(result.binary_id());
var programWithID = service.registerAnalysisForProgram(program, analysisID);
SwingUtilities.invokeLater(() -> {
tool.firePluginEvent(
new RevEngAIAnalysisStatusChangedEvent(
"Recent Analysis Dialog",
programWithID,
result.status()
)
);
close();
});
} catch (Exception ex) {
Msg.error(RecentAnalysisDialog.this, "Failed to attach to analysis: " + ex.getMessage(), ex);
}
}
}, 0);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -206,23 +206,7 @@ private Optional<TypedApiInterface.AnalysisID> getAnalysisIDFromOptions(
if (bid == ReaiPluginPackage.INVALID_ANALYSIS_ID) {
return Optional.empty();
}
var analysisID = new TypedApiInterface.AnalysisID((int) bid);
if (!statusCache.containsKey(analysisID)) {
// Check that it's really valid in the context of the currently configured API
try {
var status = api.status(analysisID);
statusCache.put(analysisID, status);
} catch (APIAuthenticationException | ApiException e) {
Msg.showError(this, null, "Invalid Analysis ID",
("The Analysis ID %s stored in the program options is invalid for the currently configured RevEng.AI server %s. "
+ "This could be an intermittent error, or you switched the servers")
.formatted(analysisID, this.apiInfo.hostURI()), e);
return Optional.empty();
}
// Now it's certain that it is a valid binary ID
}

return Optional.of(analysisID);
return Optional.of(new TypedApiInterface.AnalysisID((int) bid));
}

@Deprecated
Expand All @@ -241,7 +225,7 @@ private Optional<BinaryID> getBinaryIDfromOptions(
try {
status = api.status(binID);
} catch (APIAuthenticationException | ApiException e) {
Msg.showError(this, null, "Invalid Binary ID",
Msg.error(this,
("The Binary ID %s stored in the program options is invalid for the currently configured RevEng.AI server %s. "
+ "This could be an intermittent error, or you switched servers")
.formatted(binID, this.apiInfo.hostURI()), e);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,10 @@ public TypedApiImplementation(String baseUrl, String apiKey) {
// Set withResponseBody to true if debugging issues with the API to see the full response body
apiClient.setHttpClient(apiClient.getHttpClient().newBuilder().addInterceptor(ghidraLogger(false)).build());

apiClient.setConnectTimeout(5000);
apiClient.setReadTimeout(15000);
apiClient.setWriteTimeout(15000);

ApiKeyAuth APIKey = (ApiKeyAuth) apiClient.getAuthentication("APIKey");
APIKey.setApiKey(apiKey);

Expand Down Expand Up @@ -532,6 +536,11 @@ public AIDecompilationStatus pollAIDecompileStatus(FunctionID functionID) {
inlineCommentsStatus,
inlineComments);
} catch (ApiException e) {
if (e.getCode() == 404) {
return new AIDecompilationStatus(
DecompilationData.StatusEnum.UNINITIALISED,
null, null, null, null, null, List.of());
}
throw new RuntimeException("Failed to poll AI decompilation status", e);
}
}
Expand Down
Loading
Loading