diff --git a/quinn-h3/examples/h3.rs b/quinn-h3/examples/h3.rs index 0bc6d3f87..78d5371a8 100644 --- a/quinn-h3/examples/h3.rs +++ b/quinn-h3/examples/h3.rs @@ -228,6 +228,7 @@ fn handle_request( sender .send_response_trailers(response, trailer) + .unwrap() .map_err(|e| format_err!("failed to send response: {:?}", e)) }) } diff --git a/quinn-h3/src/headers.rs b/quinn-h3/src/headers.rs index 2033dba33..a11627913 100644 --- a/quinn-h3/src/headers.rs +++ b/quinn-h3/src/headers.rs @@ -1,5 +1,7 @@ -use futures::{Async, Future, Poll}; +use futures::{try_ready, Async, Future, Poll}; +use quinn::SendStream; use quinn_proto::StreamId; +use tokio_io::io::WriteAll; use crate::{ connection::ConnectionRef, @@ -48,3 +50,34 @@ impl Future for DecodeHeaders { } } } + +pub(crate) struct SendHeaders(WriteAll>); + +impl SendHeaders { + pub fn new( + header: Header, + conn: &ConnectionRef, + send: SendStream, + stream_id: StreamId, + ) -> Result { + let block = { + let conn = &mut conn.h3.lock().unwrap().inner; + conn.encode_header(stream_id, header)? + }; + + let mut encoded = Vec::new(); + block.encode(&mut encoded); + + Ok(Self(tokio_io::io::write_all(send, encoded))) + } +} + +impl Future for SendHeaders { + type Item = SendStream; + type Error = Error; + + fn poll(&mut self) -> Poll { + let (send, _) = try_ready!(self.0.poll()); + Ok(Async::Ready(send)) + } +} diff --git a/quinn-h3/src/server.rs b/quinn-h3/src/server.rs index 99aec5f0a..6739eb941 100644 --- a/quinn-h3/src/server.rs +++ b/quinn-h3/src/server.rs @@ -13,7 +13,7 @@ use crate::{ body::{Body, RecvBody, SendBody}, connection::{ConnectionDriver, ConnectionRef}, frame::{FrameDecoder, FrameStream}, - headers::DecodeHeaders, + headers::{DecodeHeaders, SendHeaders}, proto::{frame::HttpFrame, headers::Header}, try_take, Error, Settings, }; @@ -220,7 +220,10 @@ pub struct Sender { } impl Sender { - pub fn send_response>(self, response: Response) -> SendResponse { + pub fn send_response>( + self, + response: Response, + ) -> Result { SendResponse::new(response, self.send, self.stream_id, self.conn) } @@ -228,7 +231,7 @@ impl Sender { self, response: Response, trailer: HeaderMap, - ) -> SendResponse { + ) -> Result { SendResponse::with_trailers( response, Some(trailer), @@ -240,8 +243,7 @@ impl Sender { } enum SendResponseState { - Encoding, - SendingHeader(WriteAll>), + SendingHeader(SendHeaders), SendingBody(SendBody), SendingTrailers(WriteAll>), Closing(Shutdown), @@ -249,10 +251,8 @@ enum SendResponseState { pub struct SendResponse { state: SendResponseState, - header: Option
, body: Option, trailer: Option
, - send: Option, conn: ConnectionRef, stream_id: StreamId, } @@ -263,7 +263,7 @@ impl SendResponse { send: SendStream, stream_id: StreamId, conn: ConnectionRef, - ) -> Self { + ) -> Result { Self::with_trailers(response, None, send, stream_id, conn) } @@ -273,7 +273,7 @@ impl SendResponse { send: SendStream, stream_id: StreamId, conn: ConnectionRef, - ) -> Self { + ) -> Result { let ( response::Parts { status, headers, .. @@ -281,42 +281,29 @@ impl SendResponse { body, ) = response.into_parts(); - Self { + let headers = Header::response(status, headers); + let state = + SendResponseState::SendingHeader(SendHeaders::new(headers, &conn, send, stream_id)?); + + Ok(Self { conn, + state, stream_id, body: Some(body.into()), - send: Some(send), trailer: trailers.map(Header::trailer), - header: Some(Header::response(status, headers)), - state: SendResponseState::Encoding, - } + }) } } impl Future for SendResponse { type Item = (); type Error = Error; + fn poll(&mut self) -> Poll { loop { match self.state { - SendResponseState::Encoding => { - let header = try_take(&mut self.header, "polled after finished")?; - let block = { - let conn = &mut self.conn.h3.lock().unwrap().inner; - conn.encode_header(self.stream_id, header)? - }; - - let mut encoded = Vec::new(); - block.encode(&mut encoded); - - let send = try_take(&mut self.send, "polled after finished")?; - mem::replace( - &mut self.state, - SendResponseState::SendingHeader(tokio_io::io::write_all(send, encoded)), - ); - } SendResponseState::SendingHeader(ref mut write) => { - let (send, _) = try_ready!(write.poll()); + let send = try_ready!(write.poll()); mem::replace( &mut self.state, SendResponseState::SendingBody(SendBody::new(