monkeypatching in rust

the_good_rust.rs

trait Person {
    fn greeting(&self) -> StateOfMind;
    fn name(&self) -> &'static str;
}

struct Ixi { }
struct Katie { }

#[derive(Debug)]
enum StateOfMind {
    GoodMorning,
    HappyBirthday,
}

impl Person for Ixi {
    fn greeting(&self) -> StateOfMind {
        StateOfMind::GoodMorning
    }
    fn name(&self) -> &'static str { "ixi" }
}

impl Person for Katie {
    fn greeting(&self) -> StateOfMind {
        StateOfMind::HappyBirthday
    }
    fn name(&self) -> &'static str { "katie" }
}

#[inline(never)]
fn greet(you: &dyn Person) {
    println!("{} says {:?} to you!", you.name(), you.greeting());
}

#[inline(never)]
fn correct_ixi() {
    unsafe {
        // gotta reprogram ixi real quick
        let ixi_in_a_box: Box<dyn Person> = Box::new(Ixi {});
        #[inline(never)]
        unsafe fn inner_takes_dyn(v: Box<dyn Person>) {
            if let Some(vtable_offset) = get_greeting_offset() {
                // ok we know where in the vtable to patch, just not where the vtable itself is yet.

                // v is a pointer to (data, vtable)
                let dyn_thingie: [*const usize; 2] = std::mem::transmute(v);
                let vtable = dyn_thingie[1];
                let table_entry = vtable_offset as usize / std::mem::size_of::<usize>();
                // get a pointer to the entry in the vtable we want to fix
                let vtable_ptr = (vtable as *mut usize).offset(table_entry as isize);
                let katie_greet = get_katie_greet_address()
                    .expect("if we got here we can read ixi's vtable, and the same should work for katie");
                    
                // memory mapping will have this as read-only. gotta fix that...
                libc::mprotect(
                    (vtable_ptr as usize & !(4095)) as *mut _,
                    4096 * 2,
                    7, // READ | WRITE
                );

                // and commit a crime
                std::ptr::write_volatile(vtable_ptr, katie_greet);

                // clean up after though
                libc::mprotect(
                    (vtable_ptr as usize & !(4095)) as *mut _,
                    4096 * 2,
                    1, // READ | WRITE
                );
            } else {
                panic!("well we need a vtable offset");
            }
        }

        inner_takes_dyn(ixi_in_a_box);
    }
}

// just get the offset of `greeting` in the `Person` vtable
#[inline(never)]
unsafe fn get_greeting_offset() -> Option<u32> {
    // the function is read, so this will be compiled.
    // because `greeting` does not return a type with a destructor, it
    // will end with a tail call to the function we want. the offset
    // of that tail call is the magic number for the offset of the
    // function to replace in the vtable.
    //
    // in debug builds, this ends with `call; pop; ret`. look for both.
    // in both cases these are preceeded by `mov rdi, rax`, spelled
    // `4889c7`.
    fn do_greeting(b: &Box<dyn Person>) {
        b.greeting();
    }
    let f = do_greeting as *const u8;
    std::ptr::read_volatile(f);
    
    // find the magic sigil. 0 indicates we couldn't find it!
    let mut i = 1;
    let mut vtable_offset = None;
    while i < 100 {
        let curr = f.offset(i);
        if *curr == 0x48 && *(curr.offset(1)) == 0x89 && *(curr.offset(2)) == 0xc7 {
            // now the offset is encoded in one of two ways: either an 8-bit
            // offset if the vtable is tiny, or a 32-bit offset if it's larger.
            //
            // 8-bit offst has modrm bits of     01_rrr_mmm
            // 32-bit offset has modrm bits of   10_rrr_mmm
            // for example, "call [rcx + 0x18]" is spelled
            // ff 61 18
            //  ^  ^  ^-- offset we want
            //  |  ------ modrm 01_000_001, 8-bit displacement off reg 001 (rcx), r=0 (call)
            //  --------- opcode for call/jmp [mem]
            let modrm: u8 = *curr.offset(4);
            let offset = curr.offset(5); // we might not actually read this, depending on mod bits
            match modrm >> 6 {
                0b00 => {
                    // no offset - this is just `call [reg]`
                    vtable_offset = Some(0);
                }
                0b01 => {
                    // 8-bit offset
                    vtable_offset = Some(std::ptr::read_unaligned(offset) as u32);
                }
                0b10 => {
                    // 32-bit offset
                    vtable_offset = Some(std::ptr::read_unaligned(offset as *const u32));
                }
                0b11 => {
                    // this is actually `{call,jmp} reg`. we'd be in a bad spot here.
                }
                _ => {
                    // these are just unreachable. if everything is well-formed, anyway.
                    unreachable!();
                }
            }
            // anyway we found the instruction so lets break and move on
            break;
        }
        i += 1;
    }
    
    vtable_offset
}

fn get_katie_greet_address() -> Option<usize> {
    // this is for the most part the same logic as correcting ixi, but instead of
    // writing to ixi's vtable, we read katie's vtable
    unsafe {
        let katie_in_a_box: Box<dyn Person> = Box::new(Katie {});
        #[inline(never)]
        unsafe fn inner_takes_dyn(v: Box<dyn Person>) -> Option<usize>{
            if let Some(vtable_offset) = get_greeting_offset() {
                // ok we know where in the vtable to read, just not where the vtable itself is yet.

                // v is a pointer to (data, vtable)
                let dyn_thingie: [*const usize; 2] = std::mem::transmute(v);
                let vtable = dyn_thingie[1];
                let table_entry = vtable_offset as usize / std::mem::size_of::<usize>();
                // get a pointer to the entry in the vtable we want to read
                let vtable_ptr = (vtable as *mut usize).offset(table_entry as isize);
                // and commit a crime
                Some(*vtable_ptr)
            } else {
                None
            }
        }

        inner_takes_dyn(katie_in_a_box)
    }
}

fn main() {
    greet(&Ixi {});
    greet(&Katie {});
    println!("wait that's not right, hold on a sec...");
    correct_ixi();
    greet(&Ixi {});
}