io_uring/zcrx: add a read limit to recvzc requests

Currently multishot recvzc requests have no read limit and will remain
active so as long as the socket remains open. But, there are sometimes a
need to do a fixed length read e.g. peeking at some data in the socket.

Add a length limit to recvzc requests `len`. A value of 0 means no limit
which is the previous behaviour. A positive value N specifies how many
bytes to read from the socket.

Data will still be posted in aux completions, as before. This could be
split across multiple frags. But the primary recvzc request will now
complete once N bytes have been read. The completion of the recvzc
request will have res and cflags both set to 0.

Signed-off-by: David Wei <dw@davidwei.uk>
Reviewed-by: Pavel Begunkov <asml.silence@gmail.com>
Link: https://lore.kernel.org/r/20250224041319.2389785-2-dw@davidwei.uk
[axboe: fixup io_zcrx_recv() for !CONFIG_NET]
Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/io_uring/net.c b/io_uring/net.c
index 000dc70d08d0d1fe2d10ddcdebbf1a8abf900399..4850d4d898f9b7fcc87e99abdc972e3ff53aa4a9 100644 (file)
--- a/io_uring/net.c
+++ b/io_uring/net.c
@@ -94,6 +94,7 @@ struct io_recvzc {
        struct file                     *file;
        unsigned                        msg_flags;
        u16                             flags;
+       u32                             len;
        struct io_zcrx_ifq              *ifq;
 };
 
@@ -1241,7 +1242,7 @@ int io_recvzc_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
        unsigned ifq_idx;
 
        if (unlikely(sqe->file_index || sqe->addr2 || sqe->addr ||
-                    sqe->len || sqe->addr3))
+                    sqe->addr3))
                return -EINVAL;
 
        ifq_idx = READ_ONCE(sqe->zcrx_ifq_idx);
@@ -1250,7 +1251,7 @@ int io_recvzc_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
        zc->ifq = req->ctx->ifq;
        if (!zc->ifq)
                return -EINVAL;
-
+       zc->len = READ_ONCE(sqe->len);
        zc->flags = READ_ONCE(sqe->ioprio);
        zc->msg_flags = READ_ONCE(sqe->msg_flags);
        if (zc->msg_flags)
@@ -1270,6 +1271,7 @@ int io_recvzc(struct io_kiocb *req, unsigned int issue_flags)
 {
        struct io_recvzc *zc = io_kiocb_to_cmd(req, struct io_recvzc);
        struct socket *sock;
+       unsigned int len;
        int ret;
 
        if (!(req->flags & REQ_F_POLLED) &&
@@ -1280,8 +1282,16 @@ int io_recvzc(struct io_kiocb *req, unsigned int issue_flags)
        if (unlikely(!sock))
                return -ENOTSOCK;
 
+       len = zc->len;
        ret = io_zcrx_recv(req, zc->ifq, sock, zc->msg_flags | MSG_DONTWAIT,
-                          issue_flags);
+                          issue_flags, &zc->len);
+       if (len && zc->len == 0) {
+               io_req_set_res(req, 0, 0);
+
+               if (issue_flags & IO_URING_F_MULTISHOT)
+                       return IOU_STOP_MULTISHOT;
+               return IOU_OK;
+       }
        if (unlikely(ret <= 0) && ret != -EAGAIN) {
                if (ret == -ERESTARTSYS)
                        ret = -EINTR;

diff --git a/io_uring/zcrx.c b/io_uring/zcrx.c
index f2d326e18e67061f56954b779639f2a0228ea87a..9c95b5b6ec4eb497885c47aae37753b89c1b5b21 100644 (file)
--- a/io_uring/zcrx.c
+++ b/io_uring/zcrx.c
@@ -817,6 +817,7 @@ io_zcrx_recv_skb(read_descriptor_t *desc, struct sk_buff *skb,
        int i, copy, end, off;
        int ret = 0;
 
+       len = min_t(size_t, len, desc->count);
        if (unlikely(args->nr_skbs++ > IO_SKBS_PER_CALL_LIMIT))
                return -EAGAIN;
 
@@ -894,26 +895,30 @@ io_zcrx_recv_skb(read_descriptor_t *desc, struct sk_buff *skb,
 out:
        if (offset == start_off)
                return ret;
+       desc->count -= (offset - start_off);
        return offset - start_off;
 }
 
 static int io_zcrx_tcp_recvmsg(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
                                struct sock *sk, int flags,
-                               unsigned issue_flags)
+                               unsigned issue_flags, unsigned int *outlen)
 {
+       unsigned int len = *outlen;
        struct io_zcrx_args args = {
                .req = req,
                .ifq = ifq,
                .sock = sk->sk_socket,
        };
        read_descriptor_t rd_desc = {
-               .count = 1,
+               .count = len ? len : UINT_MAX,
                .arg.data = &args,
        };
        int ret;
 
        lock_sock(sk);
        ret = tcp_read_sock(sk, &rd_desc, io_zcrx_recv_skb);
+       if (len && ret > 0)
+               *outlen = len - ret;
        if (ret <= 0) {
                if (ret < 0 || sock_flag(sk, SOCK_DONE))
                        goto out;
@@ -942,7 +947,7 @@ out:
 
 int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
                 struct socket *sock, unsigned int flags,
-                unsigned issue_flags)
+                unsigned issue_flags, unsigned int *len)
 {
        struct sock *sk = sock->sk;
        const struct proto *prot = READ_ONCE(sk->sk_prot);
@@ -951,5 +956,5 @@ int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
                return -EPROTONOSUPPORT;
 
        sock_rps_record_flow(sk);
-       return io_zcrx_tcp_recvmsg(req, ifq, sk, flags, issue_flags);
+       return io_zcrx_tcp_recvmsg(req, ifq, sk, flags, issue_flags, len);
 }

diff --git a/io_uring/zcrx.h b/io_uring/zcrx.h
index a16bdd921f03b60a49cf053d09fed4d7668f2601..706cc7300780d6c3985fa61a8ce17a1d6660446f 100644 (file)
--- a/io_uring/zcrx.h
+++ b/io_uring/zcrx.h
@@ -46,7 +46,7 @@ void io_unregister_zcrx_ifqs(struct io_ring_ctx *ctx);
 void io_shutdown_zcrx_ifqs(struct io_ring_ctx *ctx);
 int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
                 struct socket *sock, unsigned int flags,
-                unsigned issue_flags);
+                unsigned issue_flags, unsigned int *len);
 #else
 static inline int io_register_zcrx_ifq(struct io_ring_ctx *ctx,
                                        struct io_uring_zcrx_ifq_reg __user *arg)
@@ -61,7 +61,7 @@ static inline void io_shutdown_zcrx_ifqs(struct io_ring_ctx *ctx)
 }
 static inline int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
                               struct socket *sock, unsigned int flags,
-                              unsigned issue_flags)
+                              unsigned issue_flags, unsigned int *len)
 {
        return -EOPNOTSUPP;
 }