[RFC PATCH bpf-next RESEND 12/16] bpf/crib: Introduce skb_data open-coded iterator kfuncs

[Juntong Deng]

This patch adds open-coded iterator style skb data iterator kfuncs
bpf_iter_skb_data_{new,next,destroy} that iterates over all data in
the specified skb (struct sk_buff).

This iterator is designed for checkpointing, and thus iterates over
all data in the skb (starting with head), not just application-level
data (starting with data).

Each iteration (next) copies the data to the specified buffer and
updates the offset.

The skb_data iterator has two getters
bpf_iter_skb_data_get_offset()/bpf_iter_skb_data_get_chunk_len(),
which are used to get the offset/the size of the chunk read in the
current iteration.

The skb_data iterator has a setter bpf_iter_skb_data_set_buf(),
which is used to set the buffer and the size of the buffer for
dumping the skb data during iteration.

Signed-off-by: Juntong Deng <juntong.deng@xxxxxxxxxxx>
---
 include/linux/bpf_crib.h         |  14 ++++
 kernel/bpf/crib/bpf_checkpoint.c | 116 +++++++++++++++++++++++++++++++
 kernel/bpf/crib/bpf_crib.c       |   7 ++
 3 files changed, 137 insertions(+)

diff --git a/include/linux/bpf_crib.h b/include/linux/bpf_crib.h
index e7cfa9c1ae6b..c073166e60a0 100644
--- a/include/linux/bpf_crib.h
+++ b/include/linux/bpf_crib.h
@@ -32,4 +32,18 @@ struct bpf_iter_skb_kern {
 	struct sk_buff *skb;
 } __aligned(8);
 
+struct bpf_iter_skb_data {
+	__u64 __opaque[5];
+} __aligned(8);
+
+struct bpf_iter_skb_data_kern {
+	struct sk_buff *skb;
+	char *buf;
+	unsigned int buflen;
+	int offset;
+	unsigned int headerlen;
+	unsigned int size;
+	unsigned int chunklen;
+} __aligned(8);
+
 #endif /* _BPF_CRIB_H */
diff --git a/kernel/bpf/crib/bpf_checkpoint.c b/kernel/bpf/crib/bpf_checkpoint.c
index c95844faecbc..5c56f1cbf3c8 100644
--- a/kernel/bpf/crib/bpf_checkpoint.c
+++ b/kernel/bpf/crib/bpf_checkpoint.c
@@ -241,4 +241,120 @@ __bpf_kfunc void bpf_iter_skb_destroy(struct bpf_iter_skb *it)
 		bpf_skb_release(kit->skb);
 }
 
