mirror of
https://github.com/n0-computer/noq.git
synced 2026-09-22 19:24:08 +00:00
quinn: split streams module in send/recv parts
This commit is contained in:
Dirkjan Ochtman
committed by
Benjamin Saunders
Benjamin Saunders
parent
24cf82ef83
commit
14db88562d
@@ -23,7 +23,8 @@ use tracing::info_span;
|
||||
use crate::{
|
||||
broadcast::{self, Broadcast},
|
||||
mutex::Mutex,
|
||||
streams::{RecvStream, SendStream, WriteError},
|
||||
recv_stream::RecvStream,
|
||||
send_stream::{SendStream, WriteError},
|
||||
ConnectionEvent, EndpointEvent, VarInt,
|
||||
};
|
||||
|
||||
|
||||
+6
-5
@@ -56,7 +56,8 @@ mod connection;
|
||||
mod endpoint;
|
||||
mod mutex;
|
||||
mod platform;
|
||||
mod streams;
|
||||
mod recv_stream;
|
||||
mod send_stream;
|
||||
|
||||
pub use proto::{
|
||||
crypto, ApplicationClose, Certificate, CertificateChain, Chunk, ConnectError, ConnectionClose,
|
||||
@@ -65,7 +66,8 @@ pub use proto::{
|
||||
|
||||
pub use crate::builders::EndpointError;
|
||||
pub use crate::connection::{SendDatagramError, ZeroRttAccepted};
|
||||
pub use crate::streams::{ReadError, ReadExactError, ReadToEndError, StoppedError, WriteError};
|
||||
pub use crate::recv_stream::{ReadError, ReadExactError, ReadToEndError};
|
||||
pub use crate::send_stream::{StoppedError, WriteError};
|
||||
|
||||
/// Types that are generic over the crypto protocol implementation
|
||||
pub mod generic {
|
||||
@@ -75,9 +77,8 @@ pub mod generic {
|
||||
OpenBi, OpenUni,
|
||||
};
|
||||
pub use crate::endpoint::{Endpoint, Incoming};
|
||||
pub use crate::streams::{
|
||||
Read, ReadChunk, ReadChunks, ReadExact, ReadToEnd, RecvStream, SendStream,
|
||||
};
|
||||
pub use crate::recv_stream::{Read, ReadChunk, ReadChunks, ReadExact, ReadToEnd, RecvStream};
|
||||
pub use crate::send_stream::SendStream;
|
||||
pub use proto::generic::{ClientConfig, ServerConfig};
|
||||
}
|
||||
|
||||
|
||||
@@ -6,299 +6,13 @@ use std::{
|
||||
};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::{
|
||||
channel::oneshot,
|
||||
io::{AsyncRead, AsyncWrite},
|
||||
ready, FutureExt,
|
||||
};
|
||||
use proto::{Chunk, Chunks, ConnectionError, FinishError, ReadableError, StreamId};
|
||||
use futures::{io::AsyncRead, ready};
|
||||
use proto::{Chunk, Chunks, ConnectionError, ReadableError, StreamId};
|
||||
use thiserror::Error;
|
||||
use tokio::io::ReadBuf;
|
||||
|
||||
use crate::{connection::ConnectionRef, VarInt};
|
||||
|
||||
/// A stream that can only be used to send data
|
||||
///
|
||||
/// If dropped, streams that haven't been explicitly [`reset()`] will continue to (re)transmit
|
||||
/// previously written data until it has been fully acknowledged or the connection is closed.
|
||||
///
|
||||
/// [`reset()`]: SendStream::reset
|
||||
#[derive(Debug)]
|
||||
pub struct SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
conn: ConnectionRef<S>,
|
||||
stream: StreamId,
|
||||
is_0rtt: bool,
|
||||
finishing: Option<oneshot::Receiver<Option<WriteError>>>,
|
||||
}
|
||||
|
||||
impl<S> SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
pub(crate) fn new(conn: ConnectionRef<S>, stream: StreamId, is_0rtt: bool) -> Self {
|
||||
Self {
|
||||
conn,
|
||||
stream,
|
||||
is_0rtt,
|
||||
finishing: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Write bytes to the stream
|
||||
///
|
||||
/// Yields the number of bytes written on success. Congestion and flow control may cause this to
|
||||
/// be shorter than `buf.len()`, indicating that only a prefix of `buf` was written.
|
||||
pub fn write<'a>(&'a mut self, buf: &'a [u8]) -> Write<'a, S> {
|
||||
Write { stream: self, buf }
|
||||
}
|
||||
|
||||
/// Convenience method to write an entire buffer to the stream
|
||||
pub fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> WriteAll<'a, S> {
|
||||
WriteAll { stream: self, buf }
|
||||
}
|
||||
|
||||
fn poll_write(&mut self, cx: &mut Context, buf: &[u8]) -> Poll<Result<usize, WriteError>> {
|
||||
use proto::WriteError::*;
|
||||
let mut conn = self.conn.lock("SendStream::poll_write");
|
||||
if self.is_0rtt {
|
||||
conn.check_0rtt()
|
||||
.map_err(|()| WriteError::ZeroRttRejected)?;
|
||||
}
|
||||
if let Some(ref x) = conn.error {
|
||||
return Poll::Ready(Err(WriteError::ConnectionClosed(x.clone())));
|
||||
}
|
||||
let n = match conn.inner.write(self.stream, buf) {
|
||||
Ok(n) => n,
|
||||
Err(Blocked) => {
|
||||
conn.blocked_writers.insert(self.stream, cx.waker().clone());
|
||||
return Poll::Pending;
|
||||
}
|
||||
Err(Stopped(error_code)) => {
|
||||
return Poll::Ready(Err(WriteError::Stopped(error_code)));
|
||||
}
|
||||
Err(UnknownStream) => {
|
||||
return Poll::Ready(Err(WriteError::UnknownStream));
|
||||
}
|
||||
};
|
||||
conn.wake();
|
||||
Poll::Ready(Ok(n))
|
||||
}
|
||||
|
||||
/// Shut down the send stream gracefully.
|
||||
///
|
||||
/// No new data may be written after calling this method. Completes when the peer has
|
||||
/// acknowledged all sent data, retransmitting data as needed.
|
||||
pub fn finish(&mut self) -> Finish<'_, S> {
|
||||
Finish { stream: self }
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn poll_finish(&mut self, cx: &mut Context) -> Poll<Result<(), WriteError>> {
|
||||
let mut conn = self.conn.lock("poll_finish");
|
||||
if self.is_0rtt {
|
||||
conn.check_0rtt()
|
||||
.map_err(|()| WriteError::ZeroRttRejected)?;
|
||||
}
|
||||
if self.finishing.is_none() {
|
||||
conn.inner.finish(self.stream).map_err(|e| match e {
|
||||
FinishError::UnknownStream => WriteError::UnknownStream,
|
||||
FinishError::Stopped(error_code) => WriteError::Stopped(error_code),
|
||||
})?;
|
||||
let (send, recv) = oneshot::channel();
|
||||
self.finishing = Some(recv);
|
||||
conn.finishing.insert(self.stream, send);
|
||||
conn.wake();
|
||||
}
|
||||
match self
|
||||
.finishing
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.poll_unpin(cx)
|
||||
.map(|x| x.unwrap())
|
||||
{
|
||||
Poll::Ready(None) => Poll::Ready(Ok(())),
|
||||
Poll::Ready(Some(e)) => Poll::Ready(Err(e)),
|
||||
Poll::Pending => {
|
||||
// To ensure that finished streams can be detected even after the connection is
|
||||
// closed, we must only check for connection errors after determining that the
|
||||
// stream has not yet been finished. Note that this relies on holding the connection
|
||||
// lock so that it is impossible for the stream to become finished between the above
|
||||
// poll call and this check.
|
||||
if let Some(ref x) = conn.error {
|
||||
return Poll::Ready(Err(WriteError::ConnectionClosed(x.clone())));
|
||||
}
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Close the send stream immediately.
|
||||
///
|
||||
/// No new data can be written after calling this method. Locally buffered data is dropped, and
|
||||
/// previously transmitted data will no longer be retransmitted if lost. If an attempt has
|
||||
/// already been made to finish the stream, the peer may still receive all written data.
|
||||
pub fn reset(&mut self, error_code: VarInt) -> Result<(), UnknownStream> {
|
||||
let mut conn = self.conn.lock("SendStream::reset");
|
||||
if self.is_0rtt && conn.check_0rtt().is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
conn.inner.reset(self.stream, error_code)?;
|
||||
conn.wake();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set the priority of the send stream
|
||||
///
|
||||
/// Every send stream has an initial priority of 0. Locally buffered data from streams with
|
||||
/// higher priority will be transmitted before data from streams with lower priority. Changing
|
||||
/// the priority of a stream with pending data may only take effect after that data has been
|
||||
/// transmitted. Using many different priority levels per connection may have a negative
|
||||
/// impact on performance.
|
||||
pub fn set_priority(&self, priority: i32) -> Result<(), UnknownStream> {
|
||||
let mut conn = self.conn.lock("SendStream::set_priority");
|
||||
conn.inner.set_priority(self.stream, priority)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the priority of the send stream
|
||||
pub fn priority(&self) -> Result<i32, UnknownStream> {
|
||||
let mut conn = self.conn.lock("SendStream::priority");
|
||||
Ok(conn.inner.priority(self.stream)?)
|
||||
}
|
||||
|
||||
/// Completes if/when the peer stops the stream, yielding the error code
|
||||
pub fn stopped(&mut self) -> Stopped<'_, S> {
|
||||
Stopped { stream: self }
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn poll_stopped(&mut self, cx: &mut Context) -> Poll<Result<VarInt, StoppedError>> {
|
||||
let mut conn = self.conn.lock("SendStream::poll_stopped");
|
||||
|
||||
if self.is_0rtt {
|
||||
conn.check_0rtt()
|
||||
.map_err(|()| StoppedError::ZeroRttRejected)?;
|
||||
}
|
||||
|
||||
match conn.inner.stopped(self.stream) {
|
||||
Err(_) => Poll::Ready(Err(StoppedError::UnknownStream)),
|
||||
Ok(Some(error_code)) => Poll::Ready(Ok(error_code)),
|
||||
Ok(None) => {
|
||||
conn.stopped.insert(self.stream, cx.waker().clone());
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the identity of this stream
|
||||
pub fn id(&self) -> StreamId {
|
||||
self.stream
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> AsyncWrite for SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
|
||||
SendStream::poll_write(self.get_mut(), cx, buf).map_err(Into::into)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
self.get_mut().poll_finish(cx).map_err(Into::into)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> tokio::io::AsyncWrite for SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
AsyncWrite::poll_write(self, cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
AsyncWrite::poll_close(self, cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Drop for SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
fn drop(&mut self) {
|
||||
let mut conn = self.conn.lock("SendStream::drop");
|
||||
if conn.error.is_some() || (self.is_0rtt && conn.check_0rtt().is_err()) {
|
||||
return;
|
||||
}
|
||||
if self.finishing.is_none() {
|
||||
match conn.inner.finish(self.stream) {
|
||||
Ok(()) => conn.wake(),
|
||||
Err(FinishError::Stopped(reason)) => {
|
||||
if conn.inner.reset(self.stream, reason).is_ok() {
|
||||
conn.wake();
|
||||
}
|
||||
}
|
||||
// Already finished or reset, which is fine.
|
||||
Err(FinishError::UnknownStream) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by `SendStream::finish`
|
||||
pub struct Finish<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
}
|
||||
|
||||
impl<S> Future for Finish<'_, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<(), WriteError>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
self.get_mut().stream.poll_finish(cx)
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by `SendStream::stopped`
|
||||
pub struct Stopped<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
}
|
||||
|
||||
impl<S> Future for Stopped<'_, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<VarInt, StoppedError>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
self.get_mut().stream.poll_stopped(cx)
|
||||
}
|
||||
}
|
||||
|
||||
/// A stream that can only be used to receive data
|
||||
///
|
||||
/// `stop(0)` is implicitly called on drop unless:
|
||||
@@ -761,60 +475,6 @@ impl From<ReadError> for io::Error {
|
||||
}
|
||||
}
|
||||
|
||||
/// Errors that arise from writing to a stream
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
pub enum WriteError {
|
||||
/// The peer is no longer accepting data on this stream.
|
||||
///
|
||||
/// Carries an application-defined error code.
|
||||
#[error("sending stopped by peer: error {0}")]
|
||||
Stopped(VarInt),
|
||||
/// The connection was closed.
|
||||
#[error("connection closed: {0}")]
|
||||
ConnectionClosed(#[source] ConnectionError),
|
||||
/// The stream has already been finished or reset
|
||||
#[error("unknown stream")]
|
||||
UnknownStream,
|
||||
/// This was a 0-RTT stream and the server rejected it.
|
||||
///
|
||||
/// Can only occur on clients for 0-RTT streams, which can be opened using
|
||||
/// [`Connecting::into_0rtt()`].
|
||||
///
|
||||
/// [`Connecting::into_0rtt()`]: crate::generic::Connecting::into_0rtt()
|
||||
#[error("0-RTT rejected")]
|
||||
ZeroRttRejected,
|
||||
}
|
||||
|
||||
/// Errors that arise while monitoring for a send stream stop from the peer
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
pub enum StoppedError {
|
||||
/// The connection was closed.
|
||||
#[error("connection closed: {0}")]
|
||||
ConnectionClosed(#[source] ConnectionError),
|
||||
/// The stream has already been finished or reset
|
||||
#[error("unknown stream")]
|
||||
UnknownStream,
|
||||
/// This was a 0-RTT stream and the server rejected it.
|
||||
///
|
||||
/// Can only occur on clients for 0-RTT streams, which can be opened using
|
||||
/// [`Connecting::into_0rtt()`].
|
||||
///
|
||||
/// [`Connecting::into_0rtt()`]: crate::generic::Connecting::into_0rtt()
|
||||
#[error("0-RTT rejected")]
|
||||
ZeroRttRejected,
|
||||
}
|
||||
|
||||
impl From<WriteError> for io::Error {
|
||||
fn from(x: WriteError) -> Self {
|
||||
use self::WriteError::*;
|
||||
let kind = match x {
|
||||
Stopped(_) | ZeroRttRejected => io::ErrorKind::ConnectionReset,
|
||||
ConnectionClosed(_) | UnknownStream => io::ErrorKind::NotConnected,
|
||||
};
|
||||
io::Error::new(kind, x)
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by [`RecvStream::read()`].
|
||||
///
|
||||
/// [`RecvStream::read()`]: crate::generic::RecvStream::read
|
||||
@@ -929,56 +589,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by [`SendStream::write()`].
|
||||
///
|
||||
/// [`SendStream::write()`]: crate::generic::SendStream::write
|
||||
pub struct Write<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
buf: &'a [u8],
|
||||
}
|
||||
|
||||
impl<'a, S> Future for Write<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<usize, WriteError>;
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
let this = self.get_mut();
|
||||
this.stream.poll_write(cx, this.buf)
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by [`SendStream::write_all()`].
|
||||
///
|
||||
/// [`SendStream::write_all()`]: crate::generic::SendStream::write_all
|
||||
pub struct WriteAll<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
buf: &'a [u8],
|
||||
}
|
||||
|
||||
impl<'a, S> Future for WriteAll<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<(), WriteError>;
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
let this = self.get_mut();
|
||||
loop {
|
||||
if this.buf.is_empty() {
|
||||
return Poll::Ready(Ok(()));
|
||||
}
|
||||
let n = ready!(this.stream.poll_write(cx, this.buf))?;
|
||||
this.buf = &this.buf[n..];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Error indicating that a stream has already been finished or reset
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
#[error("unknown stream")]
|
||||
@@ -0,0 +1,398 @@
|
||||
use std::{
|
||||
future::Future,
|
||||
io,
|
||||
pin::Pin,
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
use futures::{channel::oneshot, io::AsyncWrite, ready, FutureExt};
|
||||
use proto::{ConnectionError, FinishError, StreamId};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::{connection::ConnectionRef, recv_stream::UnknownStream, VarInt};
|
||||
|
||||
/// A stream that can only be used to send data
|
||||
///
|
||||
/// If dropped, streams that haven't been explicitly [`reset()`] will continue to (re)transmit
|
||||
/// previously written data until it has been fully acknowledged or the connection is closed.
|
||||
///
|
||||
/// [`reset()`]: SendStream::reset
|
||||
#[derive(Debug)]
|
||||
pub struct SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
conn: ConnectionRef<S>,
|
||||
stream: StreamId,
|
||||
is_0rtt: bool,
|
||||
finishing: Option<oneshot::Receiver<Option<WriteError>>>,
|
||||
}
|
||||
|
||||
impl<S> SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
pub(crate) fn new(conn: ConnectionRef<S>, stream: StreamId, is_0rtt: bool) -> Self {
|
||||
Self {
|
||||
conn,
|
||||
stream,
|
||||
is_0rtt,
|
||||
finishing: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Write bytes to the stream
|
||||
///
|
||||
/// Yields the number of bytes written on success. Congestion and flow control may cause this to
|
||||
/// be shorter than `buf.len()`, indicating that only a prefix of `buf` was written.
|
||||
pub fn write<'a>(&'a mut self, buf: &'a [u8]) -> Write<'a, S> {
|
||||
Write { stream: self, buf }
|
||||
}
|
||||
|
||||
/// Convenience method to write an entire buffer to the stream
|
||||
pub fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> WriteAll<'a, S> {
|
||||
WriteAll { stream: self, buf }
|
||||
}
|
||||
|
||||
fn poll_write(&mut self, cx: &mut Context, buf: &[u8]) -> Poll<Result<usize, WriteError>> {
|
||||
use proto::WriteError::*;
|
||||
let mut conn = self.conn.lock("SendStream::poll_write");
|
||||
if self.is_0rtt {
|
||||
conn.check_0rtt()
|
||||
.map_err(|()| WriteError::ZeroRttRejected)?;
|
||||
}
|
||||
if let Some(ref x) = conn.error {
|
||||
return Poll::Ready(Err(WriteError::ConnectionClosed(x.clone())));
|
||||
}
|
||||
let n = match conn.inner.write(self.stream, buf) {
|
||||
Ok(n) => n,
|
||||
Err(Blocked) => {
|
||||
conn.blocked_writers.insert(self.stream, cx.waker().clone());
|
||||
return Poll::Pending;
|
||||
}
|
||||
Err(Stopped(error_code)) => {
|
||||
return Poll::Ready(Err(WriteError::Stopped(error_code)));
|
||||
}
|
||||
Err(UnknownStream) => {
|
||||
return Poll::Ready(Err(WriteError::UnknownStream));
|
||||
}
|
||||
};
|
||||
conn.wake();
|
||||
Poll::Ready(Ok(n))
|
||||
}
|
||||
|
||||
/// Shut down the send stream gracefully.
|
||||
///
|
||||
/// No new data may be written after calling this method. Completes when the peer has
|
||||
/// acknowledged all sent data, retransmitting data as needed.
|
||||
pub fn finish(&mut self) -> Finish<'_, S> {
|
||||
Finish { stream: self }
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn poll_finish(&mut self, cx: &mut Context) -> Poll<Result<(), WriteError>> {
|
||||
let mut conn = self.conn.lock("poll_finish");
|
||||
if self.is_0rtt {
|
||||
conn.check_0rtt()
|
||||
.map_err(|()| WriteError::ZeroRttRejected)?;
|
||||
}
|
||||
if self.finishing.is_none() {
|
||||
conn.inner.finish(self.stream).map_err(|e| match e {
|
||||
FinishError::UnknownStream => WriteError::UnknownStream,
|
||||
FinishError::Stopped(error_code) => WriteError::Stopped(error_code),
|
||||
})?;
|
||||
let (send, recv) = oneshot::channel();
|
||||
self.finishing = Some(recv);
|
||||
conn.finishing.insert(self.stream, send);
|
||||
conn.wake();
|
||||
}
|
||||
match self
|
||||
.finishing
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.poll_unpin(cx)
|
||||
.map(|x| x.unwrap())
|
||||
{
|
||||
Poll::Ready(None) => Poll::Ready(Ok(())),
|
||||
Poll::Ready(Some(e)) => Poll::Ready(Err(e)),
|
||||
Poll::Pending => {
|
||||
// To ensure that finished streams can be detected even after the connection is
|
||||
// closed, we must only check for connection errors after determining that the
|
||||
// stream has not yet been finished. Note that this relies on holding the connection
|
||||
// lock so that it is impossible for the stream to become finished between the above
|
||||
// poll call and this check.
|
||||
if let Some(ref x) = conn.error {
|
||||
return Poll::Ready(Err(WriteError::ConnectionClosed(x.clone())));
|
||||
}
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Close the send stream immediately.
|
||||
///
|
||||
/// No new data can be written after calling this method. Locally buffered data is dropped, and
|
||||
/// previously transmitted data will no longer be retransmitted if lost. If an attempt has
|
||||
/// already been made to finish the stream, the peer may still receive all written data.
|
||||
pub fn reset(&mut self, error_code: VarInt) -> Result<(), UnknownStream> {
|
||||
let mut conn = self.conn.lock("SendStream::reset");
|
||||
if self.is_0rtt && conn.check_0rtt().is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
conn.inner.reset(self.stream, error_code)?;
|
||||
conn.wake();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set the priority of the send stream
|
||||
///
|
||||
/// Every send stream has an initial priority of 0. Locally buffered data from streams with
|
||||
/// higher priority will be transmitted before data from streams with lower priority. Changing
|
||||
/// the priority of a stream with pending data may only take effect after that data has been
|
||||
/// transmitted. Using many different priority levels per connection may have a negative
|
||||
/// impact on performance.
|
||||
pub fn set_priority(&self, priority: i32) -> Result<(), UnknownStream> {
|
||||
let mut conn = self.conn.lock("SendStream::set_priority");
|
||||
conn.inner.set_priority(self.stream, priority)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the priority of the send stream
|
||||
pub fn priority(&self) -> Result<i32, UnknownStream> {
|
||||
let mut conn = self.conn.lock("SendStream::priority");
|
||||
Ok(conn.inner.priority(self.stream)?)
|
||||
}
|
||||
|
||||
/// Completes if/when the peer stops the stream, yielding the error code
|
||||
pub fn stopped(&mut self) -> Stopped<'_, S> {
|
||||
Stopped { stream: self }
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn poll_stopped(&mut self, cx: &mut Context) -> Poll<Result<VarInt, StoppedError>> {
|
||||
let mut conn = self.conn.lock("SendStream::poll_stopped");
|
||||
|
||||
if self.is_0rtt {
|
||||
conn.check_0rtt()
|
||||
.map_err(|()| StoppedError::ZeroRttRejected)?;
|
||||
}
|
||||
|
||||
match conn.inner.stopped(self.stream) {
|
||||
Err(_) => Poll::Ready(Err(StoppedError::UnknownStream)),
|
||||
Ok(Some(error_code)) => Poll::Ready(Ok(error_code)),
|
||||
Ok(None) => {
|
||||
conn.stopped.insert(self.stream, cx.waker().clone());
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the identity of this stream
|
||||
pub fn id(&self) -> StreamId {
|
||||
self.stream
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> AsyncWrite for SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
|
||||
SendStream::poll_write(self.get_mut(), cx, buf).map_err(Into::into)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
self.get_mut().poll_finish(cx).map_err(Into::into)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> tokio::io::AsyncWrite for SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
AsyncWrite::poll_write(self, cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
AsyncWrite::poll_close(self, cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Drop for SendStream<S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
fn drop(&mut self) {
|
||||
let mut conn = self.conn.lock("SendStream::drop");
|
||||
if conn.error.is_some() || (self.is_0rtt && conn.check_0rtt().is_err()) {
|
||||
return;
|
||||
}
|
||||
if self.finishing.is_none() {
|
||||
match conn.inner.finish(self.stream) {
|
||||
Ok(()) => conn.wake(),
|
||||
Err(FinishError::Stopped(reason)) => {
|
||||
if conn.inner.reset(self.stream, reason).is_ok() {
|
||||
conn.wake();
|
||||
}
|
||||
}
|
||||
// Already finished or reset, which is fine.
|
||||
Err(FinishError::UnknownStream) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by `SendStream::finish`
|
||||
pub struct Finish<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
}
|
||||
|
||||
impl<S> Future for Finish<'_, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<(), WriteError>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
self.get_mut().stream.poll_finish(cx)
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by `SendStream::stopped`
|
||||
pub struct Stopped<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
}
|
||||
|
||||
impl<S> Future for Stopped<'_, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<VarInt, StoppedError>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
self.get_mut().stream.poll_stopped(cx)
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by [`SendStream::write()`].
|
||||
///
|
||||
/// [`SendStream::write()`]: crate::generic::SendStream::write
|
||||
pub struct Write<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
buf: &'a [u8],
|
||||
}
|
||||
|
||||
impl<'a, S> Future for Write<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<usize, WriteError>;
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
let this = self.get_mut();
|
||||
this.stream.poll_write(cx, this.buf)
|
||||
}
|
||||
}
|
||||
|
||||
/// Future produced by [`SendStream::write_all()`].
|
||||
///
|
||||
/// [`SendStream::write_all()`]: crate::generic::SendStream::write_all
|
||||
pub struct WriteAll<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
stream: &'a mut SendStream<S>,
|
||||
buf: &'a [u8],
|
||||
}
|
||||
|
||||
impl<'a, S> Future for WriteAll<'a, S>
|
||||
where
|
||||
S: proto::crypto::Session,
|
||||
{
|
||||
type Output = Result<(), WriteError>;
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
let this = self.get_mut();
|
||||
loop {
|
||||
if this.buf.is_empty() {
|
||||
return Poll::Ready(Ok(()));
|
||||
}
|
||||
let n = ready!(this.stream.poll_write(cx, this.buf))?;
|
||||
this.buf = &this.buf[n..];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Errors that arise from writing to a stream
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
pub enum WriteError {
|
||||
/// The peer is no longer accepting data on this stream.
|
||||
///
|
||||
/// Carries an application-defined error code.
|
||||
#[error("sending stopped by peer: error {0}")]
|
||||
Stopped(VarInt),
|
||||
/// The connection was closed.
|
||||
#[error("connection closed: {0}")]
|
||||
ConnectionClosed(#[source] ConnectionError),
|
||||
/// The stream has already been finished or reset
|
||||
#[error("unknown stream")]
|
||||
UnknownStream,
|
||||
/// This was a 0-RTT stream and the server rejected it.
|
||||
///
|
||||
/// Can only occur on clients for 0-RTT streams, which can be opened using
|
||||
/// [`Connecting::into_0rtt()`].
|
||||
///
|
||||
/// [`Connecting::into_0rtt()`]: crate::generic::Connecting::into_0rtt()
|
||||
#[error("0-RTT rejected")]
|
||||
ZeroRttRejected,
|
||||
}
|
||||
|
||||
/// Errors that arise while monitoring for a send stream stop from the peer
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
pub enum StoppedError {
|
||||
/// The connection was closed.
|
||||
#[error("connection closed: {0}")]
|
||||
ConnectionClosed(#[source] ConnectionError),
|
||||
/// The stream has already been finished or reset
|
||||
#[error("unknown stream")]
|
||||
UnknownStream,
|
||||
/// This was a 0-RTT stream and the server rejected it.
|
||||
///
|
||||
/// Can only occur on clients for 0-RTT streams, which can be opened using
|
||||
/// [`Connecting::into_0rtt()`].
|
||||
///
|
||||
/// [`Connecting::into_0rtt()`]: crate::generic::Connecting::into_0rtt()
|
||||
#[error("0-RTT rejected")]
|
||||
ZeroRttRejected,
|
||||
}
|
||||
|
||||
impl From<WriteError> for io::Error {
|
||||
fn from(x: WriteError) -> Self {
|
||||
use self::WriteError::*;
|
||||
let kind = match x {
|
||||
Stopped(_) | ZeroRttRejected => io::ErrorKind::ConnectionReset,
|
||||
ConnectionClosed(_) | UnknownStream => io::ErrorKind::NotConnected,
|
||||
};
|
||||
io::Error::new(kind, x)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user