zynq-rs/libasync/src/smoltcp/tcp_stream.rs

239 lines
7.4 KiB
Rust
Raw Normal View History

//! async TCP interface
//!
//! TODO: implement futures AsyncRead/AsyncWrite/Stream/Sink interfaces
2020-03-31 07:16:58 +08:00
use core::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
2020-04-02 05:34:53 +08:00
use alloc::vec::Vec;
2020-03-31 07:16:58 +08:00
use smoltcp::{
Error, Result,
2020-03-31 07:16:58 +08:00
socket::{
SocketHandle, SocketRef,
TcpSocketBuffer, TcpSocket, TcpState,
2020-03-31 07:16:58 +08:00
},
};
use crate::task;
2020-03-31 07:16:58 +08:00
use super::Sockets;
/// References a smoltcp TcpSocket
2020-03-31 07:16:58 +08:00
pub struct TcpStream {
handle: SocketHandle,
}
2020-04-02 05:34:53 +08:00
/// Wait while letting `$f()` poll a stream's socket
macro_rules! poll_stream {
($stream: expr, $output: ty, $f: expr) => (async {
struct Adhoc<'a> {
stream: &'a TcpStream,
}
impl<'a> Future for Adhoc<'a> {
type Output = $output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let result = self.stream.with_socket($f);
if !result.is_ready() {
Sockets::register_waker(cx.waker().clone());
}
result
}
}
Adhoc { stream: $stream }.await
})
}
2020-03-31 07:16:58 +08:00
impl TcpStream {
2020-04-02 05:34:53 +08:00
/// Allocates sockets and its buffers, registers it in the
/// SocketSet.
///
/// Not `pub` as the result can not yet be used. Use `listen()` or
/// `connect()` to obtain a valid TcpStream.
2020-03-31 07:16:58 +08:00
fn new(rx_bufsize: usize, tx_bufsize: usize) -> Self {
fn uninit_vec<T>(size: usize) -> Vec<T> {
let mut result = Vec::with_capacity(size);
unsafe {
result.set_len(size);
}
result
}
let rx_buffer = TcpSocketBuffer::new(uninit_vec(rx_bufsize));
let tx_buffer = TcpSocketBuffer::new(uninit_vec(tx_bufsize));
2020-03-31 07:16:58 +08:00
let socket = TcpSocket::new(rx_buffer, tx_buffer);
let handle = Sockets::instance().sockets.borrow_mut()
.add(socket);
TcpStream { handle }
}
/// Operate on the referenced TCP socket
2020-03-31 07:16:58 +08:00
fn with_socket<F, R>(&self, f: F) -> R
where
F: FnOnce(SocketRef<TcpSocket>) -> R,
{
let mut sockets = Sockets::instance().sockets.borrow_mut();
let socket_ref = sockets.get::<TcpSocket>(self.handle);
f(socket_ref)
}
/// Listen for the next incoming connection on a TCP
/// port. Succeeds on connection attempt.
///
/// Calling this serially in a loop will cause slow/botched
/// connection attempts stall any more new connections. Use
/// `listen()` with a backlog instead.
pub async fn accept(port: u16, rx_bufsize: usize, tx_bufsize: usize) -> Result<Self> {
2020-03-31 07:16:58 +08:00
let stream = Self::new(rx_bufsize, tx_bufsize);
2020-04-02 05:34:53 +08:00
// Set socket to listen
stream.with_socket(|mut s| s.listen(port))?;
2020-04-02 05:34:53 +08:00
// Wait for a connection
2020-04-02 05:01:48 +08:00
poll_stream!(&stream, (), |socket| {
if socket.state() != TcpState::Listen {
2020-04-02 05:01:48 +08:00
Poll::Ready(())
} else {
Poll::Pending
}
}).await;
Ok(stream)
2020-03-31 07:16:58 +08:00
}
/// Probe the receive buffer
///
/// Instead of handing you the data on the heap all at once,
/// smoltcp's read interface is wrapped so that your callback can
/// just return `Poll::Pending` if there is not enough data
/// yet. Likewise, return the amount of bytes consumed from the
/// buffer in the `Poll::Ready` result.
pub async fn recv<F, R>(&self, f: F) -> Result<R>
where
F: Fn(&[u8]) -> Poll<(usize, R)>,
{
struct Recv<'a, F: FnOnce(&[u8]) -> Poll<(usize, R)>, R> {
stream: &'a TcpStream,
f: F,
}
impl<'a, F: Fn(&[u8]) -> Poll<(usize, R)>, R> Future for Recv<'a, F, R> {
type Output = Result<R>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let result = self.stream.with_socket(|mut socket| {
if socket_is_handhshaking(&socket) {
return Ok(Poll::Pending);
}
socket.recv(|buf| {
if buf.len() > 0 {
match (self.f)(buf) {
Poll::Ready((amount, result)) =>
(amount, Poll::Ready(Ok(result))),
Poll::Pending =>
// 0 bytes consumed
(0, Poll::Pending),
}
} else {
(0, Poll::Pending)
}
})
});
match result {
Ok(Poll::Pending) => {
Sockets::register_waker(cx.waker().clone());
Poll::Pending
}
Ok(result) => {
result
}
Err(e) =>
Poll::Ready(Err(e)),
}
}
}
Recv {
stream: self,
f,
}.await
}
/// Wait until there is any space in the socket's send queue
async fn wait_can_send(&self) -> Result<()> {
poll_stream!(self, Result<()>, |socket| {
if socket_is_handhshaking(&socket) {
Poll::Pending
} else if socket.can_send() {
Poll::Ready(Ok(()))
} else if ! socket.may_send() {
Poll::Ready(Err(Error::Truncated))
} else {
Poll::Pending
}
}).await
}
/// Yields to wait for more buffer space
pub async fn send<I: IntoIterator<Item = u8>>(&self, data: I) -> Result<()> {
let mut data = data.into_iter();
let mut done = false;
while !done {
self.wait_can_send().await?;
self.with_socket(|mut socket| {
socket.send(|buf| {
for i in 0..buf.len() {
if let Some(byte) = data.next() {
buf[i] = byte;
} else {
done = true;
return (i, ())
}
}
(buf.len(), ())
})
})?;
}
Ok(())
}
/// Wait for all queued data to be sent and ACKed
///
/// **Warning:** this may not work as immediately as expected! The
/// other side may wait until it sends packets to you for
/// piggybacking the ACKs.
pub async fn flush(&self) -> Result<()> {
poll_stream!(self, Result<()>, |socket| {
if socket_is_handhshaking(&socket) {
Poll::Pending
} else if socket.may_send() && socket.send_queue() > 0 {
Poll::Pending
} else if socket.may_send() {
Poll::Ready(Ok(()))
} else {
Poll::Ready(Err(Error::Truncated))
}
}).await
2020-03-31 07:16:58 +08:00
}
}
impl Drop for TcpStream {
2020-04-02 05:34:53 +08:00
/// Free item in the socket set, which leads to deallocation of
/// the rx/tx buffers associated with this socket.
2020-03-31 07:16:58 +08:00
fn drop(&mut self) {
Sockets::instance().sockets.borrow_mut()
.remove(self.handle);
}
}
fn socket_is_handhshaking(socket: &SocketRef<TcpSocket>) -> bool {
match socket.state() {
TcpState::SynSent | TcpState::SynReceived =>
true,
_ =>
false,
}
}