diff --git a/net/ipv4/tcp_ipv4.c b/net/ipv4/tcp_ipv4.c index b8887cdd66c5..7f413f509d7d 100644 --- a/net/ipv4/tcp_ipv4.c +++ b/net/ipv4/tcp_ipv4.c @@ -3078,24 +3078,24 @@ static unsigned int bpf_iter_tcp_established_batch(struct seq_file *seq, { struct bpf_tcp_iter_state *iter = seq->private; struct hlist_nulls_node *node; - unsigned int expected = 1; - struct sock *sk; + struct sock *sk = *start_sk; + unsigned int expected = 0; - sock_hold(*start_sk); - iter->batch[iter->end_sk++].sk = *start_sk; - - sk = sk_nulls_next(*start_sk); *start_sk = NULL; sk_nulls_for_each_from(sk, node) { - if (seq_sk_match(seq, sk)) { - if (iter->end_sk < iter->max_sk) { - sock_hold(sk); - iter->batch[iter->end_sk++].sk = sk; - } else if (!*start_sk) { - /* Remember where we left off. */ - *start_sk = sk; - } - expected++; + if (!seq_sk_match(seq, sk)) + continue; + expected++; + if (iter->end_sk < iter->max_sk) { + /* reqsk_queue_hash_req() inserts with sk_refcnt == 0 + * and refcount_set()s it after the bucket lock drops. + */ + if (unlikely(!refcount_inc_not_zero(&sk->sk_refcnt))) + continue; + iter->batch[iter->end_sk++].sk = sk; + } else if (!*start_sk) { + /* Remember where we left off. */ + *start_sk = sk; } } @@ -3133,12 +3133,13 @@ static struct sock *bpf_iter_tcp_batch(struct seq_file *seq) struct sock *sk; int err; +again: sk = bpf_iter_tcp_resume(seq); if (!sk) return NULL; /* Done */ expected = bpf_iter_fill_batch(seq, &sk); - if (likely(iter->end_sk == expected)) + if (likely(!sk)) goto done; /* Batch size was too small. */ @@ -3157,7 +3158,7 @@ static struct sock *bpf_iter_tcp_batch(struct seq_file *seq) return NULL; /* Done */ expected = bpf_iter_fill_batch(seq, &sk); - if (likely(iter->end_sk == expected)) + if (likely(!sk)) goto done; /* Batch size was still too small. Hold onto the lock while we try @@ -3170,10 +3171,14 @@ static struct sock *bpf_iter_tcp_batch(struct seq_file *seq) return ERR_PTR(err); } - expected = bpf_iter_fill_batch(seq, &sk); - WARN_ON_ONCE(iter->end_sk != expected); + bpf_iter_fill_batch(seq, &sk); + WARN_ON_ONCE(sk); done: bpf_iter_tcp_unlock_bucket(seq); + if (unlikely(!iter->end_sk)) { + ++iter->state.bucket; + goto again; + } return iter->batch[0].sk; }