Skip to main content

usb_vsock/
transport.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use futures::channel::mpsc;
6use futures::sink::Sink;
7use futures::stream::Stream;
8use futures::{AsyncRead, AsyncWrite};
9use std::io::{self, Cursor, Read};
10use std::pin::Pin;
11use std::task::{Context, Poll};
12
13/// A bidirectional in-memory byte transport backed by `futures::channel::mpsc` channels.
14#[derive(Debug)]
15pub struct MpscTransport {
16    sender: mpsc::Sender<Vec<u8>>,
17    receiver: mpsc::Receiver<Vec<u8>>,
18    current_read_buf: Option<Cursor<Vec<u8>>>,
19}
20
21impl MpscTransport {
22    /// Creates a pair of connected [`MpscTransport`] endpoints with a default capacity of 128 chunks.
23    pub fn pair() -> (Self, Self) {
24        Self::pair_with_capacity(128)
25    }
26
27    /// Creates a pair of connected [`MpscTransport`] endpoints with the given channel capacity.
28    pub fn pair_with_capacity(capacity: usize) -> (Self, Self) {
29        let (tx_a, rx_b) = mpsc::channel(capacity);
30        let (tx_b, rx_a) = mpsc::channel(capacity);
31        (
32            Self { sender: tx_a, receiver: rx_a, current_read_buf: None },
33            Self { sender: tx_b, receiver: rx_b, current_read_buf: None },
34        )
35    }
36
37    /// Creates a new [`MpscTransport`] from a sender and receiver.
38    pub fn new(sender: mpsc::Sender<Vec<u8>>, receiver: mpsc::Receiver<Vec<u8>>) -> Self {
39        Self { sender, receiver, current_read_buf: None }
40    }
41}
42
43impl AsyncRead for MpscTransport {
44    fn poll_read(
45        mut self: Pin<&mut Self>,
46        cx: &mut Context<'_>,
47        buf: &mut [u8],
48    ) -> Poll<io::Result<usize>> {
49        if buf.is_empty() {
50            return Poll::Ready(Ok(0));
51        }
52
53        let this = &mut *self;
54        loop {
55            if let Some(cursor) = &mut this.current_read_buf {
56                let n = match cursor.read(buf) {
57                    Ok(n) => n,
58                    Err(e) => return Poll::Ready(Err(e)),
59                };
60                if n > 0 {
61                    if cursor.position() >= cursor.get_ref().len() as u64 {
62                        this.current_read_buf = None;
63                    }
64                    return Poll::Ready(Ok(n));
65                }
66                this.current_read_buf = None;
67            }
68
69            match Pin::new(&mut this.receiver).poll_next(cx) {
70                Poll::Ready(Some(chunk)) => {
71                    if chunk.is_empty() {
72                        continue;
73                    }
74                    this.current_read_buf = Some(Cursor::new(chunk));
75                }
76                Poll::Ready(None) => {
77                    return Poll::Ready(Ok(0));
78                }
79                Poll::Pending => {
80                    return Poll::Pending;
81                }
82            }
83        }
84    }
85}
86
87impl AsyncWrite for MpscTransport {
88    fn poll_write(
89        mut self: Pin<&mut Self>,
90        cx: &mut Context<'_>,
91        buf: &[u8],
92    ) -> Poll<io::Result<usize>> {
93        if buf.is_empty() {
94            return Poll::Ready(Ok(0));
95        }
96
97        let this = &mut *self;
98        match Pin::new(&mut this.sender).poll_ready(cx) {
99            Poll::Ready(Ok(())) => {
100                let len = buf.len();
101                match Pin::new(&mut this.sender).start_send(buf.to_vec()) {
102                    Ok(()) => Poll::Ready(Ok(len)),
103                    Err(e) if e.is_disconnected() => {
104                        Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
105                    }
106                    Err(e) => Poll::Ready(Err(io::Error::other(e))),
107                }
108            }
109            Poll::Ready(Err(e)) => {
110                if e.is_disconnected() {
111                    Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
112                } else {
113                    Poll::Ready(Err(io::Error::other(e)))
114                }
115            }
116            Poll::Pending => Poll::Pending,
117        }
118    }
119
120    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
121        let this = &mut *self;
122        match Pin::new(&mut this.sender).poll_flush(cx) {
123            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
124            Poll::Ready(Err(e)) => {
125                if e.is_disconnected() {
126                    Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
127                } else {
128                    Poll::Ready(Err(io::Error::other(e)))
129                }
130            }
131            Poll::Pending => Poll::Pending,
132        }
133    }
134
135    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
136        let this = &mut *self;
137        match Pin::new(&mut this.sender).poll_close(cx) {
138            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
139            Poll::Ready(Err(e)) => {
140                if e.is_disconnected() {
141                    Poll::Ready(Ok(()))
142                } else {
143                    Poll::Ready(Err(io::Error::other(e)))
144                }
145            }
146            Poll::Pending => Poll::Pending,
147        }
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use futures::{AsyncReadExt, AsyncWriteExt};
155
156    #[fuchsia::test]
157    async fn test_basic_transfer() {
158        let (mut a, mut b) = MpscTransport::pair();
159        a.write_all(b"hello world").await.unwrap();
160
161        let mut buf = [0u8; 11];
162        b.read_exact(&mut buf).await.unwrap();
163        assert_eq!(&buf, b"hello world");
164    }
165
166    #[fuchsia::test]
167    async fn test_fragmented_read() {
168        let (mut a, mut b) = MpscTransport::pair();
169        a.write_all(b"1234567890").await.unwrap();
170
171        let mut buf1 = [0u8; 4];
172        let mut buf2 = [0u8; 6];
173        b.read_exact(&mut buf1).await.unwrap();
174        b.read_exact(&mut buf2).await.unwrap();
175        assert_eq!(&buf1, b"1234");
176        assert_eq!(&buf2, b"567890");
177    }
178
179    #[fuchsia::test]
180    async fn test_eof_on_drop() {
181        let (a, mut b) = MpscTransport::pair();
182        drop(a);
183
184        let mut buf = [0u8; 10];
185        let n = b.read(&mut buf).await.unwrap();
186        assert_eq!(n, 0);
187    }
188
189    #[fuchsia::test]
190    async fn test_broken_pipe_on_write() {
191        let (mut a, b) = MpscTransport::pair();
192        drop(b);
193
194        let err = a.write_all(b"fail").await.unwrap_err();
195        assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);
196    }
197}