From 61b3c7884c4b26a06899428bfd740b70d1734cfe Mon Sep 17 00:00:00 2001
From: TSR Berry <20988865+TSRBerry@users.noreply.github.com>
Date: Thu, 17 Oct 2024 02:47:52 +0200
Subject: [PATCH 1/2] Mark PLT functions as Thunks

---
 .../java/adubbz/nx/loader/SwitchLoader.java   |  4 +-
 .../nx/loader/common/NXProgramBuilder.java    | 41 ++++++++++++++++---
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/src/main/java/adubbz/nx/loader/SwitchLoader.java b/src/main/java/adubbz/nx/loader/SwitchLoader.java
index 36870be..38c3fd8 100644
--- a/src/main/java/adubbz/nx/loader/SwitchLoader.java
+++ b/src/main/java/adubbz/nx/loader/SwitchLoader.java
@@ -131,8 +131,8 @@ protected List<Loaded<Program>> loadProgram(ByteProvider provider, String progra
 
     @Override
     protected void loadProgramInto(ByteProvider provider, LoadSpec loadSpec, List<Option> options,
-            MessageLog messageLog, Program program, TaskMonitor monitor) 
-                    throws IOException
+            MessageLog messageLog, Program program, TaskMonitor monitor)
+                    throws IOException, CancelledException
     {
         var space = program.getAddressFactory().getDefaultAddressSpace();
         
diff --git a/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java b/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java
index 5d99545..3dfc90a 100644
--- a/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java
+++ b/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java
@@ -25,6 +25,7 @@
 import ghidra.app.util.bin.format.elf.relocation.AARCH64_ElfRelocationType;
 import ghidra.app.util.bin.format.elf.relocation.ARM_ElfRelocationType;
 import ghidra.app.util.importer.MessageLog;
+import ghidra.program.disassemble.Disassembler;
 import ghidra.program.model.address.*;
 import ghidra.program.model.data.PointerDataType;
 import ghidra.program.model.data.TerminatedStringDataType;
@@ -35,6 +36,7 @@
 import ghidra.program.model.symbol.*;
 import ghidra.program.model.util.CodeUnitInsertionException;
 import ghidra.util.Msg;
+import ghidra.util.exception.CancelledException;
 import ghidra.util.exception.DuplicateNameException;
 import ghidra.util.exception.InvalidInputException;
 import ghidra.util.exception.NotFoundException;
@@ -69,7 +71,7 @@ public NXProgramBuilder(Program program, ByteProvider provider, NXOAdapter adapt
         this.nxo = new NXO(program, adapter, program.getImageBase().getOffset());
     }
     
-    public void load(TaskMonitor monitor)
+    public void load(TaskMonitor monitor) throws CancelledException
     {
         NXOAdapter adapter = this.nxo.getAdapter();
         ByteProvider memoryProvider = adapter.getMemoryProvider();
@@ -117,7 +119,7 @@ public void load(TaskMonitor monitor)
                 this.memBlockHelper.addSection(".dynsym", adapter.getSymbolTable(this.program).getFileOffset() - this.nxo.getBaseAddress(), adapter.getSymbolTable(this.program).getFileOffset() - this.nxo.getBaseAddress(), adapter.getSymbolTable(this.program).getLength(), true, false, false);
             }
             
-            this.setupRelocations();
+            this.setupRelocations(monitor);
             this.createGlobalOffsetTable();
             
             this.memBlockHelper.addFillerSection(".text", text.getOffset(), text.getSize(), true, false, true);
@@ -205,7 +207,7 @@ protected void setupSymbolTable()
         }
     }
     
