tcp_bpf: poll psock queues too in tcp_poll()

Signed-off-by: Cong Wang <cong.wang@bytedance.com>
congwang · Sep 18, 2021 · 5d46718 · 5d46718
1 parent 0a0020f
commit 5d46718
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/include/net/tcp.h b/include/net/tcp.h
@@ -2223,6 +2223,12 @@ struct sk_psock;
 struct proto *tcp_bpf_get_proto(struct sock *sk, struct sk_psock *psock);
 int tcp_bpf_update_proto(struct sock *sk, struct sk_psock *psock, bool restore);
 void tcp_bpf_clone(const struct sock *sk, struct sock *newsk);
+__poll_t tcp_bpf_poll(struct sock *sk);
+#else
+static inline __poll_t tcp_bpf_poll(struct sock *sk)
+{
+	return 0;
+}
 #endif /* CONFIG_BPF_SYSCALL */
 
 int tcp_bpf_sendmsg_redir(struct sock *sk, struct sk_msg *msg, u32 bytes,

diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
@@ -563,6 +563,7 @@ __poll_t tcp_poll(struct file *file, struct socket *sock, poll_table *wait)
 
 		if (tcp_stream_is_readable(sk, target))
 			mask |= EPOLLIN | EPOLLRDNORM;
+		mask |= tcp_bpf_poll(sk);
 
 		if (!(sk->sk_shutdown & SEND_SHUTDOWN)) {
 			if (__sk_stream_is_writeable(sk, 1)) {

diff --git a/net/ipv4/tcp_bpf.c b/net/ipv4/tcp_bpf.c
@@ -456,6 +456,20 @@ static int tcp_bpf_sendpage(struct sock *sk, struct page *page, int offset,
 	return copied ? copied : err;
 }
 
+__poll_t tcp_bpf_poll(struct sock *sk)
+{
+	struct sk_psock *psock;
+	__poll_t mask = 0;
+
+	psock = sk_psock_get(sk);
+	if (unlikely(!psock))
+		return 0;
+	if (!sk_psock_queue_empty(psock))
+		mask |= EPOLLIN | EPOLLRDNORM;
+	sk_psock_put(sk, psock);
+	return mask;
+}
+
 enum {
 	TCP_BPF_IPV4,
 	TCP_BPF_IPV6,