feat(map): Add GetNextKey and LookupAndDeleteElem

For BPFMapLow add the following wrapper methods: GetNextKey() -> bpf_map_get_next_key() LookupAndDeleteElem() -> bpf_map_lookup_and_delete_elem() LookupAndDeleteElemFlags() -> bpf_map_lookup_and_delete_elem_flags() In addition, introduce GetValueAndDeleteKey() and GetValueAndDeleteKeyFlags() helper methods to BPFMapLow which return a pre-allocated lookup value automatically. For BPFMap add the following wrapper methods: GetNextKey() -> bpf_map__get_next_key() LookupAndDeleteElem() -> bpf_map__lookup_and_delete_elem() As for BPFMapLow, introduce GetValueAndDeleteKey() and GetValueAndDeleteKeyFlags() helper methods to BPFMap. Signed-off-by: Tao Chen <chen.dylane@gmail.com> Co-authored-by: Geyslan Gregório <geyslan@gmail.com>
aquasecurity · Apr 6, 2024 · d1d119c · d1d119c
1 parent a55c024
commit d1d119c
Show file tree

Hide file tree

Showing 3 changed files with 210 additions and 12 deletions.
diff --git a/map-low.go b/map-low.go
@@ -279,9 +279,76 @@ func (m *BPFMapLow) GetValueFlags(key unsafe.Pointer, flags MapFlag) ([]byte, er
 	return value, nil
 }
 
