map: Introduce BatchCursor abstraction

map.go

@@ -963,32 +963,73 @@ func (m *Map) guessNonExistentKey() ([]byte, error) {
 //
 // "keysOut" and "valuesOut" must be of type slice, a pointer
 // to a slice or buffer will not work.
-// "prevKey" is the key to start the batch lookup from, it will
-// *not* be included in the results. Use nil to start at the first key.
+// "cursor" is an pointer to an opaque handle. It must be non-nil. Pass
+// "cursor" to subsequent calls of this function to continue the batching
+// operation in the case of chunking.
+//
+// Warning: This API is not very safe to use as the kernel implementation for
+// batching relies on the user to be aware of subtle details with regarding to
+// different map type implementations.
 //
 // ErrKeyNotExist is returned when the batch lookup has reached
 // the end of all possible results, even when partial results
 // are returned. It should be used to evaluate when lookup is "done".
-func (m *Map) BatchLookup(prevKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
-	return m.batchLookup(sys.BPF_MAP_LOOKUP_BATCH, prevKey, nextKeyOut, keysOut, valuesOut, opts)
+func (m *Map) BatchLookup(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	return m.batchLookup(sys.BPF_MAP_LOOKUP_BATCH, cursor, keysOut, valuesOut, opts)
 }
 
 // BatchLookupAndDelete looks up many elements in a map at once,
 //
 // It then deletes all those elements.
 // "keysOut" and "valuesOut" must be of type slice, a pointer
 // to a slice or buffer will not work.
-// "prevKey" is the key to start the batch lookup from, it will
-// *not* be included in the results. Use nil to start at the first key.
+// "cursor" is an pointer to an opaque handle. It must be non-nil. Pass
+// "cursor" to subsequent calls of this function to continue the batching
+// operation in the case of chunking.
+//
+// Warning: This API is not very safe to use as the kernel implementation for
+// batching relies on the user to be aware of subtle details with regarding to
+// different map type implementations.
 //
 // ErrKeyNotExist is returned when the batch lookup has reached
 // the end of all possible results, even when partial results
 // are returned. It should be used to evaluate when lookup is "done".
-func (m *Map) BatchLookupAndDelete(prevKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
-	return m.batchLookup(sys.BPF_MAP_LOOKUP_AND_DELETE_BATCH, prevKey, nextKeyOut, keysOut, valuesOut, opts)
-}
+func (m *Map) BatchLookupAndDelete(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	return m.batchLookup(sys.BPF_MAP_LOOKUP_AND_DELETE_BATCH, cursor, keysOut, valuesOut, opts)
+}
+
+// BatchCursor represents a starting point for a batch operation.
+type BatchCursor struct {
+	m      *Map
+	opaque []byte
+}
+
+func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	cursorLen := int(m.keySize)
+	if cursorLen < 4 {
+		// * generic_map_lookup_batch requires that batch_out is key_size bytes.
+		//   This is used by array and LPM maps.
+		//
+		// * __htab_map_lookup_and_delete_batch requires u32. This is used by the
+		//   various hash maps.
+		//
+		// Use a minimum of 4 bytes to avoid having to distinguish between the two.
+		cursorLen = 4
+	}
+
+	inBatch := cursor.opaque
+	if inBatch == nil {
+		// This is the first lookup, allocate a buffer to hold the cursor.
+		cursor.opaque = make([]byte, cursorLen)
+		cursor.m = m
+	} else if cursor.m != m {
+		// Prevent reuse of a cursor across maps. First, it's unlikely to work.
+		// Second, the maps may require different cursorLen and cursor.opaque
+		// may therefore be too short. This could lead to the kernel clobbering
+		// user space memory.
+		return 0, errors.New("a cursor may not be reused across maps")
+	}
 
-func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
 	if err := haveBatchAPI(); err != nil {
 		return 0, err
 	}
@@ -1011,40 +1052,28 @@ func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut
 	keyPtr := sys.NewSlicePointer(keyBuf)
 	valueBuf := make([]byte, count*int(m.fullValueSize))
 	valuePtr := sys.NewSlicePointer(valueBuf)
-	nextBuf := makeMapSyscallOutput(nextKeyOut, int(m.keySize))
 
 	attr := sys.MapLookupBatchAttr{
 		MapFd:    m.fd.Uint(),
 		Keys:     keyPtr,
 		Values:   valuePtr,
 		Count:    uint32(count),
-		OutBatch: nextBuf.Pointer(),
+		InBatch:  sys.NewSlicePointer(inBatch),
+		OutBatch: sys.NewSlicePointer(cursor.opaque),
 	}
 
 	if opts != nil {
 		attr.ElemFlags = opts.ElemFlags
 		attr.Flags = opts.Flags
 	}
 
-	var err error
-	if startKey != nil {
-		attr.InBatch, err = marshalMapSyscallInput(startKey, int(m.keySize))
-		if err != nil {
-			return 0, err
-		}
-	}
-
 	_, sysErr := sys.BPF(cmd, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
 	sysErr = wrapMapError(sysErr)
 	if sysErr != nil && !errors.Is(sysErr, unix.ENOENT) {
 		return 0, sysErr
 	}
 
-	err = nextBuf.Unmarshal(nextKeyOut)
-	if err != nil {
-		return 0, err
-	}
-	err = sysenc.Unmarshal(keysOut, keyBuf)
+	err := sysenc.Unmarshal(keysOut, keyBuf)
 	if err != nil {
 		return 0, err
 	}