ferron/listeners/

quic.rs

// Copyright (c) 2018 The quinn Developers
// Portions of this file are derived from Quinn (https://github.com/quinn-rs/quinn).
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//

use std::error::Error;
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use std::fmt::Debug;
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use std::future::Future;
use std::io;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use std::pin::Pin;
use std::sync::Arc;
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use std::task::{Context, Poll};
use std::time::Duration;
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use std::time::Instant;

use async_channel::{Receiver, Sender};
use ferron_common::logging::LogMessage;
#[cfg(feature = "runtime-monoio")]
use monoio::time::Sleep;
use quinn::crypto::rustls::QuicServerConfig;
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use quinn::{AsyncTimer, AsyncUdpSocket, Runtime};
use rustls::ServerConfig;
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
use send_wrapper::SendWrapper;
use tokio_util::sync::CancellationToken;
#[cfg(feature = "runtime-vibeio")]
use vibeio::time::Sleep;

use crate::listener_handler_communication::{Connection, ConnectionData};

type ListenerError = Box<dyn Error + Send + Sync>;

/// A timer for Quinn that utilizes Monoio's timer.
#[cfg(feature = "runtime-monoio")]
#[derive(Debug)]
struct MonoioTimer {
  inner: SendWrapper<Pin<Box<Sleep>>>,
}

#[cfg(feature = "runtime-monoio")]
impl AsyncTimer for MonoioTimer {
  fn reset(mut self: Pin<&mut Self>, t: Instant) {
    (*self.inner).as_mut().reset(t.into())
  }

  fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
    (*self.inner).as_mut().poll(cx)
  }
}

/// A timer for Quinn that utilizes `vibeio`'s timer.
#[cfg(feature = "runtime-vibeio")]
struct CustomAsyncTimer {
  inner: SendWrapper<Pin<Box<Sleep>>>,
}

#[cfg(feature = "runtime-vibeio")]
impl AsyncTimer for CustomAsyncTimer {
  fn reset(mut self: Pin<&mut Self>, t: Instant) {
    (*self.inner).as_mut().reset(t)
  }

  fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
    (*self.inner).as_mut().poll(cx)
  }
}

#[cfg(feature = "runtime-vibeio")]
impl Debug for CustomAsyncTimer {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.debug_struct("CustomAsyncTimer").finish()
  }
}

/// A runtime for Quinn that utilizes Tokio, if under Tokio runtime, and otherwise Monoio with async_io.
#[derive(Debug)]
#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
struct EnterTokioRuntime;

#[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
impl Runtime for EnterTokioRuntime {
  fn new_timer(&self, t: Instant) -> Pin<Box<dyn AsyncTimer>> {
    if tokio::runtime::Handle::try_current().is_ok() {
      Box::pin(tokio::time::sleep_until(t.into()))
    } else {
      #[cfg(feature = "runtime-monoio")]
      let timer = Box::pin(MonoioTimer {
        inner: SendWrapper::new(Box::pin(monoio::time::sleep_until(t.into()))),
      });
      #[cfg(feature = "runtime-vibeio")]
      let timer = Box::pin(CustomAsyncTimer {
        inner: SendWrapper::new(Box::pin(vibeio::time::sleep_until(t))),
      });
      timer
    }
  }

  fn spawn(&self, future: Pin<Box<dyn Future<Output = ()> + Send>>) {
    if let Ok(handle) = tokio::runtime::Handle::try_current() {
      handle.spawn(future);
    } else {
      #[cfg(feature = "runtime-monoio")]
      monoio::spawn(future);
      #[cfg(feature = "runtime-vibeio")]
      vibeio::spawn(future);
    }
  }

  fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>> {
    quinn::TokioRuntime::wrap_udp_socket(&quinn::TokioRuntime, sock)
  }
}

#[inline]
fn build_quic_server_config(tls_config: Arc<ServerConfig>) -> Result<quinn::ServerConfig, ListenerError> {
  let quic_server_config = QuicServerConfig::try_from(tls_config)
    .map_err(|err| anyhow::anyhow!("Cannot prepare the QUIC server configuration: {err}"))?;
  Ok(quinn::ServerConfig::with_crypto(Arc::new(quic_server_config)))
}

#[inline]
fn bind_udp_socket(address: SocketAddr) -> io::Result<std::net::UdpSocket> {
  // Create a new socket
  let listener_socket2 = socket2::Socket::new(
    if address.is_ipv6() {
      socket2::Domain::IPV6
    } else {
      socket2::Domain::IPV4
    },
    socket2::Type::DGRAM,
    Some(socket2::Protocol::UDP),
  )?;

  // Set socket options
  if address.is_ipv6() {
    listener_socket2.set_only_v6(false).unwrap_or_default();
  }

  // Bind the socket to the address
  listener_socket2.bind(&address.into())?;

  // Wrap the socket into a UdpSocket
  Ok(listener_socket2.into())
}

#[inline]
async fn log_accept_closed(logging_tx: &Option<Sender<LogMessage>>) {
  if let Some(logging_tx) = logging_tx {
    logging_tx
      .send(LogMessage::new(
        "HTTP/3 connections can't be accepted anymore".to_string(),
        true,
      ))
      .await
      .unwrap_or_default();
  }
}

