[RFC PATCH 09/22] ublk: bpf: attach bpf prog to ublk device

[Ming Lei <tom.leiming@gmail.com>]

Attach bpf program to ublk device before adding ublk disk, and detach it
after the disk is removed.

ublk device needs to provide the struct_ops ID for attaching the specific
prog, and each ublk device has to attach to only single bpf prog.

So that we can use the attached bpf prog for handling ublk IO command.

Meantime add two ublk bpf callbacks for prog to attach & detach ublk
device.

Signed-off-by: Ming Lei <tom.leiming@gmail.com>
---
 drivers/block/ublk/Makefile  |  2 +-
 drivers/block/ublk/bpf.c     | 99 ++++++++++++++++++++++++++++++++++++
 drivers/block/ublk/bpf.h     | 33 ++++++++++++
 drivers/block/ublk/bpf_ops.c | 34 +++++++++++++
 drivers/block/ublk/main.c    | 25 ++++++---
 drivers/block/ublk/ublk.h    | 16 ++++++
 6 files changed, 200 insertions(+), 9 deletions(-)
 create mode 100644 drivers/block/ublk/bpf.c

diff --git a/drivers/block/ublk/Makefile b/drivers/block/ublk/Makefile
index 7058b0fc13bf..f843a9005cdb 100644
--- a/drivers/block/ublk/Makefile
+++ b/drivers/block/ublk/Makefile
@@ -5,6 +5,6 @@ ccflags-y			+= -I$(src)
 
 ublk_drv-$(CONFIG_BLK_DEV_UBLK)	:= main.o
 ifeq ($(CONFIG_UBLK_BPF), y)
-ublk_drv-$(CONFIG_BLK_DEV_UBLK)	+= bpf_ops.o
+ublk_drv-$(CONFIG_BLK_DEV_UBLK)	+= bpf_ops.o bpf.o
 endif
 obj-$(CONFIG_BLK_DEV_UBLK)	+= ublk_drv.o
diff --git a/drivers/block/ublk/bpf.c b/drivers/block/ublk/bpf.c
new file mode 100644
index 000000000000..479045a5f0d9
--- /dev/null
+++ b/drivers/block/ublk/bpf.c
@@ -0,0 +1,99 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2024 Red Hat */
+
+#include "ublk.h"
+#include "bpf.h"
+
+static int ublk_set_bpf_ops(struct ublk_device *ub,
+		struct ublk_bpf_ops *ops)
+{
+	int i;
+
+	for (i = 0; i < ub->dev_info.nr_hw_queues; i++) {
+		if (ops && ublk_get_queue(ub, i)->bpf_ops) {
+			ublk_set_bpf_ops(ub, NULL);
+			return -EBUSY;
+		}
+		ublk_get_queue(ub, i)->bpf_ops = ops;
+	}
+	return 0;
+}
+
+static int ublk_bpf_prog_attach_cb(struct bpf_prog_consumer *consumer,
+				   struct bpf_prog_provider *provider)
+{
+	struct ublk_device *ub = container_of(consumer, struct ublk_device,
+					      prog);
+	struct ublk_bpf_ops *ops = container_of(provider,
+			struct ublk_bpf_ops, provider);
+	int ret;
+
+	if (!ublk_get_device(ub))
+		return -ENODEV;
+
+	ret = ublk_set_bpf_ops(ub, ops);
+	if (ret)
+		goto fail_put_dev;
+
+	if (ops->attach_dev) {
+		ret = ops->attach_dev(ub->dev_info.dev_id);
+		if (ret)
+			goto fail_reset_ops;
+	}
+	return 0;
+
+fail_reset_ops:
+	ublk_set_bpf_ops(ub, NULL);
+fail_put_dev:
+	ublk_put_device(ub);
+	return ret;
+}
+
+static void ublk_bpf_prog_detach_cb(struct bpf_prog_consumer *consumer,
+				    bool unreg)
+{
+	struct ublk_device *ub = container_of(consumer, struct ublk_device,
+					      prog);
+	struct ublk_bpf_ops *ops = container_of(consumer->provider,
+			struct ublk_bpf_ops, provider);
+
+	if (unreg) {
+		blk_mq_freeze_queue(ub->ub_disk->queue);
+		ublk_set_bpf_ops(ub, NULL);
+		blk_mq_unfreeze_queue(ub->ub_disk->queue);
+	} else {
+		ublk_set_bpf_ops(ub, NULL);
+	}
+	if (ops->detach_dev)
+		ops->detach_dev(ub->dev_info.dev_id);
+	ublk_put_device(ub);
+}
+
+static const struct bpf_prog_consumer_ops ublk_prog_consumer_ops = {
+	.attach_fn	= ublk_bpf_prog_attach_cb,
+	.detach_fn	= ublk_bpf_prog_detach_cb,
+};
+
+int ublk_bpf_attach(struct ublk_device *ub)
+{
+	if (!ublk_dev_support_bpf(ub))
+		return 0;
+
+	/* todo: ublk device need to provide struct_ops prog id */
+	ub->prog.prog_id = 0;
+	ub->prog.ops = &ublk_prog_consumer_ops;
+
+	return ublk_bpf_prog_attach(&ub->prog);
+}
+
+void ublk_bpf_detach(struct ublk_device *ub)
+{
+	if (!ublk_dev_support_bpf(ub))
+		return;
+	ublk_bpf_prog_detach(&ub->prog);
+}
+
+int __init ublk_bpf_init(void)
+{
+	return ublk_bpf_struct_ops_init();
+}
diff --git a/drivers/block/ublk/bpf.h b/drivers/block/ublk/bpf.h
index e3505c9ab86a..4e178cbecb74 100644
--- a/drivers/block/ublk/bpf.h
+++ b/drivers/block/ublk/bpf.h
@@ -7,6 +7,8 @@
 typedef unsigned long ublk_bpf_return_t;
 typedef ublk_bpf_return_t (*queue_io_cmd_t)(struct ublk_bpf_io *io, unsigned int);
 typedef void (*release_io_cmd_t)(struct ublk_bpf_io *io);
