From 8e6f447e9b8110d99ba9b0945a6dd1a4f5de350a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alessandro=C2=A0Decina?= <alessandro.d@gmail.com>
Date: Wed, 5 May 2021 02:47:51 +0000
Subject: [PATCH] bpf: sk_skb: add helper methods

This adds support for skb_store_bytes, skb_load_bytes and
l4_csum_replace to SkSkbContext.
---
 bpf/aya-bpf/src/maps/sock_hash.rs  |  2 +-
 bpf/aya-bpf/src/programs/sk_skb.rs | 62 +++++++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/bpf/aya-bpf/src/maps/sock_hash.rs b/bpf/aya-bpf/src/maps/sock_hash.rs
index d978690f..ddfe3d8d 100644
--- a/bpf/aya-bpf/src/maps/sock_hash.rs
+++ b/bpf/aya-bpf/src/maps/sock_hash.rs
@@ -3,7 +3,7 @@ use core::{marker::PhantomData, mem};
 use aya_bpf_cty::c_void;
 
 use crate::{
-    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKHASH, bpf_sock_ops, BPF_F_INGRESS},
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKHASH, bpf_sock_ops},
     helpers::{bpf_msg_redirect_hash, bpf_sock_hash_update},
     programs::SkMsgContext,
     BpfContext,
diff --git a/bpf/aya-bpf/src/programs/sk_skb.rs b/bpf/aya-bpf/src/programs/sk_skb.rs
index 9eac93c9..f8ce0739 100644
--- a/bpf/aya-bpf/src/programs/sk_skb.rs
+++ b/bpf/aya-bpf/src/programs/sk_skb.rs
@@ -1,4 +1,10 @@
-use core::ffi::c_void;
+use core::{
+    ffi::c_void,
+    mem::{self, MaybeUninit},
+};
+
+use aya_bpf_bindings::helpers::{bpf_l4_csum_replace, bpf_skb_load_bytes, bpf_skb_store_bytes};
+use aya_bpf_cty::c_long;
 
 use crate::{bindings::__sk_buff, BpfContext};
 
@@ -10,6 +16,60 @@ impl SkSkbContext {
     pub fn new(skb: *mut __sk_buff) -> SkSkbContext {
         SkSkbContext { skb }
     }
+
+    #[inline]
+    pub fn load<T>(&self, offset: usize) -> Result<T, c_long> {
+        unsafe {
+            let mut data = MaybeUninit::<T>::uninit();
+            let ret = bpf_skb_load_bytes(
+                self.skb as *const _,
+                offset as u32,
+                &mut data as *mut _ as *mut _,
+                mem::size_of::<T>() as u32,
+            );
+            if ret < 0 {
+                return Err(ret);
+            }
+
+            Ok(data.assume_init())
+        }
+    }
+
+    #[inline]
+    pub fn store<T>(&self, offset: usize, v: &T) -> Result<(), c_long> {
+        unsafe {
+            let ret = bpf_skb_store_bytes(
+                self.skb as *mut _,
+                offset as u32,
+                v as *const _ as *const _,
+                mem::size_of::<T>() as u32,
+                0,
+            );
+            if ret < 0 {
+                return Err(ret);
+            }
+        }
+
+        Ok(())
+    }
+
+    #[inline]
+    pub fn l4_csum_replace(
+        &self,
+        offset: usize,
+        from: u64,
+        to: u64,
+        flags: u64,
+    ) -> Result<(), c_long> {
+        unsafe {
+            let ret = bpf_l4_csum_replace(self.skb as *mut _, offset as u32, from, to, flags);
+            if ret < 0 {
+                return Err(ret);
+            }
+        }
+
+        Ok(())
+    }
 }
 
 impl BpfContext for SkSkbContext {