Compare commits

...

3 Commits

13 changed files with 282 additions and 66 deletions

View File

@ -60,60 +60,31 @@ public class ChatContextBuilder {
getFunctionCallGraph (or getFunctionAssembly + getCallers + getCallees) getFunctionCallGraph (or getFunctionAssembly + getCallers + getCallees)
to examine the function AND the functions it calls. A thorough analysis to examine the function AND the functions it calls. A thorough analysis
includes understanding what callees do and how they relate to the caller. includes understanding what callees do and how they relate to the caller.
Never just read one function in isolation and suggest analyzing callees Never just read one function in isolation analyze callees now.
"later" analyze them now as part of your response.
PRIORITY RULE: When the user asks for a specific mutation (rename, comment, [MUTATION TOOLS MANDATORY]
signature, flag), use that tool first. Do NOT respond with a different mutation When your analysis reveals findings deserving a rename, comment, signature,
type instead (e.g., don't suggest renames when the user asked for comments). flag, or label change, you MUST call the corresponding tool IMMEDIATELY.
This rule is about mutation tool selection ONLY it does NOT limit your A suggestion described only in text will NOT be reflected in the UI and is
analysis. Always do a thorough analysis including callees and callers. USELESS the user needs the tool call to see and approve it.
[FUNCTION RENAME] For proactive suggestions (requires user approval), use the suggest* tools:
Use renameFunction when the user explicitly asks to rename a specific function suggestRenameFunction, suggestAddComment, suggestUpdateSignature,
(applies immediately). Use suggestRenameFunction when you independently identify suggestPseudocode, suggestRenameLabel, suggestRenameGlobalVar,
a function's purpose and want to suggest a better name (requires user approval). suggestRenameStruct
Proactively suggest renames whenever you identify a function's purpose with high
confidence. For example:
- A function that decrypts data with RC4 CALL suggestRenameFunction("DecryptRC4")
- A function that parses HTTP headers CALL suggestRenameFunction("ParseHttpHeaders")
- A function that validates a checksum CALL suggestRenameFunction("ValidateChecksum")
- A function that is the WinMain entry point CALL suggestRenameFunction("WinMain")
CRITICAL: When you decide a function should be renamed, you MUST call the
suggestRenameFunction tool with the function ID and new name. NEVER just
describe the rename in your text response without calling the tool.
Always call the tool, then briefly explain your reasoning afterward.
[FUNCTION SIGNATURE] For changes the user explicitly asked for (applied immediately), use:
Use updateSignature (apply) or suggestUpdateSignature (pending approval) renameFunction, addComment, updateSignature,
to set a function's type signature. Infer parameter types and return type addFlag, removeFlag, renameLabel,
from assembly analysis. Example: 'int DecryptRC4(int keyLen, const uint8_t* key)'. renameGlobalVar, renameStruct
[FUNCTION COMMENTS] Rules:
Use addComment (apply) or suggestAddComment (pending approval) to document 1. Call mutation tools BEFORE writing your analysis text.
analysis findings. You can attach comments to specific addresses 2. If you have 3 suggestions, make 3 tool calls not 1 text block.
(e.g., address="0x5EB5") to document individual instructions inline. 3. For pseudocode: only suggestPseudocode when getFunctionDecompiled returned
Omit address to attach to the function. The tool appends automatically. found=false and you are confident in the assembly logic.
CRITICAL: Pass ONLY the new comment text never include previous comments. 4. When the user explicitly asks for a specific change type (e.g. "add a comment"),
CRITICAL: When the user asks to add comments, call addComment/suggestAddComment. use that tool first. Do not respond with a different mutation type instead.
Do NOT respond with renames or signature changes unless the user also asked for those.
[FUNCTION FLAGS]
Use addFlag and removeFlag to categorize functions.
Valid flags: CRYPTO, NETWORK, FILE_IO, ANTI_DEBUG, ENTRY_POINT,
CALLBACK, MEMORY, PROCESS, REGISTRY, UI.
[GLOBAL VARIABLES]
Use searchGlobalVars to find data label IDs, then renameGlobalVar (apply)
or suggestRenameGlobalVar (pending approval).
[LABELS]
Use searchLabels to find label IDs within a function, then renameLabel (apply)
or suggestRenameLabel (pending approval).
[STRUCTS]
Use searchStructs to find struct IDs, then renameStruct (apply)
or suggestRenameStruct (pending approval).
%s %s
[BINARY OVERVIEW] [BINARY OVERVIEW]
Filename: %s Filename: %s

View File

@ -2,6 +2,7 @@ package ai.decompile.ai.service.tools;
import ai.decompile.analysis.model.repository.StaticFunctionRepository; import ai.decompile.analysis.model.repository.StaticFunctionRepository;
import ai.decompile.analysis.service.MutationService; import ai.decompile.analysis.service.MutationService;
import java.util.List;
import java.util.UUID; import java.util.UUID;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.tool.annotation.Tool;
@ -217,6 +218,48 @@ public class FunctionMutationTools {
mutationService.suggestComment(funcId, comment, address, "AI", sessionId)); mutationService.suggestComment(funcId, comment, address, "AI", sessionId));
} }
@Tool(
description =
"""
Suggest C-like pseudocode for a function that currently has no decompiled code.
CRITICAL: Only call this when getFunctionDecompiled returns found=false and you
have thoroughly analyzed the function's assembly (call getFunctionAssembly first).
Be conservative only suggest pseudocode when the assembly logic is clear and
you can produce accurate, idiomatic C code. Do NOT guess or fabricate logic.
The pseudocode must faithfully reproduce the exact logic shown in the assembly,
using proper C syntax (control flow, variable names derived from stack offsets,
function calls using resolved names when available).
Provide one line of code per array element. Do NOT combine multiple statements
into a single string element. A comment "Pseudocode generated by AI" will be
automatically added.
""")
public String suggestPseudocode(
@ToolParam(description = "The function ID (UUID from searchFunctions result)")
String functionId,
@ToolParam(
description =
"Array of C-like pseudocode lines. Each element is exactly ONE line of code. "
+ "Do NOT put multiple statements in one element. "
+ "Do not include markdown code fences.")
List<String> pseudocodeLines) {
if (pseudocodeLines == null || pseudocodeLines.isEmpty()) {
return error(functionId, "Pseudocode must not be empty");
}
var pseudocode = String.join("\n", pseudocodeLines);
var funcId = UUID.fromString(functionId);
if (functionRepository.findById(funcId).isEmpty()) {
return error(functionId, "Function not found: " + functionId);
}
mutationService.applyComment(funcId, "Pseudocode generated by AI", "", "AI");
var sessionId = toolContext.getSessionId();
return helper.formatResult(
mutationService.suggest("FUNCTION", funcId, "decompiledCode", pseudocode, "AI", sessionId));
}
@Tool( @Tool(
description = description =
""" """

View File

@ -19,6 +19,7 @@ import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam; import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.data.domain.PageRequest; import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Sort;
import org.springframework.stereotype.Component; import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional; import org.springframework.transaction.annotation.Transactional;
@ -155,16 +156,18 @@ public class FunctionTools {
var func = functionRepository.findById(UUID.fromString(functionId)).orElse(null); var func = functionRepository.findById(UUID.fromString(functionId)).orElse(null);
if (Objects.isNull(func)) { if (Objects.isNull(func)) {
return new DecompiledResult("unknown", "unknown", ""); return new DecompiledResult("unknown", "unknown", "", false);
} }
var renameMap = mutationService.buildRenameMap(func.getAnalysis().getId()); var renameMap = mutationService.buildRenameMap(func.getAnalysis().getId());
var code = func.getDecompiledCode(); var code = func.getDecompiledCode();
var hasDecompiled = Objects.nonNull(code) && !code.isBlank();
return new DecompiledResult( return new DecompiledResult(
func.getName(), func.getName(),
func.getAddress(), func.getAddress(),
AnalysisService.resolveNames(Objects.nonNull(code) ? code : "", renameMap)); hasDecompiled ? AnalysisService.resolveNames(code, renameMap) : "",
hasDecompiled);
} }
@Tool( @Tool(
@ -291,7 +294,8 @@ public class FunctionTools {
var page = var page =
functionRepository.findByAnalysisId( functionRepository.findByAnalysisId(
analysis.getId(), PageRequest.of(0, MAX_LIST_FUNCTIONS)); analysis.getId(),
PageRequest.of(0, MAX_LIST_FUNCTIONS, Sort.by("address").ascending()));
return page.getContent().stream().map(this::toSummary).toList(); return page.getContent().stream().map(this::toSummary).toList();
} }
@ -345,7 +349,8 @@ public class FunctionTools {
public record FunctionSummary( public record FunctionSummary(
UUID id, String name, String address, int functionSize, String signature) {} UUID id, String name, String address, int functionSize, String signature) {}
public record DecompiledResult(String name, String address, String decompiledCode) {} public record DecompiledResult(
String name, String address, String decompiledCode, boolean found) {}
public record FunctionAssembly(String name, String address, String assembly) {} public record FunctionAssembly(String name, String address, String assembly) {}
} }

View File

@ -57,6 +57,26 @@ public class MutationToolsHelper {
if (Objects.isNull(value)) { if (Objects.isNull(value)) {
return ""; return "";
} }
return value.replace("\\", "\\\\").replace("\"", "\\\""); var sb = new StringBuilder(value.length());
for (int i = 0; i < value.length(); i++) {
var ch = value.charAt(i);
switch (ch) {
case '"' -> sb.append("\\\"");
case '\\' -> sb.append("\\\\");
case '\b' -> sb.append("\\b");
case '\f' -> sb.append("\\f");
case '\n' -> sb.append("\\n");
case '\r' -> sb.append("\\r");
case '\t' -> sb.append("\\t");
default -> {
if (ch < 0x20) {
sb.append(String.format("\\u%04x", (int) ch));
} else {
sb.append(ch);
}
}
}
}
return sb.toString();
} }
} }

