From 8f55cd728c5d02a1aa24da1092f4f2d58c62a007 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alessandro=C2=A0Decina?= <alessandro.d@gmail.com>
Date: Thu, 6 May 2021 03:17:23 +0000
Subject: [PATCH] bpf: SkSkbContext: add ::l3_csum_replace

---
 bpf/aya-bpf/src/programs/sk_skb.rs | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/bpf/aya-bpf/src/programs/sk_skb.rs b/bpf/aya-bpf/src/programs/sk_skb.rs
index f8ce0739..c44cd57d 100644
--- a/bpf/aya-bpf/src/programs/sk_skb.rs
+++ b/bpf/aya-bpf/src/programs/sk_skb.rs
@@ -3,7 +3,7 @@ use core::{
     mem::{self, MaybeUninit},
 };
 
-use aya_bpf_bindings::helpers::{bpf_l4_csum_replace, bpf_skb_load_bytes, bpf_skb_store_bytes};
+use aya_bpf_bindings::helpers::{bpf_l3_csum_replace, bpf_l4_csum_replace, bpf_skb_load_bytes, bpf_skb_store_bytes};
 use aya_bpf_cty::c_long;
 
 use crate::{bindings::__sk_buff, BpfContext};
@@ -17,6 +17,11 @@ impl SkSkbContext {
         SkSkbContext { skb }
     }
 
+    #[inline]
+    pub fn set_mark(&mut self, mark: u32) {
+        unsafe { *self.skb }.mark = mark;
+    }
+
     #[inline]
     pub fn load<T>(&self, offset: usize) -> Result<T, c_long> {
         unsafe {
@@ -36,7 +41,7 @@ impl SkSkbContext {
     }
 
     #[inline]
-    pub fn store<T>(&self, offset: usize, v: &T) -> Result<(), c_long> {
+    pub fn store<T>(&mut self, offset: usize, v: &T) -> Result<(), c_long> {
         unsafe {
             let ret = bpf_skb_store_bytes(
                 self.skb as *mut _,
@@ -53,6 +58,24 @@ impl SkSkbContext {
         Ok(())
     }
 
+    #[inline]
+    pub fn l3_csum_replace(
+        &self,
+        offset: usize,
+        from: u64,
+        to: u64,
+        size: u64,
+    ) -> Result<(), c_long> {
+        unsafe {
+            let ret = bpf_l3_csum_replace(self.skb as *mut _, offset as u32, from, to, size);
+            if ret < 0 {
+                return Err(ret);
+            }
+        }
+
+        Ok(())
+    }
+
     #[inline]
     pub fn l4_csum_replace(
         &self,