[PATCH net 1/1] udp: diag: bound bucket lock hold time

[Zihan Xi]

udp_diag_dump() currently keeps the UDP hash bucket spinlock held while
running the request's bytecode filter and filling a netlink response for
every socket in the bucket.  A large filter and a heavily populated bucket
can therefore keep bottom halves disabled for an attacker-scaled amount of
time.

Collect a bounded batch of matching sockets under the bucket lock, taking a
reference for each socket, then run the filter and fill the response after
releasing the lock.  Keep a referenced hash-list cursor for the next
unprocessed entry so each locked walk stays within the batch size.  This
remains bounded when filtered entries are skipped.  Consume the previous
resume reference after the next batch is collected so an -EMSGSIZE retry
cannot put a socket twice.  Preserve the cursor across -EMSGSIZE retries
and release it when the dump is finished.  Count every hash entry,
including foreign-network sockets in a shared table, against the walk
bound.  Take an extra module reference while a dump is active so
dump_done() can still put the cursor after inet_diag unlocks the handler.

Fixes: b6d640c2286d ("udp_diag: Implement the dump-all functionality")
Cc: stable@xxxxxxxxxxxxxxx
Reported-by: Vega <vega@xxxxxxxxxx>
Assisted-by: Codex:gpt-5.4
Signed-off-by: Zihan Xi <zihanx@xxxxxxxxxx>
---
 include/linux/inet_diag.h |   1 +
 net/ipv4/inet_diag.c      |   2 +
 net/ipv4/udp_diag.c       | 154 ++++++++++++++++++++++++++++----------
 3 files changed, 119 insertions(+), 38 deletions(-)