+/**
+ * bpf_iter_skb_data_new() - Initialize a new skb data iterator for a skb
+ * (sk_buff), used to iterates over all skb data in the specified skb
+ *
+ * @it: new bpf_iter_skb_data to be created
+ * @skb: a pointer to a sk_buff to be iterated over
+ * @buf: buffer for dumping skb data
+ * @buflen: buffer length
+ */
+__bpf_kfunc int bpf_iter_skb_data_new(struct bpf_iter_skb_data *it,
+		struct sk_buff *skb, char *buf, int buflen)
+{
+	struct bpf_iter_skb_data_kern *kit = (void *)it;
+
+	BUILD_BUG_ON(sizeof(struct bpf_iter_skb_data_kern) > sizeof(struct bpf_iter_skb_data));
+	BUILD_BUG_ON(__alignof__(struct bpf_iter_skb_data_kern) !=
+		     __alignof__(struct bpf_iter_skb_data));
+
+	int headerlen = skb_headroom(skb);
+
+	kit->skb = skb;
+	kit->headerlen = headerlen;
+	kit->offset = 0;
+	kit->chunklen = 0;
+	kit->size = headerlen + skb->len;
+	kit->buf = buf;
+	kit->buflen = buflen;
+
+	return 0;
+}
+
+/**
+ * bpf_iter_skb_data_next() - Dumps the corresponding data in skb to
+ * the buffer based on the current offset and buffer size, and updates
+ * the offset after copying the data
+ *
+ * @it: bpf_iter_skb_data to be checked
+ *
+ * @returns a pointer to the buffer if further data is available,
+ * otherwise returns NULL
+ */
+__bpf_kfunc char *bpf_iter_skb_data_next(struct bpf_iter_skb_data *it)
+{
+	struct bpf_iter_skb_data_kern *kit = (void *)it;
+
+	if (!kit->buf || kit->buflen <= 0)
+		return NULL;
+
+	if (kit->offset >= kit->size)
+		return NULL;
+
+	kit->chunklen = (kit->offset + kit->buflen > kit->size) ?
+			kit->size - kit->offset : kit->buflen;
+
+	skb_copy_bits(kit->skb, kit->offset - kit->headerlen, kit->buf, kit->chunklen);
+
+	kit->offset += kit->chunklen;
+
+	return kit->buf;
+}
+
+/**
+ * bpf_iter_skb_data_set_buf() - Set the buffer for dumping the skb data
+ * during iteration
+ *
+ * @it: bpf_iter_skb_data to be set
+ * @buf: buffer
+ * @buflen: buffer length
+ */
+__bpf_kfunc void bpf_iter_skb_data_set_buf(struct bpf_iter_skb_data *it, char *buf, int buflen)
+{
+	struct bpf_iter_skb_data_kern *kit = (void *)it;
+
+	kit->buf = buf;
+	kit->buflen = buflen;
+}
+
+/**
+ * bpf_iter_skb_data_get_chunk_len() - get the size of the chunk read
+ * in the current iteration
+ *
+ * @it: bpf_iter_skb_data to be checked
+ *
+ * @returns read size in the current iteration
+ */
+__bpf_kfunc int bpf_iter_skb_data_get_chunk_len(struct bpf_iter_skb_data *it)
+{
+	struct bpf_iter_skb_data_kern *kit = (void *)it;
+
+	return kit->chunklen;
+}
+
+/**
+ * bpf_iter_skb_data_get_offset() - get the offset of the chunk read
+ * in the current iteration
+ *
+ * @it: bpf_iter_skb_data to be checked
+ *
+ * @returns offset in the current iteration
+ */
+__bpf_kfunc int bpf_iter_skb_data_get_offset(struct bpf_iter_skb_data *it)
+{
+	struct bpf_iter_skb_data_kern *kit = (void *)it;
+
+	return kit->offset - kit->chunklen;
+}
+
+/**
+ * bpf_iter_skb_destroy() - Destroy a bpf_iter_skb_data
+ *
+ * @it: bpf_iter_skb_data to be destroyed
+ */
+__bpf_kfunc void bpf_iter_skb_data_destroy(struct bpf_iter_skb_data *it)
+{
+}
+
 __bpf_kfunc_end_defs();
diff --git a/kernel/bpf/crib/bpf_crib.c b/kernel/bpf/crib/bpf_crib.c
index fda34d8143f1..462ae1ab50e5 100644
--- a/kernel/bpf/crib/bpf_crib.c
+++ b/kernel/bpf/crib/bpf_crib.c
@@ -293,6 +293,13 @@ BTF_ID_FLAGS(func, bpf_iter_skb_new, KF_ITER_NEW | KF_TRUSTED_ARGS)
 BTF_ID_FLAGS(func, bpf_iter_skb_next, KF_ITER_NEXT | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_iter_skb_destroy, KF_ITER_DESTROY)
 
+BTF_ID_FLAGS(func, bpf_iter_skb_data_new, KF_ITER_NEW | KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_iter_skb_data_next, KF_ITER_NEXT | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_iter_skb_data_set_buf, KF_ITER_SETTER | KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_iter_skb_data_get_chunk_len, KF_ITER_GETTER)
+BTF_ID_FLAGS(func, bpf_iter_skb_data_get_offset, KF_ITER_GETTER)
+BTF_ID_FLAGS(func, bpf_iter_skb_data_destroy, KF_ITER_DESTROY)
+
 BTF_KFUNCS_END(bpf_crib_kfuncs)
 
 static int bpf_prog_run_crib(struct bpf_prog *prog,
-- 
2.39.2