+typedef int (*attach_dev_t)(int dev_id);
+typedef void (*detach_dev_t)(int dev_id);
 
 #ifdef CONFIG_UBLK_BPF
 #include <linux/filter.h>
@@ -47,6 +49,12 @@ struct ublk_bpf_ops {
 	/* called when the io command reference drops to zero, can't be sleepable */
 	release_io_cmd_t	release_io_cmd;
 
+	/* called when attaching bpf prog to this ublk dev */
+	attach_dev_t		attach_dev;
+
+	/* called when detaching bpf prog from this ublk dev */
+	detach_dev_t		detach_dev;
+
 	/* private: don't show in doc, must be the last field */
 	struct bpf_prog_provider	provider;
 };
@@ -149,7 +157,12 @@ static inline queue_io_cmd_t ublk_get_bpf_any_io_cb(struct ublk_queue *ubq)
 	return ublk_get_bpf_io_cb_daemon(ubq);
 }
 
+int ublk_bpf_init(void);
 int ublk_bpf_struct_ops_init(void);
+int ublk_bpf_prog_attach(struct bpf_prog_consumer *consumer);
+void ublk_bpf_prog_detach(struct bpf_prog_consumer *consumer);
+int ublk_bpf_attach(struct ublk_device *ub);
+void ublk_bpf_detach(struct ublk_device *ub);
 
 #else
 
@@ -176,9 +189,29 @@ static inline queue_io_cmd_t ublk_get_bpf_any_io_cb(struct ublk_queue *ubq)
 	return NULL;
 }
 
+static inline int ublk_bpf_init(void)
+{
+	return 0;
+}
+
 static inline int ublk_bpf_struct_ops_init(void)
 {
 	return 0;
 }
+
+static inline int ublk_bpf_prog_attach(struct bpf_prog_consumer *consumer)
+{
+	return 0;
+}
+static inline void ublk_bpf_prog_detach(struct bpf_prog_consumer *consumer)
+{
+}
+static inline int ublk_bpf_attach(struct ublk_device *ub)
+{
+	return 0;
+}
+static inline void ublk_bpf_detach(struct ublk_device *ub)
+{
+}
 #endif
 #endif
diff --git a/drivers/block/ublk/bpf_ops.c b/drivers/block/ublk/bpf_ops.c
index 6ac2aebd477e..05d8d415b30d 100644
--- a/drivers/block/ublk/bpf_ops.c
+++ b/drivers/block/ublk/bpf_ops.c
@@ -133,6 +133,29 @@ static void ublk_bpf_unreg(void *kdata, struct bpf_link *link)
 	kfree(curr);
 }
 
