use std::task::{Context, Poll}; use std::{cmp, future::Future, marker::PhantomData, net, pin::Pin, rc::Rc}; use actix_codec::{AsyncRead, AsyncWrite}; use actix_rt::time::{Instant, Sleep}; use actix_service::Service; use bytes::{Bytes, BytesMut}; use futures_core::ready; use h2::server::{Connection, SendResponse}; use h2::SendStream; use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING}; use log::{error, trace}; use crate::body::{BodySize, MessageBody, ResponseBody}; use crate::config::ServiceConfig; use crate::error::{DispatchError, Error}; use crate::message::ResponseHead; use crate::payload::Payload; use crate::request::Request; use crate::response::Response; use crate::service::HttpFlow; use crate::OnConnectData; const CHUNK_SIZE: usize = 16_384; /// Dispatcher for HTTP/2 protocol. #[pin_project::pin_project] pub struct Dispatcher where T: AsyncRead + AsyncWrite + Unpin, S: Service, B: MessageBody, { flow: Rc>, connection: Connection, on_connect_data: OnConnectData, config: ServiceConfig, peer_addr: Option, ka_expire: Instant, ka_timer: Option, _phantom: PhantomData, } impl Dispatcher where T: AsyncRead + AsyncWrite + Unpin, S: Service, S::Error: Into, S::Response: Into>, B: MessageBody, { pub(crate) fn new( flow: Rc>, connection: Connection, on_connect_data: OnConnectData, config: ServiceConfig, timeout: Option, peer_addr: Option, ) -> Self { // let keepalive = config.keep_alive_enabled(); // let flags = if keepalive { // Flags::KEEPALIVE | Flags::KEEPALIVE_ENABLED // } else { // Flags::empty() // }; // keep-alive timer let (ka_expire, ka_timer) = if let Some(delay) = timeout { (delay.deadline(), Some(delay)) } else if let Some(delay) = config.keep_alive_timer() { (delay.deadline(), Some(delay)) } else { (config.now(), None) }; Dispatcher { flow, config, peer_addr, connection, on_connect_data, ka_expire, ka_timer, _phantom: PhantomData, } } } impl Future for Dispatcher where T: AsyncRead + AsyncWrite + Unpin, S: Service, S::Error: Into + 'static, S::Future: 'static, S::Response: Into> + 'static, B: MessageBody + 'static, { type Output = Result<(), DispatchError>; #[inline] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); loop { match ready!(Pin::new(&mut this.connection).poll_accept(cx)) { None => return Poll::Ready(Ok(())), Some(Err(err)) => return Poll::Ready(Err(err.into())), Some(Ok((req, res))) => { // update keep-alive expire if this.ka_timer.is_some() { if let Some(expire) = this.config.keep_alive_expire() { this.ka_expire = expire; } } let (parts, body) = req.into_parts(); let pl = crate::h2::Payload::new(body); let pl = Payload::::H2(pl); let mut req = Request::with_payload(pl); let head = req.head_mut(); head.uri = parts.uri; head.method = parts.method; head.version = parts.version; head.headers = parts.headers.into(); head.peer_addr = this.peer_addr; // merge on_connect_ext data into request extensions this.on_connect_data.merge_into(&mut req); let svc = ServiceResponse:: { state: ServiceResponseState::ServiceCall( this.flow.service.call(req), Some(res), ), config: this.config.clone(), buffer: None, _phantom: PhantomData, }; actix_rt::spawn(svc); } } } } } #[pin_project::pin_project] struct ServiceResponse { #[pin] state: ServiceResponseState, config: ServiceConfig, buffer: Option, _phantom: PhantomData<(I, E)>, } #[pin_project::pin_project(project = ServiceResponseStateProj)] enum ServiceResponseState { ServiceCall(#[pin] F, Option>), SendPayload(SendStream, #[pin] ResponseBody), } impl ServiceResponse where F: Future>, E: Into, I: Into>, B: MessageBody, { fn prepare_response( &self, head: &ResponseHead, size: &mut BodySize, ) -> http::Response<()> { let mut has_date = false; let mut skip_len = size != &BodySize::Stream; let mut res = http::Response::new(()); *res.status_mut() = head.status; *res.version_mut() = http::Version::HTTP_2; // Content length match head.status { http::StatusCode::NO_CONTENT | http::StatusCode::CONTINUE | http::StatusCode::PROCESSING => *size = BodySize::None, http::StatusCode::SWITCHING_PROTOCOLS => { skip_len = true; *size = BodySize::Stream; } _ => {} } let _ = match size { BodySize::None | BodySize::Stream => None, BodySize::Empty => res .headers_mut() .insert(CONTENT_LENGTH, HeaderValue::from_static("0")), BodySize::Sized(len) => { let mut buf = itoa::Buffer::new(); res.headers_mut().insert( CONTENT_LENGTH, HeaderValue::from_str(buf.format(*len)).unwrap(), ) } }; // copy headers for (key, value) in head.headers.iter() { match *key { // TODO: consider skipping other headers according to: // https://tools.ietf.org/html/rfc7540#section-8.1.2.2 // omit HTTP/1.x only headers CONNECTION | TRANSFER_ENCODING => continue, CONTENT_LENGTH if skip_len => continue, DATE => has_date = true, _ => {} } res.headers_mut().append(key, value.clone()); } // set date header if !has_date { let mut bytes = BytesMut::with_capacity(29); self.config.set_date_header(&mut bytes); res.headers_mut().insert( DATE, // SAFETY: serialized date-times are known ASCII strings unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) }, ); } res } } impl Future for ServiceResponse where F: Future>, E: Into, I: Into>, B: MessageBody, { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut this = self.as_mut().project(); match this.state.project() { ServiceResponseStateProj::ServiceCall(call, send) => { match ready!(call.poll(cx)) { Ok(res) => { let (res, body) = res.into().replace_body(()); let mut send = send.take().unwrap(); let mut size = body.size(); let h2_res = self.as_mut().prepare_response(res.head(), &mut size); this = self.as_mut().project(); let stream = match send.send_response(h2_res, size.is_eof()) { Err(e) => { trace!("Error sending HTTP/2 response: {:?}", e); return Poll::Ready(()); } Ok(stream) => stream, }; if size.is_eof() { Poll::Ready(()) } else { this.state .set(ServiceResponseState::SendPayload(stream, body)); self.poll(cx) } } Err(e) => { let res: Response = e.into().into(); let (res, body) = res.replace_body(()); let mut send = send.take().unwrap(); let mut size = body.size(); let h2_res = self.as_mut().prepare_response(res.head(), &mut size); this = self.as_mut().project(); let stream = match send.send_response(h2_res, size.is_eof()) { Err(e) => { trace!("Error sending HTTP/2 response: {:?}", e); return Poll::Ready(()); } Ok(stream) => stream, }; if size.is_eof() { Poll::Ready(()) } else { this.state.set(ServiceResponseState::SendPayload( stream, body.into_body(), )); self.poll(cx) } } } } ServiceResponseStateProj::SendPayload(ref mut stream, ref mut body) => { loop { loop { match this.buffer { Some(ref mut buffer) => { match ready!(stream.poll_capacity(cx)) { None => return Poll::Ready(()), Some(Ok(cap)) => { let len = buffer.len(); let bytes = buffer.split_to(cmp::min(cap, len)); if let Err(e) = stream.send_data(bytes, false) { warn!("{:?}", e); return Poll::Ready(()); } else if !buffer.is_empty() { let cap = cmp::min(buffer.len(), CHUNK_SIZE); stream.reserve_capacity(cap); } else { this.buffer.take(); } } Some(Err(e)) => { warn!("{:?}", e); return Poll::Ready(()); } } } None => match ready!(body.as_mut().poll_next(cx)) { None => { if let Err(e) = stream.send_data(Bytes::new(), true) { warn!("{:?}", e); } return Poll::Ready(()); } Some(Ok(chunk)) => { stream.reserve_capacity(cmp::min( chunk.len(), CHUNK_SIZE, )); *this.buffer = Some(chunk); } Some(Err(e)) => { error!("Response payload stream error: {:?}", e); return Poll::Ready(()); } }, } } } } } } }