1use derivative::Derivative;
6use smallvec::SmallVec;
7use std::marker::PhantomData;
8use std::mem::MaybeUninit;
9use std::ops::RangeBounds;
10use zerocopy::{FromBytes, IntoBytes};
11
12#[cfg(target_arch = "aarch64")]
13mod arm64;
14
15#[cfg(target_arch = "aarch64")]
16use arm64 as arch;
17
18#[cfg(target_arch = "x86_64")]
19mod x64;
20
21#[cfg(target_arch = "x86_64")]
22use x64 as arch;
23
24#[cfg(target_arch = "riscv64")]
25mod riscv64;
26
27#[cfg(target_arch = "riscv64")]
28use riscv64 as arch;
29
30#[derive(Derivative)]
34#[derivative(Copy(bound = ""), Clone(bound = ""))]
35pub struct EbpfPtr<'a, T> {
36 ptr: *mut T,
37 phantom: PhantomData<&'a T>,
38}
39
40#[allow(clippy::undocumented_unsafe_blocks, reason = "Force documented unsafe blocks in Starnix")]
41unsafe impl<'a, T> Send for EbpfPtr<'a, T> {}
42#[allow(clippy::undocumented_unsafe_blocks, reason = "Force documented unsafe blocks in Starnix")]
43unsafe impl<'a, T> Sync for EbpfPtr<'a, T> {}
44
45impl<'a, T> EbpfPtr<'a, T>
46where
47 T: Sized,
48{
49 pub unsafe fn new(ptr: *mut T) -> Self {
55 Self { ptr, phantom: PhantomData }
56 }
57
58 pub unsafe fn deref(&self) -> &'a T {
62 #[allow(clippy::undocumented_unsafe_blocks, reason = "2024 edition migration")]
63 unsafe {
64 &*self.ptr
65 }
66 }
67
68 pub fn get_field<F, const OFFSET: usize>(&self) -> EbpfPtr<'a, F> {
69 assert!(OFFSET + std::mem::size_of::<F>() <= std::mem::size_of::<T>());
70 let field_ptr = unsafe { self.ptr.byte_offset(OFFSET as isize) } as *mut F;
73 EbpfPtr::<'a, F> { ptr: field_ptr, phantom: PhantomData }
74 }
75
76 pub fn ptr(&self) -> *mut T {
77 self.ptr
78 }
79}
80
81impl<'a, T> From<&'a mut T> for EbpfPtr<'a, T>
82where
83 T: IntoBytes + FromBytes + Sized,
84{
85 fn from(value: &'a mut T) -> Self {
86 let ptr = value.as_mut_bytes().as_mut_ptr() as *mut T;
87 unsafe { Self::new(ptr) }
91 }
92}
93
94impl EbpfPtr<'_, u64> {
95 pub fn load_relaxed(&self) -> u64 {
98 unsafe { arch::load_u64(self.ptr) }
100 }
101
102 pub fn store_relaxed(&self, value: u64) {
105 unsafe { arch::store_u64(self.ptr, value) }
107 }
108}
109
110impl EbpfPtr<'_, u32> {
111 pub fn load_relaxed(&self) -> u32 {
114 unsafe { arch::load_u32(self.ptr) }
116 }
117
118 pub fn store_relaxed(&self, value: u32) {
121 unsafe { arch::store_u32(self.ptr, value) }
123 }
124}
125
126impl EbpfPtr<'_, i32> {
127 pub fn load_relaxed(&self) -> i32 {
130 unsafe { arch::load_u32(self.ptr as *mut u32) as i32 }
132 }
133
134 pub fn store_relaxed(&self, value: i32) {
137 unsafe { arch::store_u32(self.ptr as *mut u32, value as u32) }
139 }
140}
141
142impl EbpfPtr<'_, u16> {
143 pub fn load_relaxed(&self) -> u16 {
146 unsafe { arch::load_u16(self.ptr) }
148 }
149
150 pub fn store_relaxed(&self, value: u16) {
153 unsafe { arch::store_u16(self.ptr, value) }
155 }
156}
157
158impl EbpfPtr<'_, u8> {
159 pub fn load_relaxed(&self) -> u8 {
161 unsafe { arch::load_u8(self.ptr) }
163 }
164
165 pub fn store_relaxed(&self, value: u8) {
167 unsafe { arch::store_u8(self.ptr, value) }
169 }
170}
171
172#[derive(Copy, Clone)]
180pub struct EbpfBufferPtr<'a> {
181 ptr: *mut u8,
182 size: usize,
183 phantom: PhantomData<&'a u8>,
184}
185
186macro_rules! for_each_element {
197 ($buffer:expr, $offset:ident, $ebpf_ptr:ident, $T:ident, $body:block) => {
198 let buffer = $buffer;
199 let mut $offset: usize = 0;
200 let mut cur_ptr = buffer.raw_ptr();
201 let len = buffer.len();
202 let end_ptr = unsafe { cur_ptr.add(len) };
204
205 if (cur_ptr as usize) % 8 > 0 {
206 if cur_ptr < end_ptr && (cur_ptr as usize) % 2 > 0 {
207 type $T = u8;
208 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
210 $body
211 cur_ptr = unsafe { cur_ptr.add(1) };
213 $offset += 1;
214 }
215 if (cur_ptr as usize) + 2 <= (end_ptr as usize) && (cur_ptr as usize) % 4 > 0 {
216 type $T = u16;
217 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
219 $body
220 cur_ptr = unsafe { cur_ptr.add(2) };
222 $offset += 2;
223 }
224 if (cur_ptr as usize) + 4 <= (end_ptr as usize) && (cur_ptr as usize) % 8 > 0 {
225 type $T = u32;
226 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
228 $body
229 cur_ptr = unsafe { cur_ptr.add(4) };
231 $offset += 4;
232 }
233 }
234
235 while (cur_ptr as usize) + 8 <= (end_ptr as usize) {
236 type $T = u64;
237 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
239 $body
240 cur_ptr = unsafe { cur_ptr.add(8) };
242 $offset += 8;
243 }
244
245 if cur_ptr < end_ptr {
246 if (cur_ptr as usize) + 4 <= (end_ptr as usize) {
247 type $T = u32;
248 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
250 $body
251 cur_ptr = unsafe { cur_ptr.add(4) };
253 $offset += 4;
254 }
255 if (cur_ptr as usize) + 2 <= (end_ptr as usize) {
256 type $T = u16;
257 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
259 $body
260 cur_ptr = unsafe { cur_ptr.add(2) };
262 $offset += 2;
263 }
264 if cur_ptr < end_ptr {
265 type $T = u8;
266 let $ebpf_ptr = unsafe { EbpfPtr::new(cur_ptr as *mut $T) };
268 $body
269 cur_ptr = unsafe { cur_ptr.add(1) };
271 $offset += 1;
272 }
273 }
274
275 debug_assert_eq!(cur_ptr, end_ptr);
276 debug_assert_eq!($offset, len);
277 };
278}
279
280impl<'a> EbpfBufferPtr<'a> {
281 pub const ALIGNMENT: usize = size_of::<u64>();
282
283 pub unsafe fn new(ptr: *mut u8, size: usize) -> Self {
289 Self { ptr, size, phantom: PhantomData }
290 }
291
292 pub fn len(&self) -> usize {
294 self.size
295 }
296
297 pub fn raw_ptr(&self) -> *mut u8 {
299 self.ptr
300 }
301
302 unsafe fn get_ptr_internal<T>(&self, offset: usize) -> EbpfPtr<'a, T> {
305 unsafe { EbpfPtr::new(self.ptr.byte_offset(offset as isize) as *mut T) }
308 }
309
310 pub fn get_ptr<T>(&self, offset: usize) -> Option<EbpfPtr<'a, T>> {
312 if offset + std::mem::size_of::<T>() <= self.size {
313 Some(unsafe { self.get_ptr_internal(offset) })
315 } else {
316 None
317 }
318 }
319
320 pub fn slice(&self, range: impl RangeBounds<usize>) -> Option<Self> {
322 let start = match range.start_bound() {
323 std::ops::Bound::Included(&start) => start,
324 std::ops::Bound::Excluded(&start) => start + 1,
325 std::ops::Bound::Unbounded => 0,
326 };
327 let end = match range.end_bound() {
328 std::ops::Bound::Included(&end) => end + 1,
329 std::ops::Bound::Excluded(&end) => end,
330 std::ops::Bound::Unbounded => self.size,
331 };
332
333 assert!(start <= end);
334 (end <= self.size).then(|| {
335 unsafe {
339 Self {
340 ptr: self.ptr.byte_offset(start as isize),
341 size: end - start,
342 phantom: PhantomData,
343 }
344 }
345 })
346 }
347
348 pub fn load_to_slice(&self, dst: &mut [MaybeUninit<u8>]) {
351 assert_eq!(dst.len(), self.size);
352
353 for_each_element!(self, offset, ptr, T, {
354 let value = ptr.load_relaxed();
355 let value_bytes = value.as_bytes();
356
357 unsafe {
359 std::ptr::copy_nonoverlapping(
360 value_bytes.as_ptr(),
361 dst[offset].as_mut_ptr(),
362 std::mem::size_of::<T>(),
363 )
364 }
365 });
366 }
367
368 pub fn load<const N: usize>(&self) -> SmallVec<[u8; N]> {
370 if self.size <= N {
371 let mut buf = MaybeUninit::<[u8; N]>::uninit();
372 self.load_to_slice(&mut AsMut::<[MaybeUninit<u8>]>::as_mut(&mut buf)[..self.size]);
373 unsafe { SmallVec::from_buf_and_len_unchecked(buf, self.size) }
375 } else {
376 let mut vec = Vec::<u8>::with_capacity(self.size);
377 self.load_to_slice(vec.spare_capacity_mut());
378 unsafe { vec.set_len(self.size) };
380 SmallVec::from_vec(vec)
381 }
382 }
383
384 pub fn store(&self, data: &[u8]) {
386 assert!(data.len() <= self.size);
387
388 let buffer = self.slice(..data.len()).unwrap();
389 for_each_element!(buffer, offset, ptr, T, {
390 let value =
391 T::read_from_bytes(&data[offset..offset + std::mem::size_of::<T>()]).unwrap();
392 ptr.store_relaxed(value);
393 });
394 }
395
396 pub fn copy(&self, src: &EbpfBufferPtr<'_>) {
399 assert!(src.len() <= self.size);
400
401 let mut dst_ptr = self.ptr;
402 let dst_end = unsafe { dst_ptr.add(src.len()) };
404
405 let mut src_ptr = src.ptr;
406 let src_end = unsafe { src_ptr.add(src.len()) };
408
409 if (src_ptr as usize) % 8 == 0 && (dst_ptr as usize) % 8 == 0 {
411 while (src_ptr as usize) + 8 <= src_end as usize {
412 unsafe {
414 let value: u64 = arch::load_u64(src_ptr as *const u64);
415 arch::store_u64(dst_ptr as *mut u64, value);
416 src_ptr = src_ptr.add(8);
417 dst_ptr = dst_ptr.add(8);
418 }
419 }
420
421 if src_ptr < src_end {
422 if (src_ptr as usize) + 4 <= src_end as usize {
423 unsafe {
425 let value: u32 = arch::load_u32(src_ptr as *const u32);
426 arch::store_u32(dst_ptr as *mut u32, value);
427 src_ptr = src_ptr.add(4);
428 dst_ptr = dst_ptr.add(4);
429 }
430 }
431
432 if (src_ptr as usize) + 2 <= src_end as usize {
433 unsafe {
435 let value: u16 = arch::load_u16(src_ptr as *const u16);
436 arch::store_u16(dst_ptr as *mut u16, value);
437 src_ptr = src_ptr.add(2);
438 dst_ptr = dst_ptr.add(2);
439 }
440 }
441
442 if src_ptr < src_end {
443 unsafe {
445 let value: u8 = arch::load_u8(src_ptr as *const u8);
446 arch::store_u8(dst_ptr as *mut u8, value);
447 src_ptr = src_ptr.add(1);
448 dst_ptr = dst_ptr.add(1);
449 }
450 }
451 }
452
453 debug_assert_eq!(src_ptr, src_end);
454 debug_assert_eq!(dst_ptr, dst_end);
455 } else {
456 self.store(&src.load::<128>());
460 }
461 }
462
463 pub fn eq_slice(&self, slice: &[u8]) -> bool {
466 if self.size != slice.len() {
467 return false;
468 }
469
470 for_each_element!(self, offset, ptr, T, {
471 let val = ptr.load_relaxed();
472 let slice_val =
473 T::read_from_bytes(&slice[offset..offset + std::mem::size_of::<T>()]).unwrap();
474 if val != slice_val {
475 return false;
476 }
477 });
478
479 true
480 }
481}
482
483impl<'a> From<&'a mut [u8]> for EbpfBufferPtr<'a> {
484 fn from(value: &'a mut [u8]) -> Self {
485 let ptr = value.as_mut_ptr() as *mut u8;
486 unsafe { Self::new(ptr, value.len()) }
490 }
491}
492
493impl<'a> From<&'a mut Vec<u8>> for EbpfBufferPtr<'a> {
494 fn from(value: &'a mut Vec<u8>) -> Self {
495 let ptr = value.as_mut_ptr() as *mut u8;
496 unsafe { Self::new(ptr, value.len()) }
500 }
501}
502impl<'a, const N: usize> From<&'a mut [u8; N]> for EbpfBufferPtr<'a> {
503 fn from(value: &'a mut [u8; N]) -> Self {
504 let ptr = value.as_mut_ptr() as *mut u8;
505 unsafe { Self::new(ptr, N) }
509 }
510}
511
512#[cfg(test)]
513mod test {
514 use super::*;
515 use fuchsia_runtime::vmar_root_self;
516 use std::sync::Barrier;
517 use std::sync::atomic::{AtomicU32, Ordering};
518 use std::thread;
519
520 #[test]
521 fn test_u64_atomicity() {
522 let vmo_size = zx::system_get_page_size() as usize;
523 let vmo = zx::Vmo::create(vmo_size as u64).unwrap();
524 let addr = vmar_root_self()
525 .map(0, &vmo, 0, vmo_size, zx::VmarFlags::PERM_READ | zx::VmarFlags::PERM_WRITE)
526 .unwrap();
527 #[allow(
528 clippy::undocumented_unsafe_blocks,
529 reason = "Force documented unsafe blocks in Starnix"
530 )]
531 let shared_ptr = unsafe { EbpfPtr::new(addr as *mut u64) };
532
533 const NUM_THREADS: usize = 10;
534
535 let barrier = Barrier::new(NUM_THREADS * 2);
537
538 let finished_writers = AtomicU32::new(0);
539
540 thread::scope(|scope| {
541 let mut threads = Vec::new();
542
543 for _ in 0..10 {
544 threads.push(scope.spawn(|| {
545 barrier.wait();
546 for _ in 0..1000 {
547 for i in 0..255 {
548 let v = i << 8 | i;
550 let v = v << 16 | v;
551 let v = v << 32 | v;
552 shared_ptr.store_relaxed(v);
553 }
554 }
555 finished_writers.fetch_add(1, Ordering::Relaxed);
556 }));
557
558 threads.push(scope.spawn(|| {
559 barrier.wait();
560 loop {
561 for _ in 0..1000 {
562 let v = shared_ptr.load_relaxed();
563 assert!(v >> 32 == v & 0xffff_ffff);
565 assert!((v >> 16) & 0xffff == v & 0xffff);
566 assert!((v >> 8) & 0xff == v & 0xff);
567 }
568 if finished_writers.load(Ordering::Relaxed) == NUM_THREADS as u32 {
569 break;
570 }
571 }
572 }));
573 }
574
575 for t in threads.into_iter() {
576 t.join().expect("failed to join a test thread");
577 }
578 });
579
580 #[allow(
581 clippy::undocumented_unsafe_blocks,
582 reason = "Force documented unsafe blocks in Starnix"
583 )]
584 unsafe {
585 vmar_root_self().unmap(addr, vmo_size).unwrap()
586 };
587 }
588
589 #[test]
590 fn test_buffer_slice() {
591 const SIZE: usize = 32;
592
593 let mut buf = [0; SIZE];
594 #[allow(
595 clippy::undocumented_unsafe_blocks,
596 reason = "Force documented unsafe blocks in Starnix"
597 )]
598 let buf_ptr = unsafe { EbpfBufferPtr::new(buf.as_mut_ptr(), SIZE) };
599
600 buf_ptr.slice(8..16).unwrap().store(&[1, 2, 3, 4, 5, 6, 7, 8]);
601 let value = buf_ptr.slice(0..24).unwrap().load::<16>();
602 assert_eq!(
603 &value[..],
604 &[0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0]
605 );
606
607 assert!(buf_ptr.slice(8..40).is_none());
608 }
609
610 #[test]
611 fn test_buffer_load() {
612 const FULL_SIZE: usize = 40;
613 let mut buf = (0..(FULL_SIZE as u8)).map(|v| v as u8).collect::<Vec<_>>();
614 let buf_ptr = unsafe { EbpfBufferPtr::new(buf.as_mut_ptr(), FULL_SIZE) };
616
617 for start in 0..FULL_SIZE {
618 for end in start..=FULL_SIZE {
619 let slice = buf_ptr.slice(start..end).unwrap();
620 let loaded = slice.load::<16>();
621
622 let expected = (start..end).map(|v| v as u8).collect::<Vec<_>>();
623 assert_eq!(&loaded[..], &expected[..], "failed for range {}..{}", start, end);
624 }
625 }
626 }
627
628 #[test]
629 fn test_buffer_store() {
630 const FULL_SIZE: usize = 40;
631 let mut buf = [0u8; FULL_SIZE];
632 let buf_ptr = unsafe { EbpfBufferPtr::new(buf.as_mut_ptr(), FULL_SIZE) };
634
635 for start in 0..FULL_SIZE {
636 for end in start..=FULL_SIZE {
637 let slice = buf_ptr.slice(start..end).unwrap();
638 let data_to_store = (start..end).map(|v| v as u8).collect::<Vec<_>>();
639 slice.store(&data_to_store);
640
641 let loaded = slice.load::<16>();
642 assert_eq!(&loaded[..], &data_to_store[..], "failed for range {}..{}", start, end);
643 }
644 }
645 }
646
647 #[test]
648 fn test_buffer_copy() {
649 const BASE_SIZE: usize = 48;
650 let mut src_buf = (0..(BASE_SIZE as u8)).map(|v| v as u8).collect::<Vec<_>>();
651 let mut dst_buf = [0u8; BASE_SIZE];
652
653 let src_base = unsafe { EbpfBufferPtr::new(src_buf.as_mut_ptr(), BASE_SIZE) };
655
656 let dst_base = unsafe { EbpfBufferPtr::new(dst_buf.as_mut_ptr(), BASE_SIZE) };
658
659 for src_align in 0..8 {
660 for dst_align in 0..8 {
661 for len in 0..=32 {
662 dst_buf.fill(0);
663
664 let src_slice = src_base.slice(src_align..(src_align + len)).unwrap();
665 let dst_slice = dst_base.slice(dst_align..(dst_align + len)).unwrap();
666
667 dst_slice.copy(&src_slice);
668
669 let loaded = dst_slice.load::<16>();
670 let expected =
671 (src_align..(src_align + len)).map(|v| v as u8).collect::<Vec<_>>();
672 assert_eq!(
673 &loaded[..],
674 &expected[..],
675 "copy failed for length {} with src align {} and dst align {}",
676 len,
677 src_align,
678 dst_align
679 );
680
681 for i in 0..BASE_SIZE {
682 if i < dst_align || i >= dst_align + len {
683 assert_eq!(
684 dst_buf[i], 0,
685 "out-of-bounds memory modified at index {} (dst_align={}, len={})",
686 i, dst_align, len
687 );
688 }
689 }
690 }
691 }
692 }
693 }
694
695 #[test]
696 fn test_buffer_eq_slice() {
697 const BASE_SIZE: usize = 48;
698 let mut src_buf = (0..(BASE_SIZE as u8)).map(|v| v as u8).collect::<Vec<_>>();
699 let src_base = unsafe { EbpfBufferPtr::new(src_buf.as_mut_ptr(), BASE_SIZE) };
701
702 for src_align in 0..8 {
703 for len in 0..=32 {
704 let src_slice = src_base.slice(src_align..(src_align + len)).unwrap();
705 let expected = (src_align..(src_align + len)).map(|v| v as u8).collect::<Vec<_>>();
706
707 assert!(src_slice.eq_slice(&expected));
708
709 if len > 0 {
711 let mut modified = expected.clone();
712 modified[0] ^= 0xff;
713 assert!(!src_slice.eq_slice(&modified));
714
715 let mut modified_end = expected.clone();
716 let last = len - 1;
717 modified_end[last] ^= 0xff;
718 assert!(!src_slice.eq_slice(&modified_end));
719 }
720
721 let mut longer = expected.clone();
723 longer.push(0);
724 assert!(!src_slice.eq_slice(&longer));
725
726 if len > 0 {
727 let shorter = &expected[..len - 1];
728 assert!(!src_slice.eq_slice(shorter));
729 }
730 }
731 }
732 }
733}