-// TODO: implement `bpf_map__lookup_and_delete_elem`
-// func (m *BPFMapLow) GetValueAndDeleteKey(key unsafe.Pointer) ([]byte, error) {
-// }
+func (m *BPFMapLow) LookupAndDeleteElem(
+	key unsafe.Pointer,
+	value unsafe.Pointer,
+) error {
+	retC := C.bpf_map_lookup_and_delete_elem(
+		C.int(m.FileDescriptor()),
+		key,
+		value,
+	)
+	if retC < 0 {
+		return fmt.Errorf("failed to lookup and delete value %v in map %s: %w", key, m.Name(), syscall.Errno(-retC))
+	}
+
+	return nil
+}
+
+func (m *BPFMapLow) LookupAndDeleteElemFlags(
+	key unsafe.Pointer,
+	value unsafe.Pointer,
+	flags MapFlag,
+) error {
+	retC := C.bpf_map_lookup_and_delete_elem_flags(
+		C.int(m.FileDescriptor()),
+		key,
+		value,
+		C.ulonglong(flags),
+	)
+	if retC < 0 {
+		return fmt.Errorf("failed to lookup and delete value %v in map %s: %w", key, m.Name(), syscall.Errno(-retC))
+	}
+
+	return nil
+}
+
+func (m *BPFMapLow) GetValueAndDeleteKey(key unsafe.Pointer) ([]byte, error) {
+	valueSize, err := calcMapValueSize(m.ValueSize(), m.Type())
+	if err != nil {
+		return nil, fmt.Errorf("map %s %w", m.Name(), err)
+	}
+
+	value := make([]byte, valueSize)
+	err = m.LookupAndDeleteElem(
+		key,
+		unsafe.Pointer(&value[0]),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return value, nil
+}
+
+func (m *BPFMapLow) GetValueAndDeleteKeyFlags(key unsafe.Pointer, flags MapFlag) ([]byte, error) {
+	valueSize, err := calcMapValueSize(m.ValueSize(), m.Type())
+	if err != nil {
+		return nil, fmt.Errorf("map %s %w", m.Name(), err)
+	}
+
+	value := make([]byte, valueSize)
+	err = m.LookupAndDeleteElemFlags(
+		key,
+		unsafe.Pointer(&value[0]),
+		flags,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return value, nil
+}
 
 func (m *BPFMapLow) Update(key, value unsafe.Pointer) error {
 	return m.UpdateValueFlags(key, value, MapFlagUpdateAny)
@@ -310,9 +377,18 @@ func (m *BPFMapLow) DeleteKey(key unsafe.Pointer) error {
 	return nil
 }
 
-// TODO: implement `bpf_map__get_next_key`
-// func (m *BPFMapLow) GetNextKey(key unsafe.Pointer) (unsafe.Pointer, error) {
-// }
+func (m *BPFMapLow) GetNextKey(key unsafe.Pointer, nextKey unsafe.Pointer) error {
+	retC := C.bpf_map_get_next_key(
+		C.int(m.FileDescriptor()),
+		key,
+		nextKey,
+	)
+	if retC < 0 {
+		return fmt.Errorf("failed to get next key in map %s: %w", m.Name(), syscall.Errno(-retC))
+	}
+
+	return nil
+}
 
 //
 // BPFMapLow Batch Operations

diff --git a/map.go b/map.go
@@ -405,9 +405,53 @@ func (m *BPFMap) GetValueFlags(key unsafe.Pointer, flags MapFlag) ([]byte, error
 	return value, nil
 }
 
-// TODO: implement `bpf_map__lookup_and_delete_elem` wrapper
-// func (m *BPFMap) GetValueAndDeleteKey(key unsafe.Pointer) ([]byte, error) {
-// }
+// LookupAndDeleteElem stores the value associated with a given key into the
+// provided unsafe.Pointer and deletes the key from the BPFMap.
+func (m *BPFMap) LookupAndDeleteElem(
+	key unsafe.Pointer,
+	value unsafe.Pointer,
+	valueSize uint64,
+	flags MapFlag,
+) error {
+	retC := C.bpf_map__lookup_and_delete_elem(
+		m.bpfMap,
+		key,
+		C.ulong(m.KeySize()),
+		value,
+		C.ulong(valueSize),
+		C.ulonglong(flags),
+	)
+	if retC < 0 {
+		return fmt.Errorf("failed to lookup and delete value %v in map %s: %w", key, m.Name(), syscall.Errno(-retC))
+	}
+
+	return nil
+}
+
+// GetValueAndDeleteKey retrieves the value associated with a given key
+// and delete the key in the BPFMap.
+// It returns the value as a slice of bytes.
+func (m *BPFMap) GetValueAndDeleteKey(key unsafe.Pointer) ([]byte, error) {
+	return m.GetValueAndDeleteKeyFlags(key, MapFlagUpdateAny)
+}
+
+// GetValueAndDeleteKeyFlags retrieves the value associated with a given key
+// and delete the key in the BPFMap, with the specified flags.
+// It returns the value as a slice of bytes.
+func (m *BPFMap) GetValueAndDeleteKeyFlags(key unsafe.Pointer, flags MapFlag) ([]byte, error) {
+	valueSize, err := calcMapValueSize(m.ValueSize(), m.Type())
+	if err != nil {
+		return nil, fmt.Errorf("map %s %w", m.Name(), err)
+	}
+
+	value := make([]byte, valueSize)
+	err = m.LookupAndDeleteElem(key, unsafe.Pointer(&value[0]), uint64(valueSize), flags)
+	if err != nil {
+		return nil, err
+	}
+
+	return value, nil
+}
 
 // Deprecated: use BPFMap.GetValue() or BPFMap.GetValueFlags() instead, since
 // they already calculate the value size for per-cpu maps.
@@ -480,9 +524,20 @@ func (m *BPFMap) DeleteKey(key unsafe.Pointer) error {
 	return nil
 }
 
-// TODO: implement `bpf_map__get_next_key` wrapper
-// func (m *BPFMap) GetNextKey(key unsafe.Pointer) (unsafe.Pointer, error) {
-// }
+// GetNextKey allows to iterate BPF map keys by fetching next key that follows current key.
+func (m *BPFMap) GetNextKey(key unsafe.Pointer, nextKey unsafe.Pointer) error {
+	retC := C.bpf_map__get_next_key(
+		m.bpfMap,
+		key,
+		nextKey,
+		C.ulong(m.KeySize()),
+	)
+	if retC < 0 {
+		return fmt.Errorf("failed to get next key %d in map %s: %w", key, m.Name(), syscall.Errno(-retC))
+	}
+
+	return nil
+}
 
 //
 // BPFMap Batch Operations (low-level API)

diff --git a/selftest/map-batch/main.go b/selftest/map-batch/main.go
@@ -247,6 +247,73 @@ func main() {
 	if count != uint32(fewer) {
 		log.Fatalf("testerMap.DeleteKeyBatch failed: count=%d", count)
 	}
+
+	// map contains only 1 key-value pair.
+
+	// Re-add deleted entries.
+	_, err = testerMap.UpdateBatch(
+		unsafe.Pointer(&keys[0]),
+		unsafe.Pointer(&values[0]),
+		uint32(len(keys)),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	//
+	// GetNextKey
+	//
+
+	// Populate the map again.
+	_, err = testerMap.UpdateBatch(
+		unsafe.Pointer(&keys[0]),
+		unsafe.Pointer(&values[0]),
+		uint32(len(keys)),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Test GetNextKey.
+	key := uint32(0)
+	keyPtr := unsafe.Pointer(&key)
+	keyCnt := 0
+	for {
+		err := testerMap.GetNextKey(keyPtr, keyPtr)
+		if err != nil {
+			if !errors.Is(err, syscall.ENOENT) {
+				log.Fatalf("testerMap.GetNextKey failed: err=%v", err)
+			}
+			break
+		}
+		keyCnt++
+	}
+	if keyCnt != len(keys) {
+		log.Fatalf("testerMap.GetNextKey failed: count=%d", keyCnt)
+	}
+
+	//
+	// GetValueAndDeleteKey
+	//
+
+	// Test GetValueAndDeleteKey.
+	for i, key := range keys {
+		val, err := testerMap.GetValueAndDeleteKey(unsafe.Pointer(&key))
+		if err != nil {
+			log.Fatalf("testerMap.GetValueAndDelete failed: err=%v", err)
+		}
+
+		if endian().Uint32(val) != values[i] {
+			log.Fatalf("testerMpa.GetValueAndDetele failed: val=%d", endian().Uint32(val))
+		}
+	}
+
+	// Check if all keys are deleted.
+	key = 0
+	err = testerMap.GetNextKey(keyPtr, keyPtr)
+	if !errors.Is(err, syscall.ENOENT) {
+		log.Fatalf("testerMap.GetValueAndDeleteKey failed: err=%v", err)
+	}
 }
 
 func endian() binary.ByteOrder {