[PATCH bpf v3 2/2] bpf: Fix reference leak in bpf_sk_assign()

sk_is_refcounted() is mutable; it depends on sk_state and SOCK_RCU_FREE.
TC bpf_sk_assign() uses it at assign time to take a reference and lets
sock_pfree() re-evaluate it at free time. If the socket becomes
non-refcounted in between, e.g. connect(AF_UNSPEC) + listen() sets
SOCK_RCU_FREE, the reference is leaked.

Freeze the decision at assign time. If the socket is refcounted (i.e. can
become non-refcounted), take a reference and set sock_pfree_refcounted()
as the destructor; otherwise stick with sock_pfree(), which effectively
becomes a no-op destructor. Next, extend skb_sk_is_prefetched() to handle
both dtors and, for the sake of CONFIG_INET=n, add a specialized
skb_sk_is_prefetched_noref(). Adapt prefetched-skb treatment in
__nf_queue() and skb_steal_sock(). The latter changes its contract with
callers: inet{,6}_steal_sock() must now expect a refcounted TCP_LISTEN
socket.

Fixes: 7ae215d23c12 ("bpf: Don't refcount LISTEN sockets in sk_assign()")
Signed-off-by: Michal Luczaj <[email protected]>
---
 include/net/inet6_hashtables.h |  9 +++++----
 include/net/inet_hashtables.h  |  9 +++++----
 include/net/request_sock.h     |  2 +-
 include/net/sock.h             | 11 +++++++++++
 net/core/filter.c              | 19 ++++++++++++++++---
 net/netfilter/nf_queue.c       | 18 +++++++++---------
 6 files changed, 47 insertions(+), 21 deletions(-)

diff --git a/include/net/inet6_hashtables.h b/include/net/inet6_hashtables.h
index 2cc5d416bbb5..a310cb91611c 100644
--- a/include/net/inet6_hashtables.h
+++ b/include/net/inet6_hashtables.h
@@ -134,10 +134,11 @@ struct sock *inet6_steal_sock(struct net *net, struct sk_buff *skb, int doff,
 	if (!reuse_sk)
 		return sk;
 
-	/* We've chosen a new reuseport sock which is never refcounted. This
-	 * implies that sk also isn't refcounted.
-	 */
-	WARN_ON_ONCE(*refcounted);
+	/* New reuseport sock is never refcounted; drop the old sk's ref. */
+	if (*refcounted) {
+		sock_put(sk);
+		*refcounted = false;
+	}
 
 	return reuse_sk;
 }
diff --git a/include/net/inet_hashtables.h b/include/net/inet_hashtables.h
index 6e2fe186d0dc..1589e9e7ecb2 100644
--- a/include/net/inet_hashtables.h
+++ b/include/net/inet_hashtables.h
@@ -474,10 +474,11 @@ struct sock *inet_steal_sock(struct net *net, struct sk_buff *skb, int doff,
 	if (!reuse_sk)
 		return sk;
 
-	/* We've chosen a new reuseport sock which is never refcounted. This
-	 * implies that sk also isn't refcounted.
-	 */
-	WARN_ON_ONCE(*refcounted);
+	/* New reuseport sock is never refcounted; drop the old sk's ref. */
+	if (*refcounted) {
+		sock_put(sk);
+		*refcounted = false;
+	}
 
 	return reuse_sk;
 }
diff --git a/include/net/request_sock.h b/include/net/request_sock.h
index 5a9c826a7092..3fb114fa77de 100644
--- a/include/net/request_sock.h
+++ b/include/net/request_sock.h
@@ -113,7 +113,7 @@ static inline struct sock *skb_steal_sock(struct sk_buff *skb,
 			return sk;
 		}
 #endif
-		*refcounted = sk_is_refcounted(sk);
+		*refcounted = !skb_sk_is_prefetched_noref(skb);
 	} else {
 		*refcounted = true;
 	}
diff --git a/include/net/sock.h b/include/net/sock.h
index 51185222aac2..fb553697db83 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -1870,6 +1870,7 @@ void sock_efree(struct sk_buff *skb);
 #ifdef CONFIG_INET
 void sock_edemux(struct sk_buff *skb);
 void sock_pfree(struct sk_buff *skb);