/// Creates a QUIC listener
#[allow(clippy::type_complexity)]
pub fn create_quic_listener(
  address: SocketAddr,
  tls_config: Arc<ServerConfig>,
  tx: Sender<ConnectionData>,
  logging_tx: Option<Sender<LogMessage>>,
  first_startup: bool,
) -> Result<(CancellationToken, Sender<Arc<ServerConfig>>), Box<dyn Error + Send + Sync>> {
  let shutdown_tx = CancellationToken::new();
  let shutdown_rx = shutdown_tx.clone();
  let (rustls_config_tx, rustls_config_rx) = async_channel::unbounded();
  let (listen_error_tx, listen_error_rx) = async_channel::unbounded();
  std::thread::Builder::new()
    .name(format!("QUIC listener for {address}"))
    .spawn(move || {
      let mut rt = match crate::runtime::Runtime::new_runtime_tokio_only() {
        Ok(rt) => rt,
        Err(error) => {
          listen_error_tx
            .send_blocking(Some(
              anyhow::anyhow!("Can't create async runtime: {error}").into_boxed_dyn_error(),
            ))
            .unwrap_or_default();
          return;
        }
      };
      rt.run(async move {
        if let Err(error) = quic_listener_fn(
          address,
          tls_config,
          tx,
          &listen_error_tx,
          logging_tx,
          first_startup,
          shutdown_rx,
          rustls_config_rx,
        )
        .await
        {
          listen_error_tx.send(Some(error)).await.unwrap_or_default();
        }
      });
    })?;

  if let Some(error) = listen_error_rx.recv_blocking()? {
    Err(error)?;
  }

  Ok((shutdown_tx, rustls_config_tx))
}

/// QUIC listener function
#[allow(clippy::too_many_arguments)]
async fn quic_listener_fn(
  address: SocketAddr,
  tls_config: Arc<ServerConfig>,
  tx: Sender<ConnectionData>,
  listen_error_tx: &Sender<Option<ListenerError>>,
  logging_tx: Option<Sender<LogMessage>>,
  first_startup: bool,
  shutdown_rx: CancellationToken,
  rustls_config_rx: Receiver<Arc<ServerConfig>>,
) -> Result<(), ListenerError> {
  let server_config = build_quic_server_config(tls_config)?;
  let udp_port = address.port();
  let mut udp_socket_result;
  let mut tries: u64 = 0;
  loop {
    udp_socket_result = bind_udp_socket(address);
    if first_startup || udp_socket_result.is_ok() {
      break;
    }
    tries += 1;
    let duration = Duration::from_millis(1000);
    if tries >= 10 {
      println!("HTTP/3 port is used at try #{tries}, skipping...");
      listen_error_tx.send(None).await.unwrap_or_default();
      break;
    }
    println!("HTTP/3 port is used at try #{tries}, retrying in {duration:?}...");
    if shutdown_rx.is_cancelled() {
      break;
    }
    tokio::time::sleep(duration).await;
  }
  let udp_socket = match udp_socket_result {
    Ok(socket) => socket,
    Err(err) => Err(anyhow::anyhow!("Cannot listen to HTTP/3 port: {err}"))?,
  };
  let endpoint = match quinn::Endpoint::new(quinn::EndpointConfig::default(), Some(server_config), udp_socket, {
    #[cfg(any(feature = "runtime-vibeio", feature = "runtime-monoio"))]
    let runtime = Arc::new(EnterTokioRuntime);
    #[cfg(feature = "runtime-tokio")]
    let runtime = Arc::new(quinn::TokioRuntime);

    runtime
  }) {
    Ok(endpoint) => endpoint,
    Err(err) => Err(anyhow::anyhow!("Cannot listen to HTTP/3 port: {err}"))?,
  };
  println!("HTTP/3 server is listening on {address}...");
  listen_error_tx.send(None).await.unwrap_or_default();

  loop {
    let rustls_receive_future = async { rustls_config_rx.recv().await.ok() };

    let connection = tokio::select! {
      result = endpoint.accept() => {
        match result {
          Some(conn) => conn,
          None => {
            log_accept_closed(&logging_tx).await;
            break;
          }
        }
      }
      tls_config = rustls_receive_future => {
        let Some(tls_config) = tls_config else {
          futures_util::future::pending::<()>().await;
          unreachable!();
        };

        if let Ok(server_config) = build_quic_server_config(tls_config) {
          endpoint.set_server_config(Some(server_config));
        }
        continue;
      }
      _ = shutdown_rx.cancelled() => {
        break;
      }
    };
    let remote_address = connection.remote_address();
    let local_address = SocketAddr::new(
      connection.local_ip().unwrap_or(IpAddr::V6(Ipv6Addr::UNSPECIFIED)),
      udp_port,
    );
    let quic_data = ConnectionData {
      connection: Connection::Quic(connection),
      client_address: remote_address,
      server_address: local_address,
    };
    let quic_tx = tx.clone();
    tokio::spawn(async move {
      quic_tx.send(quic_data).await.unwrap_or_default();
    });
  }

  endpoint.wait_idle().await;

  Ok(())
}