View File

@ -96,7 +96,7 @@ public record FunctionDetailResponse(
f.getFunctionSize() != null ? f.getFunctionSize() : 0, f.getFunctionSize() != null ? f.getFunctionSize() : 0,
f.getFlags(), f.getFlags(),
AnalysisService.resolveNames(f.getAssembly(), renameMap), AnalysisService.resolveNames(f.getAssembly(), renameMap),
AnalysisService.resolveNames(f.getDecompiledCode(), renameMap), escapeNewlines(AnalysisService.resolveNames(f.getDecompiledCode(), renameMap)),
f.getSignature(), f.getSignature(),
callerCount, callerCount,
calleeCount, calleeCount,
@ -154,4 +154,11 @@ public record FunctionDetailResponse(
return Collections.emptyList(); return Collections.emptyList();
} }
} }
private static String escapeNewlines(String value) {
if (Objects.isNull(value)) {
return null;
}
return value.replace("\n", "\\n");
}
} }

View File

@ -47,6 +47,7 @@ public class StaticAnalysis {
cascade = CascadeType.ALL, cascade = CascadeType.ALL,
orphanRemoval = true, orphanRemoval = true,
fetch = FetchType.LAZY) fetch = FetchType.LAZY)
@OrderBy("address ASC")
@Builder.Default @Builder.Default
private List<StaticFunction> functions = new ArrayList<>(); private List<StaticFunction> functions = new ArrayList<>();

