From fed6ae454f659f4d2b36520474d1998b526a27dd Mon Sep 17 00:00:00 2001 From: Johannes Stoelp Date: Mon, 27 Feb 2023 22:00:11 +0100 Subject: Add JMP, JZ, and Label --- src/insn.rs | 10 ++++++++ src/insn/jmp.rs | 7 ++++++ src/insn/jz.rs | 7 ++++++ src/label.rs | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 48 ++++++++++++++++++++++++++++++++++++-- src/prelude.rs | 8 ++++--- 6 files changed, 147 insertions(+), 5 deletions(-) create mode 100644 src/insn/jmp.rs create mode 100644 src/insn/jz.rs create mode 100644 src/label.rs diff --git a/src/insn.rs b/src/insn.rs index 2bfaf18..be2e689 100644 --- a/src/insn.rs +++ b/src/insn.rs @@ -1,5 +1,7 @@ mod add; mod dec; +mod jmp; +mod jz; mod mov; mod ret; mod test; @@ -12,6 +14,14 @@ pub trait Dec { fn dec(&mut self, op1: T); } +pub trait Jmp { + fn jmp(&mut self, op1: T); +} + +pub trait Jz { + fn jz(&mut self, op1: T); +} + pub trait Mov { fn mov(&mut self, op1: T, op2: U); } diff --git a/src/insn/jmp.rs b/src/insn/jmp.rs new file mode 100644 index 0000000..71d1dbc --- /dev/null +++ b/src/insn/jmp.rs @@ -0,0 +1,7 @@ +use crate::prelude::*; + +impl Jmp<&mut Label> for Asm { + fn jmp(&mut self, op1: &mut Label) { + self.encode_jmp_label(&[0xe9], op1); + } +} diff --git a/src/insn/jz.rs b/src/insn/jz.rs new file mode 100644 index 0000000..6563ca2 --- /dev/null +++ b/src/insn/jz.rs @@ -0,0 +1,7 @@ +use crate::prelude::*; + +impl Jz<&mut Label> for Asm { + fn jz(&mut self, op1: &mut Label) { + self.encode_jmp_label(&[0x0f, 0x84], op1); + } +} diff --git a/src/label.rs b/src/label.rs new file mode 100644 index 0000000..b1f1133 --- /dev/null +++ b/src/label.rs @@ -0,0 +1,72 @@ +use std::collections::HashSet; + +/// A label which is used as target for jump instructions. +/// +/// ```rust +/// use juicebox_asm::prelude::*; +/// +/// let mut lbl = Label::new(); +/// let mut asm = Asm::new(); +/// +/// // Skip the mov instruction. +/// asm.jmp(&mut lbl); +/// asm.mov(Reg64::rax, Reg64::rax); +/// asm.bind(&mut lbl); +/// ``` +/// +/// # Panics +/// +/// Panics if the label is dropped while not yet bound, or having unresolved relocations. +/// This is mainly a safety-guard to detect wrong usage. +pub struct Label { + /// Location of the label. Will be set after the label is bound, else None. + location: Option, + + /// Offsets that must be patched with the label location. + offsets: HashSet, +} + +impl Label { + /// Create a new `unbound` [Label]. + pub fn new() -> Label { + Label { + location: None, + offsets: HashSet::new(), + } + } + + /// Bind the label to the `location`. + pub(crate) fn bind(&mut self, loc: usize) { + // A label can only be bound once! + assert!(!self.is_bound()); + + self.location = Some(loc); + } + + /// Record an offset that must be patched with the label location. + pub(crate) fn record_offset(&mut self, off: usize) { + self.offsets.insert(off); + } + + pub(crate) fn location(&self) -> Option { + self.location + } + + pub(crate) fn offsets_mut(&mut self) -> &mut HashSet { + &mut self.offsets + } + + /// Check whether the label is bound to a location. + const fn is_bound(&self) -> bool { + self.location.is_some() + } +} + +impl Drop for Label { + fn drop(&mut self) { + // Ensure the label was bound when it is dropped. + assert!(self.is_bound()); + // Ensure all offsets have been patched when the label is dropped. + assert!(self.offsets.is_empty()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 27164ad..3368e3d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,12 +2,13 @@ pub mod prelude; mod imm; mod insn; +mod label; mod reg; use imm::Imm; -pub use imm::{Imm16, Imm32, Imm64, Imm8}; +use label::Label; use reg::Reg; -pub use reg::{Reg16, Reg32, Reg64, Reg8}; +use reg::{Reg16, Reg32, Reg64, Reg8}; pub enum MemOp { Indirect(Reg64), @@ -69,6 +70,35 @@ impl Asm { } } + /// Bind the [Label] to the current location. + pub fn bind(&mut self, label: &mut Label) { + // Bind the label to the current offset. + label.bind(self.buf.len()); + + // Resolve any pending relocations for the label. + self.resolve(label); + } + + /// If the [Label] is bound, patch any pending relocation. + pub fn resolve(&mut self, label: &mut Label) { + if let Some(loc) = label.location() { + let loc = i32::try_from(loc).expect("Label location did not fit into i32."); + + // Resolve any pending relocations for the label. + for off in label.offsets_mut().drain() { + // Displacement is relative to the next instruction following the jump. + // We record the offset to patch at the first byte of the disp32 therefore we need + // to account for that in the disp computation. + let disp32 = loc - i32::try_from(off).expect("Label offset did not fit into i32") - 4 /* account for the disp32 */; + + // Patch the relocation with the disp32. + self.emit_at(off, &disp32.to_ne_bytes()); + } + } + } + + // -- Encode utilities. + fn encode_rr(&mut self, opc: u8, op1: T, op2: T) where Self: EncodeRR, @@ -185,6 +215,20 @@ impl Asm { // op2 -> modrm.rm self.encode_mr(opc, op2, op1); } + + fn encode_jmp_label(&mut self, opc: &[u8], op1: &mut Label) { + // Emit the opcode. + self.emit(opc); + + // Record relocation offset starting at the first byte of the disp32. + op1.record_offset(self.buf.len()); + + // Emit a zeroed disp32, which serves as placeholder for the relocation. + self.emit(&[0u8; 4]); + + // Resolve any pending relocations for the label. + self.resolve(op1); + } } // -- Encoder helper. diff --git a/src/prelude.rs b/src/prelude.rs index bca972f..29e7265 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -1,6 +1,8 @@ pub use crate::Asm; pub use crate::MemOp; -pub use crate::{Imm16, Imm32, Imm64, Imm8}; -pub use crate::{Reg16, Reg32, Reg64, Reg8}; -pub use crate::insn::{Add, Dec, Mov, Test}; +pub use crate::imm::{Imm16, Imm32, Imm64, Imm8}; +pub use crate::label::Label; +pub use crate::reg::{Reg16, Reg32, Reg64, Reg8}; + +pub use crate::insn::{Add, Dec, Jmp, Jz, Mov, Test}; -- cgit v1.2.3