Merge pull request #461 from FallingSnow/main

Add ability to iterate over LpmTrie keys and matches

aya/src/maps/lpm_trie.rs

@@ -3,12 +3,11 @@ use std::{
     borrow::Borrow,
     convert::{AsMut, AsRef},
     marker::PhantomData,
-    mem,
 };
 
 use crate::{
-    maps::{check_kv_size, IterableMap, MapData, MapError},
-    sys::{bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem},
+    maps::{check_kv_size, IterableMap, MapData, MapError, MapIter, MapKeys},
+    sys::{bpf_map_delete_elem, bpf_map_get_next_key, bpf_map_lookup_elem, bpf_map_update_elem},
     Pod,
 };
 
@@ -125,6 +124,24 @@ impl<T: AsRef<MapData>, K: Pod, V: Pod> LpmTrie<T, K, V> {
         })?;
         value.ok_or(MapError::KeyNotFound)
     }
+
+    /// An iterator visiting all key-value pairs in arbitrary order. The
+    /// iterator item type is `Result<(K, V), MapError>`.
+    pub fn iter(&self) -> MapIter<'_, Key<K>, V, Self> {
+        MapIter::new(self)
+    }
+
+    /// An iterator visiting all keys in arbitrary order. The iterator element
+    /// type is `Result<Key<K>, MapError>`.
+    pub fn keys(&self) -> MapKeys<'_, Key<K>> {
+        MapKeys::new(self.inner.as_ref())
+    }
+
+    /// An iterator visiting all keys matching key. The
+    /// iterator item type is `Result<Key<K>, MapError>`.
+    pub fn iter_key(&self, key: Key<K>) -> LpmTrieKeys<'_, K> {
+        LpmTrieKeys::new(self.inner.as_ref(), key)
+    }
 }
 
 impl<T: AsMut<MapData>, K: Pod, V: Pod> LpmTrie<T, K, V> {
@@ -160,14 +177,63 @@ impl<T: AsMut<MapData>, K: Pod, V: Pod> LpmTrie<T, K, V> {
     }
 }
 
-impl<T: AsRef<MapData>, K: Pod, V: Pod> IterableMap<K, V> for LpmTrie<T, K, V> {
+impl<T: AsRef<MapData>, K: Pod, V: Pod> IterableMap<Key<K>, V> for LpmTrie<T, K, V> {
     fn map(&self) -> &MapData {
         self.inner.as_ref()
     }
 
-    fn get(&self, key: &K) -> Result<V, MapError> {
-        let lookup = Key::new(mem::size_of::<K>() as u32, *key);
-        self.get(&lookup, 0)
+    fn get(&self, key: &Key<K>) -> Result<V, MapError> {
+        self.get(key, 0)
+    }
+}
+
+/// Iterator returned by `LpmTrie::iter_key()`.
+pub struct LpmTrieKeys<'coll, K: Pod> {
+    map: &'coll MapData,
+    err: bool,
+    key: Key<K>,
+}
+
+impl<'coll, K: Pod> LpmTrieKeys<'coll, K> {
+    fn new(map: &'coll MapData, key: Key<K>) -> LpmTrieKeys<'coll, K> {
+        LpmTrieKeys {
+            map,
+            err: false,
+            key,
+        }
+    }
+}
+
+impl<K: Pod> Iterator for LpmTrieKeys<'_, K> {
+    type Item = Result<Key<K>, MapError>;
+
+    fn next(&mut self) -> Option<Result<Key<K>, MapError>> {
+        if self.err {
+            return None;
+        }
+
+        let fd = match self.map.fd_or_err() {
+            Ok(fd) => fd,
+            Err(e) => {
+                self.err = true;
+                return Some(Err(e));
+            }
+        };
+
+        match bpf_map_get_next_key(fd, Some(&self.key)) {
+            Ok(Some(key)) => {
+                self.key = key;
+                Some(Ok(key))
+            }
+            Ok(None) => None,
+            Err((_, io_error)) => {
+                self.err = true;
+                Some(Err(MapError::SyscallError {
+                    call: "bpf_map_get_next_key".to_owned(),
+                    io_error,
+                }))
+            }
+        }
     }
 }