-    protected void setupRelocations() throws AddressOutOfBoundsException, NotFoundException, IOException {
+    protected void setupRelocations(TaskMonitor monitor) throws AddressOutOfBoundsException, NotFoundException, IOException, CancelledException {
         NXOAdapter adapter = this.nxo.getAdapter();
         ByteProvider memoryProvider = adapter.getMemoryProvider();
         BinaryReader memoryReader = adapter.getMemoryReader();
@@ -274,6 +276,8 @@ protected void setupRelocations() throws AddressOutOfBoundsException, NotFoundEx
             long pltStart = this.pltEntries.get(0).off;
             long pltEnd = this.pltEntries.get(this.pltEntries.size() - 1).off + 0x10;
             this.memBlockHelper.addSection(".plt", pltStart, pltStart, pltEnd - pltStart, true, false, false);
+            // Disassemble the entire section, so AARCH64PltThunkAnalyzer works.
+            disassembleRange(program.getImageBase().add(pltStart), program.getImageBase().add(pltEnd), program, monitor);
         }
         else {
             // TODO: Find a way to locate the plt in CFI-enabled binaries.
@@ -382,12 +386,15 @@ else if (reloc.r_type == R_FAKE_RELR) {
         {
             if (gotNameLookup.containsKey(entry.target))
             {
-                long addr = this.nxo.getBaseAddress() + entry.off;
+                Address addr = this.aSpace.getAddress(this.nxo.getBaseAddress() + entry.off);
                 String name = gotNameLookup.get(entry.target);
-                // TODO: Mark as func
                 if (name != null && !name.isEmpty())
                 {
-                    this.createSymbol(this.aSpace.getAddress(addr), name, false, false, null);
+                    Function func = createOneByteFunction(name, addr, false);
+                    ExternalLocation extLoc = program.getExternalManager().getUniqueExternalLocation(Library.UNKNOWN, name);
+                    if (extLoc != null) {
+                        func.setThunkedFunction(extLoc.getFunction());
+                    }
                 }
             }
         }
@@ -642,6 +649,28 @@ public Symbol createSymbol(Address addr, String name, boolean isPrimary, boolean
         }
         return sym;
     }
+
+    // Source: https://github.com/NationalSecurityAgency/ghidra/blob/de7c3eaee2a4bc993a402e371b039c2bb2d6c545/Ghidra/Features/Base/src/main/java/ghidra/app/util/bin/format/elf/ElfDefaultGotPltMarkup.java#L545
+    private void disassembleRange(Address start, Address end, Program program, TaskMonitor monitor)
+            throws CancelledException {
+        AddressSet set = new AddressSet(start, end);
+        Disassembler disassembler = Disassembler.getDisassembler(program, monitor, m -> {
+            /* silent */});
+        while (!set.isEmpty()) {
+            monitor.checkCancelled();
+            AddressSet disset = disassembler.disassemble(set.getMinAddress(), null, true);
+            if (disset.isEmpty()) {
+                // Stop on first error but discard error bookmark since
+                // some plt sections are partly empty and must rely
+                // on normal flow disassembly during analysis
+                program.getBookmarkManager()
+                        .removeBookmarks(set, BookmarkType.ERROR,
+                                Disassembler.ERROR_BOOKMARK_CATEGORY, monitor);
+                break;//we did not disassemble anything...
+            }
+            set.delete(disset);
+        }
+    }
     
     private Symbol checkPrimary(Symbol sym) 
     {

From 882fd718068f35938431376954b00d10e17e87bb Mon Sep 17 00:00:00 2001
From: TSR Berry <20988865+TSRBerry@users.noreply.github.com>
Date: Fri, 1 Nov 2024 22:51:30 +0100
Subject: [PATCH 2/2] Improve thunk creation for external references

This also makes AARCH64PltThunkAnalyzer work for references within the same binary.
---
 .../nx/loader/common/NXProgramBuilder.java    | 28 +++++++++++++++++--
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java b/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java
index 3dfc90a..2bff236 100644
--- a/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java
+++ b/src/main/java/adubbz/nx/loader/common/NXProgramBuilder.java
@@ -25,6 +25,7 @@
 import ghidra.app.util.bin.format.elf.relocation.AARCH64_ElfRelocationType;
 import ghidra.app.util.bin.format.elf.relocation.ARM_ElfRelocationType;
 import ghidra.app.util.importer.MessageLog;
+import ghidra.program.database.function.OverlappingFunctionException;
 import ghidra.program.disassemble.Disassembler;
 import ghidra.program.model.address.*;
 import ghidra.program.model.data.PointerDataType;
@@ -276,7 +277,7 @@ protected void setupRelocations(TaskMonitor monitor) throws AddressOutOfBoundsEx
             long pltStart = this.pltEntries.get(0).off;
             long pltEnd = this.pltEntries.get(this.pltEntries.size() - 1).off + 0x10;
             this.memBlockHelper.addSection(".plt", pltStart, pltStart, pltEnd - pltStart, true, false, false);
-            // Disassemble the entire section, so AARCH64PltThunkAnalyzer works.
+            // Disassemble the entire section, so AARCH64PltThunkAnalyzer works for functions within this binary.
             disassembleRange(program.getImageBase().add(pltStart), program.getImageBase().add(pltEnd), program, monitor);
         }
         else {
@@ -390,10 +391,31 @@ else if (reloc.r_type == R_FAKE_RELR) {
                 String name = gotNameLookup.get(entry.target);
                 if (name != null && !name.isEmpty())
                 {
-                    Function func = createOneByteFunction(name, addr, false);
                     ExternalLocation extLoc = program.getExternalManager().getUniqueExternalLocation(Library.UNKNOWN, name);
+                    // AARCH64PltThunkAnalyzer won't be able to find a valid destination address for entries referencing external functions.
                     if (extLoc != null) {
-                        func.setThunkedFunction(extLoc.getFunction());
+                        Address lastAddr = addr.add(0x10 - 0x4);
+                        Instruction lastInstruction = program.getListing().getInstructionAt(lastAddr);
+
+                        // Make sure the last instruction of this function is a branch.
+                        // This check is also performed by AARCH64PltThunkAnalyzer.
+                        if (lastInstruction != null && lastInstruction.getMnemonicString().equals("br")) {
+                            try {
+                                program.getFunctionManager().createThunkFunction(name, null, addr, new AddressSet(addr, lastAddr), extLoc.getFunction(), SourceType.IMPORTED);
+                                lastInstruction.setFlowOverride(FlowOverride.CALL_RETURN);
+                            } catch (OverlappingFunctionException e) {
+                                Msg.error(this, String.format("Couldn't create thunk function at %s: %s", addr, e));
+                            }
+                        }
+                        else {
+                            // This shouldn't happen.
+                            Msg.warn(this, String.format("Couldn't find a valid last instruction for thunk function %s at: %s", addr, lastAddr));
+                            createOneByteFunction(name, addr, false).setThunkedFunction(extLoc.getFunction());
+                        }
+                    }
+                    else {
+                        // Already defined functions are skipped by AARCH64PltThunkAnalyzer, so we only create a symbol.
+                        this.createSymbol(addr, name, false, false, null);
                     }
                 }
             }