View File

@ -6,12 +6,28 @@ import java.util.UUID;
import org.springframework.data.domain.Page; import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable; import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.repository.JpaRepository; import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.Query;
import org.springframework.data.repository.query.Param;
public interface StaticFunctionRepository extends JpaRepository<StaticFunction, UUID> { public interface StaticFunctionRepository extends JpaRepository<StaticFunction, UUID> {
Page<StaticFunction> findByAnalysisId(UUID analysisId, Pageable pageable); Page<StaticFunction> findByAnalysisId(UUID analysisId, Pageable pageable);
List<StaticFunction> findByNameIgnoreCaseContainingAndAnalysisId(String name, UUID analysisId); @Query(
"SELECT f FROM StaticFunction f WHERE LOWER(f.name) LIKE LOWER(CONCAT('%', :name, '%'))"
+ " AND f.analysis.id = :analysisId ORDER BY f.address ASC")
List<StaticFunction> findByNameIgnoreCaseContainingAndAnalysisId(
@Param("name") String name, @Param("analysisId") UUID analysisId);
@Query(
"SELECT f FROM StaticFunction f WHERE LOWER(f.address) LIKE LOWER(CONCAT('%', :address, '%'))"
+ " AND f.analysis.id = :analysisId ORDER BY f.address ASC")
List<StaticFunction> findByAddressIgnoreCaseContainingAndAnalysisId( List<StaticFunction> findByAddressIgnoreCaseContainingAndAnalysisId(
String address, UUID analysisId); @Param("address") String address, @Param("analysisId") UUID analysisId);
@Query(
"SELECT f FROM StaticFunction f WHERE f.analysis.id = :analysisId "
+ "AND (f.decompiledCode LIKE CONCAT('%', :name, '%') "
+ "OR f.assembly LIKE CONCAT('%', :name, '%')) ORDER BY f.address ASC")
List<StaticFunction> findByAnalysisIdAndNameInCode(
@Param("analysisId") UUID analysisId, @Param("name") String name);
} }

