From c6e1d566849be92e4ae36e6ae9d2ad007d252858 Mon Sep 17 00:00:00 2001
From: Michal Rostecki <vadorovsky@gmail.com>
Date: Tue, 19 Jul 2022 11:59:37 +0200
Subject: [PATCH] bpf: Add `load_bytes` method to SKB context

This new method allows to load bytes into the given bytes slice, not
requiring to alloate the memory on stack. It can be used with
PerCpuArrays.

Example:

https://github.com/vadorovsky/aya-examples/blob/main/tc-bytes/tc-bytes-ebpf/src/main.rs

Signed-off-by: Michal Rostecki <vadorovsky@gmail.com>
---
 bpf/aya-bpf/src/programs/sk_buff.rs | 76 +++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/bpf/aya-bpf/src/programs/sk_buff.rs b/bpf/aya-bpf/src/programs/sk_buff.rs
index a442e163..57f9d127 100644
--- a/bpf/aya-bpf/src/programs/sk_buff.rs
+++ b/bpf/aya-bpf/src/programs/sk_buff.rs
@@ -1,4 +1,5 @@
 use core::{
+    cmp,
     ffi::c_void,
     mem::{self, MaybeUninit},
 };
@@ -65,6 +66,81 @@ impl SkBuffContext {
         }
     }
 
+    /// Reads some bytes from the packet into the specified buffer, returning
+    /// how many bytes were read.
+    ///
+    /// Starts reading at `offset` and reads at most `dst.len()` or
+    /// `self.len() - offset` bytes, depending on which one is smaller.
+    ///
+    /// # Examples
+    ///
+    /// Read into a `PerCpuArray`.
+    ///
+    /// ```no_run
+    /// use core::mem;
+    ///
+    /// use aya_bpf::{bindings::TC_ACT_PIPE, macros::map, maps::PerCpuArray, programs::SkBuffContext};
+    /// # #[allow(non_camel_case_types)]
+    /// # struct ethhdr {};
+    /// # #[allow(non_camel_case_types)]
+    /// # struct iphdr {};
+    /// # #[allow(non_camel_case_types)]
+    /// # struct tcphdr {};
+    ///
+    /// const ETH_HDR_LEN: usize = mem::size_of::<ethhdr>();
+    /// const IP_HDR_LEN: usize = mem::size_of::<iphdr>();
+    /// const TCP_HDR_LEN: usize = mem::size_of::<tcphdr>();
+    ///
+    /// #[repr(C)]
+    /// pub struct Buf {
+    ///    pub buf: [u8; 1500],
+    /// }
+    ///
+    /// #[map]
+    /// pub static mut BUF: PerCpuArray<Buf> = PerCpuArray::with_max_entries(1, 0);
+    ///
+    /// fn try_classifier(ctx: SkBuffContext) -> Result<i32, i32> {
+    ///     let buf = unsafe {
+    ///         let ptr = BUF.get_ptr_mut(0).ok_or(TC_ACT_PIPE)?;
+    ///         &mut *ptr
+    ///     };
+    ///     let offset = ETH_HDR_LEN + IP_HDR_LEN + TCP_HDR_LEN;
+    ///     ctx.load_bytes(offset, &mut buf.buf).map_err(|_| TC_ACT_PIPE)?;
+    ///
+    ///     // do something with `buf`
+    ///
+    ///     Ok(TC_ACT_PIPE)
+    /// }
+    /// ```
+    #[inline(always)]
+    pub fn load_bytes(&self, offset: usize, dst: &mut [u8]) -> Result<usize, c_long> {
+        if offset >= self.len() as usize {
+            return Err(-1);
+        }
+        let len = cmp::min(self.len() as isize - offset as isize, dst.len() as isize);
+        // The verifier rejects the program if it can't see that `len > 0`.
+        if len <= 0 {
+            return Err(-1);
+        }
+        // This is only needed to ensure the verifier can see the upper bound.
+        if len > dst.len() as isize {
+            return Err(-1);
+        }
+        let ret = unsafe {
+            bpf_skb_load_bytes(
+                self.skb as *const _,
+                offset as u32,
+                dst.as_mut_ptr() as *mut _,
+                len as u32,
+            )
+        };
+        if ret == 0 {
+            Ok(len as usize)
+        } else {
+            Err(ret)
+        }
+    }
+
     #[inline]
     pub fn store<T>(&mut self, offset: usize, v: &T, flags: u64) -> Result<(), c_long> {
         unsafe {