+void sock_pfree_refcounted(struct sk_buff *skb);
 
 static inline void skb_set_owner_edemux(struct sk_buff *skb, struct sock *sk)
 {
@@ -3015,6 +3016,16 @@ static inline void sk_eat_skb(struct sock *sk, struct sk_buff *skb)
 static inline bool
 skb_sk_is_prefetched(struct sk_buff *skb)
 {
+#ifdef CONFIG_INET
+	return skb->destructor == sock_pfree ||
+	       skb->destructor == sock_pfree_refcounted;
+#else
+	return false;
+#endif /* CONFIG_INET */
+}
+
+static inline bool skb_sk_is_prefetched_noref(struct sk_buff *skb)
+{
 #ifdef CONFIG_INET
 	return skb->destructor == sock_pfree;
 #else
diff --git a/net/core/filter.c b/net/core/filter.c
index 66d83a198ea2..e473ad8c9b3c 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -7989,8 +7989,20 @@ static const struct bpf_func_proto bpf_tcp_gen_syncookie_proto = {
 	.arg5_type	= ARG_MEM_SIZE,
 };
 
+/*
+ * skb destructor set by TC bpf_sk_assign(), refcounted path only.
+ * Acts as a "prefetched by bpf, ref taken" marker.
+ */
+void sock_pfree_refcounted(struct sk_buff *skb)
+{
+	sock_gen_put(skb->sk);
+}
+EXPORT_SYMBOL(sock_pfree_refcounted);
+
 BPF_CALL_3(bpf_sk_assign, struct sk_buff *, skb, struct sock *, sk, u64, flags)
 {
+	bool refcounted;
+
 	if (!sk || flags != 0)
 		return -EINVAL;
 	if (!skb_at_tc_ingress(skb))
@@ -7999,13 +8011,14 @@ BPF_CALL_3(bpf_sk_assign, struct sk_buff *, skb, struct sock *, sk, u64, flags)
 		return -ENETUNREACH;
 	if (sk_unhashed(sk))
 		return -EOPNOTSUPP;
-	if (sk_is_refcounted(sk) &&
-	    unlikely(!refcount_inc_not_zero(&sk->sk_refcnt)))
+
+	refcounted = sk_is_refcounted(sk);
+	if (refcounted && unlikely(!refcount_inc_not_zero(&sk->sk_refcnt)))
 		return -ENOENT;
 
 	skb_orphan(skb);
 	skb->sk = sk;
-	skb->destructor = sock_pfree;
+	skb->destructor = refcounted ? sock_pfree_refcounted : sock_pfree;
 
 	return 0;
 }
diff --git a/net/netfilter/nf_queue.c b/net/netfilter/nf_queue.c
index 73363ceedebe..e2e9739c00ac 100644
--- a/net/netfilter/nf_queue.c
+++ b/net/netfilter/nf_queue.c
@@ -195,16 +195,16 @@ static int __nf_queue(struct sk_buff *skb, const struct nf_hook_state *state,
 		break;
 	}
 
-	if (skb_sk_is_prefetched(skb)) {
-		struct sock *sk = skb->sk;
-
-		if (!sk_is_refcounted(sk)) {
-			if (!refcount_inc_not_zero(&sk->sk_refcnt))
-				return -ENOTCONN;
+	/*
+	 * If the skb was prefetched without taking sock's ref, bump it.
+	 * Skip sock_pfree-dtor'ed skbs coming from bpf_sk_assign_tcp_reqsk().
+	 */
+	if (skb_sk_is_prefetched_noref(skb) && sk_fullsock(skb->sk)) {
+		if (!refcount_inc_not_zero(&skb->sk->sk_refcnt))
+			return -ENOTCONN;
 
-			/* drop refcount on skb_orphan */
-			skb->destructor = sock_edemux;
-		}
+		/* drop refcount on skb_orphan */
+		skb->destructor = sock_edemux;
 	}
 
 	entry = kmalloc(sizeof(*entry) + route_key_size, GFP_ATOMIC);

-- 
2.55.0