View File

@ -9,17 +9,23 @@ import org.springframework.data.jpa.repository.Query;
import org.springframework.data.repository.query.Param; import org.springframework.data.repository.query.Param;
public interface StaticXrefRepository extends JpaRepository<StaticXref, UUID> { public interface StaticXrefRepository extends JpaRepository<StaticXref, UUID> {
List<StaticXref> findByCallerId(UUID callerId); @Query("SELECT x FROM StaticXref x WHERE x.caller.id = :callerId ORDER BY x.callee.address ASC")
List<StaticXref> findByCallerId(@Param("callerId") UUID callerId);
List<StaticXref> findByCalleeId(UUID calleeId); @Query("SELECT x FROM StaticXref x WHERE x.callee.id = :calleeId ORDER BY x.caller.address ASC")
List<StaticXref> findByCalleeId(@Param("calleeId") UUID calleeId);
int countByCallerId(UUID callerId); int countByCallerId(UUID callerId);
int countByCalleeId(UUID calleeId); int countByCalleeId(UUID calleeId);
@Query("SELECT DISTINCT x.caller FROM StaticXref x WHERE x.callee.id = :functionId") @Query(
"SELECT DISTINCT x.caller FROM StaticXref x WHERE x.callee.id = :functionId"
+ " ORDER BY x.caller.address ASC")
List<StaticFunction> findCallersByCalleeId(@Param("functionId") UUID functionId); List<StaticFunction> findCallersByCalleeId(@Param("functionId") UUID functionId);
@Query("SELECT DISTINCT x.callee FROM StaticXref x WHERE x.caller.id = :functionId") @Query(
"SELECT DISTINCT x.callee FROM StaticXref x WHERE x.caller.id = :functionId"
+ " ORDER BY x.callee.address ASC")
List<StaticFunction> findCalleesByCallerId(@Param("functionId") UUID functionId); List<StaticFunction> findCalleesByCallerId(@Param("functionId") UUID functionId);
} }

View File

