Improve ZydisRegisterGetLargestEnclosing (#456)

zyantific · Oct 9, 2023 · 6ab4463 · 6ab4463
1 parent 8a121be
commit 6ab4463
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 18 deletions.
diff --git a/include/Zydis/Register.h b/include/Zydis/Register.h
@@ -282,7 +282,7 @@ ZYDIS_EXPORT ZydisRegisterWidth ZydisRegisterGetWidth(ZydisMachineMode mode, Zyd
  * @param   reg     The register.
  *
  * @return  The largest enclosing register of the given register, or `ZYDIS_REGISTER_NONE` if the
- *          register is invalid for the active machine-mode or does not have an enclosing-register.
+ *          register is invalid for the active machine-mode.
  */
 ZYDIS_EXPORT ZydisRegister ZydisRegisterGetLargestEnclosing(ZydisMachineMode mode,
     ZydisRegister reg);

diff --git a/src/Register.c b/src/Register.c
@@ -156,6 +156,52 @@ ZydisRegister ZydisRegisterGetLargestEnclosing(ZydisMachineMode mode, ZydisRegis
         return ZYDIS_REGISTER_NONE;
     }
 
+    if (mode > ZYDIS_MACHINE_MODE_MAX_VALUE)
+    {
+        return ZYDIS_REGISTER_NONE;
+    }
+
+    const ZydisRegisterClass reg_class = REG_LOOKUP[reg].class;
+
+    if ((reg_class == ZYDIS_REGCLASS_INVALID) ||
+        ((reg_class == ZYDIS_REGCLASS_GPR64) && (mode != ZYDIS_MACHINE_MODE_LONG_64)))
+    {
+        return ZYDIS_REGISTER_NONE;
+    }
+
+    static const ZydisRegister STATIC_MAPPING[ZYDIS_REGCLASS_MAX_VALUE + 1][3] =
+    {
+                                 /* 16              */ /* 32               */ /* 64                  */
+        [ZYDIS_REGCLASS_FLAGS] = { ZYDIS_REGISTER_FLAGS, ZYDIS_REGISTER_EFLAGS, ZYDIS_REGISTER_RFLAGS },
+        [ZYDIS_REGCLASS_IP   ] = { ZYDIS_REGISTER_IP   , ZYDIS_REGISTER_EIP   , ZYDIS_REGISTER_RIP    },
+    };
+    ZYAN_ASSERT(reg_class < ZYAN_ARRAY_LENGTH(STATIC_MAPPING));
+
+    ZyanU8 mode_bits;
+    switch (mode)
+    {
+    case ZYDIS_MACHINE_MODE_LONG_64:
+        mode_bits = 2;
+        break;
+    case ZYDIS_MACHINE_MODE_LONG_COMPAT_32:
+    case ZYDIS_MACHINE_MODE_LEGACY_32:
+        mode_bits = 1;
+        break;
+    case ZYDIS_MACHINE_MODE_LONG_COMPAT_16:
+    case ZYDIS_MACHINE_MODE_LEGACY_16:
+    case ZYDIS_MACHINE_MODE_REAL_16:
+        mode_bits = 0;
+        break;
+    default:
+        ZYAN_UNREACHABLE;
+    }
+
+    const ZydisRegister static_reg = STATIC_MAPPING[reg_class][mode_bits];
+    if (static_reg != ZYDIS_REGISTER_NONE)
+    {
+        return static_reg;
+    }
+
     static const ZyanU8 GPR8_MAPPING[20] =
     {
         /* AL   */  0,
@@ -177,16 +223,9 @@ ZydisRegister ZydisRegisterGetLargestEnclosing(ZydisMachineMode mode, ZydisRegis
         /* R12B */ 12,
         /* R13B */ 13,
         /* R14B */ 14,
-        /* R15B */ 15,
+        /* R15B */ 15
     };
 
-    const ZydisRegisterClass reg_class = REG_LOOKUP[reg].class;
-    if ((reg_class == ZYDIS_REGCLASS_INVALID) ||
-        ((reg_class == ZYDIS_REGCLASS_GPR64) && (mode != ZYDIS_MACHINE_MODE_LONG_64)))
-    {
-        return ZYDIS_REGISTER_NONE;
-    }
-
     ZyanU8 reg_id = REG_LOOKUP[reg].id;
     switch (reg_class)
     {
@@ -196,19 +235,16 @@ ZydisRegister ZydisRegisterGetLargestEnclosing(ZydisMachineMode mode, ZydisRegis
     case ZYDIS_REGCLASS_GPR16:
     case ZYDIS_REGCLASS_GPR32:
     case ZYDIS_REGCLASS_GPR64:
-        switch (mode)
+        switch (mode_bits)
         {
-        case ZYDIS_MACHINE_MODE_LONG_64:
+        case 2:
             return REG_CLASS_LOOKUP[ZYDIS_REGCLASS_GPR64].lo + reg_id;
-        case ZYDIS_MACHINE_MODE_LONG_COMPAT_32:
-        case ZYDIS_MACHINE_MODE_LEGACY_32:
+        case 1:
             return REG_CLASS_LOOKUP[ZYDIS_REGCLASS_GPR32].lo + reg_id;
-        case ZYDIS_MACHINE_MODE_LONG_COMPAT_16:
-        case ZYDIS_MACHINE_MODE_LEGACY_16:
-        case ZYDIS_MACHINE_MODE_REAL_16:
+        case 0:
             return REG_CLASS_LOOKUP[ZYDIS_REGCLASS_GPR16].lo + reg_id;
         default:
-            return ZYDIS_REGISTER_NONE;
+            ZYAN_UNREACHABLE;
         }
     case ZYDIS_REGCLASS_XMM:
     case ZYDIS_REGCLASS_YMM:
@@ -219,7 +255,7 @@ ZydisRegister ZydisRegisterGetLargestEnclosing(ZydisMachineMode mode, ZydisRegis
         return REG_CLASS_LOOKUP[ZYDIS_REGCLASS_ZMM].lo + reg_id;
 #endif
     default:
-        return ZYDIS_REGISTER_NONE;
+        return reg;
     }
 }