diff --git a/include/linux/inet_diag.h b/include/linux/inet_diag.h
index 704fd415c2b4..8dfcc37cc29d 100644
--- a/include/linux/inet_diag.h
+++ b/include/linux/inet_diag.h
@@ -38,6 +38,7 @@ struct inet_diag_dump_data {
 #define inet_diag_nla_bpf_stgs req_nlas[INET_DIAG_REQ_SK_BPF_STORAGES]
 
 	struct bpf_sk_storage_diag *bpf_stg_diag;
+	void (*dump_done)(struct netlink_callback *cb);
 	bool mark_needed;	/* INET_DIAG_BC_MARK_COND present. */
 #ifdef CONFIG_SOCK_CGROUP_DATA
 	bool cgroup_needed;	/* INET_DIAG_BC_CGROUP_COND present. */
diff --git a/net/ipv4/inet_diag.c b/net/ipv4/inet_diag.c
index 34b77aa87d0a..1a5ea36348e7 100644
--- a/net/ipv4/inet_diag.c
+++ b/net/ipv4/inet_diag.c
@@ -895,6 +895,8 @@ static int inet_diag_dump_done(struct netlink_callback *cb)
 {
 	struct inet_diag_dump_data *cb_data = cb->data;
 
+	if (cb_data->dump_done)
+		cb_data->dump_done(cb);
 	bpf_sk_storage_diag_free(cb_data->bpf_stg_diag);
 	kfree(cb->data);
 
diff --git a/net/ipv4/udp_diag.c b/net/ipv4/udp_diag.c
index f4b24e628cf8..18cfe0df654d 100644
--- a/net/ipv4/udp_diag.c
+++ b/net/ipv4/udp_diag.c
@@ -24,6 +24,36 @@ static int sk_diag_dump(struct sock *sk, struct sk_buff *skb,
 				 net_admin);
 }
 
+/* Process a maximum of SKARR_SZ hash entries at a time when walking hash
+ * buckets with bh disabled.
+ */
+#define SKARR_SZ 16
+
+static bool udp_diag_cursor_valid(struct udp_table *table,
+				  struct udp_hslot *hslot,
+				  struct sock *sk)
+{
+	if (!sk || hlist_unhashed_lockless(&sk->sk_node))
+		return false;
+
+	return sock_net(sk)->ipv4.udp_table == table &&
+	       udp_hashslot(table, sock_net(sk),
+			    udp_sk(sk)->udp_port_hash) == hslot;
+}
+
+static void udp_diag_dump_done(struct netlink_callback *cb)
+{
+	struct inet_diag_dump_data *cb_data = cb->data;
+	struct sock *sk = (struct sock *)cb->args[2];
+
+	if (sk) {
+		cb->args[2] = 0;
+		sock_put(sk);
+	}
+	cb_data->dump_done = NULL;
+	module_put(THIS_MODULE);
+}
+
 static int udp_diag_dump_one(struct netlink_callback *cb,
 			     const struct inet_diag_req_v2 *req)
 {
@@ -90,55 +120,103 @@ static void udp_diag_dump(struct sk_buff *skb, struct netlink_callback *cb,
 			  const struct inet_diag_req_v2 *r)
 {
 	bool net_admin = netlink_net_capable(cb->skb, CAP_NET_ADMIN);
+	struct sock *cursor = (struct sock *)cb->args[2];
+	struct inet_diag_dump_data *cb_data = cb->data;
 	struct net *net = sock_net(skb->sk);
-	int num, s_num, slot, s_slot;
+	unsigned int slot = cb->args[0];
 	struct udp_table *table;
 
 	table = net->ipv4.udp_table;
-	s_slot = cb->args[0];
-	num = s_num = cb->args[1];
+	/* Keep this module loaded until dump_done() drops the cursor. */
+	if (!cb_data->dump_done) {
+		__module_get(THIS_MODULE);
+		cb_data->dump_done = udp_diag_dump_done;
+	}
 
-	for (slot = s_slot; slot <= table->mask; s_num = 0, slot++) {
+	for (; slot <= table->mask; slot++) {
 		struct udp_hslot *hslot = &table->hash[slot];
-		struct sock *sk;
-
-		num = 0;
-
-		if (hlist_empty(&hslot->head))
-			continue;
-
-		spin_lock_bh(&hslot->lock);
-		sk_for_each(sk, &hslot->head) {
-			struct inet_sock *inet = inet_sk(sk);
-
-			if (!net_eq(sock_net(sk), net))
-				continue;
-			if (num < s_num)
-				goto next;
-			if (!(r->idiag_states & (1 << sk->sk_state)))
-				goto next;
-			if (r->sdiag_family != AF_UNSPEC &&
-					sk->sk_family != r->sdiag_family)
-				goto next;
-			if (r->id.idiag_sport != inet->inet_sport &&
-			    r->id.idiag_sport)
-				goto next;
-			if (r->id.idiag_dport != inet->inet_dport &&
-			    r->id.idiag_dport)
-				goto next;
-
-			if (sk_diag_dump(sk, skb, cb, r, net_admin) < 0) {
-				spin_unlock_bh(&hslot->lock);
-				goto done;
+
+		for (;;) {
+			struct sock *sk, *next_cursor = NULL;
+			int idx, accum = 0, walked = 0, res;
+			struct sock *old_cursor = NULL;
+			struct sock *sk_arr[SKARR_SZ];
+
+			spin_lock_bh(&hslot->lock);
+			sk = cursor;
+			if (sk && !udp_diag_cursor_valid(table, hslot, sk)) {
+				old_cursor = sk;
+				cursor = NULL;
+				sk = NULL;
+			}
+			if (!sk)
+				sk = hlist_entry_safe(hslot->head.first,
+						      struct sock, sk_node);
+
+			while (sk && walked < SKARR_SZ) {
+				struct inet_sock *inet = inet_sk(sk);
+				struct sock *next;
+
+				next = hlist_entry_safe(sk->sk_node.next,
+							struct sock, sk_node);
+				if (net_eq(sock_net(sk), net) &&
+				    (r->idiag_states & (1 << sk->sk_state)) &&
+				    (r->sdiag_family == AF_UNSPEC ||
+				     sk->sk_family == r->sdiag_family) &&
+				    (r->id.idiag_sport == inet->inet_sport ||
+				     !r->id.idiag_sport) &&
+				    (r->id.idiag_dport == inet->inet_dport ||
+				     !r->id.idiag_dport)) {
+					sock_hold(sk);
+					sk_arr[accum++] = sk;
+				}
+
+				walked++;
+				if (walked == SKARR_SZ) {
+					if (next) {
+						sock_hold(next);
+						next_cursor = next;
+					}
+					break;
+				}
+				sk = next;
 			}
-next:
-			num++;
+			spin_unlock_bh(&hslot->lock);
+			/* Consume the resume ref; remaining refs are in
+			 * sk_arr / next_cursor.
+			 */
+			if (old_cursor)
+				sock_put(old_cursor);
+			else if (cursor)
+				sock_put(cursor);
+			cursor = NULL;
+
+			for (idx = 0; idx < accum; idx++) {
+				res = sk_diag_dump(sk_arr[idx], skb, cb, r,
+						   net_admin);
+				if (res < 0) {
+					cursor = sk_arr[idx];
+					while (++idx < accum)
+						sock_put(sk_arr[idx]);
+					if (next_cursor)
+						sock_put(next_cursor);
+					goto done;
+				}
+				sock_put(sk_arr[idx]);
+			}
+
+			cursor = next_cursor;
+			if (!cursor)
+				break;
+
+			cond_resched();
 		}
-		spin_unlock_bh(&hslot->lock);
 	}
+
 done:
 	cb->args[0] = slot;
-	cb->args[1] = num;
+	cb->args[1] = 0;
+	cb->args[2] = (unsigned long)cursor;
 }
 
 static void udp_diag_get_info(struct sock *sk, struct inet_diag_msg *r,
-- 
2.43.0