@ -49,7 +49,9 @@ import java.util.regex.Pattern;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2; import lombok.extern.log4j.Log4j2;
import org.springframework.data.domain.Page; import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable; import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.stereotype.Service; import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional; import org.springframework.transaction.annotation.Transactional;
@ -106,6 +108,13 @@ public class AnalysisService {
public Page<FunctionSummaryResponse> listFunctions(UUID analysisId, Pageable pageable) { public Page<FunctionSummaryResponse> listFunctions(UUID analysisId, Pageable pageable) {
assertAnalysisExists(analysisId); assertAnalysisExists(analysisId);
var resolution = buildNameResolutionMap(analysisId); var resolution = buildNameResolutionMap(analysisId);
if (!pageable.getSort().isSorted()) {
pageable =
PageRequest.of(
pageable.getPageNumber(), pageable.getPageSize(), Sort.by("address").ascending());
}
var page = functionRepository.findByAnalysisId(analysisId, pageable); var page = functionRepository.findByAnalysisId(analysisId, pageable);
var functionIds = page.getContent().stream().map(StaticFunction::getId).toList(); var functionIds = page.getContent().stream().map(StaticFunction::getId).toList();
var renameMap = mutationService.loadApprovedOldValues(functionIds); var renameMap = mutationService.loadApprovedOldValues(functionIds);

View File

@ -30,7 +30,12 @@ public class MutationService {
public record ResolvedEntity(UUID binaryId, UUID projectId) {} public record ResolvedEntity(UUID binaryId, UUID projectId) {}
private static final Map<String, Integer> FIELD_MAX_LENGTH = private static final Map<String, Integer> FIELD_MAX_LENGTH =
Map.of("name", 500, "signature", 1000, "comments", 10000, "flags", 100); Map.of(
"name", 500,
"signature", 1000,
"comments", 10000,
"flags", 100,
"decompiledCode", 100000);
private final MutationSuggestionRepository mutationRepository; private final MutationSuggestionRepository mutationRepository;
private final List<MutationTargetHandler> handlers; private final List<MutationTargetHandler> handlers;
@ -50,6 +55,10 @@ public class MutationService {
validateField(handler, fieldName); validateField(handler, fieldName);
validateNewValue(fieldName, newValue); validateNewValue(fieldName, newValue);
if ("decompiledCode".equals(fieldName)) {
newValue = normalizePseudocode(newValue);
}
var oldValue = handler.getCurrentValue(targetId, fieldName); var oldValue = handler.getCurrentValue(targetId, fieldName);
var mutation = var mutation =
@ -93,6 +102,10 @@ public class MutationService {
validateField(handler, fieldName); validateField(handler, fieldName);
validateNewValue(fieldName, newValue); validateNewValue(fieldName, newValue);
if ("decompiledCode".equals(fieldName)) {
newValue = normalizePseudocode(newValue);
}
var oldValue = handler.getCurrentValue(targetId, fieldName); var oldValue = handler.getCurrentValue(targetId, fieldName);
handler.applyChange(targetId, fieldName, newValue.trim(), operation); handler.applyChange(targetId, fieldName, newValue.trim(), operation);
@ -212,8 +225,11 @@ public class MutationService {
} }
var handler = handlerFor(mutation.getTargetType()); var handler = handlerFor(mutation.getTargetType());
handler.applyChange( var approvedValue = mutation.getNewValue();
mutation.getTargetId(), mutation.getFieldName(), mutation.getNewValue(), "SET"); if ("decompiledCode".equals(mutation.getFieldName())) {
approvedValue = normalizePseudocode(approvedValue);
}
handler.applyChange(mutation.getTargetId(), mutation.getFieldName(), approvedValue, "SET");
mutation.setStatus("APPROVED"); mutation.setStatus("APPROVED");
mutation.setResolvedAt(Instant.now()); mutation.setResolvedAt(Instant.now());
@ -425,4 +441,33 @@ public class MutationService {
new NotFoundException( new NotFoundException(
"Mutation suggestion with id '" + mutationId + "' not found")); "Mutation suggestion with id '" + mutationId + "' not found"));
} }
private String normalizePseudocode(String value) {
if (Objects.isNull(value) || value.isBlank()) {
return value;
}
if (value.contains("\n")) {
return value;
}
log.warn(
"Pseudocode has no line breaks — splitting on statement boundaries ({} chars)",
value.length());
return formatSingleLinePseudocode(value);
}
private String formatSingleLinePseudocode(String code) {
var result =
code.replaceAll(";\\s*", ";\n")
.replaceAll("\\{\\s*", "{\n")
.replaceAll("\\s*\\}", "\n}")
.replaceAll("\n\n+", "\n")
.trim();
result = result.replaceAll("(?m)^\\\\n([ \\t]*)", "$1");
result = result.replaceAll("\n[ \\t]*\n", "\n");
return result;
}
} }

