From ecd6a6e96bf2a316a62f92adb910b04c6198a000 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alessandro=C2=A0Decina?= <alessandro.d@gmail.com>
Date: Sat, 19 Feb 2022 10:07:24 +0000
Subject: [PATCH] bpf: hash_map: add get_mut() method

---
 bpf/aya-bpf/src/maps/hash_map.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/bpf/aya-bpf/src/maps/hash_map.rs b/bpf/aya-bpf/src/maps/hash_map.rs
index e77b4135..90fb467e 100644
--- a/bpf/aya-bpf/src/maps/hash_map.rs
+++ b/bpf/aya-bpf/src/maps/hash_map.rs
@@ -40,6 +40,11 @@ impl<K, V> HashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -85,6 +90,11 @@ impl<K, V> LruHashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -135,6 +145,11 @@ impl<K, V> PerCpuHashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -185,6 +200,11 @@ impl<K, V> LruPerCpuHashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -217,6 +237,15 @@ fn get<'a, K, V>(def: &mut bpf_map_def, key: &K) -> Option<&'a V> {
     }
 }
 
+#[inline]
+fn get_mut<'a, K, V>(def: &mut bpf_map_def, key: &K) -> Option<&'a mut V> {
+    unsafe {
+        let value = bpf_map_lookup_elem(def as *mut _ as *mut _, key as *const _ as *const c_void);
+        // FIXME: alignment
+        NonNull::new(value as *mut V).map(|mut p| p.as_mut())
+    }
+}
+
 #[inline]
 fn insert<K, V>(def: &mut bpf_map_def, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
     let ret = unsafe {