aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/send.rs205
-rw-r--r--src/stream.rs29
2 files changed, 68 insertions, 166 deletions
diff --git a/src/send.rs b/src/send.rs
index 59805cf..a8cf966 100644
--- a/src/send.rs
+++ b/src/send.rs
@@ -8,7 +8,6 @@ use bytes::Bytes;
use log::trace;
use futures::AsyncWriteExt;
-use futures::Stream;
use kuska_handshake::async_std::BoxStreamWrite;
use tokio::sync::mpsc;
@@ -30,152 +29,14 @@ pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x3FF0;
pub(crate) const ERROR_MARKER: ChunkLength = 0x4000;
pub(crate) const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000;
+struct SendQueue {
+ items: VecDeque<(u8, VecDeque<SendQueueItem>)>,
+}
+
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>)>,
+ data: ByteStreamReader,
}
impl SendQueue {
@@ -232,35 +93,69 @@ impl<'a> futures::Future for SendQueuePollNextReady<'a> {
let mut ready_item = None;
for (j, item) in items_at_prio.iter_mut().enumerate() {
- match Pin::new(&mut item.data).poll_next(ctx) {
+ 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));
+ ready_item = Some((j, ready_v, item.data.eos()));
break;
}
}
}
- if let Some((j, ready_v)) = ready_item {
+ if let Some((j, bytes_or_err, eos)) = ready_item {
+ let data_frame = match bytes_or_err {
+ Ok(bytes) => DataFrame::Data(bytes, !eos),
+ Err(e) => DataFrame::Error(match e {
+ ReadExactError::Stream(code) => code,
+ _ => unreachable!(),
+ }),
+ };
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)
- {
+ if !eos {
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)));
+ return Poll::Ready((id, data_frame));
}
}
- // TODO what do we do if self.queue is empty? We won't get scheduled again.
+ // 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(u8),
+}
+
+impl DataFrame {
+ 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(e) => *e as u16 | ERROR_MARKER,
+ };
+ ChunkLength::to_be_bytes(header_u16)
+ }
+
+ fn data(&self) -> &[u8] {
+ match self {
+ DataFrame::Data(ref data, _) => &data[..],
+ DataFrame::Error(_) => &[],
+ }
+ }
+}
+
/// 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,
@@ -295,7 +190,7 @@ pub(crate) trait SendLoop: Sync {
sending.push(SendQueueItem {
id,
prio,
- data: data.into(),
+ data: ByteStreamReader::new(data),
});
} else {
should_exit = true;
diff --git a/src/stream.rs b/src/stream.rs
index 6c23f4a..ae57d62 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -82,6 +82,23 @@ impl ByteStreamReader {
}
}
+ pub fn take_buffer(&mut self) -> Bytes {
+ let bytes = Bytes::from(
+ self .buf
+ .iter()
+ .map(|x| &x[..])
+ .collect::<Vec<_>>()
+ .concat(),
+ );
+ self.buf.clear();
+ self.buf_len = 0;
+ bytes
+ }
+
+ pub fn eos(&self) -> bool {
+ self.buf.is_empty() && self.eos
+ }
+
fn try_get(&mut self, read_len: usize) -> Option<Bytes> {
if self.buf_len >= read_len {
let mut slices = Vec::with_capacity(self.buf.len());
@@ -144,17 +161,7 @@ impl<'a> Future for ByteStreamReadExact<'a> {
if *this.fail_on_eos {
return Poll::Ready(Err(ReadExactError::UnexpectedEos));
} else {
- let bytes = Bytes::from(
- this.reader
- .buf
- .iter()
- .map(|x| &x[..])
- .collect::<Vec<_>>()
- .concat(),
- );
- this.reader.buf.clear();
- this.reader.buf_len = 0;
- return Poll::Ready(Ok(bytes));
+ return Poll::Ready(Ok(this.reader.take_buffer()));
}
}
generated by cgit v1.2.3 (git 2.25.1) at 2024-11-22 17:28:54 +0000