1use smallvec::SmallVec;
2
3use crate::{
4 AsmError,
5 core::{
6 arch_traits::Arch,
7 buffer::{CodeBufferFinalized, CodeOffset, LabelUse},
8 },
9};
10
11#[cfg(feature = "jit")]
12use crate::core::jit_allocator::{JitAllocator, Span};
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
14pub struct PatchBlockId(u32);
15
16impl PatchBlockId {
17 pub(crate) const fn from_index(index: usize) -> Self {
18 Self(index as u32)
19 }
20
21 pub const fn index(self) -> usize {
22 self.0 as usize
23 }
24}
25
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
27pub struct PatchSiteId(u32);
28
29impl PatchSiteId {
30 pub(crate) const fn from_index(index: usize) -> Self {
31 Self(index as u32)
32 }
33
34 pub const fn index(self) -> usize {
35 self.0 as usize
36 }
37}
38
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub struct PatchBlock {
41 pub offset: CodeOffset,
42 pub size: CodeOffset,
43 pub align: CodeOffset,
44}
45
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub struct PatchSite {
48 pub offset: CodeOffset,
49 pub kind: LabelUse,
50 pub current_target: CodeOffset,
51 pub addend: i64,
52}
53
54#[derive(Clone, Debug, PartialEq, Eq)]
55pub struct PatchCatalog {
56 arch: Arch,
57 blocks: SmallVec<[PatchBlock; 4]>,
58 sites: SmallVec<[PatchSite; 8]>,
59}
60
61impl PatchCatalog {
62 pub(crate) fn with_parts(
63 arch: Arch,
64 blocks: SmallVec<[PatchBlock; 4]>,
65 sites: SmallVec<[PatchSite; 8]>,
66 ) -> Self {
67 Self {
68 arch,
69 blocks,
70 sites,
71 }
72 }
73
74 pub fn arch(&self) -> Arch {
75 self.arch
76 }
77
78 pub fn is_empty(&self) -> bool {
79 self.blocks.is_empty() && self.sites.is_empty()
80 }
81
82 pub fn blocks(&self) -> &[PatchBlock] {
83 &self.blocks
84 }
85
86 pub fn sites(&self) -> &[PatchSite] {
87 &self.sites
88 }
89
90 pub fn block(&self, id: PatchBlockId) -> Option<&PatchBlock> {
91 self.blocks.get(id.index())
92 }
93
94 pub fn site(&self, id: PatchSiteId) -> Option<&PatchSite> {
95 self.sites.get(id.index())
96 }
97
98 pub fn site_mut(&mut self, id: PatchSiteId) -> Option<&mut PatchSite> {
99 self.sites.get_mut(id.index())
100 }
101}
102
103pub fn minimum_patch_alignment(arch: Arch) -> CodeOffset {
104 match arch {
105 Arch::AArch64 | Arch::RISCV32 | Arch::RISCV64 => 4,
106 _ => 1,
107 }
108}
109
110pub fn fill_with_nops(arch: Arch, buffer: &mut [u8]) -> Result<(), AsmError> {
111 let pattern: &[u8] = match arch {
112 Arch::X86 | Arch::X64 => &[0x90],
113 Arch::AArch64 => &[0x1f, 0x20, 0x03, 0xd5],
114 Arch::RISCV32 | Arch::RISCV64 => &[0x13, 0x00, 0x00, 0x00],
115 _ => return Err(AsmError::InvalidArgument),
116 };
117
118 if pattern.len() > 1 && buffer.len() % pattern.len() != 0 {
119 return Err(AsmError::InvalidArgument);
120 }
121
122 for chunk in buffer.chunks_mut(pattern.len()) {
123 chunk.copy_from_slice(pattern);
124 }
125
126 Ok(())
127}
128
129#[cfg(feature = "jit")]
130pub struct LoadedPatchableCode {
131 catalog: PatchCatalog,
132 span: Span,
133}
134
135#[cfg(feature = "jit")]
136impl LoadedPatchableCode {
137 pub(crate) fn new(span: Span, catalog: PatchCatalog) -> Self {
138 Self { catalog, span }
139 }
140
141 pub fn patch_catalog(&self) -> &PatchCatalog {
142 &self.catalog
143 }
144
145 pub const fn rx(&self) -> *const u8 {
146 self.span.rx()
147 }
148
149 pub const fn rw(&self) -> *mut u8 {
150 self.span.rw()
151 }
152
153 pub const fn span(&self) -> &Span {
154 &self.span
155 }
156
157 pub fn retarget_site(
158 &mut self,
159 jit_allocator: &mut JitAllocator,
160 id: PatchSiteId,
161 target_offset: CodeOffset,
162 ) -> Result<(), AsmError> {
163 let site = *self.catalog.site(id).ok_or(AsmError::InvalidArgument)?;
164 if !site.kind.can_reach(site.offset, target_offset) {
165 return Err(AsmError::TooLarge);
166 }
167 let patch_size = site.kind.patch_size();
168 let patch_end = (site.offset as usize)
169 .checked_add(patch_size)
170 .ok_or(AsmError::InvalidState)?;
171 if patch_end > self.span.size() {
172 return Err(AsmError::InvalidState);
173 }
174
175 unsafe {
176 jit_allocator.write(&mut self.span, |span| {
177 let patch_ptr = span.rw().add(site.offset as usize);
178 let patch_slice = core::slice::from_raw_parts_mut(patch_ptr, patch_size);
179 site.kind
180 .patch_with_addend(patch_slice, site.offset, target_offset, site.addend);
181 })?;
182 }
183
184 self.catalog.site_mut(id).unwrap().current_target = target_offset;
185 Ok(())
186 }
187
188 pub fn rewrite_block(
189 &mut self,
190 jit_allocator: &mut JitAllocator,
191 id: PatchBlockId,
192 bytes: &[u8],
193 ) -> Result<(), AsmError> {
194 let block = *self.catalog.block(id).ok_or(AsmError::InvalidArgument)?;
195 if bytes.len() > block.size as usize {
196 return Err(AsmError::TooLarge);
197 }
198 let instruction_alignment = minimum_patch_alignment(self.catalog.arch()) as usize;
199 if bytes.len() % instruction_alignment != 0 {
200 return Err(AsmError::InvalidArgument);
201 }
202 let block_end = (block.offset as usize)
203 .checked_add(block.size as usize)
204 .ok_or(AsmError::InvalidState)?;
205 if block_end > self.span.size() {
206 return Err(AsmError::InvalidState);
207 }
208
209 let mut fill_result = Ok(());
210 unsafe {
211 jit_allocator.write(&mut self.span, |span| {
212 let block_ptr = span.rw().add(block.offset as usize);
213 block_ptr.copy_from_nonoverlapping(bytes.as_ptr(), bytes.len());
214 let tail = core::slice::from_raw_parts_mut(
215 block_ptr.add(bytes.len()),
216 block.size as usize - bytes.len(),
217 );
218 fill_result = fill_with_nops(self.catalog.arch(), tail);
219 })?;
220 }
221 fill_result?;
222
223 Ok(())
224 }
225}
226
227impl CodeBufferFinalized {
228 pub fn patch_catalog(&self) -> &PatchCatalog {
229 &self.patch_catalog
230 }
231
232 #[cfg(feature = "jit")]
233 pub fn allocate_patched(
234 &self,
235 jit_allocator: &mut JitAllocator,
236 ) -> Result<LoadedPatchableCode, AsmError> {
237 let span = self.allocate(jit_allocator)?;
238 Ok(LoadedPatchableCode::new(span, self.patch_catalog.clone()))
239 }
240}
241
242#[cfg(all(test, feature = "jit"))]
243mod tests {
244 use super::*;
245 use crate::core::jit_allocator::JitAllocatorOptions;
246
247 #[test]
248 fn loaded_patch_operations_reject_ranges_outside_span() {
249 let mut allocator = JitAllocator::new(JitAllocatorOptions::default());
250 let span = allocator.alloc(64).unwrap();
251 let span_size = span.size() as CodeOffset;
252
253 let mut blocks = SmallVec::new();
254 blocks.push(PatchBlock {
255 offset: span_size,
256 size: 1,
257 align: 1,
258 });
259 let mut sites = SmallVec::new();
260 sites.push(PatchSite {
261 offset: span_size - 3,
262 kind: LabelUse::X86JmpRel32,
263 current_target: 0,
264 addend: 0,
265 });
266 let catalog = PatchCatalog::with_parts(Arch::X64, blocks, sites);
267 let mut loaded = LoadedPatchableCode::new(span, catalog);
268
269 assert_eq!(
270 loaded
271 .rewrite_block(&mut allocator, PatchBlockId::from_index(0), &[0x90])
272 .unwrap_err(),
273 AsmError::InvalidState
274 );
275 assert_eq!(
276 loaded
277 .retarget_site(&mut allocator, PatchSiteId::from_index(0), 0)
278 .unwrap_err(),
279 AsmError::InvalidState
280 );
281
282 let span = allocator.alloc(64).unwrap();
283 let mut blocks = SmallVec::new();
284 blocks.push(PatchBlock {
285 offset: 0,
286 size: 4,
287 align: 4,
288 });
289 let catalog = PatchCatalog::with_parts(Arch::AArch64, blocks, SmallVec::new());
290 let mut loaded = LoadedPatchableCode::new(span, catalog);
291 assert_eq!(
292 loaded
293 .rewrite_block(&mut allocator, PatchBlockId::from_index(0), &[0])
294 .unwrap_err(),
295 AsmError::InvalidArgument
296 );
297 }
298}