use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use async_trait::async_trait; use bytes::Bytes; use log::*; use futures::AsyncWriteExt; use kuska_handshake::async_std::BoxStreamWrite; use tokio::sync::mpsc; use crate::error::*; use crate::message::*; use crate::stream::*; // Messages are sent by chunks // Chunk format: // - u32 BE: request id (same for request and response) // - u16 BE: chunk length + flags: // CHUNK_HAS_CONTINUATION when this is not the last chunk of the stream // ERROR_MARKER if this chunk denotes an error // (these two flags are exclusive, an error denotes the end of the stream) // - [u8; chunk_length] chunk data / error message pub(crate) type RequestID = u32; pub(crate) type ChunkLength = u16; pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x3FF0; pub(crate) const ERROR_MARKER: ChunkLength = 0x4000; pub(crate) const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000; pub(crate) const CHUNK_LENGTH_MASK: ChunkLength = 0x3FFF; struct SendQueue { items: Vec<(u8, VecDeque)>, } struct SendQueueItem { id: RequestID, prio: RequestPriority, data: ByteStreamReader, } impl SendQueue { fn new() -> Self { Self { items: Vec::with_capacity(64), } } fn push(&mut self, item: SendQueueItem) { let prio = item.prio; let pos_prio = match self.items.binary_search_by(|(p, _)| p.cmp(&prio)) { Ok(i) => i, Err(i) => { self.items.insert(i, (prio, VecDeque::new())); i } }; self.items[pos_prio].1.push_back(item); } fn is_empty(&self) -> bool { self.items.iter().all(|(_k, v)| v.is_empty()) } // this is like an async fn, but hand implemented fn next_ready(&mut self) -> SendQueuePollNextReady<'_> { SendQueuePollNextReady { queue: self } } } struct SendQueuePollNextReady<'a> { queue: &'a mut SendQueue, } impl<'a> futures::Future for SendQueuePollNextReady<'a> { type Output = (RequestID, DataFrame); fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { for (i, (_prio, items_at_prio)) in self.queue.items.iter_mut().enumerate() { let mut ready_item = None; for (j, item) in items_at_prio.iter_mut().enumerate() { let mut item_reader = item.data.read_exact_or_eos(MAX_CHUNK_LENGTH as usize); match Pin::new(&mut item_reader).poll(ctx) { Poll::Pending => (), Poll::Ready(ready_v) => { ready_item = Some((j, ready_v)); break; } } } if let Some((j, bytes_or_err)) = ready_item { let item = items_at_prio.remove(j).unwrap(); let id = item.id; let eos = item.data.eos(); let packet = bytes_or_err.map_err(|e| match e { ReadExactError::Stream(err) => err, _ => unreachable!(), }); let data_frame = DataFrame::from_packet(packet, !eos); if !eos && !matches!(data_frame, DataFrame::Error(_)) { items_at_prio.push_back(item); } else if items_at_prio.is_empty() { self.queue.items.remove(i); } return Poll::Ready((id, data_frame)); } } // If the queue is empty, this futures is eternally pending. // This is ok because we use it in a select with another future // that can interrupt it. Poll::Pending } } enum DataFrame { /// a fixed size buffer containing some data + a boolean indicating whether /// there may be more data comming from this stream. Can be used for some /// optimization. It's an error to set it to false if there is more data, but it is correct /// (albeit sub-optimal) to set it to true if there is nothing coming after Data(Bytes, bool), /// An error code automatically signals the end of the stream Error(Bytes), } impl DataFrame { fn from_packet(p: Packet, has_cont: bool) -> Self { match p { Ok(bytes) => { assert!(bytes.len() <= MAX_CHUNK_LENGTH as usize); Self::Data(bytes, has_cont) } Err(e) => { let msg = format!("{}", e); let mut msg = Bytes::from(msg.into_bytes()); if msg.len() > MAX_CHUNK_LENGTH as usize { msg = msg.slice(..MAX_CHUNK_LENGTH as usize); } Self::Error(msg) } } } fn header(&self) -> [u8; 2] { let header_u16 = match self { DataFrame::Data(data, false) => data.len() as u16, DataFrame::Data(data, true) => data.len() as u16 | CHUNK_HAS_CONTINUATION, DataFrame::Error(msg) => msg.len() as u16 | ERROR_MARKER, }; ChunkLength::to_be_bytes(header_u16) } fn data(&self) -> &[u8] { match self { DataFrame::Data(ref data, _) => &data[..], DataFrame::Error(ref msg) => &msg[..], } } } /// The SendLoop trait, which is implemented both by the client and the server /// connection objects (ServerConna and ClientConn) adds a method `.send_loop()` /// that takes a channel of messages to send and an asynchronous writer, /// and sends messages from the channel to the async writer, putting them in a queue /// before being sent and doing the round-robin sending strategy. /// /// The `.send_loop()` exits when the sending end of the channel is closed, /// or if there is an error at any time writing to the async writer. #[async_trait] pub(crate) trait SendLoop: Sync { async fn send_loop( self: Arc, msg_recv: mpsc::UnboundedReceiver<(RequestID, RequestPriority, ByteStream)>, mut write: BoxStreamWrite, ) -> Result<(), Error> where W: AsyncWriteExt + Unpin + Send + Sync, { let mut sending = SendQueue::new(); let mut msg_recv = Some(msg_recv); while msg_recv.is_some() || !sending.is_empty() { debug!( "Sending: {:?}", sending .items .iter() .map(|(_, i)| i.iter().map(|x| x.id)) .flatten() .collect::>() ); let recv_fut = async { if let Some(chan) = &mut msg_recv { chan.recv().await } else { futures::future::pending().await } }; let send_fut = sending.next_ready(); // recv_fut is cancellation-safe according to tokio doc, // send_fut is cancellation-safe as implemented above? tokio::select! { sth = recv_fut => { if let Some((id, prio, data)) = sth { trace!("send_loop: add stream {} to send", id); sending.push(SendQueueItem { id, prio, data: ByteStreamReader::new(data), }); } else { msg_recv = None; }; } (id, data) = send_fut => { trace!( "send_loop: id {}, send {} bytes, header_size {}", id, data.data().len(), hex::encode(data.header()) ); let header_id = RequestID::to_be_bytes(id); write.write_all(&header_id[..]).await?; write.write_all(&data.header()).await?; write.write_all(data.data()).await?; write.flush().await?; } } } let _ = write.goodbye().await; Ok(()) } }