asmkit/core/
patch.rs

1//! Post-emit code patching (JSC MacroAssembler–style).
2//!
3//! # Workflow
4//!
5//! 1. During emission, call arch `patchable_*` helpers (or
6//!    [`CodeBuffer::reserve_patch_block`](crate::CodeBuffer::reserve_patch_block)).
7//!    These return self-describing [`PatchableSite`] / [`PatchableBlock`] handles.
8//! 2. Finalize with [`CodeBuffer::finish_patched`](crate::CodeBuffer::finish_patched)
9//!    so the [`PatchCatalog`] validates reachability (and the linker can rebase it).
10//! 3. Apply patches with **`unsafe`** methods on a JIT [`Span`] or `&mut [u8]`.
11//!    Patching does **not** go through [`CodeBufferFinalized`].
12//!
13//! # JSC correspondence
14//!
15//! | JSC | asmkit |
16//! |---|---|
17//! | `PatchableJump` / near call | [`PatchableSite`] + [`PatchableSite::retarget`] |
18//! | `DataLabel32` / `DataLabelPtr` | [`PatchableBlock`] + [`PatchableBlock::repatch_u32`] / [`repatch_u64`](PatchableBlock::repatch_u64) |
19//! | `padBeforePatch` + custom stub | [`reserve_patch_block`](crate::CodeBuffer::reserve_patch_block) → [`PatchableBlock::rewrite`] |
20//! | `repatchJump` on a code pointer | `unsafe` apply on [`Span`] / `&mut [u8]` |
21//!
22//! # Site vs block vs custom region
23//!
24//! - **Site** — a fixed-width displacement field (`LabelUse`) for jumps/calls; retarget by
25//!   code offset within the same image.
26//! - **Block** — a reserved byte range (immediate field or whole insn sequence); rewrite or
27//!   repatch integer payloads; shorter rewrites are nop-padded.
28//! - **Custom** — [`reserve_patch_block`](crate::CodeBuffer::reserve_patch_block) plants a
29//!   nop island you fill later with [`PatchableBlock::rewrite`].
30//!
31//! # Safety
32//!
33//! Handles from `patchable_*` describe locations in the buffer that produced them. Applying a
34//! handle is still `unsafe`: the bytes/`Span` must be that image (after finalize/link with
35//! stable offsets), and the caller must synchronize concurrent execution of the patched region.
36//! [`PatchableSite::new`] / [`PatchableBlock::new`] are `unsafe` for the same reason when
37//! constructing handles by hand.
38//!
39//! [`PatchCatalog`] remains for `finish_patched` validation and linker rebasing. After rebase,
40//! convert entries with [`PatchSite::to_patchable`] / [`PatchBlock::to_patchable`].
41
42use smallvec::SmallVec;
43
44use crate::{
45    AsmError,
46    core::{
47        arch_traits::Arch,
48        buffer::{CodeBufferFinalized, CodeOffset, LabelUse},
49    },
50};
51
52#[cfg(feature = "jit")]
53use crate::core::jit_allocator::{JitAllocator, Span};
54
55/// Catalog index for a patch block (linker / introspection).
56#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
57pub struct PatchBlockId(u32);
58
59impl PatchBlockId {
60    pub(crate) const fn from_index(index: usize) -> Self {
61        Self(index as u32)
62    }
63
64    pub const fn index(self) -> usize {
65        self.0 as usize
66    }
67}
68
69/// Catalog index for a patch site (linker / introspection).
70#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
71pub struct PatchSiteId(u32);
72
73impl PatchSiteId {
74    pub(crate) const fn from_index(index: usize) -> Self {
75        Self(index as u32)
76    }
77
78    pub const fn index(self) -> usize {
79        self.0 as usize
80    }
81}
82
83/// Catalog entry for a rewritable byte range.
84#[derive(Clone, Copy, Debug, PartialEq, Eq)]
85pub struct PatchBlock {
86    pub offset: CodeOffset,
87    pub size: CodeOffset,
88    pub align: CodeOffset,
89}
90
91impl PatchBlock {
92    /// Build a patch handle from a (possibly rebased) catalog entry.
93    pub const fn to_patchable(self, arch: Arch) -> PatchableBlock {
94        // SAFETY: catalog entries describe blocks recorded during emission.
95        unsafe { PatchableBlock::new(self.offset, self.size, arch) }
96    }
97}
98
99/// Catalog entry for a displacement field.
100#[derive(Clone, Copy, Debug, PartialEq, Eq)]
101pub struct PatchSite {
102    pub offset: CodeOffset,
103    pub kind: LabelUse,
104    pub current_target: CodeOffset,
105    pub addend: i64,
106}
107
108impl PatchSite {
109    /// Build a patch handle from a (possibly rebased) catalog entry.
110    pub const fn to_patchable(self) -> PatchableSite {
111        // SAFETY: catalog entries describe sites recorded during emission.
112        unsafe { PatchableSite::new(self.offset, self.kind, self.addend) }
113    }
114}
115
116/// Finalized patch metadata for validation and linker rebasing.
117#[derive(Clone, Debug, PartialEq, Eq)]
118pub struct PatchCatalog {
119    arch: Arch,
120    blocks: SmallVec<[PatchBlock; 4]>,
121    sites: SmallVec<[PatchSite; 8]>,
122}
123
124impl PatchCatalog {
125    pub(crate) fn with_parts(
126        arch: Arch,
127        blocks: SmallVec<[PatchBlock; 4]>,
128        sites: SmallVec<[PatchSite; 8]>,
129    ) -> Self {
130        Self {
131            arch,
132            blocks,
133            sites,
134        }
135    }
136
137    pub fn arch(&self) -> Arch {
138        self.arch
139    }
140
141    pub fn is_empty(&self) -> bool {
142        self.blocks.is_empty() && self.sites.is_empty()
143    }
144
145    pub fn blocks(&self) -> &[PatchBlock] {
146        &self.blocks
147    }
148
149    pub fn sites(&self) -> &[PatchSite] {
150        &self.sites
151    }
152
153    pub fn block(&self, id: PatchBlockId) -> Option<&PatchBlock> {
154        self.blocks.get(id.index())
155    }
156
157    pub fn site(&self, id: PatchSiteId) -> Option<&PatchSite> {
158        self.sites.get(id.index())
159    }
160
161    pub fn site_mut(&mut self, id: PatchSiteId) -> Option<&mut PatchSite> {
162        self.sites.get_mut(id.index())
163    }
164}
165
166/// Self-describing handle for a patchable displacement (jump/call).
167#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
168pub struct PatchableSite {
169    offset: CodeOffset,
170    kind: LabelUse,
171    addend: i64,
172}
173
174impl PatchableSite {
175    /// Construct a site handle manually.
176    ///
177    /// # Safety
178    ///
179    /// `offset` in any image you later patch must be a valid displacement field of `kind`
180    /// with the same layout as when the site was emitted or recorded.
181    pub const unsafe fn new(offset: CodeOffset, kind: LabelUse, addend: i64) -> Self {
182        Self {
183            offset,
184            kind,
185            addend,
186        }
187    }
188
189    pub const fn offset(self) -> CodeOffset {
190        self.offset
191    }
192
193    pub const fn kind(self) -> LabelUse {
194        self.kind
195    }
196
197    pub const fn addend(self) -> i64 {
198        self.addend
199    }
200
201    /// Retarget this site in a mutable code image.
202    ///
203    /// # Safety
204    ///
205    /// `bytes` must be the code image this site was recorded against (same layout). The caller
206    /// synchronizes concurrent execution of the patched region.
207    pub unsafe fn retarget(
208        self,
209        bytes: &mut [u8],
210        target_offset: CodeOffset,
211    ) -> Result<(), AsmError> {
212        if !self.kind.can_reach(self.offset, target_offset) {
213            return Err(AsmError::TooLarge);
214        }
215        let patch_size = self.kind.patch_size();
216        let patch_end = (self.offset as usize)
217            .checked_add(patch_size)
218            .ok_or(AsmError::InvalidState)?;
219        if patch_end > bytes.len() {
220            return Err(AsmError::InvalidState);
221        }
222        let patch_slice = &mut bytes[self.offset as usize..patch_end];
223        self.kind
224            .patch_with_addend(patch_slice, self.offset, target_offset, self.addend);
225        Ok(())
226    }
227
228    /// Retarget this site in executable memory.
229    ///
230    /// # Safety
231    ///
232    /// `span` must be the loaded image this site was recorded against. The caller synchronizes
233    /// concurrent execution of the patched region.
234    #[cfg(feature = "jit")]
235    pub unsafe fn retarget_span(
236        self,
237        jit_allocator: &mut JitAllocator,
238        span: &mut Span,
239        target_offset: CodeOffset,
240    ) -> Result<(), AsmError> {
241        if !self.kind.can_reach(self.offset, target_offset) {
242            return Err(AsmError::TooLarge);
243        }
244        let patch_size = self.kind.patch_size();
245        let patch_end = (self.offset as usize)
246            .checked_add(patch_size)
247            .ok_or(AsmError::InvalidState)?;
248        if patch_end > span.size() {
249            return Err(AsmError::InvalidState);
250        }
251
252        // SAFETY: per this method's contract, `span` is the loaded image this
253        // site was recorded against; the range `offset..offset + patch_size`
254        // was bounds-checked against `span.size()` above. `JitAllocator::write`
255        // toggles JIT write access and flushes the instruction cache, and the
256        // caller synchronizes concurrent execution.
257        unsafe {
258            jit_allocator.write(span, |span| {
259                let patch_ptr = span.rw().add(self.offset as usize);
260                let patch_slice = core::slice::from_raw_parts_mut(patch_ptr, patch_size);
261                self.kind
262                    .patch_with_addend(patch_slice, self.offset, target_offset, self.addend);
263            })?;
264        }
265        Ok(())
266    }
267}
268
269/// Self-describing handle for a rewritable code region (immediate or custom block).
270#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
271pub struct PatchableBlock {
272    offset: CodeOffset,
273    size: CodeOffset,
274    arch: Arch,
275}
276
277impl PatchableBlock {
278    /// Construct a block handle manually.
279    ///
280    /// # Safety
281    ///
282    /// `offset..offset+size` in any image you later patch must be a reserved patch region for
283    /// `arch` (instruction alignment and nop-fill rules apply).
284    pub const unsafe fn new(offset: CodeOffset, size: CodeOffset, arch: Arch) -> Self {
285        Self { offset, size, arch }
286    }
287
288    pub const fn offset(self) -> CodeOffset {
289        self.offset
290    }
291
292    pub const fn size(self) -> CodeOffset {
293        self.size
294    }
295
296    pub const fn arch(self) -> Arch {
297        self.arch
298    }
299
300    /// Overwrite this block; shorter payloads are padded with architecture nops.
301    ///
302    /// # Safety
303    ///
304    /// `bytes` must be the code image this block was recorded against. The caller synchronizes
305    /// concurrent execution of the patched region.
306    pub unsafe fn rewrite(self, bytes: &mut [u8], new_bytes: &[u8]) -> Result<(), AsmError> {
307        if new_bytes.len() > self.size as usize {
308            return Err(AsmError::TooLarge);
309        }
310        let instruction_alignment = minimum_patch_alignment(self.arch) as usize;
311        if new_bytes.len() % instruction_alignment != 0 {
312            return Err(AsmError::InvalidArgument);
313        }
314        let block_end = (self.offset as usize)
315            .checked_add(self.size as usize)
316            .ok_or(AsmError::InvalidState)?;
317        if block_end > bytes.len() {
318            return Err(AsmError::InvalidState);
319        }
320
321        let block = &mut bytes[self.offset as usize..block_end];
322        block[..new_bytes.len()].copy_from_slice(new_bytes);
323        fill_with_nops(self.arch, &mut block[new_bytes.len()..])?;
324        Ok(())
325    }
326
327    /// Write a little-endian `u32` into a 4-byte block (x86 `patchable_mov` on `Gp32`, etc.).
328    ///
329    /// # Safety
330    ///
331    /// See [`rewrite`](Self::rewrite).
332    pub unsafe fn repatch_u32(self, bytes: &mut [u8], value: u32) -> Result<(), AsmError> {
333        if self.size != 4 {
334            return Err(AsmError::InvalidArgument);
335        }
336        // SAFETY: forwarding the caller's `rewrite` contract; a 4-byte payload
337        // is aligned for every supported architecture.
338        unsafe { self.rewrite(bytes, &value.to_le_bytes()) }
339    }
340
341    /// Write a little-endian `u64` into an 8-byte block (x86 `patchable_mov` on `Gp64`, etc.).
342    ///
343    /// # Safety
344    ///
345    /// See [`rewrite`](Self::rewrite).
346    pub unsafe fn repatch_u64(self, bytes: &mut [u8], value: u64) -> Result<(), AsmError> {
347        if self.size != 8 {
348            return Err(AsmError::InvalidArgument);
349        }
350        // SAFETY: forwarding the caller's `rewrite` contract; an 8-byte payload
351        // is aligned for every supported architecture.
352        unsafe { self.rewrite(bytes, &value.to_le_bytes()) }
353    }
354
355    /// Overwrite this block in executable memory.
356    ///
357    /// # Safety
358    ///
359    /// `span` must be the loaded image this block was recorded against.
360    #[cfg(feature = "jit")]
361    pub unsafe fn rewrite_span(
362        self,
363        jit_allocator: &mut JitAllocator,
364        span: &mut Span,
365        new_bytes: &[u8],
366    ) -> Result<(), AsmError> {
367        if new_bytes.len() > self.size as usize {
368            return Err(AsmError::TooLarge);
369        }
370        let instruction_alignment = minimum_patch_alignment(self.arch) as usize;
371        if new_bytes.len() % instruction_alignment != 0 {
372            return Err(AsmError::InvalidArgument);
373        }
374        let block_end = (self.offset as usize)
375            .checked_add(self.size as usize)
376            .ok_or(AsmError::InvalidState)?;
377        if block_end > span.size() {
378            return Err(AsmError::InvalidState);
379        }
380
381        let mut fill_result = Ok(());
382        // SAFETY: per this method's contract, `span` is the loaded image this
383        // block was recorded against; `offset..offset + size` was bounds-checked
384        // against `span.size()` above and the payload length was checked to be
385        // at most `size`. `JitAllocator::write` synchronizes JIT access and the
386        // instruction cache; `fill_result` is written by the closure and read
387        // only after `write` returns.
388        unsafe {
389            jit_allocator.write(span, |span| {
390                let block_ptr = span.rw().add(self.offset as usize);
391                block_ptr.copy_from_nonoverlapping(new_bytes.as_ptr(), new_bytes.len());
392                let tail = core::slice::from_raw_parts_mut(
393                    block_ptr.add(new_bytes.len()),
394                    self.size as usize - new_bytes.len(),
395                );
396                fill_result = fill_with_nops(self.arch, tail);
397            })?;
398        }
399        fill_result
400    }
401
402    /// Write a little-endian `u32` into a 4-byte block in executable memory.
403    ///
404    /// # Safety
405    ///
406    /// See [`rewrite_span`](Self::rewrite_span).
407    #[cfg(feature = "jit")]
408    pub unsafe fn repatch_u32_span(
409        self,
410        jit_allocator: &mut JitAllocator,
411        span: &mut Span,
412        value: u32,
413    ) -> Result<(), AsmError> {
414        if self.size != 4 {
415            return Err(AsmError::InvalidArgument);
416        }
417        // SAFETY: forwarding the caller's `rewrite_span` contract; a 4-byte
418        // payload satisfies the minimum patch alignment for every arch.
419        unsafe { self.rewrite_span(jit_allocator, span, &value.to_le_bytes()) }
420    }
421
422    /// Write a little-endian `u64` into an 8-byte block in executable memory.
423    ///
424    /// # Safety
425    ///
426    /// See [`rewrite_span`](Self::rewrite_span).
427    #[cfg(feature = "jit")]
428    pub unsafe fn repatch_u64_span(
429        self,
430        jit_allocator: &mut JitAllocator,
431        span: &mut Span,
432        value: u64,
433    ) -> Result<(), AsmError> {
434        if self.size != 8 {
435            return Err(AsmError::InvalidArgument);
436        }
437        // SAFETY: forwarding the caller's `rewrite_span` contract; an 8-byte
438        // payload satisfies the minimum patch alignment for every arch.
439        unsafe { self.rewrite_span(jit_allocator, span, &value.to_le_bytes()) }
440    }
441}
442
443pub fn minimum_patch_alignment(arch: Arch) -> CodeOffset {
444    match arch {
445        Arch::AArch64 | Arch::RISCV32 | Arch::RISCV64 => 4,
446        _ => 1,
447    }
448}
449
450pub fn fill_with_nops(arch: Arch, buffer: &mut [u8]) -> Result<(), AsmError> {
451    let pattern: &[u8] = match arch {
452        Arch::X86 | Arch::X64 => &[0x90],
453        Arch::AArch64 => &[0x1f, 0x20, 0x03, 0xd5],
454        Arch::RISCV32 | Arch::RISCV64 => &[0x13, 0x00, 0x00, 0x00],
455        _ => return Err(AsmError::InvalidArgument),
456    };
457
458    if pattern.len() > 1 && buffer.len() % pattern.len() != 0 {
459        return Err(AsmError::InvalidArgument);
460    }
461
462    for chunk in buffer.chunks_mut(pattern.len()) {
463        chunk.copy_from_slice(pattern);
464    }
465
466    Ok(())
467}
468
469impl CodeBufferFinalized {
470    pub fn patch_catalog(&self) -> &PatchCatalog {
471        &self.patch_catalog
472    }
473}
474
475#[cfg(test)]
476mod tests {
477    use super::*;
478
479    #[test]
480    fn retarget_rejects_out_of_range_slice() {
481        let site = unsafe { PatchableSite::new(62, LabelUse::X86JmpRel32, 0) };
482        let mut bytes = [0u8; 64];
483        assert_eq!(
484            unsafe { site.retarget(&mut bytes, 0) }.unwrap_err(),
485            AsmError::InvalidState
486        );
487    }
488
489    #[test]
490    fn rewrite_rejects_misaligned_payload_for_a64() {
491        let block = unsafe { PatchableBlock::new(0, 4, Arch::AArch64) };
492        let mut bytes = [0u8; 4];
493        assert_eq!(
494            unsafe { block.rewrite(&mut bytes, &[0]) }.unwrap_err(),
495            AsmError::InvalidArgument
496        );
497    }
498
499    #[test]
500    fn repatch_u32_round_trips() {
501        let block = unsafe { PatchableBlock::new(1, 4, Arch::X64) };
502        let mut bytes = [0xB8, 0, 0, 0, 0];
503        unsafe { block.repatch_u32(&mut bytes, 0x11223344).unwrap() };
504        assert_eq!(&bytes[1..], &[0x44, 0x33, 0x22, 0x11]);
505    }
506
507    #[cfg(feature = "jit")]
508    #[test]
509    fn span_patch_rejects_ranges_outside_span() {
510        use crate::core::jit_allocator::JitAllocatorOptions;
511
512        let mut allocator = JitAllocator::new(JitAllocatorOptions::default());
513        let mut span = allocator.alloc(64).unwrap();
514        let span_size = span.size() as CodeOffset;
515
516        let block = unsafe { PatchableBlock::new(span_size, 1, Arch::X64) };
517        assert_eq!(
518            unsafe { block.rewrite_span(&mut allocator, &mut span, &[0x90]) }.unwrap_err(),
519            AsmError::InvalidState
520        );
521
522        let site = unsafe { PatchableSite::new(span_size - 3, LabelUse::X86JmpRel32, 0) };
523        assert_eq!(
524            unsafe { site.retarget_span(&mut allocator, &mut span, 0) }.unwrap_err(),
525            AsmError::InvalidState
526        );
527    }
528}
asmkit/core/patch.rs

asmkit/core/
patch.rs
