Skip to content

Instantly share code, notes, and snippets.

@lifthrasiir

lifthrasiir/disfilter.rs

Created June 18, 2024 01:25
Show Gist options
  • Save lifthrasiir/df47509caac2f065032ef72e70f2ec05 to your computer and use it in GitHub Desktop.
Save lifthrasiir/df47509caac2f065032ef72e70f2ec05 to your computer and use it in GitHub Desktop.
Download ZIP
Disfilter: Fabian Giesen's x86-32 transformer, reworked for x86-64 in Rust (WIP)
Raw
disfilter.rs
// TODO:
// - Better error checking
// - Better code structure
// - Detect function boundary (skip CC or multi-byte nops)
use std::array;
use std::cell::Cell;
use std::fmt;
use std::fs::File;
use std::io::{self, BufReader, Read, Seek, SeekFrom, Write};
use std::mem;
use bytemuck::must_cast_slice;
macro_rules! log_encode {
() => {
cfg!(trace_encode)
};
}
macro_rules! log_decode {
() => {
cfg!(trace_decode)
};
}
fn to_u16(s: &[u8]) -> u16 {
u16::from_le_bytes(s.try_into().unwrap())
}
fn to_u32(s: &[u8]) -> u32 {
u32::from_le_bytes(s.try_into().unwrap())
}
fn to_u64(s: &[u8]) -> u64 {
u64::from_le_bytes(s.try_into().unwrap())
}
const N_: u8 = 0b0000; // no immediate
const N1: u8 = 0b0001; // 8-bit immediate
const N2: u8 = 0b0010; // 16-bit immediate
const N4: u8 = 0b0011; // 32/16-bit immediate, depending on operand size prefix
const NZ: u8 = 0b0100; // 32/64-bit immediate, depending on REX.W
const J1: u8 = 0b0101; // 8-bit relative jump target
const J4: u8 = 0b0110; // 32-bit relative jump target
const JA: u8 = 0b0111; // 32-bit absolute jump target
const A_: u8 = 0b1000; // 16/32/64-bit absolute address that is not a jump target,
// exact bit size depends on the instruction and context
const R1: u8 = 0b1001; // modR/M + 8-bit immediate
const XX: u8 = 0b1010; // invalid, has to be escaped
const R4: u8 = 0b1011; // modR/M + 32/16-bit immediate, depending on operand size prefix
const R_: u8 = 0b1100; // modR/M + no immediate
const BP: u8 = 0b1101; // 1-byte prefix that will be marked
const M2: u8 = 0b1110; // opcode byte + modR/M + no immediate (map 2)
const M3: u8 = 0b1111; // opcode byte + modR/M + 8-bit immediate (map 3)
#[inline(always)]
const fn op3_followed(flags: u8) -> bool {
flags >= M2
}
#[inline(always)]
const fn has_modrm(flags: u8) -> bool {
flags >= R1
}
#[inline(always)]
const fn modrm_to_imm(flags: u8) -> u8 {
flags & 3
}
const _: () = {
assert!(N_ + 1 == N1 && N1 + 1 == N2 && N2 + 1 == N4 && N4 + 1 == NZ);
assert!(modrm_to_imm(R_) == N_ && modrm_to_imm(R1) == N1 && modrm_to_imm(R4) == N4);
assert!(M2 + 1 == M3);
};
// https://sandpile.org/
//
// https://github.com/torvalds/linux/blob/master/arch/x86/lib/x86-opcode-map.txt
// - Placeholders are an argument made of one uppercase letter and one lowercase letter.
// - Immediate or address is present if some placeholder starts with AIJLO.
// - For -b/-w/-d/-z placeholders, its size is fixed to 8/16/32/32 bits.
// - For -v placeholders, its size is normally 32 bits but becomes 16 bits with 66 prefixed.
// - -p placeholder introduces an additional 16-bit segment selector before address.
// - L- placeholder introduces an additional 8-bit register selector.
// - ModR/M byte is present if some placeholder starts with CDEGMNPQRSTUVW.
// 1-byte opcodes (legacy map 0)
const OPCODES0: [u8; 256] = [
R_, R_, R_, R_, N1, N4, N_, N_, R_, R_, R_, R_, N1, N4, N_, XX, // 0
R_, R_, R_, R_, N1, N4, N_, N_, R_, R_, R_, R_, N1, N4, N_, N_, // 1
R_, R_, R_, R_, N1, N4, BP, N_, R_, R_, R_, R_, N1, N4, BP, N_, // 2
R_, R_, R_, R_, N1, N4, BP, N_, R_, R_, R_, R_, N1, N4, BP, N_, // 3
N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, // 4
N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, // 5
N_, N_, R_, R_, BP, BP, BP, BP, N4, R4, N1, R1, N_, N_, N_, N_, // 6
J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, J1, // 7
R1, R4, R1, R1, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 8
N_, N_, N_, N_, N_, N_, N_, N_, N_, N_, JA, N_, N_, N_, N_, N_, // 9
A_, A_, A_, A_, N_, N_, N_, N_, N1, N4, N_, N_, N_, N_, N_, N_, // A
N1, N1, N1, N1, N1, N1, N1, N1, NZ, NZ, NZ, NZ, NZ, NZ, NZ, NZ, // B
R1, R1, N2, N_, R_, R_, R1, R4, N2, N_, N2, N_, N_, N1, N_, N_, // C
R_, R_, R_, R_, N1, N1, N_, N_, R_, R_, R_, R_, R_, R_, R_, R_, // D
J1, J1, J1, J1, N1, N1, N1, N1, J4, J4, A_, J1, N_, N_, N_, N_, // E
BP, N_, BP, BP, N_, N_, R1, R4, N_, N_, N_, N_, N_, N_, R_, R_, // F
];
// Some 1-byte opcodes are outright invalid in x86-64.
const OPCODES0_I64: [u32; 8] = [
//FEDCBA9876543210 FEDCBA9876543210
0b0100000011000000_1100000011000000, // 1x, 0x
0b1000000010000000_1000000010000000, // 3x, 2x
0b0000000000000000_1111111111111111, // 5x, 4x
0b0000000000000000_0000000000000111, // 7x, 6x
0b0000010000000000_0000000000000100, // 9x, 8x
0b0000000000000000_0000000000000000, // Bx, Ax
0b0000000001110000_0100000000110000, // Dx, Cx
0b0000000000000000_0000010000000000, // Fx, Ex
];
const PRE_VEX3: u8 = 0xc4;
const PRE_VEX2: u8 = 0xc5;
const PRE_EVEX: u8 = 0x62;
const PRE_REX2: u8 = 0xd5;
const PRE_2BYTE: u8 = 0x0f;
const PRE_OSIZE: u8 = 0x66;
const PRE_REPNE: u8 = 0xf2;
const PRE_REP: u8 = 0xf3;
const OP_CALLN: u16 = 0x0_e8; // CALL Jz
// Opcodes in the map 0 that need an additional 16-bit immediate.
const OP_CALLF: u16 = 0x0_9a; // CALL Ap (32-bit only)
const OP_JMPF: u16 = 0x0_ea; // JMP Ap (32-bit only)
const OP_ENTER: u16 = 0x0_c8; // ENTER Iw,Ib
// Opcodes in the map 0 that have immediates only with /0 or /1.
const OP_GRP3_1: u16 = 0x0_f6; // TEST Eb,Ib; NOT/NEG/[I]MUL/[I]DIV Eb
const OP_GRP3_2: u16 = 0x0_f7; // TEST Ev,Iv; NOT/NEG/[I]MUL/[I]DIV Ev
// 2-byte opcodes, starting with 0F (legacy map 1)
const OPCODES1: [u8; 256] = [
R_, R_, N_, N_, XX, N_, N_, N_, N_, N_, XX, XX, XX, R_, N_, R1, // 0F 0
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, XX, R_, R_, // 0F 1
R_, R_, R_, R_, XX, XX, XX, XX, R_, R_, R_, R_, R_, R_, R_, R_, // 0F 2
N_, N_, N_, N_, N_, N_, XX, N_, M2, XX, M3, XX, XX, XX, XX, XX, // 0F 3
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 0F 4
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 0F 5
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 0F 6
R1, R1, R1, R1, R_, R_, R_, N_, N_, R_, R_, R_, R_, R_, R_, R_, // 0F 7
J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, J4, // 0F 8
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 0F 9
N_, N_, N_, R_, R1, R_, R_, R_, N_, N_, N_, R_, R1, R_, R_, R_, // 0F A
R_, R_, R_, R_, R_, R_, R_, R_, N_, R_, R1, R_, R_, R_, R_, R_, // 0F B
R_, R_, R1, R_, R1, R1, R1, R_, N_, N_, N_, N_, N_, N_, N_, N_, // 0F C
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 0F D
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 0F E
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, XX, // 0F F
];
// Opcode in the map 1 that has a ModR/M byte only with some prefixes.
const OP_JMPE_POPCNT: u16 = 0x1_b8; // JMPE Jz (IA-64 only); POPCNT Gv,Ev (with F3 prefix)
// EVEX opcode map 4 (largely a subset of the opcode map 0)
const OPCODES4: [u8; 256] = [
R_, R_, R_, R_, XX, XX, XX, XX, R_, R_, R_, R_, XX, XX, XX, XX, // 0
R_, R_, R_, R_, XX, XX, XX, XX, R_, R_, R_, R_, XX, XX, XX, XX, // 1
R_, R_, R_, R_, R1, XX, XX, XX, R_, R_, R_, R_, R1, XX, XX, XX, // 2
R_, R_, R_, R_, XX, XX, XX, XX, R_, R_, R_, R_, XX, XX, XX, XX, // 3
R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, R_, // 4
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 5
R_, R_, XX, XX, XX, R_, R_, XX, XX, R4, XX, R1, XX, XX, XX, XX, // 6
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 7
R1, R4, XX, R1, N_, N_, XX, XX, R_, XX, XX, XX, XX, XX, XX, R_, // 8
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 9
XX, XX, XX, XX, XX, R_, XX, XX, XX, XX, XX, XX, XX, R_, XX, R_, // A
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // B
R1, R1, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // C
R_, R_, R_, R_, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // D
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // E
R_, R_, R_, XX, R_, R_, R1, R4, R_, R_, XX, XX, XX, XX, R_, R_, // F
];
const OP_MAP4_GRP3_1: u16 = 0x4_f6;
const OP_MAP4_GRP3_2: u16 = 0x4_f7;
// EVEX opcode map 7
const OPCODES7: [u8; 256] = [
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 0
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 1
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 2
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 3
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 4
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 5
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 6
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 7
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 8
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // 9
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // A
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // B
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // C
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // D
XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, XX, // E
XX, XX, XX, XX, XX, XX, XX, XX, R4, XX, XX, XX, XX, XX, XX, XX, // F
];
// Opcode in the map 7 which immediate isn't affected by the operand size prefix.
const OP_URDMSR_UWRMSR: u16 = 0x7_f8; // URDMSR Rq,Id; UWRMSR Id,Rq
// Special bytes in the ST_OP stream. They should be available in any operating modes
// and decode without any more operand, but yet have to be extremely unlikely to occur.
//
// - Since we regularize any *VEX & REX prefixes, a single byte REX prefix is
// also rewritten as a uniform marker followed by a normalized REX data byte.
// REX bytes themselves (40-4F) are used as an opcode map prefix instead,
// so the original 2-byte prefix (0F) is repurposed for the REX marker instead.
//
// - Jump tables are indicated by INT1 (F1), which is meant to be used for hardware debugging
// and therefore extremely unlikely to occur in the regular x86 opcode stream.
//
// - Verbatim bytes are indicated by HLT (F4), which is privileged and can only occur rarely
// due to its semantics, making it a good opcode to steal.
//
// The original disfilter used to use INTO (CE) instead of HLT, but it is now invalid in
// the long mode and has a chance to be repurposed in the future.
const REX_MARKER: u8 = 0x0f;
const JUMPTAB: u8 = 0xf1;
const ESC: u8 = 0xf4;
#[inline(always)]
fn lookup_opcode(op: u8, map: u8, is64: bool) -> u8 {
const OPCODES: [u16; 256] = {
let mut tab = [0u16; 256];
let mut i = 0;
while i < 256 {
assert!(OPCODES0[i] | OPCODES1[i] | OPCODES4[i] | OPCODES7[i] < 16);
tab[i] = OPCODES0[i] as u16
| (OPCODES1[i] as u16) << 4
| (OPCODES4[i] as u16) << 8
| (OPCODES7[i] as u16) << 12;
i += 1;
}
// Override escape codes.
assert!(OPCODES0[ESC as usize] == N_);
assert!(OPCODES0[JUMPTAB as usize] == N_);
tab[ESC as usize] = tab[ESC as usize] & !0xf | XX as u16;
tab[JUMPTAB as usize] = tab[JUMPTAB as usize] & !0xf | XX as u16;
tab
};
const MAP_FLAGS: [u8; 16] = [
0x20, 0x24, R_, R1, 0x28, R_, R_, 0x2c, XX, XX, XX, XX, XX, XX, XX, XX,
];
if map == 0 && is64 && OPCODES0_I64[op as usize >> 5] >> (op as u32 & 31) & 1 == 1 {
XX
} else {
let map_flags = MAP_FLAGS[map as usize];
if map_flags < 0x20 {
map_flags
} else {
(OPCODES[op as usize] >> (map_flags & 0x1f)) as u8 & 15
}
}
}
#[inline(always)]
const fn prefix_hash(b: u8) -> u32 {
let b = b as u32;
((b << 2) ^ (b >> 2)) & 31
}
#[inline(always)]
const fn has_osize_prefix(prefixes: u32) -> bool {
(prefixes >> prefix_hash(PRE_OSIZE)) & 1 != 0
}
#[inline(always)]
const fn has_rep_prefix(prefixes: u32) -> bool {
prefixes & ((1 << prefix_hash(PRE_REPNE)) | (1 << prefix_hash(PRE_REP))) != 0
}
const _: () = {
let prefixes = [
0x26, 0x2e, 0x36, 0x3e, 0x64, 0x65, 0x66, 0x67, 0xf0, 0xf2, 0xf3,
];
let mut bitset = 0;
let mut i = 0;
while i < prefixes.len() {
let b = prefixes[i];
bitset |= 1u32 << prefix_hash(b);
assert!(OPCODES0[b as usize] == BP);
i += 1;
}
assert!(bitset.count_ones() == 11);
};
macro_rules! define_streams {
($($i:ident $s:literal),* $(,)?) => (
define_streams! { @0 $($i)* }
const STREAM_NAMES: [&str; NUM_STREAMS] = [$(stringify!($i)),*];
const STREAM_SIZES: [usize; NUM_STREAMS] = [$($s / 8),*];
);
(@$c:tt $i:ident $($t:tt)*) => (
#[allow(dead_code)] const $i: usize = $c;
define_streams! { @(1 + $c) $($t)* }
);
(@$c:tt) => (
const NUM_STREAMS: usize = $c;
);
}
define_streams! {
ST_OP 8,
ST_EVEX 8, ST_VEX 8, ST_REX 8, ST_SIB 8,
ST_CALL_IDX 8,
ST_DISP8_R0 8, ST_DISP8_R1 8, ST_DISP8_R2 8, ST_DISP8_R3 8,
ST_DISP8_R4 8, ST_DISP8_R5 8, ST_DISP8_R6 8, ST_DISP8_R7 8,
ST_DISP32 32,
ST_JUMP8 8, ST_JUMP32 32, ST_JUMP64 64,
ST_IMM8 8, ST_IMM16 16, ST_IMM32 32, ST_IMM64 64,
ST_ADDR16 16, ST_ADDR32 32, ST_ADDR64 64,
ST_CALL32 32, ST_CALL64 64,
ST_JUMPTAB64 64,
ST_PAD0 0, ST_PAD1 0, ST_PAD2 0, ST_PAD3 0, ST_PAD4 0, ST_PAD5 0, ST_PAD6 0, ST_PAD7 0,
ST_PAD8 0, ST_PAD9 0, ST_PAD10 0, ST_PAD11 0, ST_PAD12 0, ST_PAD13 0, ST_PAD14 0, ST_PAD15 0,
}
const ST_MODRM: usize = ST_OP;
const ST_AJUMP32: usize = ST_JUMP32;
const ST_JUMPTAB_COUNT: usize = ST_OP;
#[derive(Debug)]
pub struct Streams {
origin: u64,
streams: [Vec<u8>; NUM_STREAMS],
}
impl Streams {
fn new(origin: u64) -> Self {
Self {
origin,
streams: array::from_fn(|_| Vec::new()),
}
}
fn check(&self, st: usize, size: usize) {
let expected = STREAM_SIZES[st];
if expected > 0 {
debug_assert_eq!(size, expected);
}
}
fn put8(&mut self, st: usize, v: u8) {
if log_encode!() {
print!("({}:{v:02X})", &STREAM_NAMES[st][3..]);
}
self.check(st, 1);
self.streams[st].push(v);
}
fn put32(&mut self, st: usize, v: u32) {
if log_encode!() {
print!("({}:{v:08X})", &STREAM_NAMES[st][3..]);
}
self.check(st, 4);
self.streams[st].extend_from_slice(&v.to_le_bytes());
}
fn put64(&mut self, st: usize, v: u64) {
if log_encode!() {
print!("({}:{v:016X})", &STREAM_NAMES[st][3..]);
}
self.check(st, 8);
self.streams[st].extend_from_slice(&v.to_le_bytes());
}
fn copy(&mut self, st: usize, v: &[u8]) {
if log_encode!() {
print!("({}:{})", &STREAM_NAMES[st][3..], {
struct Hex<'a>(&'a [u8]);
impl fmt::Display for Hex<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for b in self.0 {
write!(f, "{b:02X}")?;
}
Ok(())
}
}
Hex(v)
});
}
self.check(st, v.len());
self.streams[st].extend_from_slice(v);
}
fn op(&mut self, map: u8, op: u8) {
if map > 0 || op & 0xf0 == 0x40 {
self.op_map(map);
}
self.put8(ST_OP, op);
}
fn op_esc(&mut self, b: u8) {
if log_encode!() {
//print!("(OP:{ESC:02X}=ESC)(OP:{b:02X})");
print!("(OP:{ESC:02X})(OP:{b:02X})");
}
self.streams[ST_OP].extend_from_slice(&[ESC, b]);
}
fn op_map(&mut self, m: u8) {
let b = 0x40 + m;
if log_encode!() {
//print!("(OP:{b:02X}=MAP{m:X})");
print!("(OP:{b:02X})");
}
assert!(m < 16);
self.streams[ST_OP].push(b);
}
fn call32(&mut self, idx: u8, target: u32) {
if idx == 0xff {
if log_encode!() {
print!("(CALL_IDX:{idx:02X})(CALL32:{target:08X})");
}
self.streams[ST_CALL_IDX].push(idx);
self.streams[ST_CALL32].extend_from_slice(&target.to_le_bytes());
} else {
if log_encode!() {
print!("(CALL_IDX:{idx:02X})");
}
self.streams[ST_CALL_IDX].push(idx);
}
}
fn call64(&mut self, idx: u8, target: u64) {
if idx == 0xff {
if log_encode!() {
print!("(CALL_IDX:{idx:02X})(CALL64:{target:016X})");
}
self.streams[ST_CALL_IDX].push(idx);
self.streams[ST_CALL64].extend_from_slice(&target.to_le_bytes());
} else {
if log_encode!() {
print!("(CALL_IDX:{idx:02X})");
}
self.streams[ST_CALL_IDX].push(idx);
}
}
fn jumptab(&mut self, count: u8) {
if log_encode!() {
//print!("(OP:{JUMPTAB:02X}=JUMPTAB)(JUMPTAB_COUNT:{count:02X})");
print!("(OP:{JUMPTAB:02X})(OP:{count:02X})");
}
self.streams[ST_OP].push(JUMPTAB);
self.streams[ST_JUMPTAB_COUNT].push(count);
}
pub fn write_to(&self, w: &mut impl Write) -> io::Result<()> {
let mut mask = 0u64;
for (i, stream) in self.streams.iter().enumerate() {
if !stream.is_empty() {
mask |= 1 << i as u64;
}
}
assert!(NUM_STREAMS <= 60);
let mut header = vec![];
header.extend_from_slice(&self.origin.to_le_bytes());
header.extend_from_slice(&mask.to_le_bytes());
for stream in &self.streams {
if !stream.is_empty() {
header.extend_from_slice(&(stream.len() as u32).to_le_bytes());
}
}
w.write_all(must_cast_slice(&header))?;
for stream in &self.streams {
w.write_all(stream)?;
}
Ok(())
}
}
#[inline(always)]
const fn rex_has_w(rex: u8) -> bool {
rex & 0x08 != 0
}
// ___ ____ ____
// VEX3 RXB0mmmm WvvvvLpp -> 1vvvvLpp 0000WRXB + map mmmm
// ^
// The bit 4 of the second byte is technically the fifth map bit,
// which gets ignored by the current encoding scheme
// because it is currently completely unused.
#[inline(always)]
fn shuffle_vex3([x, y]: [u8; 2]) -> Option<([u8; 2], u8)> {
if x & 0x10 != 0 {
return None;
}
let map = x & 0x0f;
let rex = (y >> 4) & 0x08 | (!x >> 5);
let vex = 0x80 | (y & 0x7f);
Some(([vex, rex], map))
}
#[inline(always)]
fn unshuffle_vex3([vex, rex]: [u8; 2], map: u8) -> Option<[u8; 2]> {
if vex & 0x80 == 0 || rex & 0xf0 != 0 || map >= 16 {
return None;
}
let x = (!rex & 0x07) << 5 | map;
let y = (rex & 0x08) << 4 | (vex & 0x7f);
Some([x, y])
}
// _____ ____
// VEX2 RvvvvLpp -> 1vvvvLpp 00000R00 + map 1 (implied)
#[inline(always)]
fn shuffle_vex2([x]: [u8; 1]) -> ([u8; 2], u8) {
let rex = (!x >> 5) & 0x04;
let vex = 0x80 | (x & 0x7f);
([vex, rex], 1)
}
#[inline(always)]
fn unshuffle_vex2([vex, rex]: [u8; 2], map: u8) -> Option<[u8; 1]> {
if vex & 0x80 == 0 || rex & 0xfb != 0 || map != 1 {
return None;
}
let x = (!rex & 0x04) << 5 | (vex & 0x7f);
Some([x])
}
// _____ _____ _ _____
// EVEX RXBrbmmm Wvvvvxpp **L*V*** -> **0*0*** VvvvvLpp 0rxbWRXB + map mmm
#[inline(always)]
fn shuffle_evex([x, y, z]: [u8; 3]) -> ([u8; 3], u8) {
let map = x & 0x07;
let rex = (!x & 0x10) << 2 | (!x & 0x08) << 1 | (!x >> 5) | (!y & 0x04) << 3 | (y >> 4) & 0x08;
let vex = (y & 0x7b) | (z >> 3) & 0x04 | (z & 0x08) << 4;
let evex = z & 0xd7;
([evex, vex, rex], map)
}
#[inline(always)]
fn unshuffle_evex([evex, vex, rex]: [u8; 3], map: u8) -> Option<[u8; 3]> {
if evex & 0x28 != 0 || rex & 0x80 != 0 || map >= 8 {
return None;
}
let x = (!rex & 0x07) << 5 | (!rex >> 2) & 0x10 | (!rex >> 1) & 0x08 | map;
let y = (rex & 0x08) << 4 | (vex & 0x7b) | (!rex >> 3) & 0x04;
let z = evex | (vex & 0x04) << 3 | (vex >> 4) & 0x08;
Some([x, y, z])
}
#[inline(always)]
const fn parse_modrm(modrm: u8) -> (u8 /*mode*/, u8 /*base*/) {
(modrm >> 6, modrm & 0b111)
}
#[inline(always)]
const fn modrm_is_reg_only((mode, _base): (u8, u8)) -> bool {
mode == 0b11
}
#[inline(always)]
const fn modrm_reg(modrm: u8) -> u8 {
modrm >> 3 & 7
}
#[inline(always)]
const fn modrm_has_sib((mode, base): (u8, u8)) -> bool {
mode < 0b11 && base == 0b100
}
fn range_chunks(count: usize, chunk_size: usize) -> impl Iterator<Item = std::ops::Range<usize>> {
let remainder = count % chunk_size;
(0..count - remainder)
.step_by(chunk_size)
.map(move |start| start..start + chunk_size)
.chain(if remainder > 0 {
Some(count - remainder..count)
} else {
None
})
}
// Try to recognize common function boundary padding starting from `code[0]`:
// - `00` (ADD Eb,Gb)
// - `90` (NOP)
// - `0F 1F /0` (NOP E[bv])
// `0F 1F 00'000'reg`
// `0F 1F 00'000'100 zz'zzz'zzz`
// `0F 1F 00'000'100 zz'zzz'101 XX XX XX XX`
// `0F 1F 00'000'101 XX XX XX XX`
// `0F 1F 01'000'reg XX`
// `0F 1F 01'000'100 ZZ XX`
// `0F 1F 10'000'reg XX XX XX XX`
// `0F 1F 10'000'100 ZZ XX XX XX XX`
// - `CC` (INT3)
// - `89 11'reg'reg` (MOV R#,R#)
// - `8D 00'reg'reg` (LEA R#,[R#]) where reg != 101
// - `8D 01'reg'reg 00` (LEA R#,[R#+00h])
// - `8D 10'reg'reg 00 00 00 00` (LEA R#,[R#+00000000h])
// - `8D 00'reg'100 zz'100'reg (LEA R#,[R#*1]) where reg != 101
// - `8D 01'reg'100 zz'100'reg 00` (LEA R#,[R#*1+00h]) where reg != 101
// - `8D 10'reg'100 zz'100'reg 00 00 00 00` (LEA R#,[R#*1+00000000h]) where reg != 101
// - Any 1-byte prefix besides from LOCK and REX: `26 2E 36 3E 64 65 66 67 F2 F3`
fn scan_pad(mut code: &[u8]) -> usize {
let len = code.len();
// Skip any 00 bytes only at the very beginning of possible padding.
let n = code.iter().position(|&b| b != 0x00).unwrap_or(len);
code = &code[n..];
loop {
code = match code {
[0xcc, rest @ ..]
| [0x90, rest @ ..]
| [0x66, 0x90, rest @ ..]
| [0x0f, 0x1f, 0x00, rest @ ..]
| [0x0f, 0x1f, 0x40, _, rest @ ..]
| [0x0f, 0x1f, 0x44, 0x00, _, rest @ ..]
| [0x66, 0x0f, 0x1f, 0x44, 0x00, _, rest @ ..]
| [0x0f, 0x1f, 0x80, _, _, _, _, rest @ ..]
| [0x0f, 0x1f, 0x84, 0x00, _, _, _, _, rest @ ..]
| [0x66, 0x0f, 0x1f, 0x84, 0x00, _, _, _, _, rest @ ..]
| [0x66, 0x66, 0x0f, 0x1f, 0x84, 0x00, _, _, _, _, rest @ ..]
| [0x66, 0x66, 0x66, 0x0f, 0x1f, 0x84, 0x00, _, _, _, _, rest @ ..] => rest,
_ => break,
};
}
len - code.len()
}
/*
https://stackoverflow.com/questions/25545470/long-multi-byte-nops-commonly-understood-macros-or-other-notation
https://gist.github.com/stevemk14ebr/d117e8d0fd1432fb2a92354a034ce5b9
if(code == OP_RETNI || code == OP_RETN || code == OP_INT3) // return/padding
nextIsFunc = sTRUE; // next opcode is likely to be first of a new function
CALL 9A E8 FF/2-3
JMP E9 EA EB FF/4-5
JMPE 0FB8(w/o F2/F3) 0F00/5
RET C2 C3 CA CB CF
JMPABS REX2+A1
REX2 invalid: 7x Ax Ex 13x
JMPABS transfers program control to the 64-bit absolute address target64 given as a quadword
immediate. JMPABS is in legacy map 0 and requires a REX2 prefix with REX2.M0 = 0 and REX2.W = 0. All
other REX2 payload bits are ignored, and code-generators should set these bits to 0. JMPABS does not
have a ModRM byte and target64 is placed immediately after the opcode byte, so the entire instruction is
11 bytes long. Prefixing JMPABS with 0x66, 0x67, 0xF0, 0xF2, or 0xF3 triggers #UD. Segment overrides are
allowed but ignored by JMPABS.
padding + function target
*/
const fn pad_followed(op: u16) -> bool {
// TODO
matches!(
op,
0x9a | 0xe8 | 0xe9 | 0xea | 0xeb | 0xc2 | 0xc3 | 0xca | 0xcb | 0xcc | 0xcf
)
}
struct CallCache {
cache: [u64; 0x100],
}
impl CallCache {
fn new() -> Self {
Self {
cache: [0u64; 0x100],
}
}
#[inline(always)]
fn find_index(&self, target: u64) -> u8 {
self.cache
.iter()
.position(|&cached| cached == target)
.unwrap_or(0xff) as u8
}
#[inline(always)]
fn find_target(&self, idx: u8) -> Option<u64> {
if idx < 0xff {
Some(self.cache[idx as usize])
} else {
None
}
}
#[inline(always)]
fn record(&mut self, idx: u8, target: u64) {
self.cache.copy_within(0..(idx as usize), 1);
self.cache[0] = target;
}
}
pub fn encode(mut code: Vec<u8>, mut origin: u64, is64: bool) -> io::Result<Streams> {
let mut st = Streams::new(origin);
let mut call_cache = CallCache::new();
let mut call_idx = |target: u64| {
let idx = call_cache.find_index(target);
call_cache.record(idx, target);
idx
};
let code_len = code.len();
let code_start = origin;
let code_end = origin + code_len as u64;
// should be enough for this encoding scheme
const SENTINEL: usize = 15;
code.extend_from_slice(&[0u8; SENTINEL]);
let mut prefixes = 0;
let mut pad = false;
let mut code = &code[..];
while code.len() > SENTINEL {
if log_encode!() {
println!();
print!("{origin:06X}: ");
//print!("{:02X?} ", &code[..SENTINEL]);
}
// Try to skip any padding.
if pad {
let stream = ST_PAD0 + (origin & 0xf) as usize;
let pad_size = scan_pad(&code[..code.len() - SENTINEL]);
let mut padding;
(padding, code) = code.split_at(pad_size);
origin += pad_size as u64;
while padding.len() >= 0xff {
let chunk;
(chunk, padding) = padding.split_at(0xff);
st.put8(stream, 0xff);
st.copy(stream, chunk);
}
st.put8(stream, padding.len() as u8);
st.copy(stream, padding);
if false {
// Speculatively record the next instruction as a call target.
call_idx(origin);
}
pad = false;
continue;
}
// Detect a possible jump table of at least 3 entries.
const MIN_JUMPTAB: usize = 3;
if !is64 {
// TBW
let mut i = 0;
let min_addr = code_start as u32;
let max_addr = (code_end - 1) as u32;
while i < code.len() - SENTINEL - 4 {
let addr = to_u32(&code[i..i + 4]);
if addr < min_addr || max_addr < addr {
break;
}
i += 4;
}
if i >= MIN_JUMPTAB * 4 {
for range in range_chunks(i / 4, 0x100) {
st.jumptab((range.len() - 1) as u8);
for j in range {
let addr = to_u32(&code[j * 4..j * 4 + 4]);
st.call32(call_idx(addr as u64), addr);
}
}
code = &code[i..];
origin += i as u64;
continue;
}
} else if origin % 8 == 0 {
// In x86-64, jump tables are typically 64-bit aligned offsets.
// Since we don't know where they will be used,
// we assume that they may be used anywhere within this `code`.
let min_offset = -(code_len as i64 - 1);
let max_offset = code_len as i64 - 1;
let mut i = 0;
while i < code.len() - SENTINEL - 8 {
let offset = to_u64(&code[i..i + 8]) as i64;
if offset < min_offset || max_offset < offset {
break;
}
i += 8;
}
if i >= MIN_JUMPTAB * 8 {
for range in range_chunks(i / 8, 0x100) {
st.jumptab((range.len() - 1) as u8);
for j in range {
st.put64(ST_JUMPTAB64, to_u64(&code[j * 8..j * 8 + 8]));
}
}
code = &code[i..];
origin += i as u64;
continue;
}
}
let mut i = 0;
let mut pre = 0;
let mut evex = 0;
let mut vex = 0;
let mut rex = 0;
let mut map = 0;
let mut op = 0;
// Handle prefixes that cannot be combined first.
// They are all followed by ModR/M, where mode = 0b11 is required in x86-32.
if is64 || modrm_is_reg_only(parse_modrm(code[1])) {
match code[0] {
PRE_VEX3 => {
if let Some(ret) = shuffle_vex3([code[1], code[2]]) {
pre = PRE_VEX3;
([vex, rex], map) = ret;
op = code[3];
i = 4;
}
}
PRE_VEX2 => {
pre = PRE_VEX2;
([vex, rex], map) = shuffle_vex2([code[1]]);
op = code[2];
i = 3;
}
PRE_EVEX => {
pre = PRE_EVEX;
([evex, vex, rex], map) = shuffle_evex([code[1], code[2], code[3]]);
op = code[4];
i = 5;
}
_ => {}
}
}
// *VEX cannot coexist with 0F or REX.
let has_vex = i > 0;
if !has_vex {
if is64 {
let c = code[0];
if c & 0xf0 == 0x40 {
// REX (40..4F)
pre = REX_MARKER;
rex = c & 0x0f;
i = 1;
} else if c == PRE_REX2 {
pre = PRE_REX2;
rex = code[1];
i = 2;
}
}
if code[i] == PRE_2BYTE {
map = 1;
op = code[i + 1];
i += 2;
} else {
op = code[i];
i += 1;
}
}
let i = Cell::new(i);
let fetch8 = || {
let ret = code[i.get()];
i.set(i.get() + 1);
ret
};
let fetch32 = || {
let ret = to_u32(&code[i.get()..i.get() + 4]);
i.set(i.get() + 4);
ret
};
let copy = |n: usize, streams: &mut Streams, st: usize| {
debug_assert_eq!(STREAM_SIZES[st], n);
streams.copy(st, &code[i.get()..i.get() + n]);
i.set(i.get() + n);
};
let rel_to_abs = |addr: u32, delta: usize| {
(origin + (i.get() + delta) as u64).wrapping_add(addr as i32 as u64)
};
let mut flags = lookup_opcode(op, map, is64);
if flags == BP {
// 1-byte prefixes can't have any *VEX and REX prefix.
if i.get() != 1 {
flags = XX;
} else {
assert_eq!(map, 0);
prefixes |= 1 << prefix_hash(op);
st.op(map, op);
code = &code[1..];
origin += 1;
continue;
}
} else if op3_followed(flags) {
// 3-byte opcode prefixes can't have any *VEX prefix which has its own map index.
if has_vex {
flags = XX;
} else {
assert_eq!(map, 1);
map = (flags - M2) + 2;
flags = (flags - M2) + R_;
op = fetch8();
}
}
let mut prefixes = mem::replace(&mut prefixes, 0);
if flags == XX {
// Do NOT commit the current position if this instruction is invalid.
st.op_esc(code[0]);
code = &code[1..];
origin += 1;
continue;
}
// Now we can commit any prefixes and opcode.
if pre != 0 {
st.put8(ST_OP, pre);
if pre == PRE_EVEX {
st.put8(ST_EVEX, evex);
st.put8(ST_VEX, vex);
} else if pre == PRE_VEX2 || pre == PRE_VEX3 {
st.put8(ST_VEX, vex);
}
st.put8(ST_REX, rex);
}
st.op(map, op);
let op = (map as u16) << 8 | op as u16;
match op {
// Parse an additional 16-bit immediate for these:
//
// 9A/EA: CALL/JMP Ap (16-bit segment + 32-bit address)
// C8: ENTER Iw,Ib (16-bit immediate + 8-bit immediate)
OP_CALLF | OP_JMPF | OP_ENTER => {
copy(16 / 8, &mut st, ST_IMM16);
}
// F6/F7: TEST E,I (/0-1) vs. NOT/NEG/[I]MUL/[I]DIV E (/2-7)
OP_GRP3_1 | OP_GRP3_2 | OP_MAP4_GRP3_1 | OP_MAP4_GRP3_2
if modrm_reg(code[i.get()]) >= 2 =>
{
flags = R_;
}
// 0F B8: JMPE Jz (IA-64 only) vs. POPCNT Gv,Ev (F3)
OP_JMPE_POPCNT if has_rep_prefix(prefixes) => {
flags = R_;
}
// MAP7 F8: URDMSR Rq,Id; UWRMSR Id,Rq (immediate size doesn't depend on 66)
OP_URDMSR_UWRMSR => {
prefixes &= !(1 << prefix_hash(PRE_OSIZE));
}
_ => {}
}
pad = pad_followed(op);
// ModR/M present
if has_modrm(flags) {
flags = modrm_to_imm(flags);
let modrm = fetch8();
st.put8(ST_MODRM, modrm);
let (mode, base) = parse_modrm(modrm);
let sib;
if modrm_has_sib((mode, base)) {
sib = fetch8();
st.put8(ST_SIB, sib);
} else {
sib = 0;
}
match mode {
0 if base == 5 => {
let addr = fetch32();
if is64 {
// [eip+disp32] or [rip+disp32]
// Note that we haven't fully decoded operands yet, hence a delta.
let delta = [0, 1, 2, 4][flags as usize];
st.put64(ST_ADDR64, rel_to_abs(addr, delta));
} else {
st.put32(ST_ADDR32, addr); // [disp32]
}
}
0 if sib & 7 == 5 => copy(32 / 8, &mut st, ST_DISP32), // [reg*scale+disp32]
// [reg+disp8] or [reg*scale+disp8]
1 => copy(8 / 8, &mut st, ST_DISP8_R0 + base as usize),
2 => copy(32 / 8, &mut st, ST_DISP32), // [reg+disp32]
_ => {}
}
}
match flags {
J4 => {
let target = rel_to_abs(fetch32(), 0);
if op == OP_CALLN {
let idx = call_idx(target);
if is64 {
st.call64(idx, target);
} else {
st.call32(idx, target as u32);
}
} else {
if is64 {
st.put64(ST_JUMP64, target);
} else {
st.put32(ST_JUMP32, target as u32);
}
}
}
A_ => {
// EA: 32-bit only, 16-bit if 66 ("Ap" = w:z)
// Ax: 32-bit or 64-bit, fixed per operating mode ("Ov")
let lgn = if is64 {
3
} else if op == OP_JMPF && has_osize_prefix(prefixes) {
1
} else {
2
};
copy(1 << lgn, &mut st, (ST_ADDR16 - 1) + lgn);
}
JA => copy(32 / 8, &mut st, ST_AJUMP32),
J1 => copy(8 / 8, &mut st, ST_JUMP8),
N_ => {}
_ => {
assert!(matches!(flags, N1 | N2 | N4 | NZ));
if flags == N4 && has_osize_prefix(prefixes) {
flags = N2;
}
if flags == NZ && !rex_has_w(rex) {
flags = N4;
}
let lgn = (flags - N1) as usize;
copy(1 << lgn, &mut st, ST_IMM8 + lgn);
}
}
let i = i.get();
code = &code[i..];
origin += i as u64;
}
Ok(st)
}
pub fn decode(streams: &Streams, is64: bool) -> Option<Vec<u8>> {
let origin = streams.origin;
let streams = streams
.streams
.each_ref()
.map(|stream| Cell::new(&stream[..]));
let pc = Cell::new(origin);
let mut code = Vec::new();
let read8 = |st: usize| {
if log_decode!() {
print!("({}:", &STREAM_NAMES[st][3..]);
}
let (&[head], tail) = streams[st].get().split_first_chunk::<1>()?;
if log_decode!() {
print!("{head:02X})");
}
streams[st].set(tail);
Some(head)
};
let read32 = |st: usize| {
if log_decode!() {
print!("({}:", &STREAM_NAMES[st][3..]);
}
let (&head, tail) = streams[st].get().split_first_chunk::<4>()?;
let head = u32::from_le_bytes(head);
if log_decode!() {
print!("{head:08X})");
}
streams[st].set(tail);
Some(head)
};
let read64 = |st: usize| {
if log_decode!() {
print!("({}:", &STREAM_NAMES[st][3..]);
}
let (&head, tail) = streams[st].get().split_first_chunk::<8>()?;
let head = u64::from_le_bytes(head);
if log_decode!() {
print!("{head:016X})");
}
streams[st].set(tail);
Some(head)
};
let mut call_cache = CallCache::new();
let read_call = |cache: &mut CallCache| {
let idx = read8(ST_CALL_IDX)?;
let target = if let Some(target) = cache.find_target(idx) {
target
} else if is64 {
read64(ST_CALL64)?
} else {
read32(ST_CALL32)? as u64
};
cache.record(idx, target);
Some(target)
};
let copy = |n: usize, code: &mut Vec<u8>, st: usize| {
if log_decode!() {
print!("({}:", &STREAM_NAMES[st][3..]);
}
let (head, tail) = streams[st].get().split_at(n);
if log_decode!() {
for &b in head {
print!("{:02X}", b);
}
print!(")");
}
code.extend_from_slice(head);
streams[st].set(tail);
Some(())
};
macro_rules! fatal {
($fmt:tt) => {
panic!(
concat!($fmt, " @ {:06X} {:02X?}"),
pc.get(),
&code[code.len().max(15) - 15..]
)
};
}
let mut prefixes = 0;
let mut pad = false;
while !streams[ST_OP].get().is_empty() {
pc.set(origin + code.len() as u64);
if log_decode!() {
println!();
print!("{:06X}: ", pc.get());
}
if pad {
let stream = ST_PAD0 + ((origin as usize + code.len()) & 0xf);
loop {
let pad_size = read8(stream)?;
copy(pad_size as usize, &mut code, stream)?;
if pad_size < 0xff {
break;
}
}
if false {
// Speculatively record the next instruction as a call target.
let target = origin + code.len() as u64;
let idx = call_cache.find_index(target);
call_cache.record(idx, target);
}
pad = false;
continue;
}
let mut op = read8(ST_OP).unwrap();
if op == ESC {
code.push(read8(ST_OP)?);
continue;
}
if op == JUMPTAB {
let count = read8(ST_JUMPTAB_COUNT)? as usize + 1;
if is64 {
for _ in 0..count {
code.extend_from_slice(&read64(ST_JUMPTAB64)?.to_le_bytes());
}
} else {
for _ in 0..count {
code.extend_from_slice(&(read_call(&mut call_cache)? as u32).to_le_bytes());
}
}
continue;
}
let (pre, evex, vex, rex) = match op {
PRE_VEX3 => (op, 0, read8(ST_VEX)?, read8(ST_REX)?),
PRE_VEX2 => (op, 0, read8(ST_VEX)?, read8(ST_REX)?),
PRE_EVEX => (op, read8(ST_EVEX)?, read8(ST_VEX)?, read8(ST_REX)?),
REX_MARKER | PRE_REX2 if !is64 => fatal!("unsupported REX prefixes in x86-32"),
REX_MARKER => (op, 0, 0, read8(ST_REX)?),
PRE_REX2 => (op, 0, 0, read8(ST_REX)?),
_ => (0, 0, 0, 0),
};
if pre != 0 {
op = read8(ST_OP)?;
}
let map;
if op & 0xf0 == 0x40 {
map = op & 0x0f;
op = read8(ST_OP)?;
} else {
map = 0;
}
let mut flags = lookup_opcode(op, map, is64);
if flags == XX {
fatal!("invalid opcode");
} else if flags == BP {
assert_eq!(map, 0);
prefixes |= 1 << prefix_hash(op);
code.push(op);
continue;
}
'prefix: {
match pre {
PRE_VEX3 => {
let Some([x, y]) = unshuffle_vex3([vex, rex], map) else {
fatal!("bad VEX3 prefix");
};
code.extend_from_slice(&[PRE_VEX3, x, y, op]);
break 'prefix;
}
PRE_VEX2 => {
let Some([x]) = unshuffle_vex2([vex, rex], map) else {
fatal!("bad VEX2 prefix");
};
code.extend_from_slice(&[PRE_VEX2, x, op]);
break 'prefix;
}
PRE_EVEX => {
let Some([x, y, z]) = unshuffle_evex([evex, vex, rex], map) else {
fatal!("bad EVEX prefix");
};
code.extend_from_slice(&[PRE_EVEX, x, y, z, op]);
break 'prefix;
}
REX_MARKER => {
if rex & 0xf0 != 0 {
fatal!("bad REX prefix");
}
code.push(0x40 | rex);
}
PRE_REX2 => code.extend_from_slice(&[PRE_REX2, rex]),
0 => {}
_ => unreachable!(),
}
// Only applicable with non-*VEX prefixes.
match map {
0 => code.push(op),
1 => code.extend_from_slice(&[0x0f, op]),
2 => code.extend_from_slice(&[0x0f, 0x38, op]),
3 => code.extend_from_slice(&[0x0f, 0x3a, op]),
_ => fatal!("bad opcode map"),
}
}
let mut prefixes = mem::replace(&mut prefixes, 0);
let op = (map as u16) << 8 | op as u16;
match op {
// Parse an additional 16-bit immediate for these:
//
// 9A/EA: CALL/JMP Ap (16-bit segment + 32-bit address)
// C8: ENTER Iw,Ib (16-bit immediate + 8-bit immediate)
OP_CALLF | OP_JMPF | OP_ENTER => {
copy(16 / 8, &mut code, ST_IMM16)?;
}
// F6/F7: TEST E,I (/0-1) vs. NOT/NEG/[I]MUL/[I]DIV E (/2-7)
OP_GRP3_1 | OP_GRP3_2 | OP_MAP4_GRP3_1 | OP_MAP4_GRP3_2
if modrm_reg(streams[ST_MODRM].get()[0]) >= 2 =>
{
flags = R_;
}
// 0F B8: JMPE Jz (IA-64 only) vs. POPCNT Gv,Ev (F3)
OP_JMPE_POPCNT if has_rep_prefix(prefixes) => {
flags = R_;
}
// MAP7 F8: URDMSR Rq,Id; UWRMSR Id,Rq (immediate size doesn't depend on 66)
OP_URDMSR_UWRMSR => {
prefixes &= !(1 << prefix_hash(PRE_OSIZE));
}
_ => {}
}
pad = pad_followed(op);
let abs_to_rel = |addr: u64, code: &[u8], delta: usize| {
addr.wrapping_sub(origin + (code.len() + delta) as u64)
};
// ModR/M present
if has_modrm(flags) {
flags = modrm_to_imm(flags);
let modrm = read8(ST_MODRM)?;
code.push(modrm);
let (mode, base) = parse_modrm(modrm);
let sib;
if modrm_has_sib((mode, base)) {
sib = read8(ST_SIB)?;
code.push(sib);
} else {
sib = 0;
}
match mode {
0 if base == 5 => {
let addr = if is64 {
// [eip+disp32] or [rip+disp32]
// Note that we haven't fully decoded operands yet, hence a delta.
let delta = [0, 1, 2, 4][flags as usize];
abs_to_rel(read64(ST_ADDR64)?, &code, delta + 4) as u32
} else {
read32(ST_ADDR32)? // [disp32]
};
code.extend_from_slice(&addr.to_le_bytes());
}
// [reg*scale+disp32]
0 if sib & 7 == 5 => copy(32 / 8, &mut code, ST_DISP32)?,
// [reg+disp8] or [reg*scale+disp8]
1 => copy(8 / 8, &mut code, ST_DISP8_R0 + base as usize)?,
2 => copy(32 / 8, &mut code, ST_DISP32)?, // [reg+disp32]
_ => {}
}
}
match flags {
J4 => {
let target = if op == OP_CALLN {
read_call(&mut call_cache)?
} else if is64 {
read64(ST_JUMP64)?
} else {
read32(ST_JUMP32)? as u64
};
let target = abs_to_rel(target as u64, &code, 4) as u32;
code.extend_from_slice(&target.to_le_bytes());
}
A_ => {
// EA: 32-bit only, 16-bit if 66 ("Ap" = w:z)
// Ax: 32-bit or 64-bit, fixed per operating mode ("Ov")
let lgn = if is64 {
3
} else if op == OP_JMPF && has_osize_prefix(prefixes) {
1
} else {
2
};
copy(1 << lgn, &mut code, (ST_ADDR16 - 1) + lgn)?;
}
JA => copy(32 / 8, &mut code, ST_AJUMP32)?,
J1 => copy(8 / 8, &mut code, ST_JUMP8)?,
N_ => {}
_ => {
assert!(matches!(flags, N1 | N2 | N4 | NZ));
if flags == N4 && has_osize_prefix(prefixes) {
flags = N2;
}
if flags == NZ && !rex_has_w(rex) {
flags = N4;
}
let lgn = (flags - N1) as usize;
copy(1 << lgn, &mut code, ST_IMM8 + lgn)?;
}
}
}
Some(code)
}
pub fn locate_code(f: &mut File) -> io::Result<Vec<(u64, u64, usize)>> {
f.seek(SeekFrom::Start(0))?;
let mut r = BufReader::new(f);
// MZ header
let mut buf = [0u8; 0x40];
r.read_exact(&mut buf)?;
let sig = to_u16(&buf[0..2]);
if sig != 0x5a4d {
return Err(io::Error::other(format!("bad MZ signature {sig:#x}")));
}
let pe_offset = to_u32(&buf[0x3c..0x40]);
if pe_offset < 0x40 {
return Err(io::Error::other(format!(
"too low offset to PE header {pe_offset:#x}"
)));
}
r.seek_relative(pe_offset as i64 - 0x40)?;
// PE header
let mut buf = [0u8; 0x18];
r.read_exact(&mut buf)?;
let sig = to_u32(&buf[0..4]);
if sig != 0x4550 {
return Err(io::Error::other(format!("bad PE signature {sig:#x}")));
}
let num_sections = to_u16(&buf[6..8]);
let opt_header_size = to_u16(&buf[0x14..0x16]) as usize;
// PE optional header
let mut opt_header = vec![0u8; opt_header_size];
r.read_exact(&mut opt_header)?;
let magic = to_u16(&opt_header[0..2]);
match magic {
0x10b => {
// IMAGE_OPTIONAL_HEADER32
todo!();
}
0x20b => {
// IMAGE_OPTIONAL_HEADER64
if opt_header_size < 0x60 {
return Err(io::Error::other(format!(
"PE64 optional header too small ({opt_header_size:#x} < 0x60)"
)));
}
let num_data_dirs = to_u32(&opt_header[0x5c..0x60]) as usize;
let min_size = 0x60 + num_data_dirs * 0x10;
if opt_header_size < min_size {
return Err(io::Error::other(format!(
"PE64 optional header too small ({opt_header_size:#x} < {min_size:#x})"
)));
}
// data directories:
// EXPORT, IMPORT, RESOURCE, EXCEPTION, SECURITY, BASERELOC, DEBUG, COPYRIGHT,
// GLOBALPTR, TLS, LOAD_CONFIG, BOUND_IMPORT, IAT, DELAY_IMPORT, COM_DESCRIPTOR, -
}
_ => {
return Err(io::Error::other(format!(
"bad PE optional header magic {magic:#x}"
)))
}
}
// section headers
let mut exec_sections = Vec::new();
for _ in 0..num_sections {
let mut buf = [0u8; 40];
r.read_exact(&mut buf)?;
let name = &buf[..8];
let _name = String::from_utf8_lossy(name);
let rva = to_u32(&buf[12..16]);
let stored_size = to_u32(&buf[16..20]);
let stored_offset = to_u32(&buf[20..24]);
let flags = to_u32(&buf[36..40]);
//println!("section {_name:?}: rva {rva:#x}, stored {stored_offset:#x} + {stored_size:#x}, flags {flags:#x}");
if flags & 0x20 != 0 {
exec_sections.push((rva as u64, stored_offset as u64, stored_size as usize));
}
}
Ok(exec_sections)
}
#[test]
fn test_disfilter() -> io::Result<()> {
std::env::set_var("RUST_LOG", "trace");
env_logger::init();
let mut f = File::open(r"c:\Program Files\ImageMagick-7.1.1-Q16-HDRI\ffmpeg.exe")?;
for (origin, offset, size) in locate_code(&mut f)? {
f.seek(SeekFrom::Start(offset))?;
let mut input = vec![0u8; size];
f.read_exact(&mut input)?;
File::create(r"x:\unfiltered.bin")?.write_all(&input)?;
let start = std::time::Instant::now();
let streams = encode(input.clone(), origin, true)?;
let enc_rate = size as f64 / start.elapsed().as_secs_f64() / 1e6;
streams.write_to(&mut File::create(r"x:\filtered.bin")?)?;
let start = std::time::Instant::now();
let recons = decode(&streams, true).expect("round trip failed");
let dec_rate = size as f64 / start.elapsed().as_secs_f64() / 1e6;
if input != recons {
let mismatch = std::iter::zip(&input, &recons)
.position(|(a, b)| a != b)
.unwrap();
let lo = mismatch.max(15) - 15;
let hi = mismatch + 15;
panic!(
"input != recons\n \
Input: {mismatch}/{} {:02X?}\n \
Recons: {mismatch}/{} {:02X?}",
input.len(),
&input[lo..hi.min(input.len())],
recons.len(),
&recons[lo..hi.min(recons.len())],
);
}
eprintln!("Disfilter: encoding {enc_rate:.2} MB/s, decoding {dec_rate:.2} MB/s");
break;
}
Ok(())
}
#[test]
fn test_call_cache() {
let mut expected = vec![
(1234, 0xff),
(1234, 0),
(5678, 0xff),
(1234, 1),
(1234, 0),
(5678, 1),
];
for i in 9000..9256 {
expected.push((i, 0xff));
}
expected.push((1234, 0xff));
expected.push((5678, 0xff));
let mut cache = CallCache::new();
for &(target, idx) in &expected {
assert_eq!(cache.find_index(target), idx);
cache.record(idx, target);
}
let mut cache = CallCache::new();
for &(target, idx) in &expected {
if idx == 0xff {
assert_eq!(cache.find_target(idx), None);
} else {
assert_eq!(cache.find_target(idx), Some(target));
}
cache.record(idx, target);
}
}
#[cfg(test)]
fn test_shuffle<const IN: usize, const OUT: usize>(
shuffle: impl Fn([u8; IN]) -> Option<([u8; OUT], u8)>,
unshuffle: impl Fn([u8; OUT], u8) -> Option<[u8; IN]>,
map_range: impl Iterator<Item = u8> + Clone,
) {
#[inline(always)]
fn generate<const N: usize>() -> impl Iterator<Item = [u8; N]> {
assert!(N <= 3);
(0..1u32 << (N * 8)).map(|n| {
let mut b = [0u8; N];
for i in 0..N {
b[i] = (n >> (i * 8) & 0xff) as u8;
}
b
})
}
let mut seen = std::collections::HashSet::new();
for i in generate::<IN>() {
if let Some((o, map)) = shuffle(i) {
assert_eq!(
unshuffle(o, map),
Some(i),
"{i:02X?} -> {o:02X?} + {map} -> (roundtrip failed)"
);
assert!(seen.insert((o, map)), "{i:02X?} -> {o:02X?} + {map} (dupe)");
}
}
for o in generate::<OUT>() {
for map in map_range.clone() {
if !seen.contains(&(o, map)) {
assert_eq!(
unshuffle(o, map),
None,
"? <- {o:02X?} + {map} (didn't fail)"
);
}
}
}
}
#[test]
fn test_shuffle_vex3() {
test_shuffle(shuffle_vex3, unshuffle_vex3, 0..=16);
}
#[test]
fn test_shuffle_vex2() {
test_shuffle(|i| Some(shuffle_vex2(i)), unshuffle_vex2, 0..=16);
}
#[test]
fn test_shuffle_evex() {
test_shuffle(
|i| Some(shuffle_evex(i)),
unshuffle_evex,
[5, 8].into_iter(),
);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment