[RFC PATCH v7 5/8] net_sched: Add init program

[Amery Hung <ameryhung@gmail.com>]

This patch adds another optional program to be called during
the creation of a qdisc for initializating data in the bpf world.
The program takes bpf_qdisc_ctx as context, but cannot not access
any field.

Signed-off-by: Amery Hung <amery.hung@bytedance.com>
---
 include/uapi/linux/bpf.h       |  1 +
 include/uapi/linux/pkt_sched.h |  4 ++++
 kernel/bpf/syscall.c           |  1 +
 net/core/filter.c              |  3 ++-
 net/sched/sch_bpf.c            | 23 ++++++++++++++++++++++-
 tools/include/uapi/linux/bpf.h |  1 +
 6 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 84669886a493..cad0788bef99 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1060,6 +1060,7 @@ enum bpf_attach_type {
 	BPF_QDISC_ENQUEUE,
 	BPF_QDISC_DEQUEUE,
 	BPF_QDISC_RESET,
+	BPF_QDISC_INIT,
 	__MAX_BPF_ATTACH_TYPE
 };
 
diff --git a/include/uapi/linux/pkt_sched.h b/include/uapi/linux/pkt_sched.h
index e9e1a83c22f7..61f0cf4a088c 100644
--- a/include/uapi/linux/pkt_sched.h
+++ b/include/uapi/linux/pkt_sched.h
@@ -1332,6 +1332,10 @@ enum {
 	TCA_SCH_BPF_RESET_PROG_FD,	/* u32 */
 	TCA_SCH_BPF_RESET_PROG_ID,	/* u32 */
 	TCA_SCH_BPF_RESET_PROG_TAG,	/* data */
+	TCA_SCH_BPF_INIT_PROG_NAME,	/* string */
+	TCA_SCH_BPF_INIT_PROG_FD,	/* u32 */
+	TCA_SCH_BPF_INIT_PROG_ID,	/* u32 */
+	TCA_SCH_BPF_INIT_PROG_TAG,	/* data */
 	__TCA_SCH_BPF_MAX,
 };
 
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 9af6fa542f2e..0959905044b9 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2507,6 +2507,7 @@ bpf_prog_load_check_attach(enum bpf_prog_type prog_type,
 		case BPF_QDISC_ENQUEUE:
 		case BPF_QDISC_DEQUEUE:
 		case BPF_QDISC_RESET:
+		case BPF_QDISC_INIT:
 			return 0;
 		default:
 			return -EINVAL;
diff --git a/net/core/filter.c b/net/core/filter.c
index f8e17465377f..5619a12c0d06 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -8905,7 +8905,8 @@ static bool tc_qdisc_is_valid_access(int off, int size,
 {
 	struct btf *btf;
 
-	if (prog->expected_attach_type == BPF_QDISC_RESET)
+	if (prog->expected_attach_type == BPF_QDISC_RESET ||
+	    prog->expected_attach_type == BPF_QDISC_INIT)
 		return false;
 
 	if (off < 0 || off >= sizeof(struct bpf_qdisc_ctx))
diff --git a/net/sched/sch_bpf.c b/net/sched/sch_bpf.c
index 3f0f809dced6..925a131016f0 100644
--- a/net/sched/sch_bpf.c
+++ b/net/sched/sch_bpf.c
@@ -43,6 +43,7 @@ struct bpf_sched_data {
 	struct sch_bpf_prog __rcu enqueue_prog;
 	struct sch_bpf_prog __rcu dequeue_prog;
 	struct sch_bpf_prog __rcu reset_prog;
+	struct sch_bpf_prog __rcu init_prog;
 
 	struct qdisc_watchdog watchdog;
 };
@@ -88,6 +89,9 @@ static int sch_bpf_dump(struct Qdisc *sch, struct sk_buff *skb)
 	if (sch_bpf_dump_prog(&q->reset_prog, skb, TCA_SCH_BPF_RESET_PROG_NAME,
 			      TCA_SCH_BPF_RESET_PROG_ID, TCA_SCH_BPF_RESET_PROG_TAG))
 		goto nla_put_failure;
+	if (sch_bpf_dump_prog(&q->init_prog, skb, TCA_SCH_BPF_INIT_PROG_NAME,
+			      TCA_SCH_BPF_INIT_PROG_ID, TCA_SCH_BPF_INIT_PROG_TAG))
+		goto nla_put_failure;
 
 	return nla_nest_end(skb, opts);
 
@@ -269,6 +273,9 @@ static const struct nla_policy sch_bpf_policy[TCA_SCH_BPF_MAX + 1] = {
 	[TCA_SCH_BPF_RESET_PROG_FD]	= { .type = NLA_U32 },
 	[TCA_SCH_BPF_RESET_PROG_NAME]	= { .type = NLA_NUL_STRING,
 					    .len = ACT_BPF_NAME_LEN },
+	[TCA_SCH_BPF_INIT_PROG_FD]	= { .type = NLA_U32 },
+	[TCA_SCH_BPF_INIT_PROG_NAME]	= { .type = NLA_NUL_STRING,
+					    .len = ACT_BPF_NAME_LEN },
 };
 
 static int bpf_init_prog(struct nlattr *fd, struct nlattr *name,
@@ -348,6 +355,10 @@ static int sch_bpf_change(struct Qdisc *sch, struct nlattr *opt,
 		goto failure;
 	err = bpf_init_prog(tb[TCA_SCH_BPF_RESET_PROG_FD],
 			    tb[TCA_SCH_BPF_RESET_PROG_NAME], &q->reset_prog, true);
+	if (err)
+		goto failure;
+	err = bpf_init_prog(tb[TCA_SCH_BPF_INIT_PROG_FD],
+			    tb[TCA_SCH_BPF_INIT_PROG_NAME], &q->init_prog, true);
 failure:
 	sch_tree_unlock(sch);
 	return err;
@@ -357,6 +368,8 @@ static int sch_bpf_init(struct Qdisc *sch, struct nlattr *opt,
 			struct netlink_ext_ack *extack)
 {
 	struct bpf_sched_data *q = qdisc_priv(sch);
+	struct bpf_qdisc_ctx ctx = {};
+	struct bpf_prog *init;
 	int err;
 
 	qdisc_watchdog_init(&q->watchdog, sch);
@@ -370,7 +383,14 @@ static int sch_bpf_init(struct Qdisc *sch, struct nlattr *opt,
 	if (err)
 		return err;
 
-	return qdisc_class_hash_init(&q->clhash);
+	err = qdisc_class_hash_init(&q->clhash);
+	if (err < 0)
+		return err;
+
+	init = rcu_dereference(q->init_prog.prog);
+	if (init)
+		bpf_prog_run(init, &ctx);
+	return 0;
 }
 
 static void sch_bpf_reset(struct Qdisc *sch)
@@ -420,6 +440,7 @@ static void sch_bpf_destroy(struct Qdisc *sch)
 	bpf_cleanup_prog(&q->enqueue_prog);
 	bpf_cleanup_prog(&q->dequeue_prog);
 	bpf_cleanup_prog(&q->reset_prog);
+	bpf_cleanup_prog(&q->init_prog);
 	sch_tree_unlock(sch);
 }
 
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index 84669886a493..cad0788bef99 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -1060,6 +1060,7 @@ enum bpf_attach_type {
 	BPF_QDISC_ENQUEUE,
 	BPF_QDISC_DEQUEUE,
 	BPF_QDISC_RESET,
+	BPF_QDISC_INIT,
 	__MAX_BPF_ATTACH_TYPE
 };
 
-- 
2.20.1