View File

@ -19,7 +19,8 @@ import org.springframework.stereotype.Component;
@Component @Component
@RequiredArgsConstructor @RequiredArgsConstructor
public class FunctionTargetHandler implements MutationTargetHandler { public class FunctionTargetHandler implements MutationTargetHandler {
private static final Set<String> FIELDS = Set.of("name", "signature", "comments", "flags"); private static final Set<String> FIELDS =
Set.of("name", "signature", "comments", "flags", "decompiledCode");
private final StaticFunctionRepository functionRepository; private final StaticFunctionRepository functionRepository;
private final ObjectMapper objectMapper; private final ObjectMapper objectMapper;
@ -42,6 +43,8 @@ public class FunctionTargetHandler implements MutationTargetHandler {
case "signature" -> func.getSignature(); case "signature" -> func.getSignature();
case "comments" -> func.getComments(); case "comments" -> func.getComments();
case "flags" -> func.getFlags(); case "flags" -> func.getFlags();
case "decompiledCode" ->
Objects.nonNull(func.getDecompiledCode()) ? func.getDecompiledCode() : "";
default -> throw new IllegalArgumentException("Unknown field: " + fieldName); default -> throw new IllegalArgumentException("Unknown field: " + fieldName);
}; };
} }
@ -60,6 +63,7 @@ public class FunctionTargetHandler implements MutationTargetHandler {
} }
} }
case "flags" -> func.setFlags(applyFlagOperation(func.getFlags(), newValue, operation)); case "flags" -> func.setFlags(applyFlagOperation(func.getFlags(), newValue, operation));
case "decompiledCode" -> func.setDecompiledCode(normalizePseudocode(newValue));
default -> throw new IllegalArgumentException("Unknown field: " + fieldName); default -> throw new IllegalArgumentException("Unknown field: " + fieldName);
} }
functionRepository.save(func); functionRepository.save(func);
@ -154,4 +158,32 @@ public class FunctionTargetHandler implements MutationTargetHandler {
.orElseThrow( .orElseThrow(
() -> new NotFoundException("Function with id '" + functionId + "' not found")); () -> new NotFoundException("Function with id '" + functionId + "' not found"));
} }
private String normalizePseudocode(String value) {
if (Objects.isNull(value) || value.isBlank()) {
return value;
}
if (value.contains("\n")) {
return value;
}
log.warn(
"decompiledCode has no line breaks — normalizing ({} chars)", value.length());
return formatMultiline(value);
}
private String formatMultiline(String code) {
var result =
code.replaceAll(";\\s*", ";\n")
.replaceAll("\\{\\s*", "{\n")
.replaceAll("\\s*\\}", "\n}")
.replaceAll("\n\n+", "\n")
.trim();
result = result.replaceAll("(?m)^\\\\n([ \\t]*)", "$1");
result = result.replaceAll("\n[ \\t]*\n", "\n");
return result;
}
} }

View File

@ -9,6 +9,8 @@
"analysis::events", "analysis::events",
"analysis::services", "analysis::services",
"analysis::dto", "analysis::dto",
"analysis::repositories",
"analysis::entities",
"engine::services", "engine::services",
"engine::model", "engine::model",
"ai", "ai",

