From 39f6445c232bb5caf5c4a2a996de91dfa20c48e8 Mon Sep 17 00:00:00 2001
From: Mike Hunhoff <mehunhoff@google.com>
Date: Tue, 23 Apr 2024 11:37:37 -0700
Subject: [PATCH] BinExport: Ghidra: Export module mapping for imported
 functions

PiperOrigin-RevId: 627454882
Change-Id: I9fb6e24ef6897b3d9d351caacecc2fe1aff7a19b
---
 .../security/binexport/BinExport2Builder.java | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/java/src/main/java/com/google/security/binexport/BinExport2Builder.java b/java/src/main/java/com/google/security/binexport/BinExport2Builder.java
index 8704f1c2..1a574c24 100644
--- a/java/src/main/java/com/google/security/binexport/BinExport2Builder.java
+++ b/java/src/main/java/com/google/security/binexport/BinExport2Builder.java
@@ -34,6 +34,7 @@
 import ghidra.program.model.listing.Function;
 import ghidra.program.model.listing.FunctionManager;
 import ghidra.program.model.listing.Instruction;
+import ghidra.program.model.listing.Library;
 import ghidra.program.model.listing.Listing;
 import ghidra.program.model.listing.Parameter;
 import ghidra.program.model.listing.Program;
@@ -368,7 +369,7 @@ private void buildFlowGraphs(Map<Long, Integer> basicBlockIndices) throws Cancel
     }
   }
 
-  private void buildCallGraph() throws CancelledException {
+  private void buildCallGraphAndModuleList() throws CancelledException {
     var callGraph = builder.getCallGraphBuilder();
     FunctionManager funcManager = program.getFunctionManager();
     monitor.setIndeterminate(false);
@@ -377,6 +378,7 @@ private void buildCallGraph() throws CancelledException {
     int id = 0;
     TreeMap<Address, Function> orderedFunctions = new TreeMap<>();
     Map<Long, Integer> vertexIndices = new HashMap<>();
+    Map<String, Integer> moduleIndices = new HashMap<>();
 
     // First round, create ordered function mapping because Ghidra does not guarantee the order of
     // external (imported) functions and BinExport requires ordered vertices
@@ -410,6 +412,21 @@ private void buildCallGraph() throws CancelledException {
       }
       if (func.isExternal()) {
         vertex.setType(BinExport2.CallGraph.Vertex.Type.IMPORTED);
+        // Add module mapping
+        String moduleName = func.getParentNamespace().getName();
+        if (!moduleName.equals(Library.UNKNOWN)) {
+          if (moduleName.contains(".")) {
+            moduleName = moduleName.substring(0, moduleName.lastIndexOf('.'));
+          }
+          Integer moduleId =
+              moduleIndices.computeIfAbsent(
+                  moduleName,
+                  (String k) -> {
+                    builder.addModuleBuilder().setName(k);
+                    return builder.getModuleCount() - 1;
+                  });
+          vertex.setModuleIndex(moduleId);
+        }
       }
       if (!func.getName().equals(SymbolUtilities.getDefaultFunctionName(func.getEntryPoint()))) {
         // Ghidra does not seem to provide both mangled and demangled names
@@ -701,7 +718,7 @@ public BinExport2 build(TaskMonitor taskMonitor) throws CancelledException {
     monitor.setMessage("Exporting flow graphs");
     buildFlowGraphs(basicBlockIndices);
     monitor.setMessage("Exporting call graph");
-    buildCallGraph();
+    buildCallGraphAndModuleList();
     buildSections();
 
     return builder.build();