Preserve callee saved registers.

src/main/java/com/bai/solver/PcodeVisitor.java

@@ -738,6 +738,12 @@ public void visit_CALL(PcodeOp pcode, AbsEnv inOutEnv, AbsEnv tmpEnv) {
                 }
             } else {
                 for (JImmutableMap.Entry<ALoc, KSet> entry : exit) {
+                    if (GlobalState.config.getPreserveCalleeSavedReg()) {
+                        // do not overwrite callee saved register
+                        if (GlobalState.arch.isCalleeSavedRegister(entry.getKey())) {
+                            continue;
+                        }
+                    }
                     inOutEnv.set(entry.getKey(), entry.getValue(), true);
                 }
             }
@@ -877,6 +883,12 @@ public void visit_CALLIND(PcodeOp pcode, AbsEnv inOutEnv, AbsEnv tmpEnv) {
 
         if (isFinished) {
             for (JImmutableTreeMap.Entry<ALoc, KSet> entry : resEnv.getEnvMap()) {
+                if (GlobalState.config.getPreserveCalleeSavedReg()) {
+                    // do not overwrite callee saved register
+                    if (GlobalState.arch.isCalleeSavedRegister(entry.getKey())) {
+                        continue;
+                    }
+                }
                 inOutEnv.set(entry.getKey(), entry.getValue(), true);
             }
         }

src/main/java/com/bai/util/Architecture.java

@@ -1,6 +1,6 @@
 package com.bai.util;
 
-
+import com.bai.env.ALoc;
 import com.bai.env.AbsVal;
 import com.bai.env.KSet;
 import com.bai.env.region.Global;
@@ -132,6 +132,24 @@ public boolean isX86() {
         return processor.equalsIgnoreCase("x86");
     }
 
+    /**
+     * Checks if current program is ARM architecture (32 bits).
+     * 
+     * @return true if it is ARM 32 bits, false otherwise
+     */
+    public boolean isArm32() {
+        return processor.equalsIgnoreCase("ARM");
+    }
+
+    /**
+     * Checks if current program is AArch64 architecture (64 bits).
+     * 
+     * @return true if it is AArch64, false otherwise
+     */
+    public boolean isAArch64() {
+        return processor.equalsIgnoreCase("AARCH64");
+    }
+
     /**
      * Get the KSet of pc register.
      * @param currentAddress address of current instruction.
@@ -164,4 +182,72 @@ public KSet getPcKSet(Address currentAddress) {
         pcKSet = pcKSet.insert(AbsVal.getPtr(Global.getInstance(), pcValue));
         return pcKSet;
     }
+
+    /**
+     * Checks if the register is a callee saved register.
+     * 
+     * @param aLoc abstract location of the register.
+     * @return the KSet.
+     */
+    public boolean isCalleeSavedRegister(ALoc aLoc) {
+        switch (processor) {
+            case "ARM":
+                return isArm32SavedRegister(aLoc);
+            case "AARCH64":
+                return isAArch64SavedRegister(aLoc);
+            case "x86":
+            default:
+                return false;
+        }
+    }
+
+    /**
+     * Checks if the register is a callee saved register of Arm 32 architecture.
+     * Checks if the register is in range r4 - r11, but not r9.
+     * See
+     * https://github.com/NationalSecurityAgency/ghidra/blob/ef3e74b6d3cb9e7ce1ccef20c8660dcb5346066e/Ghidra/Processors/ARM/data/languages/ARM.sinc#L10
+     * See also:
+     * https://github.com/ARM-software/abi-aa/blob/60a8eb8c55e999d74dac5e368fc9d7e36e38dda4/aapcs32/aapcs32.rst
+     * 
+     * @param aLoc abstract location of the register.
+     * @return the KSet.
+     */
+    public static boolean isArm32SavedRegister(ALoc aLoc) {
+        if (!aLoc.getRegion().isReg()) {
+            return false;
+        }
+        long start = aLoc.getBegin();
+        long end = start + aLoc.getLen();
+        // r9
+        if (start == 0x44 && end == 0x48) {
+            return false;
+        }
+        // r4 (0x30) - r11
+        if (start >= 0x30 && end <= 0x50) {
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Checks if the register is a callee saved register of AArch64 architecture.
+     * Checks if the register is in range x19 - x29.
+     * See
+     * https://github.com/NationalSecurityAgency/ghidra/blob/master/Ghidra/Processors/AARCH64/data/languages/AARCH64instructions.sinc#L448
+     * 
+     * @param aLoc abstract location of the register.
+     * @return the KSet.
+     */
+    public static boolean isAArch64SavedRegister(ALoc aLoc) {
+        if (!aLoc.getRegion().isReg()) {
+            return false;
+        }
+        long start = aLoc.getBegin();
+        long end = start + aLoc.getLen();
+        // r19 - r29
+        if (start >= 0x4098 && end <= 0x40F0) {
+            return true;
+        }
+        return false;
+    }
 }

src/main/java/com/bai/util/Config.java

@@ -56,6 +56,7 @@ private static void usage() {
             System.out.println("        [-disableZ3]");
             System.out.println("        [-all]");
             System.out.println("        [-debug]");
+            System.out.println("        [-PreserveCalleeSavedReg]");
             System.out.println("        [-check \"<cweNo1>[;<cweNo2>...]\"]");
         }
 
@@ -93,6 +94,8 @@ public static Config parseConfig(String fullArgs) {
                         CheckerManager.loadAllCheckers(config);
                     } else if (arg.equalsIgnoreCase("-debug")) {
                         config.setDebug(true);
+                    } else if (arg.equalsIgnoreCase("-preserveCalleeSavedReg")) {
+                        config.setPreserveCalleeSavedReg(true);
                     } else if (checkArgument("-check", args, argi)) {
                         String[] checkers = getSubArguments(args, argi);
                         Arrays.stream(checkers)
@@ -142,6 +145,8 @@ public static Config parseConfig(String fullArgs) {
 
     private boolean isGUI;
 
+    private boolean preserveCalleeSavedReg;
+
     // for tactic tuning, see:
     // http://www.cs.tau.ac.il/~msagiv/courses/asv/z3py/strategies-examples.htm
     private List<String> z3Tactics = new ArrayList<>();
@@ -157,6 +162,7 @@ public Config() {
         this.entryAddress = null;
         this.isEnableZ3 = true;
         this.externalMapPath = null;
+        this.preserveCalleeSavedReg = false;
     }
 
     /**
@@ -350,6 +356,22 @@ public void setGUI(boolean isGUI) {
         this.isGUI = isGUI;
     }
 
+    /**
+     * Preserve the callee saved registers.
+     * @param preserveCalleeSavedReg preserve or not.
+     */
+    public void setPreserveCalleeSavedReg(boolean preserveCalleeSavedReg) {
+        this.preserveCalleeSavedReg = preserveCalleeSavedReg;
+    }
+
+    /**
+     *  Get if the callee saved registers are preserved.
+     * @return true if preserved, false otherwise.
+     */
+    public boolean getPreserveCalleeSavedReg() {
+        return preserveCalleeSavedReg;
+    }
+
     @Override
     public String toString() {
         return "Config{"