View File

@ -3,6 +3,8 @@ package ai.decompile.job.service;
import ai.decompile.ai.service.EmbeddingService; import ai.decompile.ai.service.EmbeddingService;
import ai.decompile.analysis.event.MutationAppliedEvent; import ai.decompile.analysis.event.MutationAppliedEvent;
import ai.decompile.analysis.event.StaticAnalysisCompletedEvent; import ai.decompile.analysis.event.StaticAnalysisCompletedEvent;
import ai.decompile.analysis.model.repository.StaticAnalysisRepository;
import ai.decompile.analysis.model.repository.StaticFunctionRepository;
import ai.decompile.analysis.model.repository.StaticLabelRepository; import ai.decompile.analysis.model.repository.StaticLabelRepository;
import ai.decompile.job.messaging.JobMessagePublisher; import ai.decompile.job.messaging.JobMessagePublisher;
import ai.decompile.job.model.entity.Job; import ai.decompile.job.model.entity.Job;
@ -10,6 +12,8 @@ import ai.decompile.job.model.enums.JobStatus;
import ai.decompile.job.model.enums.JobType; import ai.decompile.job.model.enums.JobType;
import ai.decompile.job.model.repository.JobRepository; import ai.decompile.job.model.repository.JobRepository;
import java.util.Objects; import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2; import lombok.extern.log4j.Log4j2;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
@ -27,6 +31,8 @@ public class EmbeddingEventSubscriber {
private final JobMessagePublisher jobMessagePublisher; private final JobMessagePublisher jobMessagePublisher;
private final EmbeddingService embeddingService; private final EmbeddingService embeddingService;
private final StaticLabelRepository labelRepository; private final StaticLabelRepository labelRepository;
private final StaticFunctionRepository functionRepository;
private final StaticAnalysisRepository analysisRepository;
@ApplicationModuleListener @ApplicationModuleListener
public void on(StaticAnalysisCompletedEvent event) { public void on(StaticAnalysisCompletedEvent event) {
@ -71,6 +77,59 @@ public class EmbeddingEventSubscriber {
default: default:
break; break;
} }
updateCodeReferences(event);
}
private void updateCodeReferences(MutationAppliedEvent event) {
var analysisOpt = analysisRepository.findFirstByBinaryIdOrderByCreatedAtDesc(event.binaryId());
if (analysisOpt.isEmpty()) {
return;
}
var analysis = analysisOpt.get();
var affected =
functionRepository.findByAnalysisIdAndNameInCode(analysis.getId(), event.oldValue());
var pattern = Pattern.compile("\\b" + Pattern.quote(event.oldValue()) + "\\b");
var replacement = Matcher.quoteReplacement(event.newValue());
for (var fn : affected) {
if ("FUNCTION".equals(event.targetType()) && fn.getId().equals(event.targetId())) {
continue;
}
var updated = false;
if (Objects.nonNull(fn.getDecompiledCode()) && !fn.getDecompiledCode().isBlank()) {
var newDecompiled = pattern.matcher(fn.getDecompiledCode()).replaceAll(replacement);
if (!newDecompiled.equals(fn.getDecompiledCode())) {
fn.setDecompiledCode(newDecompiled);
updated = true;
}
}
if (Objects.nonNull(fn.getAssembly()) && !fn.getAssembly().isBlank()) {
var newAssembly = pattern.matcher(fn.getAssembly()).replaceAll(replacement);
if (!newAssembly.equals(fn.getAssembly())) {
fn.setAssembly(newAssembly);
updated = true;
}
}
if (updated) {
functionRepository.save(fn);
log.debug(
"Updated name reference {} → {} in function {} ({})",
event.oldValue(),
event.newValue(),
fn.getName(),
fn.getAddress());
}
}
} }
@Transactional(propagation = Propagation.REQUIRES_NEW) @Transactional(propagation = Propagation.REQUIRES_NEW)