Use cilent time out for h2 handshake timeout. (#2483)

actix-http/src/h2/dispatcher.rs

@@ -10,7 +10,7 @@ use std::{
 };
 
 use actix_codec::{AsyncRead, AsyncWrite};
-use actix_rt::time::Sleep;
+use actix_rt::time::{sleep, Sleep};
 use actix_service::Service;
 use actix_utils::future::poll_fn;
 use bytes::{Bytes, BytesMut};
@@ -55,9 +55,16 @@ where
         on_connect_data: OnConnectData,
         config: ServiceConfig,
         peer_addr: Option<net::SocketAddr>,
+        timer: Option<Pin<Box<Sleep>>>,
     ) -> Self {
-        let ping_pong = config.keep_alive_timer().map(|timer| H2PingPong {
-            timer: Box::pin(timer),
+        let ping_pong = config.keep_alive().map(|dur| H2PingPong {
+            timer: timer
+                .map(|mut timer| {
+                    // reset timer if it's received from new function.
+                    timer.as_mut().reset(config.now() + dur);
+                    timer
+                })
+                .unwrap_or_else(|| Box::pin(sleep(dur))),
             on_flight: false,
             ping_pong: connection.ping_pong().unwrap(),
         });

actix-http/src/h2/mod.rs

@@ -1,20 +1,30 @@
 //! HTTP/2 protocol.
 
 use std::{
+    future::Future,
     pin::Pin,
     task::{Context, Poll},
 };
 
+use actix_codec::{AsyncRead, AsyncWrite};
+use actix_rt::time::Sleep;
 use bytes::Bytes;
 use futures_core::{ready, Stream};
-use h2::RecvStream;
+use h2::{
+    server::{handshake, Connection, Handshake},
+    RecvStream,
+};
 
 mod dispatcher;
 mod service;
 
 pub use self::dispatcher::Dispatcher;
 pub use self::service::H2Service;
-use crate::error::PayloadError;
+
+use crate::{
+    config::ServiceConfig,
+    error::{DispatchError, PayloadError},
+};
 
 /// HTTP/2 peer stream.
 pub struct Payload {
@@ -50,3 +60,44 @@ impl Stream for Payload {
         }
     }
 }
+
+pub(crate) fn handshake_with_timeout<T>(
+    io: T,
+    config: &ServiceConfig,
+) -> HandshakeWithTimeout<T>
+where
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    HandshakeWithTimeout {
+        handshake: handshake(io),
+        timer: config.client_timer().map(Box::pin),
+    }
+}
+
+pub(crate) struct HandshakeWithTimeout<T: AsyncRead + AsyncWrite + Unpin> {
+    handshake: Handshake<T>,
+    timer: Option<Pin<Box<Sleep>>>,
+}
+
+impl<T> Future for HandshakeWithTimeout<T>
+where
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = Result<(Connection<T, Bytes>, Option<Pin<Box<Sleep>>>), DispatchError>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+
+        match Pin::new(&mut this.handshake).poll(cx)? {
+            // return the timer on success handshake. It can be re-used for h2 ping-pong.
+            Poll::Ready(conn) => Poll::Ready(Ok((conn, this.timer.take()))),
+            Poll::Pending => match this.timer.as_mut() {
+                Some(timer) => {
+                    ready!(timer.as_mut().poll(cx));
+                    Poll::Ready(Err(DispatchError::SlowRequestTimeout))
+                }
+                None => Poll::Pending,
+            },
+        }
+    }
+}

actix-http/src/h2/service.rs

@@ -15,9 +15,7 @@ use actix_service::{
     ServiceFactoryExt as _,
 };
 use actix_utils::future::ready;
-use bytes::Bytes;
 use futures_core::{future::LocalBoxFuture, ready};
-use h2::server::{handshake as h2_handshake, Handshake as H2Handshake};
 use log::error;
 
 use crate::{
@@ -28,7 +26,7 @@ use crate::{
     ConnectCallback, OnConnectData, Request, Response,
 };
 
-use super::dispatcher::Dispatcher;
+use super::{dispatcher::Dispatcher, handshake_with_timeout, HandshakeWithTimeout};
 
 /// `ServiceFactory` implementation for HTTP/2 transport
 pub struct H2Service<T, S, B> {
@@ -297,7 +295,7 @@ where
                 Some(self.cfg.clone()),
                 addr,
                 on_connect_data,
-                h2_handshake(io),
+                handshake_with_timeout(io, &self.cfg),
             ),
         }
     }
@@ -314,7 +312,7 @@ where
         Option<ServiceConfig>,
         Option<net::SocketAddr>,
         OnConnectData,
-        H2Handshake<T, Bytes>,
+        HandshakeWithTimeout<T>,
     ),
 }
 
@@ -352,7 +350,7 @@ where
                 ref mut on_connect_data,
                 ref mut handshake,
             ) => match ready!(Pin::new(handshake).poll(cx)) {
-                Ok(conn) => {
+                Ok((conn, timer)) => {
                     let on_connect_data = std::mem::take(on_connect_data);
                     self.state = State::Incoming(Dispatcher::new(
                         srv.take().unwrap(),
@@ -360,12 +358,13 @@ where
                         on_connect_data,
                         config.take().unwrap(),
                         *peer_addr,
+                        timer,
                     ));
                     self.poll(cx)
                 }
                 Err(err) => {
                     trace!("H2 handshake error: {}", err);
-                    Poll::Ready(Err(err.into()))
+                    Poll::Ready(Err(err))
                 }
             },
         }