[PATCH bpf v2 2/2] bpf: Unconditionally take socket references in lookup helpers

Lookup helpers gate whether to acquire a socket reference on
sk_is_refcounted(), a check re-evaluated at release. An established socket
refcounted at acquire time can gain SOCK_RCU_FREE via
connect(AF_UNSPEC)+listen() before release runs; the release-side re-check
then reads sk_is_refcounted() == false and skips the put. The reference
leaks.

Make acquire and release unconditional and symmetric: always take a
reference, always put it. Adapt sk_select_reuseport().

Fixes: 6acc9b432e67 ("bpf: Add helper to retrieve socket in BPF")
Fixes: 64d85290d79c ("bpf: Allow bpf_map_lookup_elem for SOCKMAP and SOCKHASH")
Reported-by: Sashiko <[email protected]>
Closes: https://lore.kernel.org/bpf/[email protected]/
Signed-off-by: Michal Luczaj <[email protected]>
Reviewed-by: Emil Tsalapatis <[email protected]>
---
TC bpf_sk_assign() has the same issue; it takes a reference only when
sk_is_refcounted() is true at assign time, but sock_pfree() (the skb
destructor it installs) re-checks sk_is_refcounted() independently at
release time. The same connect(AF_UNSPEC)+listen() transition leaks the
socket here too. I'd welcome suggestions on the right way to handle this.
---
 net/core/filter.c   | 27 ++++++++++++++++++---------
 net/core/sock_map.c |  4 ++--
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/net/core/filter.c b/net/core/filter.c
index fede810ef37f..d71e069f669a 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -7032,6 +7032,14 @@ static struct sock *sk_lookup(struct net *net, struct bpf_sock_tuple *tuple,
 		WARN_ONCE(1, "Found non-RCU, unreferenced socket!");
 		sk = NULL;
 	}
+
+	/*
+	 * Always take a reference, even if the lookup skipped one;
+	 * bpf_sk_release() always puts one.
+	 */
+	if (sk && !refcounted && !refcount_inc_not_zero(&sk->sk_refcnt))
+		sk = NULL;
+
 	return sk;
 }
 
@@ -7090,11 +7098,16 @@ bpf_sk_lookup_full_sk(struct sock *sk)
 	 */
 	if (sk2 != sk) {
 		sock_gen_put(sk);
-		/* Ensure there is no need to bump sk2 refcnt. */
 		if (unlikely(sk2 && !sock_flag(sk2, SOCK_RCU_FREE))) {
 			WARN_ONCE(1, "Found non-RCU, unreferenced socket!");
 			return NULL;
 		}
+		/*
+		 * sk2 is RCU-free, but take a reference anyway;
+		 * bpf_sk_release() puts.
+		 */
+		if (sk2 && !refcount_inc_not_zero(&sk2->sk_refcnt))
+			sk2 = NULL;
 		sk = sk2;
 	}
 
@@ -7279,7 +7292,7 @@ static const struct bpf_func_proto bpf_tc_sk_lookup_udp_proto = {
 
 BPF_CALL_1(bpf_sk_release, struct sock *, sk)
 {
-	if (sk && sk_is_refcounted(sk))
+	if (sk)
 		sock_gen_put(sk);
 	return 0;
 }
@@ -11571,11 +11584,13 @@ BPF_CALL_4(sk_select_reuseport, struct sk_reuseport_kern *, reuse_kern,
 	bool is_sockarray = map->map_type == BPF_MAP_TYPE_REUSEPORT_SOCKARRAY;
 	struct sock_reuseport *reuse;
 	struct sock *selected_sk;
-	int err;
+	int err = 0;
 
 	selected_sk = map->ops->map_lookup_elem(map, key);
 	if (!selected_sk)
 		return -ENOENT;
+	if (!is_sockarray)
+		sock_put(selected_sk);
 
 	reuse = rcu_dereference(selected_sk->sk_reuseport_cb);
 	if (!reuse) {
@@ -11605,13 +11620,7 @@ BPF_CALL_4(sk_select_reuseport, struct sk_reuseport_kern *, reuse_kern,
 	}
 
 	reuse_kern->selected_sk = selected_sk;
-
-	return 0;
 error:
-	/* Lookup in sock_map can return TCP ESTABLISHED sockets. */
-	if (sk_is_refcounted(selected_sk))
-		sock_put(selected_sk);
-
 	return err;
 }
 
diff --git a/net/core/sock_map.c b/net/core/sock_map.c
index 9efbd8ca7db8..92a006fd3368 100644
--- a/net/core/sock_map.c
+++ b/net/core/sock_map.c
@@ -392,7 +392,7 @@ static void *sock_map_lookup(struct bpf_map *map, void *key)
 	sk = __sock_map_lookup_elem(map, *(u32 *)key);
 	if (!sk)
 		return NULL;
-	if (sk_is_refcounted(sk) && !refcount_inc_not_zero(&sk->sk_refcnt))
+	if (!refcount_inc_not_zero(&sk->sk_refcnt))
 		return NULL;
 	return sk;
 }
@@ -1218,7 +1218,7 @@ static void *sock_hash_lookup(struct bpf_map *map, void *key)
 	sk = __sock_hash_lookup_elem(map, key);
 	if (!sk)
 		return NULL;
-	if (sk_is_refcounted(sk) && !refcount_inc_not_zero(&sk->sk_refcnt))
+	if (!refcount_inc_not_zero(&sk->sk_refcnt))
 		return NULL;
 	return sk;
 }

-- 
2.55.0