forked from lx/netapp
112 lines
3.1 KiB
Rust
112 lines
3.1 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use async_trait::async_trait;
|
|
use bytes::Bytes;
|
|
use log::trace;
|
|
|
|
use futures::channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
|
|
use futures::AsyncReadExt;
|
|
|
|
use crate::error::*;
|
|
use crate::send::*;
|
|
use crate::stream::*;
|
|
|
|
/// Structure to warn when the sender is dropped before end of stream was reached, like when
|
|
/// connection to some remote drops while transmitting data
|
|
struct Sender {
|
|
inner: UnboundedSender<Packet>,
|
|
closed: bool,
|
|
}
|
|
|
|
impl Sender {
|
|
fn new(inner: UnboundedSender<Packet>) -> Self {
|
|
Sender {
|
|
inner,
|
|
closed: false,
|
|
}
|
|
}
|
|
|
|
fn send(&self, packet: Packet) {
|
|
let _ = self.inner.unbounded_send(packet);
|
|
}
|
|
|
|
fn end(&mut self) {
|
|
self.closed = true;
|
|
}
|
|
}
|
|
|
|
impl Drop for Sender {
|
|
fn drop(&mut self) {
|
|
if !self.closed {
|
|
self.send(Err(255));
|
|
}
|
|
self.inner.close_channel();
|
|
}
|
|
}
|
|
|
|
/// The RecvLoop trait, which is implemented both by the client and the server
|
|
/// connection objects (ServerConn and ClientConn) adds a method `.recv_loop()`
|
|
/// and a prototype of a handler for received messages `.recv_handler()` that
|
|
/// must be filled by implementors. `.recv_loop()` receives messages in a loop
|
|
/// according to the protocol defined above: chunks of message in progress of being
|
|
/// received are stored in a buffer, and when the last chunk of a message is received,
|
|
/// the full message is passed to the receive handler.
|
|
#[async_trait]
|
|
pub(crate) trait RecvLoop: Sync + 'static {
|
|
fn recv_handler(self: &Arc<Self>, id: RequestID, stream: UnboundedReceiver<Packet>);
|
|
|
|
async fn recv_loop<R>(self: Arc<Self>, mut read: R) -> Result<(), Error>
|
|
where
|
|
R: AsyncReadExt + Unpin + Send + Sync,
|
|
{
|
|
let mut streams: HashMap<RequestID, Sender> = HashMap::new();
|
|
loop {
|
|
trace!("recv_loop: reading packet");
|
|
let mut header_id = [0u8; RequestID::BITS as usize / 8];
|
|
match read.read_exact(&mut header_id[..]).await {
|
|
Ok(_) => (),
|
|
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break,
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
let id = RequestID::from_be_bytes(header_id);
|
|
trace!("recv_loop: got header id: {:04x}", id);
|
|
|
|
let mut header_size = [0u8; ChunkLength::BITS as usize / 8];
|
|
read.read_exact(&mut header_size[..]).await?;
|
|
let size = ChunkLength::from_be_bytes(header_size);
|
|
trace!("recv_loop: got header size: {:04x}", size);
|
|
|
|
let has_cont = (size & CHUNK_HAS_CONTINUATION) != 0;
|
|
let is_error = (size & ERROR_MARKER) != 0;
|
|
let packet = if is_error {
|
|
Err(size as u8)
|
|
} else {
|
|
let size = size & !CHUNK_HAS_CONTINUATION;
|
|
let mut next_slice = vec![0; size as usize];
|
|
read.read_exact(&mut next_slice[..]).await?;
|
|
trace!("recv_loop: read {} bytes", next_slice.len());
|
|
Ok(Bytes::from(next_slice))
|
|
};
|
|
|
|
let mut sender = if let Some(send) = streams.remove(&(id)) {
|
|
send
|
|
} else {
|
|
let (send, recv) = unbounded();
|
|
self.recv_handler(id, recv);
|
|
Sender::new(send)
|
|
};
|
|
|
|
// if we get an error, the receiving end is disconnected. We still need to
|
|
// reach eos before dropping this sender
|
|
sender.send(packet);
|
|
|
|
if has_cont {
|
|
streams.insert(id, sender);
|
|
} else {
|
|
sender.end();
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|