This patch adds a subflow pointers array in struct mptcp_sched_data. Set
the array before invoking get_subflow(), then get it in get_subflow() in
the BPF contexts.

Signed-off-by: Geliang Tang <geliang.tang@suse.com>
---
 include/net/mptcp.h                           |  2 ++
 net/mptcp/sched.c                             | 29 +++++++++++++++++++
 tools/testing/selftests/bpf/bpf_tcp_helpers.h |  7 +++++
 3 files changed, 38 insertions(+)

diff --git a/include/net/mptcp.h b/include/net/mptcp.h
index 33a44ec21701..10eb96ed2ecc 100644
--- a/include/net/mptcp.h
+++ b/include/net/mptcp.h
@@ -97,9 +97,11 @@ struct mptcp_out_options {
 };
 
 #define MPTCP_SCHED_NAME_MAX	16
+#define MPTCP_SUBFLOWS_MAX	8
 
 struct mptcp_sched_data {
 	unsigned long	bitmap;
+	struct mptcp_subflow_context *contexts[MPTCP_SUBFLOWS_MAX];
 };
 
 struct mptcp_sched_ops {
diff --git a/net/mptcp/sched.c b/net/mptcp/sched.c
index 207ab422ac5d..5e224e3a17e8 100644
--- a/net/mptcp/sched.c
+++ b/net/mptcp/sched.c
@@ -91,8 +91,22 @@ void mptcp_release_sched(struct mptcp_sock *msk)
 static int mptcp_sched_data_init(struct mptcp_sock *msk,
 				 struct mptcp_sched_data *data)
 {
+	struct mptcp_subflow_context *subflow;
+	int i = 0;
+
 	data->bitmap = 0;
 
+	mptcp_for_each_subflow(msk, subflow) {
+		if (i == MPTCP_SUBFLOWS_MAX) {
+			pr_warn_once("too many subflows");
+			break;
+		}
+		data->contexts[i++] = subflow;
+	}
+
+	for (; i < MPTCP_SUBFLOWS_MAX; i++)
+		data->contexts[i++] = NULL;
+
 	return 0;
 }
 
@@ -101,6 +115,7 @@ int mptcp_sched_get_send(struct mptcp_sock *msk)
 	struct mptcp_subflow_context *subflow;
 	struct mptcp_sched_data data;
 	struct sock *ssk;
+	int i;
 
 	sock_owned_by_me((struct sock *)msk);
 
@@ -133,6 +148,12 @@ int mptcp_sched_get_send(struct mptcp_sock *msk)
 	mptcp_sched_data_init(msk, &data);
 	msk->sched->get_subflow(msk, false, &data);
 
+	for (i = 0; i < MPTCP_SUBFLOWS_MAX; i++) {
+		if (test_bit(i, &data.bitmap) && data.contexts[i]) {
+			data.contexts[i]->scheduled = 1;
+			msk->last_snd = data.contexts[i]->tcp_sock;
+		}
+	}
 	return 0;
 
 err:
@@ -149,6 +170,7 @@ int mptcp_sched_get_retrans(struct mptcp_sock *msk)
 	struct mptcp_subflow_context *subflow;
 	struct mptcp_sched_data data;
 	struct sock *ssk;
+	int i;
 
 	sock_owned_by_me((const struct sock *)msk);
 
@@ -175,6 +197,13 @@ int mptcp_sched_get_retrans(struct mptcp_sock *msk)
 	mptcp_sched_data_init(msk, &data);
 	msk->sched->get_subflow(msk, true, &data);
 
+	for (i = 0; i < MPTCP_SUBFLOWS_MAX; i++) {
+		if (test_bit(i, &data.bitmap) && data.contexts[i]) {
+			data.contexts[i]->scheduled = 1;
+			msk->last_snd = data.contexts[i]->tcp_sock;
+		}
+	}
+
 	return 0;
 
 err:
diff --git a/tools/testing/selftests/bpf/bpf_tcp_helpers.h b/tools/testing/selftests/bpf/bpf_tcp_helpers.h
index 60c8239f95ff..4c7192cb6134 100644
--- a/tools/testing/selftests/bpf/bpf_tcp_helpers.h
+++ b/tools/testing/selftests/bpf/bpf_tcp_helpers.h
@@ -231,9 +231,16 @@ extern __u32 tcp_slow_start(struct tcp_sock *tp, __u32 acked) __ksym;
 extern void tcp_cong_avoid_ai(struct tcp_sock *tp, __u32 w, __u32 acked) __ksym;
 
 #define MPTCP_SCHED_NAME_MAX	16
+#define MPTCP_SUBFLOWS_MAX	8
+
+struct mptcp_subflow_context {
+	__u32	token;
+	struct	sock *tcp_sock;	    /* tcp sk backpointer */
+} __attribute__((preserve_access_index));
 
 struct mptcp_sched_data {
 	unsigned long	bitmap;
+	struct mptcp_subflow_context *contexts[MPTCP_SUBFLOWS_MAX];
 };
 
 struct mptcp_sched_ops {
-- 
2.34.1