Supporting two-phase Listeners

[SabrinaJewson]

• I have looked for existing issues (including closed) about this

Feature Request

Motivation

Some kinds of listener, like TLS listeners, will want to perform work after the connection has been accepted, but before HTTP requests can actually be served. This logic can be moved into the listener itself via complex multiplexing logic – for example, the TlsListener crate uses FuturesUnordered to do this – however I think that ideally it would be performed on the same task that serves the connection.

Proposal

Modify the Listener trait as follows:

pub trait Listener: Send + 'static {
    type Io: AsyncRead + AsyncWrite + Unpin + Send + 'static;
    type Addr: Send;
    fn accept(&mut self) -> impl Future<Output = impl Future<Output = Option<(Self::Io, Self::Addr)>> + Send + 'static> + Send;
    fn local_addr(&self) -> io::Result<Self::Addr>;
}

The outer future of accept will resolve when a connection has been accepted, while the inner future runs on the spawned task. Option<(Self::Io, Self::Addr)> is used to account for errors: if None is returned, then the task is aborted.

Alternatives

• Use a type Error: Into<Box<dyn Error + Send + Sync>>; and Result<(Self::Io, Self::Addr), Self::Error> instead of Option. This might be more convenient for the user, as they would have to handle the errors less themselves.
  • A significant disadvantage of this is it makes composability much more difficult, as now errors have to compose. So I’m leaning toward Option.
• Say that users should define their own “lazy TLS stream” type to handle this case. It would be an enum over the handshake future and the underlying TLS stream, and its implementation of AsyncRead/AsyncWrite would first drive the handshake to completion before performing I/O (or report EOF should the handshake fail). I can’t particularly think of practical reasons to avoid this other than “it feels weird”.
• Say that users who want TLS in this way should not use axum::serve. That function is, however, a very convenient abstraction, taking care of graceful shutdown and constructing the server::conn::auto::Builder with the right configuration.

[SabrinaJewson]

To prove the utility of this API, here’s pseudocode for what a TLS listener API might look like:

TLS listener pseudocode

pub trait TlsAcceptor<C> {
    fn begin_handshake(&mut self) -> impl TlsHandshake<C>;
}

pub trait TlsHandshake<C>: Send + 'static {
    type Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static;
    fn run(self, stream: C) -> impl Future<Output = Option<Self::Stream>> + Send;
}

pub fn tls_handshake<F, Fut, C, S>(f: F) -> impl TlsHandshake<C>
where
    F: FnOnce(C) -> Fut + Send + 'static,
    Fut: Future<Output = Option<S>> + Send + 'static,
    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{ /* … */ }

// Errors from the TLS handshake are logged and then discarded
impl<C> TlsAcceptor<C> for tokio_rustls::TlsAcceptor { /* … */ }
impl<C> TlsAcceptor<C> for tokio_native_tls::TlsAcceptor { /* … */ }

#[non_exhaustive]
pub struct TlsListener<L, A> {
    pub listener: L,
    pub acceptor: A,
    pub timeout: Duration,
}

impl<L, A> TlsListener<L, A> {
    // choose a sensible default timeout
    pub fn new(listener: L, acceptor: A) -> Self { /* … */ }
    pub fn with_timeout(self, timeout: Duration) -> Self { /* … */ }
    pub fn with_replacements<F>(self, replacements: F) -> TlsListener<L, WithReplacements<A, F>>
    where
        F: FnMut() -> Option<A>,
    { /* … */ }
}

impl<L: Listener, A: TlsAcceptor<L::Io>> Listener for TlsListener<L, A> {
    type Io = A::Stream;
    type Addr = L::Addr;
    async fn accept(&mut self) -> impl Future<Output = Option<(Self::Io, Self::Addr)>> + Send + 'static {
        let future = self.listener.accept().await;
        let handshake = self.acceptor.begin_handshake();
        let timeout = self.timeout;
        async move {
            let (io, addr) = future.await?;
            tokio::time::timeout(handshake.run(io), timeout)
                .await
                .unwrap_or_else(|_| {
                    error!("TLS handshake timeout after {timeout:?}");
                    None
                })
                .map(|io| (io, addr))
        }
    }
}

// An example combinator that allows dynamically replacing TLS certificates.
pub struct WithReplacements<A, F> {
    pub current: A,
    pub replacements: F,
}

impl<A, F, C> TlsAcceptor<C> for WithReplacements<A, F>
where
    A: TlsAcceptor<C>,
    F: FnMut() -> Option<A>,
{
    fn begin_handshake(&mut self) -> impl TlsHandshake<C> {
        if let Some(replacement) = (self.replacements)() {
            self.current = replacement;
        }
        self.current.begin_handshake()
    }
}

[mladedav]

Personally, I don't like futures returning futures as it's easy to misuse and explaining well why there are two futures and what to put in which might be hard.

For example here someone could do

async fn accept(...) {
  std::future::ready(async { ... })
}

instead of

async fn accept(...) {
  async { std::future::ready(...) }
}

where the first version would spawn an infinite number of tasks, each trying to accept a new connection. It's a minor oversight with potentially big consequences.

I think that in most cases people should do what the TLS listener crate is doing or writing a wrapper like you described in the second alternative.

We could maybe provide helpers for either or both of those cases.

[SabrinaJewson]

Personally, I don't like futures returning futures as it's easy to misuse and explaining well why there are two futures and what to put in which might be hard.

We could alternatively have an explicit trait implementation:

pub trait Listener: Send + 'static {
    type Io: AsyncRead + AsyncWrite + Unpin + Send + 'static;
    type Addr: Send;
    fn accept(&mut self) -> impl Future<Output = impl Handshake<Io = Self::Io, Addr = Self::Addr>> + Send;
    fn local_addr(&self) -> io::Result<Self::Addr>;
}
pub trait Handshake: Send + 'static {
    type Io: AsyncRead + AsyncWrite + Unpin + Send + 'static;
    type Addr: Send;
    pub fn handshake(self) -> impl Future<Output = Option<(Self::Io, Self::Addr)>> + Send;
}

thus making it more clear what the purpose of the second handshake is. But I’m also not massively concerned about footguns, since I wouldn’t expect this trait to be implemented by users to begin with – I would assume the target audience is very much people who know what they’re doing and could easily diagnose an issue like this.

I think that in most cases people should do what the TLS listener crate is doing or writing a wrapper like you described in the second alternative.

The TlsListener approach feels very much like a hack to me – I guess I just find it weird to be forcing an inherently parallel operation through a serial interface. The lazy wrapper approach is something that feels like a very creative use of the API – so I guess I would like confirmation from someone knowledgable that it’s a somewhat sane thing to do.