feat: tls handshake timeout (#368)

This commit is contained in:
sigoden 2024-03-08 10:29:12 +08:00 committed by GitHub
parent 7c0fa3dab7
commit d66c9de8c8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -29,6 +29,8 @@ use std::sync::{
atomic::{AtomicBool, Ordering}, atomic::{AtomicBool, Ordering},
Arc, Arc,
}; };
use std::time::Duration;
use tokio::time::timeout;
use tokio::{net::TcpListener, task::JoinHandle}; use tokio::{net::TcpListener, task::JoinHandle};
#[cfg(feature = "tls")] #[cfg(feature = "tls")]
use tokio_rustls::{rustls::ServerConfig, TlsAcceptor}; use tokio_rustls::{rustls::ServerConfig, TlsAcceptor};
@ -91,12 +93,19 @@ fn serve(args: Args, running: Arc<AtomicBool>) -> Result<Vec<JoinHandle<()>>> {
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
let config = Arc::new(config); let config = Arc::new(config);
let tls_accepter = TlsAcceptor::from(config); let tls_accepter = TlsAcceptor::from(config);
let handshake_timeout = Duration::from_secs(10);
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
loop { loop {
let (cnx, addr) = listener.accept().await.unwrap(); let Ok((stream, addr)) = listener.accept().await else {
let Ok(stream) = tls_accepter.accept(cnx).await else { continue;
warn!("During tls handshake connection from {}", addr); };
let Some(stream) =
timeout(handshake_timeout, tls_accepter.accept(stream))
.await
.ok()
.and_then(|v| v.ok())
else {
continue; continue;
}; };
let stream = TokioIo::new(stream); let stream = TokioIo::new(stream);
@ -113,8 +122,10 @@ fn serve(args: Args, running: Arc<AtomicBool>) -> Result<Vec<JoinHandle<()>>> {
(None, None) => { (None, None) => {
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
loop { loop {
let (cnx, addr) = listener.accept().await.unwrap(); let Ok((stream, addr)) = listener.accept().await else {
let stream = TokioIo::new(cnx); continue;
};
let stream = TokioIo::new(stream);
tokio::spawn(handle_stream( tokio::spawn(handle_stream(
server_handle.clone(), server_handle.clone(),
stream, stream,
@ -139,8 +150,10 @@ fn serve(args: Args, running: Arc<AtomicBool>) -> Result<Vec<JoinHandle<()>>> {
.with_context(|| format!("Failed to bind `{}`", path.display()))?; .with_context(|| format!("Failed to bind `{}`", path.display()))?;
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
loop { loop {
let (cnx, _) = listener.accept().await.unwrap(); let Ok((stream, _addr)) = listener.accept().await else {
let stream = TokioIo::new(cnx); continue;
};
let stream = TokioIo::new(stream);
tokio::spawn(handle_stream(server_handle.clone(), stream, None)); tokio::spawn(handle_stream(server_handle.clone(), stream, None));
} }
}); });
@ -160,18 +173,15 @@ where
let hyper_service = let hyper_service =
service_fn(move |request: Request<Incoming>| handle.clone().call(request, addr)); service_fn(move |request: Request<Incoming>| handle.clone().call(request, addr));
let ret = Builder::new(TokioExecutor::new()) match Builder::new(TokioExecutor::new())
.serve_connection_with_upgrades(stream, hyper_service) .serve_connection_with_upgrades(stream, hyper_service)
.await; .await
{
if let Err(err) = ret { Ok(()) => {}
let scope = match addr { Err(_err) => {
Some(addr) => format!(" from {}", addr), // This error only appears when the client doesn't send a request and terminate the connection.
None => String::new(), //
}; // If client sends one request then terminate connection whenever, it doesn't appear.
match err.downcast_ref::<std::io::Error>() {
Some(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => {}
_ => warn!("Serving connection{}: {}", scope, err),
} }
} }
} }