BPF streams are only valid for the main programs, to make it easier to
access streams from subprogs, introduce main_prog_aux in struct
bpf_prog_aux.

prog->aux->main_prog_aux = prog->aux, for main programs and
prog->aux->main_prog_aux = main_prog->aux, for subprograms.

This makes it easy to access streams like:
stream = bpf_stream_get(stream_id, prog->main_prog_aux);

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
---
 include/linux/bpf.h   | 1 +
 kernel/bpf/core.c     | 3 +--
 kernel/bpf/stream.c   | 6 +++---
 kernel/bpf/verifier.c | 1 +
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 8f6e87f0f3a89..d133171c4d2a9 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1633,6 +1633,7 @@ struct bpf_prog_aux {
 	/* function name for valid attach_btf_id */
 	const char *attach_func_name;
 	struct bpf_prog **func;
+	struct bpf_prog_aux *main_prog_aux;
 	void *jit_data; /* JIT specific data. arch dependent */
 	struct bpf_jit_poke_descriptor *poke_tab;
 	struct bpf_kfunc_desc_tab *kfunc_tab;
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index ef01cc644a965..dbbf8e4b6e4c2 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -120,6 +120,7 @@ struct bpf_prog *bpf_prog_alloc_no_stats(unsigned int size, gfp_t gfp_extra_flag
 
 	fp->pages = size / PAGE_SIZE;
 	fp->aux = aux;
+	fp->aux->main_prog_aux = aux;
 	fp->aux->prog = fp;
 	fp->jit_requested = ebpf_jit_enabled();
 	fp->blinding_requested = bpf_jit_blinding_enabled(fp);
@@ -3292,8 +3293,6 @@ static bool find_from_stack_cb(void *cookie, u64 ip, u64 sp, u64 bp)
 	rcu_read_unlock();
 	if (!prog)
 		return true;
-	if (bpf_is_subprog(prog))
-		return true;
 	ctxp->prog = prog;
 	return false;
 }
diff --git a/kernel/bpf/stream.c b/kernel/bpf/stream.c
index ab592db4a4bf6..a36dee4a95d59 100644
--- a/kernel/bpf/stream.c
+++ b/kernel/bpf/stream.c
@@ -343,7 +343,7 @@ int bpf_prog_stream_read(struct bpf_prog *prog, enum bpf_stream_id stream_id, vo
 {
 	struct bpf_stream *stream;
 
-	stream = bpf_stream_get(stream_id, prog->aux);
+	stream = bpf_stream_get(stream_id, prog->aux->main_prog_aux);
 	if (!stream)
 		return -ENOENT;
 	return bpf_stream_read(stream, buf, len);
@@ -367,7 +367,7 @@ __bpf_kfunc int bpf_stream_vprintk(int stream_id, const char *fmt__str, const vo
 	u32 data_len = len__sz;
 	int ret, num_args;
 
-	stream = bpf_stream_get(stream_id, aux);
+	stream = bpf_stream_get(stream_id, aux->main_prog_aux);
 	if (!stream)
 		return -ENOENT;
 
@@ -457,7 +457,7 @@ int bpf_stream_stage_commit(struct bpf_stream_stage *ss, struct bpf_prog *prog,
 	struct bpf_stream *stream;
 	int ret;
 
-	stream = bpf_stream_get(stream_id, prog->aux);
+	stream = bpf_stream_get(stream_id, prog->aux->main_prog_aux);
 	if (!stream)
 		return -EINVAL;
 
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 5c9dd16b2c56b..fa110656099c4 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -21597,6 +21597,7 @@ static int jit_subprogs(struct bpf_verifier_env *env)
 		func[i]->aux->func_info_cnt = prog->aux->func_info_cnt;
 		func[i]->aux->poke_tab = prog->aux->poke_tab;
 		func[i]->aux->size_poke_tab = prog->aux->size_poke_tab;
+		func[i]->aux->main_prog_aux = prog->aux;
 
 		for (j = 0; j < prog->aux->size_poke_tab; j++) {
 			struct bpf_jit_poke_descriptor *poke;
-- 
2.47.3