+int ublk_bpf_prog_attach(struct bpf_prog_consumer *consumer)
+{
+	unsigned id = consumer->prog_id;
+	struct ublk_bpf_ops *ops;
+	int ret = -EINVAL;
+
+	mutex_lock(&ublk_bpf_ops_lock);
+	ops = xa_load(&ublk_ops, id);
+	if (ops && ops->id == id)
+		ret = bpf_prog_consumer_attach(consumer, &ops->provider);
+	mutex_unlock(&ublk_bpf_ops_lock);
+
+	return ret;
+}
+
+void ublk_bpf_prog_detach(struct bpf_prog_consumer *consumer)
+{
+	mutex_lock(&ublk_bpf_ops_lock);
+	bpf_prog_consumer_detach(consumer, false);
+	mutex_unlock(&ublk_bpf_ops_lock);
+}
+
+
 static void ublk_bpf_prep_io(struct ublk_bpf_io *io,
 		const struct ublksrv_io_desc *iod)
 {
@@ -231,10 +254,21 @@ static void ublk_bpf_release_io_cmd(struct ublk_bpf_io *io)
 {
 }
 
+static int ublk_bpf_attach_dev(int dev_id)
+{
+	return 0;
+}
+
+static void ublk_bpf_detach_dev(int dev_id)
+{
+}
+
 static struct ublk_bpf_ops __bpf_ublk_bpf_ops = {
 	.queue_io_cmd = ublk_bpf_queue_io_cmd,
 	.queue_io_cmd_daemon = ublk_bpf_run_io_task,
 	.release_io_cmd = ublk_bpf_release_io_cmd,
+	.attach_dev	= ublk_bpf_attach_dev,
+	.detach_dev	= ublk_bpf_detach_dev,
 };
 
 static struct bpf_struct_ops bpf_ublk_bpf_ops = {
diff --git a/drivers/block/ublk/main.c b/drivers/block/ublk/main.c
index 29d3e7f656a7..0b136bc5247f 100644
--- a/drivers/block/ublk/main.c
+++ b/drivers/block/ublk/main.c
@@ -486,7 +486,7 @@ static inline bool ublk_need_get_data(const struct ublk_queue *ubq)
 }
 
 /* Called in slow path only, keep it noinline for trace purpose */
-static noinline struct ublk_device *ublk_get_device(struct ublk_device *ub)
+struct ublk_device *ublk_get_device(struct ublk_device *ub)
 {
 	if (kobject_get_unless_zero(&ub->cdev_dev.kobj))
 		return ub;
@@ -499,12 +499,6 @@ void ublk_put_device(struct ublk_device *ub)
 	put_device(&ub->cdev_dev);
 }
 
-static inline struct ublk_queue *ublk_get_queue(struct ublk_device *dev,
-		int qid)
-{
-       return (struct ublk_queue *)&(dev->__queues[qid * dev->queue_size]);
-}
-
 static inline bool ublk_rq_has_data(const struct request *rq)
 {
 	return bio_has_data(rq->bio);
@@ -1492,6 +1486,8 @@ static struct gendisk *ublk_detach_disk(struct ublk_device *ub)
 {
 	struct gendisk *disk;
 
+	ublk_bpf_detach(ub);
+
 	/* Sync with ublk_abort_queue() by holding the lock */
 	spin_lock(&ub->lock);
 	disk = ub->ub_disk;
@@ -2206,12 +2202,19 @@ static int ublk_ctrl_start_dev(struct ublk_device *ub, struct io_uring_cmd *cmd)
 			goto out_put_cdev;
 	}
 
-	ret = add_disk(disk);
+	ret = ublk_bpf_attach(ub);
 	if (ret)
 		goto out_put_cdev;
 
+	ret = add_disk(disk);
+	if (ret)
+		goto out_put_bpf;
+
 	set_bit(UB_STATE_USED, &ub->state);
 
+out_put_bpf:
+	if (ret)
+		ublk_bpf_detach(ub);
 out_put_cdev:
 	if (ret) {
 		ublk_detach_disk(ub);
@@ -2967,8 +2970,14 @@ static int __init ublk_init(void)
 	if (ret)
 		goto free_chrdev_region;
 
+	ret = ublk_bpf_init();
+	if (ret)
+		goto unregister_class;
+
 	return 0;
 
+unregister_class:
+	class_unregister(&ublk_chr_class);
 free_chrdev_region:
 	unregister_chrdev_region(ublk_chr_devt, UBLK_MINORS);
 unregister_mis:
diff --git a/drivers/block/ublk/ublk.h b/drivers/block/ublk/ublk.h
index e9ceadbc616d..7579b0032a3c 100644
--- a/drivers/block/ublk/ublk.h
+++ b/drivers/block/ublk/ublk.h
@@ -7,6 +7,8 @@
 #include <linux/cdev.h>
 #include <uapi/linux/ublk_cmd.h>
 
+#include "bpf_reg.h"
+
 #define UBLK_MINORS		(1U << MINORBITS)
 
 /* private ioctl command mirror */
@@ -153,6 +155,9 @@ struct ublk_device {
 	unsigned long		state;
 	int			ub_number;
 
+#ifdef CONFIG_UBLK_BPF
+	struct bpf_prog_consumer prog;
+#endif
 	struct mutex		mutex;
 
 	spinlock_t		lock;
@@ -173,6 +178,11 @@ struct ublk_params_header {
 	__u32	types;
 };
 
+static inline struct ublk_queue *ublk_get_queue(struct ublk_device *dev,
+		int qid)
+{
+       return (struct ublk_queue *)&(dev->__queues[qid * dev->queue_size]);
+}
 
 static inline struct ublksrv_io_desc *ublk_get_iod(struct ublk_queue *ubq,
 		int tag)
@@ -186,6 +196,12 @@ static inline bool ublk_support_bpf(const struct ublk_queue *ubq)
 	return false;
 }
 
+static inline bool ublk_dev_support_bpf(const struct ublk_device *ub)
+{
+	return false;
+}
+
+struct ublk_device *ublk_get_device(struct ublk_device *ub);
 struct ublk_device *ublk_get_device_from_id(int idx);
 void ublk_put_device(struct ublk_device *ub);
 void __ublk_complete_rq(struct request *req);
-- 
2.47.0