Cover the two loads which used to lose the BPF_PROBE_MEM rewrite, both
reached from an RCU read-side critical section.

  # LDLIBS=-static PKG_CONFIG='pkg-config --static' ./vmtest.sh -- ./test_progs -t rcu_read_lock
  [...]
  #332/1   rcu_read_lock/success:OK
  #332/2   rcu_read_lock/rcuptr_acquire:OK
  #332/3   rcu_read_lock/negative_tests_inproper_region:OK
  #332/4   rcu_read_lock/negative_tests_rcuptr_misuse:OK
  #332     rcu_read_lock:OK
  Summary: 1/4 PASSED, 0 SKIPPED, 0/0 FAILED

Signed-off-by: Daniel Borkmann <daniel@iogearbox.net>
---
 v1 -> v2:
   - new patch

 .../selftests/bpf/prog_tests/rcu_read_lock.c  |  2 +
 .../selftests/bpf/progs/rcu_read_lock.c       | 76 +++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/tools/testing/selftests/bpf/prog_tests/rcu_read_lock.c b/tools/testing/selftests/bpf/prog_tests/rcu_read_lock.c
index 246eb259c08a..6a07b2b418d1 100644
--- a/tools/testing/selftests/bpf/prog_tests/rcu_read_lock.c
+++ b/tools/testing/selftests/bpf/prog_tests/rcu_read_lock.c
@@ -34,6 +34,8 @@ static void test_success(void)
 	bpf_program__set_autoload(skel->progs.rcu_read_lock_global_subprog, true);
 	bpf_program__set_autoload(skel->progs.rcu_read_lock_subprog_lock, true);
 	bpf_program__set_autoload(skel->progs.rcu_read_lock_subprog_unlock, true);
+	bpf_program__set_autoload(skel->progs.non_own_ref_untrusted_ld, true);
+	bpf_program__set_autoload(skel->progs.rcu_untrusted_union_ld, true);
 	err = rcu_read_lock__load(skel);
 	if (!ASSERT_OK(err, "skel_load"))
 		goto out;
diff --git a/tools/testing/selftests/bpf/progs/rcu_read_lock.c b/tools/testing/selftests/bpf/progs/rcu_read_lock.c
index b4e073168fb1..31d4081c3a9f 100644
--- a/tools/testing/selftests/bpf/progs/rcu_read_lock.c
+++ b/tools/testing/selftests/bpf/progs/rcu_read_lock.c
@@ -549,3 +549,79 @@ int rcu_read_lock_sleepable_global_subprog_indirect(void *ctx)
 	bpf_rcu_read_unlock();
 	return 0;
 }
+
+struct rcu_node_data {
+	long key;
+	struct bpf_rb_node node;
+};
+
+struct rcu_node_stash {
+	struct rcu_node_data __kptr *node;
+};
+
+/*
+ * Necessary so that LLVM emits BTF for rcu_node_data rather than just a
+ * fwd reference to it, same as in progs/local_kptr_stash.c.
+ */
+struct rcu_node_data *just_here_because_btf_bug;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, 1);
+	__type(key, int);
+	__type(value, struct rcu_node_stash);
+} node_stash SEC(".maps");
+
+long non_own_ref_key;
+
+SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
+int non_own_ref_untrusted_ld(void *ctx)
+{
+	struct rcu_node_stash *stash;
+	struct rcu_node_data *node;
+	int key = 0;
+
+	stash = bpf_map_lookup_elem(&node_stash, &key);
+	if (!stash)
+		return 0;
+	bpf_rcu_read_lock();
+	node = stash->node;
+	if (!node) {
+		bpf_rcu_read_unlock();
+		return 0;
+	}
+	bpf_rcu_read_unlock();
+	/*
+	 * The unlock leaves node as PTR_TO_BTF_ID | MEM_ALLOC | PTR_UNTRUSTED
+	 * | NON_OWN_REF, and the load below has to get the BPF_PROBE_MEM
+	 * rewrite for it, otherwise a bad address panics the kernel.
+	 */
+	non_own_ref_key = node->key;
+	return 0;
+}
+
+long rcu_untrusted_wq_flags;
+
+SEC("?tp_btf/tcp_probe")
+int BPF_PROG(rcu_untrusted_union_ld, struct sock *sk)
+{
+	struct socket_wq *wq;
+
+	/*
+	 * sk_wq sits in a two member union, so btf_struct_walk() marks the
+	 * pointer PTR_UNTRUSTED, and the __rcu tag on the member adds MEM_RCU
+	 * on top of it. struct sock is not on the __safe_rcu_or_null allow
+	 * list, hence the two stay combined and the load below has to get the
+	 * BPF_PROBE_MEM rewrite for PTR_TO_BTF_ID | PTR_UNTRUSTED | MEM_RCU,
+	 * otherwise a bad address panics the kernel.
+	 *
+	 * The __rcu tag only reaches BTF on a clang built kernel, that is, one
+	 * with CONFIG_PAHOLE_HAS_BTF_TAG. On a gcc built kernel the walk yields
+	 * a plain untrusted pointer, which is rewritten either way.
+	 */
+	wq = sk->sk_wq;
+	if (!wq)
+		return 0;
+	rcu_untrusted_wq_flags = wq->flags;
+	return 0;
+}
-- 
2.43.0