path: root/src/send.rs

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::trace;

use futures::AsyncWriteExt;
use futures::Stream;
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, possibly with CHUNK_HAS_CONTINUATION flag
//					when this is not the last chunk of the message
// - [u8; chunk_length] chunk data

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;

struct SendQueueItem {
	id: RequestID,
	prio: RequestPriority,
	data: DataReader,
}

#[pin_project::pin_project]
struct DataReader {
	#[pin]
	reader: ByteStream,
	packet: Packet,
	pos: usize,
	buf: Vec<u8>,
	eos: bool,
}

impl From<ByteStream> for DataReader {
	fn from(data: ByteStream) -> DataReader {
		DataReader {
			reader: data,
			packet: Ok(Bytes::new()),
			pos: 0,
			buf: Vec::with_capacity(MAX_CHUNK_LENGTH as usize),
			eos: false,
		}
	}
}

enum DataFrame {
	Data {
		/// a fixed size buffer containing some data, possibly padded with 0s
		data: [u8; MAX_CHUNK_LENGTH as usize],
		/// actual lenght of data
		len: usize,
		/// whethere 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
		may_have_more: bool,
	},
	/// An error code automatically signals the end of the stream
	Error(u8),
}

impl DataFrame {
	fn empty_last() -> Self {
		DataFrame::Data {
			data: [0; MAX_CHUNK_LENGTH as usize],
			len: 0,
			may_have_more: false,
		}
	}

	fn header(&self) -> [u8; 2] {
		let header_u16 = match self {
			DataFrame::Data {
				len,
				may_have_more: false,
				..
			} => *len as u16,
			DataFrame::Data {
				len,
				may_have_more: true,
				..
			} => *len as u16 | CHUNK_HAS_CONTINUATION,
			DataFrame::Error(e) => *e as u16 | ERROR_MARKER,
		};
		ChunkLength::to_be_bytes(header_u16)
	}

	fn data(&self) -> &[u8] {
		match self {
			DataFrame::Data { ref data, len, .. } => &data[..*len],
			DataFrame::Error(_) => &[],
		}
	}

	fn may_have_more(&self) -> bool {
		match self {
			DataFrame::Data { may_have_more, .. } => *may_have_more,
			DataFrame::Error(_) => false,
		}
	}
}

impl Stream for DataReader {
	type Item = DataFrame;

	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
		let mut this = self.project();

		if *this.eos {
			// eos was reached at previous call to poll_next, where a partial packet
			// was returned. Now return None
			return Poll::Ready(None);
		}

		loop {
			let packet = match this.packet {
				Ok(v) => v,
				Err(e) => {
					let e = *e;
					*this.packet = Ok(Bytes::new());
					*this.eos = true;
					return Poll::Ready(Some(DataFrame::Error(e)));
				}
			};
			let packet_left = packet.len() - *this.pos;
			let buf_left = MAX_CHUNK_LENGTH as usize - this.buf.len();
			let to_read = std::cmp::min(buf_left, packet_left);
			this.buf
				.extend_from_slice(&packet[*this.pos..*this.pos + to_read]);
			*this.pos += to_read;
			if this.buf.len() == MAX_CHUNK_LENGTH as usize {
				// we have a full buf, ready to send
				break;
			}

			// we don't have a full buf, packet is empty; try receive more
			if let Some(p) = futures::ready!(this.reader.as_mut().poll_next(cx)) {
				*this.packet = p;
				*this.pos = 0;
				// if buf is empty, we will loop and return the error directly. If buf
				// isn't empty, send it before by breaking.
				if this.packet.is_err() && !this.buf.is_empty() {
					break;
				}
			} else {
				*this.eos = true;
				break;
			}
		}

		let mut body = [0; MAX_CHUNK_LENGTH as usize];
		let len = this.buf.len();
		body[..len].copy_from_slice(this.buf);
		this.buf.clear();
		Poll::Ready(Some(DataFrame::Data {
			data: body,
			len,
			may_have_more: !*this.eos,
		}))
	}
}

struct SendQueue {
	items: VecDeque<(u8, VecDeque<SendQueueItem>)>,
}

impl SendQueue {
	fn new() -> Self {
		Self {
			items: VecDeque::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);
	}
	// used only in tests. They should probably be rewriten
	#[allow(dead_code)]
	fn pop(&mut self) -> Option<SendQueueItem> {
		match self.items.pop_front() {
			None => None,
			Some((prio, mut items_at_prio)) => {
				let ret = items_at_prio.pop_front();
				if !items_at_prio.is_empty() {
					self.items.push_front((prio, items_at_prio));
				}
				ret.or_else(|| self.pop())
			}
		}
	}
	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<Self::Output> {
		for i in 0..self.queue.items.len() {
			let (_prio, items_at_prio) = &mut self.queue.items[i];

			let mut ready_item = None;
			for (j, item) in items_at_prio.iter_mut().enumerate() {
				match Pin::new(&mut item.data).poll_next(ctx) {
					Poll::Pending => (),
					Poll::Ready(ready_v) => {
						ready_item = Some((j, ready_v));
						break;
					}
				}
			}

			if let Some((j, ready_v)) = ready_item {
				let item = items_at_prio.remove(j).unwrap();
				let id = item.id;
				if ready_v
					.as_ref()
					.map(|data| data.may_have_more())
					.unwrap_or(false)
				{
					items_at_prio.push_back(item);
				} else if items_at_prio.is_empty() {
					self.queue.items.remove(i);
				}
				return Poll::Ready((id, ready_v.unwrap_or_else(DataFrame::empty_last)));
			}
		}
		// TODO what do we do if self.queue is empty? We won't get scheduled again.
		Poll::Pending
	}
}

/// 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<W>(
		self: Arc<Self>,
		mut msg_recv: mpsc::UnboundedReceiver<(RequestID, RequestPriority, ByteStream)>,
		mut write: BoxStreamWrite<W>,
	) -> Result<(), Error>
	where
		W: AsyncWriteExt + Unpin + Send + Sync,
	{
		let mut sending = SendQueue::new();
		let mut should_exit = false;
		while !should_exit || !sending.is_empty() {
			let recv_fut = msg_recv.recv();
			futures::pin_mut!(recv_fut);
			let send_fut = sending.next_ready();

			// recv_fut is cancellation-safe according to tokio doc,
			// send_fut is cancellation-safe as implemented above?
			use futures::future::Either;
			match futures::future::select(recv_fut, send_fut).await {
				Either::Left((sth, _send_fut)) => {
					if let Some((id, prio, data)) = sth {
						sending.push(SendQueueItem {
							id,
							prio,
							data: data.into(),
						});
					} else {
						should_exit = true;
					};
				}
				Either::Right(((id, data), _recv_fut)) => {
					trace!("send_loop: sending bytes for {}", id);

					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(())
	}
}

#[cfg(test)]
mod test {
	use super::*;

	fn empty_data() -> DataReader {
		type Item = Packet;
		let stream: Pin<Box<dyn futures::Stream<Item = Item> + Send + 'static>> =
			Box::pin(futures::stream::empty::<Packet>());
		stream.into()
	}

	#[test]
	fn test_priority_queue() {
		let i1 = SendQueueItem {
			id: 1,
			prio: PRIO_NORMAL,
			data: empty_data(),
		};
		let i2 = SendQueueItem {
			id: 2,
			prio: PRIO_HIGH,
			data: empty_data(),
		};
		let i2bis = SendQueueItem {
			id: 20,
			prio: PRIO_HIGH,
			data: empty_data(),
		};
		let i3 = SendQueueItem {
			id: 3,
			prio: PRIO_HIGH | PRIO_SECONDARY,
			data: empty_data(),
		};
		let i4 = SendQueueItem {
			id: 4,
			prio: PRIO_BACKGROUND | PRIO_SECONDARY,
			data: empty_data(),
		};
		let i5 = SendQueueItem {
			id: 5,
			prio: PRIO_BACKGROUND | PRIO_PRIMARY,
			data: empty_data(),
		};

		let mut q = SendQueue::new();

		q.push(i1); // 1
		let a = q.pop().unwrap(); // empty -> 1
		assert_eq!(a.id, 1);
		assert!(q.pop().is_none());

		q.push(a); // 1
		q.push(i2); // 2 1
		q.push(i2bis); // [2 20] 1
		let a = q.pop().unwrap(); // 20 1 -> 2
		assert_eq!(a.id, 2);
		let b = q.pop().unwrap(); // 1 -> 20
		assert_eq!(b.id, 20);
		let c = q.pop().unwrap(); // empty -> 1
		assert_eq!(c.id, 1);
		assert!(q.pop().is_none());

		q.push(a); // 2
		q.push(b); // [2 20]
		q.push(c); // [2 20] 1
		q.push(i3); // [2 20] 3 1
		q.push(i4); // [2 20] 3 1 4
		q.push(i5); // [2 20] 3 1 5 4

		let a = q.pop().unwrap(); // 20 3 1 5 4 -> 2
		assert_eq!(a.id, 2);
		q.push(a); // [20 2] 3 1 5 4

		let a = q.pop().unwrap(); // 2 3 1 5 4 -> 20
		assert_eq!(a.id, 20);
		let b = q.pop().unwrap(); // 3 1 5 4 -> 2
		assert_eq!(b.id, 2);
		q.push(b); // 2 3 1 5 4
		let b = q.pop().unwrap(); // 3 1 5 4 -> 2
		assert_eq!(b.id, 2);
		let c = q.pop().unwrap(); // 1 5 4 -> 3
		assert_eq!(c.id, 3);
		q.push(b); // 2 1 5 4
		let b = q.pop().unwrap(); // 1 5 4 -> 2
		assert_eq!(b.id, 2);
		let e = q.pop().unwrap(); // 5 4 -> 1
		assert_eq!(e.id, 1);
		let f = q.pop().unwrap(); // 4 -> 5
		assert_eq!(f.id, 5);
		let g = q.pop().unwrap(); // empty -> 4
		assert_eq!(g.id, 4);
		assert!(q.pop().is_none());
	}
}