Skip to main content

microvms_edges/session/
serve.rs

1// SPDX-License-Identifier: Apache-2.0
2//! The serving loops behind `microvm tunnel` and `microvm port-forward`, and behind the SDKs'
3//! tunnel and port-forward handles (#263).
4//!
5//! Each loop accepts a local connection, serves it on a task of its own, and counts how it
6//! ended. The three things the CLI used to own are parameters: a stop future (the CLI's Ctrl-C,
7//! a handle's `stop()`), a limit on how many connections to accept, and a callback that hears
8//! about each connection's end, so a caller can print a warning or keep a list. A connection
9//! that fails never ends the loop; only the stop, the limit, or a listener that stops accepting
10//! does. The report's counts are computed here, once, for every surface.
11//!
12//! # Every connection gets a task, both loops alike
13//!
14//! One WebSocket per tunnel connection matches the daemon's no-multiplexing decision, and one
15//! task per forwarded connection means a slow request (a long poll, a first compile) doesn't
16//! hold the browser's next one behind it, which the CLI's port-forward used to do. When the
17//! loop stops accepting, it waits for the connections already open to end: a task still
18//! relaying holds bytes the local client is waiting for, and dropping it would cut them short.
19//! Only an SDK handle cuts them, when its caller's `stop` gives a grace that runs out or the
20//! handle is dropped.
21//!
22//! # A direct session serves too
23//!
24//! A session with no proxy credential reaches the daemon without the endpoint proxy. The
25//! tunnel then offers no proxy subprotocols, and the daemon's bearer check is the gate; the
26//! forwarder mints no proxy headers and sends each request to the endpoint's host at the guest
27//! port ([`super::forward::direct_url`]), since there is no proxy to route it by header.
28
29use std::future::Future;
30use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
31use std::sync::{Arc, Mutex, PoisonError};
32use std::time::Duration;
33
34use tokio::net::TcpListener;
35use tokio::sync::{oneshot, watch};
36use tokio::task::JoinSet;
37
38use microvms_app::error::{Error, ErrorKind};
39use microvms_app::identity::TunnelIdentity;
40use microvms_app::session::Session;
41use microvms_app::session::proxy::ProxyAuth;
42
43use super::forward::{ForwardClient, ForwardEvent, ForwardSpec, serve_connection_via};
44use super::tunnel::{TunnelEnd, relay_connection_inner};
45
46/// When a loop stops on its own, besides its caller's stop.
47#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
48pub struct ServeLimits {
49    /// Stop accepting after this many connections. `None` serves until stopped.
50    pub max_connections: Option<u32>,
51}
52
53/// Why a loop stopped accepting.
54#[derive(Clone, Debug, Eq, PartialEq)]
55pub enum StopReason {
56    /// The caller's stop future resolved.
57    Requested,
58    /// [`ServeLimits::max_connections`] connections were accepted.
59    Limit,
60    /// The listener stopped accepting, with why. A listener's failure isn't one connection's,
61    /// so the loop ends rather than retrying a socket that will keep failing.
62    ListenerFailed(String),
63}
64
65impl StopReason {
66    /// Whether the caller stopped the loop, which the CLI reports as `interrupted`: a stop is
67    /// a success, and a loop that ended at its limit or on its listener wasn't interrupted.
68    pub fn was_requested(&self) -> bool {
69        matches!(self, StopReason::Requested)
70    }
71}
72
73/// What a tunnel reaches: the daemon, the guest port, and the credentials to present.
74#[derive(Clone)]
75pub struct TunnelTarget {
76    pub endpoint: String,
77    pub agent_token: String,
78    /// The session's proxy credential; `None` for a direct session.
79    pub auth: Option<Arc<ProxyAuth>>,
80    pub guest_port: u16,
81    /// The pair to verify the far end with (`--verify-identity`), when asked for.
82    pub identity: Option<TunnelIdentity>,
83}
84
85impl TunnelTarget {
86    /// The target for `guest_port` over `session`'s endpoint and credentials.
87    pub fn for_session(
88        session: &Session,
89        guest_port: u16,
90        identity: Option<TunnelIdentity>,
91    ) -> Self {
92        Self {
93            endpoint: session.endpoint().to_string(),
94            agent_token: session.agent_token().to_string(),
95            auth: session.proxy_auth().cloned(),
96            guest_port,
97            identity,
98        }
99    }
100}
101
102/// How one tunnel connection ended.
103#[derive(Debug)]
104pub struct TunnelConnection {
105    pub peer: SocketAddr,
106    pub end: Result<TunnelEnd, Error>,
107}
108
109/// What a tunnel loop did.
110#[derive(Clone, Debug, Eq, PartialEq)]
111pub struct TunnelReport {
112    /// Connections accepted.
113    pub served: u32,
114    /// Connections the daemon refused, or that failed with an error.
115    pub refused: u32,
116    /// Verified connections that ended without the daemon's end of stream (#342).
117    pub truncated: u32,
118    /// Verified connections whose daemon predates the end of stream (#342).
119    pub unproven: u32,
120    /// Proxy tokens the session minted by the time the loop ended.
121    pub proxy_token_mints: u64,
122    pub stopped: StopReason,
123}
124
125impl TunnelReport {
126    /// Counts one connection's end in its column. A clean end is only served.
127    fn count(&mut self, end: &Result<TunnelEnd, Error>) {
128        match end {
129            Ok(TunnelEnd::Closed) => {}
130            Ok(TunnelEnd::ClosedUnproven) => self.unproven += 1,
131            Ok(TunnelEnd::Truncated { .. }) => self.truncated += 1,
132            Ok(TunnelEnd::Refused { .. }) | Err(_) => self.refused += 1,
133        }
134    }
135}
136
137/// Serves `listener` as a TCP tunnel to `target` until `stop` resolves, the limit is reached,
138/// or the listener fails, then waits for the open connections to end.
139///
140/// `on_end` hears about each connection as it ends, on the loop's own task.
141pub async fn serve_tunnel<F>(
142    listener: TcpListener,
143    target: TunnelTarget,
144    limits: ServeLimits,
145    stop: impl Future<Output = ()>,
146    on_end: F,
147) -> TunnelReport
148where
149    F: FnMut(TunnelConnection),
150{
151    tunnel_loop(
152        listener,
153        target,
154        limits,
155        stop,
156        std::future::pending(),
157        on_end,
158    )
159    .await
160}
161
162/// [`serve_tunnel`], cutting the connections still open once `cut` resolves after the stop.
163async fn tunnel_loop<F>(
164    listener: TcpListener,
165    target: TunnelTarget,
166    limits: ServeLimits,
167    stop: impl Future<Output = ()>,
168    cut: impl Future<Output = ()>,
169    mut on_end: F,
170) -> TunnelReport
171where
172    F: FnMut(TunnelConnection),
173{
174    let target = Arc::new(target);
175    let mut report = TunnelReport {
176        served: 0,
177        refused: 0,
178        truncated: 0,
179        unproven: 0,
180        proxy_token_mints: 0,
181        stopped: StopReason::Requested,
182    };
183    let (stopped, served) = serve(
184        &listener,
185        limits,
186        stop,
187        cut,
188        |stream, _| {
189            let target = Arc::clone(&target);
190            async move {
191                relay_connection_inner(
192                    stream,
193                    &target.endpoint,
194                    target.guest_port,
195                    &target.agent_token,
196                    target.auth.as_ref(),
197                    target.identity.as_ref(),
198                )
199                .await
200            }
201        },
202        |peer, end| {
203            report.count(&end);
204            on_end(TunnelConnection { peer, end });
205        },
206    )
207    .await;
208    report.served = served;
209    report.stopped = stopped;
210    report.proxy_token_mints = target.auth.as_ref().map_or(0, |auth| auth.mint_count());
211    report
212}
213
214/// Something one forwarded connection did, for a caller's warnings.
215#[derive(Debug)]
216pub enum ForwardNotice {
217    /// An event the forwarder reported for the connection from `peer`.
218    Event {
219        peer: SocketAddr,
220        event: ForwardEvent,
221    },
222    /// The connection from `peer` ended early, with why.
223    Failed { peer: SocketAddr, error: Error },
224}
225
226/// What a port-forward loop did.
227#[derive(Clone, Debug, Eq, PartialEq)]
228pub struct ForwardReport {
229    /// Connections accepted.
230    pub served: u32,
231    /// Exchanges the proxy refused ([`ForwardEvent::Refused`]).
232    pub refused: u32,
233    /// Exchanges that upgraded, a WebSocket among them.
234    pub upgrades: u32,
235    /// Proxy tokens the session minted by the time the loop ended.
236    pub proxy_token_mints: u64,
237    pub stopped: StopReason,
238}
239
240impl ForwardReport {
241    /// Counts one of a connection's events in its column.
242    fn count(&mut self, event: &ForwardEvent) {
243        match event {
244            ForwardEvent::Refused { .. } => self.refused += 1,
245            ForwardEvent::Forwarded { upgraded: true, .. } => self.upgrades += 1,
246            _ => {}
247        }
248    }
249}
250
251/// Serves `listener` as an HTTP and WebSocket forward to `spec`'s guest port until `stop`
252/// resolves, the limit is reached, or the listener fails, then waits for the open connections
253/// to end.
254///
255/// `auth` is the session's proxy credential, `None` for a direct session. `on_notice` hears
256/// about each connection's events and failure as it ends, on the loop's own task.
257pub async fn serve_forward<F>(
258    listener: TcpListener,
259    spec: ForwardSpec,
260    auth: Option<Arc<ProxyAuth>>,
261    limits: ServeLimits,
262    stop: impl Future<Output = ()>,
263    on_notice: F,
264) -> Result<ForwardReport, Error>
265where
266    F: FnMut(ForwardNotice),
267{
268    forward_loop(
269        listener,
270        spec,
271        auth,
272        limits,
273        stop,
274        std::future::pending(),
275        on_notice,
276    )
277    .await
278}
279
280/// [`serve_forward`], cutting the connections still open once `cut` resolves after the stop.
281async fn forward_loop<F>(
282    listener: TcpListener,
283    spec: ForwardSpec,
284    auth: Option<Arc<ProxyAuth>>,
285    limits: ServeLimits,
286    stop: impl Future<Output = ()>,
287    cut: impl Future<Output = ()>,
288    mut on_notice: F,
289) -> Result<ForwardReport, Error>
290where
291    F: FnMut(ForwardNotice),
292{
293    // One client for the whole forward: connection reuse to the endpoint is what keeps a
294    // page-load of thirty assets from paying thirty TLS handshakes.
295    let client = Arc::new(ForwardClient::new()?);
296    let spec = Arc::new(spec);
297    let mut report = ForwardReport {
298        served: 0,
299        refused: 0,
300        upgrades: 0,
301        proxy_token_mints: 0,
302        stopped: StopReason::Requested,
303    };
304    let (stopped, served) = serve(
305        &listener,
306        limits,
307        stop,
308        cut,
309        |stream, _| {
310            let (spec, auth, client) = (Arc::clone(&spec), auth.clone(), Arc::clone(&client));
311            async move {
312                let mut events = Vec::new();
313                let outcome =
314                    serve_connection_via(stream, &spec, auth.as_ref(), &client, |event| {
315                        events.push(event);
316                    })
317                    .await;
318                outcome.map(|()| events)
319            }
320        },
321        |peer, done| {
322            let (events, failure) = match done {
323                Ok(events) => (events, None),
324                Err(error) => (Vec::new(), Some(error)),
325            };
326            for event in events {
327                report.count(&event);
328                on_notice(ForwardNotice::Event { peer, event });
329            }
330            if let Some(error) = failure {
331                on_notice(ForwardNotice::Failed { peer, error });
332            }
333        },
334    )
335    .await;
336    report.served = served;
337    report.stopped = stopped;
338    report.proxy_token_mints = auth.as_ref().map_or(0, |auth| auth.mint_count());
339    Ok(report)
340}
341
342// ── the SDKs' handles ────────────────────────────────────────────────────────
343
344/// Where an SDK's tunnel or forward listens unless told otherwise: loopback, on a port the OS
345/// picks. Loopback because a forward carries the VM's credentials to whoever connects.
346pub const DEFAULT_BIND: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0));
347
348/// A local address to listen on: `bind` when given, [`DEFAULT_BIND`] otherwise.
349pub fn bind_address(bind: Option<&str>) -> Result<SocketAddr, Error> {
350    match bind {
351        None => Ok(DEFAULT_BIND),
352        Some(bind) => bind.parse().map_err(|err| {
353            Error::new(
354                ErrorKind::InvalidArg,
355                format!("{bind:?} is not a local address to listen on, such as 127.0.0.1:0: {err}"),
356            )
357        }),
358    }
359}
360
361/// How a connection that didn't end clean ended, as a handle's report lists it.
362#[derive(Clone, Copy, Debug, Eq, PartialEq)]
363pub enum EndKind {
364    /// The daemon refused the tunnel, or the proxy refused the exchange.
365    Refused,
366    /// A verified tunnel ended without the daemon's end of stream (#342).
367    Truncated,
368    /// A verified tunnel ended into a daemon from before the end of stream (#342).
369    Unproven,
370    /// The connection failed with an error.
371    Failed,
372}
373
374impl EndKind {
375    /// The wire spelling the bindings report.
376    pub fn as_str(self) -> &'static str {
377        match self {
378            EndKind::Refused => "refused",
379            EndKind::Truncated => "truncated",
380            EndKind::Unproven => "unproven",
381            EndKind::Failed => "failed",
382        }
383    }
384}
385
386/// One connection that didn't end clean. A clean end is only counted, so a long-lived handle's
387/// report grows with its failures, not with its traffic.
388#[derive(Clone, Debug, Eq, PartialEq)]
389pub struct ConnectionEnd {
390    pub peer: SocketAddr,
391    pub kind: EndKind,
392    /// The close code or the HTTP status, when the end carried one.
393    pub code: Option<u16>,
394    /// The daemon's reason, the forwarder's explanation, or the error.
395    pub detail: String,
396}
397
398impl ConnectionEnd {
399    /// A tunnel connection's end, unless it was clean.
400    pub fn of_tunnel(connection: &TunnelConnection) -> Option<Self> {
401        let (kind, code, detail) = match &connection.end {
402            Ok(TunnelEnd::Closed) => return None,
403            Ok(TunnelEnd::ClosedUnproven) => (EndKind::Unproven, None, String::new()),
404            Ok(TunnelEnd::Truncated { code }) => (EndKind::Truncated, *code, String::new()),
405            Ok(TunnelEnd::Refused { code, reason }) => {
406                (EndKind::Refused, Some(*code), reason.clone())
407            }
408            Err(error) => (EndKind::Failed, None, error.to_string()),
409        };
410        Some(Self {
411            peer: connection.peer,
412            kind,
413            code,
414            detail,
415        })
416    }
417
418    /// A forwarded connection's refusal or failure; its other events aren't ends worth listing.
419    pub fn of_forward(notice: &ForwardNotice) -> Option<Self> {
420        match notice {
421            ForwardNotice::Event {
422                peer,
423                event:
424                    ForwardEvent::Refused {
425                        status,
426                        explanation,
427                    },
428            } => Some(Self {
429                peer: *peer,
430                kind: EndKind::Refused,
431                code: Some(*status),
432                detail: explanation.clone(),
433            }),
434            ForwardNotice::Event { .. } => None,
435            ForwardNotice::Failed { peer, error } => Some(Self {
436                peer: *peer,
437                kind: EndKind::Failed,
438                code: None,
439                detail: error.to_string(),
440            }),
441        }
442    }
443}
444
445/// What a tunnel handle did: the loop's report, and each connection that didn't end clean.
446#[derive(Clone, Debug, Eq, PartialEq)]
447pub struct TunnelSummary {
448    pub report: TunnelReport,
449    pub ended: Vec<ConnectionEnd>,
450}
451
452/// What a port-forward handle did: the loop's report, and each connection that didn't end
453/// clean.
454#[derive(Clone, Debug, Eq, PartialEq)]
455pub struct ForwardSummary {
456    pub report: ForwardReport,
457    pub ended: Vec<ConnectionEnd>,
458}
459
460type Outcome<S> = Result<S, (ErrorKind, String)>;
461
462/// A serving loop on a background task: the SDKs' tunnel and port-forward handles.
463///
464/// Dropping it stops the loop and cuts the connections still open, so no relay outlives the
465/// handle that owns it: a relay carries the VM's credentials to whoever connected.
466pub struct Serving<S> {
467    local: SocketAddr,
468    stop: Mutex<Option<oneshot::Sender<()>>>,
469    cut: Mutex<Option<oneshot::Sender<()>>>,
470    outcome: Arc<Mutex<Option<Outcome<S>>>>,
471    done: watch::Receiver<bool>,
472}
473
474impl<S: Clone + Send + 'static> Serving<S> {
475    /// Runs `serving` on the current tokio runtime, handing it the stop and the cut futures.
476    /// Panics outside a runtime, like `tokio::spawn`.
477    fn spawn<Fut>(
478        local: SocketAddr,
479        serving: impl FnOnce(oneshot::Receiver<()>, oneshot::Receiver<()>) -> Fut,
480    ) -> Self
481    where
482        Fut: Future<Output = Result<S, Error>> + Send + 'static,
483    {
484        let (stop_tx, stop_rx) = oneshot::channel();
485        let (cut_tx, cut_rx) = oneshot::channel();
486        let (done_tx, done_rx) = watch::channel(false);
487        let outcome = Arc::new(Mutex::new(None));
488        let written = Arc::clone(&outcome);
489        let serving = serving(stop_rx, cut_rx);
490        tokio::spawn(async move {
491            let result = serving
492                .await
493                .map_err(|error| (error.kind(), error.to_string()));
494            *written.lock().unwrap_or_else(PoisonError::into_inner) = Some(result);
495            done_tx.send_replace(true);
496        });
497        Self {
498            local,
499            stop: Mutex::new(Some(stop_tx)),
500            cut: Mutex::new(Some(cut_tx)),
501            outcome,
502            done: done_rx,
503        }
504    }
505
506    /// The address the loop listens on, with the port the OS picked for a port 0.
507    pub fn local_address(&self) -> SocketAddr {
508        self.local
509    }
510
511    /// Whether the loop is still serving, or waiting for its open connections to end.
512    pub fn is_running(&self) -> bool {
513        !*self.done.borrow()
514    }
515
516    /// Asks the loop to stop accepting. Idempotent; [`Self::finished`] waits for it, and for the
517    /// connections still open.
518    pub fn request_stop(&self) {
519        signal(&self.stop);
520    }
521
522    /// Stops accepting and waits for the connections still open to end, cutting the ones
523    /// still open after `grace` when it's given, and returns what the loop did.
524    ///
525    /// Without a grace, a client that keeps its connection open keeps this waiting: the
526    /// relay holds bytes that client may still be reading, so the loop won't cut it unasked.
527    pub async fn stop(&self, grace: Option<Duration>) -> Result<S, Error> {
528        self.request_stop();
529        if let Some(grace) = grace {
530            if let Ok(outcome) = tokio::time::timeout(grace, self.finished()).await {
531                return outcome;
532            }
533            signal(&self.cut);
534        }
535        self.finished().await
536    }
537
538    /// Waits for the loop to end and returns what it did. Callable any number of times.
539    pub async fn finished(&self) -> Result<S, Error> {
540        let mut done = self.done.clone();
541        // The sender drops only after writing the outcome, so an error here still means done.
542        let _ = done.wait_for(|finished| *finished).await;
543        self.outcome
544            .lock()
545            .unwrap_or_else(PoisonError::into_inner)
546            .clone()
547            .expect("the loop records its outcome before signalling done")
548            .map_err(|(kind, message)| Error::new(kind, message))
549    }
550}
551
552/// Sends a handle's stop or cut, once.
553fn signal(sender: &Mutex<Option<oneshot::Sender<()>>>) {
554    if let Some(sender) = sender.lock().unwrap_or_else(PoisonError::into_inner).take() {
555        let _ = sender.send(());
556    }
557}
558
559/// A handle's stop or cut future: the sender's `send`, or its drop.
560async fn signalled(signal: oneshot::Receiver<()>) {
561    let _ = signal.await;
562}
563
564/// Binds `bind` and starts a tunnel to `target` on the current tokio runtime: the SDKs'
565/// `Session.tunnel`.
566pub async fn start_tunnel(
567    bind: SocketAddr,
568    target: TunnelTarget,
569    limits: ServeLimits,
570) -> Result<Serving<TunnelSummary>, Error> {
571    let listener =
572        super::forward::bind(&ForwardSpec::new(bind, target.guest_port, &target.endpoint)).await?;
573    let local = local_address(&listener)?;
574    Ok(Serving::spawn(local, move |stop, cut| async move {
575        let mut ended = Vec::new();
576        let report = tunnel_loop(
577            listener,
578            target,
579            limits,
580            signalled(stop),
581            signalled(cut),
582            |connection| ended.extend(ConnectionEnd::of_tunnel(&connection)),
583        )
584        .await;
585        Ok(TunnelSummary { report, ended })
586    }))
587}
588
589/// Binds `bind` and starts a port-forward to `guest_port` behind `endpoint` on the current
590/// tokio runtime: the SDKs' `Session.port_forward`. `auth` is `None` for a direct session.
591pub async fn start_forward(
592    bind: SocketAddr,
593    endpoint: &str,
594    guest_port: u16,
595    auth: Option<Arc<ProxyAuth>>,
596    limits: ServeLimits,
597) -> Result<Serving<ForwardSummary>, Error> {
598    let spec = ForwardSpec::new(bind, guest_port, endpoint);
599    let listener = super::forward::bind(&spec).await?;
600    let local = local_address(&listener)?;
601    Ok(Serving::spawn(local, move |stop, cut| async move {
602        let mut ended = Vec::new();
603        let report = forward_loop(
604            listener,
605            spec,
606            auth,
607            limits,
608            signalled(stop),
609            signalled(cut),
610            |notice| ended.extend(ConnectionEnd::of_forward(&notice)),
611        )
612        .await?;
613        Ok(ForwardSummary { report, ended })
614    }))
615}
616
617fn local_address(listener: &TcpListener) -> Result<SocketAddr, Error> {
618    listener.local_addr().map_err(|err| {
619        Error::new(
620            ErrorKind::Unexpected,
621            format!("the listener has no address: {err}"),
622        )
623    })
624}
625
626/// The accept loop both serving loops share: accept until stopped, a task per connection, then
627/// wait for the open ones, cutting them if `cut` resolves first. Answers why it stopped and how
628/// many connections it accepted.
629///
630/// `spawn` makes a connection's task. `ended` takes each connection's result as its task ends,
631/// on this task; a task that panicked or was cut reads as an error.
632async fn serve<T, Fut>(
633    listener: &TcpListener,
634    limits: ServeLimits,
635    stop: impl Future<Output = ()>,
636    cut: impl Future<Output = ()>,
637    mut spawn: impl FnMut(tokio::net::TcpStream, SocketAddr) -> Fut,
638    mut ended: impl FnMut(SocketAddr, Result<T, Error>),
639) -> (StopReason, u32)
640where
641    T: Send + 'static,
642    Fut: Future<Output = Result<T, Error>> + Send + 'static,
643{
644    let mut stop = std::pin::pin!(stop);
645    let mut tasks = JoinSet::new();
646    // Peers by task, so a task that panics is still reported against its connection.
647    let mut peers = Peers::new();
648    type Peers = std::collections::HashMap<tokio::task::Id, SocketAddr>;
649    let mut finish =
650        |peers: &mut Peers,
651         joined: Result<(tokio::task::Id, Result<T, Error>), tokio::task::JoinError>| {
652            let (id, result) = match joined {
653                Ok((id, result)) => (id, result),
654                Err(error) if error.is_cancelled() => (
655                    error.id(),
656                    Err(Error::new(
657                        ErrorKind::Unexpected,
658                        "the connection was cut when its loop stopped",
659                    )),
660                ),
661                Err(error) => (
662                    error.id(),
663                    Err(Error::new(
664                        ErrorKind::Unexpected,
665                        format!("the connection's task ended abnormally: {error}"),
666                    )),
667                ),
668            };
669            let peer = peers
670                .remove(&id)
671                .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
672            ended(peer, result);
673        };
674    let mut served = 0_u32;
675    let stopped = loop {
676        if limits.max_connections.is_some_and(|max| served >= max) {
677            break StopReason::Limit;
678        }
679        tokio::select! {
680            // Biased so a stop wins over a connection that arrived in the same wakeup: the
681            // caller asked to stop, and serving one more first looks like the stop did nothing.
682            biased;
683            () = &mut stop => break StopReason::Requested,
684            Some(joined) = tasks.join_next_with_id(), if !tasks.is_empty() => finish(&mut peers, joined),
685            next = listener.accept() => match next {
686                Ok((stream, peer)) => {
687                    served += 1;
688                    let handle = tasks.spawn(spawn(stream, peer));
689                    peers.insert(handle.id(), peer);
690                }
691                Err(error) => {
692                    break StopReason::ListenerFailed(format!(
693                        "the local listener stopped accepting: {error}"
694                    ));
695                }
696            },
697        }
698    };
699    // Drained rather than abandoned (see the module docs), until the caller cuts what's left.
700    let mut cut = std::pin::pin!(cut);
701    let mut cutting = false;
702    loop {
703        tokio::select! {
704            joined = tasks.join_next_with_id() => match joined {
705                Some(joined) => finish(&mut peers, joined),
706                None => break,
707            },
708            () = &mut cut, if !cutting => {
709                cutting = true;
710                tasks.abort_all();
711            }
712        }
713    }
714    (stopped, served)
715}
716
717#[cfg(test)]
718mod tests {
719    use super::*;
720
721    fn tunnel_report() -> TunnelReport {
722        TunnelReport {
723            served: 0,
724            refused: 0,
725            truncated: 0,
726            unproven: 0,
727            proxy_token_mints: 0,
728            stopped: StopReason::Requested,
729        }
730    }
731
732    /// Each way a tunnel connection ends lands in its own column, once.
733    #[test]
734    fn each_tunnel_end_is_counted_once_in_its_column() {
735        let mut report = tunnel_report();
736        for end in [
737            Ok(TunnelEnd::Closed),
738            Ok(TunnelEnd::ClosedUnproven),
739            Ok(TunnelEnd::ClosedUnproven),
740            Ok(TunnelEnd::Truncated { code: None }),
741            Ok(TunnelEnd::Truncated { code: Some(1000) }),
742            Ok(TunnelEnd::Truncated { code: Some(1000) }),
743            Ok(TunnelEnd::Refused {
744                code: 4502,
745                reason: String::new(),
746            }),
747            Err(Error::new(ErrorKind::Unexpected, "a transport error")),
748        ] {
749            report.count(&end);
750        }
751        assert_eq!(
752            (report.refused, report.truncated, report.unproven),
753            (2, 3, 2),
754            "{report:?}"
755        );
756    }
757
758    /// Only the caller's own stop reads as requested.
759    #[test]
760    fn only_the_callers_stop_was_requested() {
761        assert!(StopReason::Requested.was_requested());
762        assert!(!StopReason::Limit.was_requested());
763        assert!(!StopReason::ListenerFailed("gone".to_string()).was_requested());
764    }
765
766    /// A refusal and an upgrade each land in their column; other events count nowhere.
767    #[test]
768    fn each_forward_event_is_counted_once_in_its_column() {
769        let mut report = ForwardReport {
770            served: 0,
771            refused: 0,
772            upgrades: 0,
773            proxy_token_mints: 0,
774            stopped: StopReason::Requested,
775        };
776        for event in [
777            ForwardEvent::Refused {
778                status: 502,
779                explanation: String::new(),
780            },
781            ForwardEvent::Forwarded {
782                status: 101,
783                upgraded: true,
784            },
785            ForwardEvent::Forwarded {
786                status: 101,
787                upgraded: true,
788            },
789            ForwardEvent::Forwarded {
790                status: 200,
791                upgraded: false,
792            },
793            ForwardEvent::ConnectionError {
794                detail: String::new(),
795            },
796        ] {
797            report.count(&event);
798        }
799        assert_eq!((report.refused, report.upgrades), (1, 2), "{report:?}");
800    }
801
802    /// Each end but a clean one is listed with its kind, its code and its reason; a forward's
803    /// events other than a refusal aren't ends.
804    #[test]
805    fn each_end_but_a_clean_one_is_listed_with_its_kind() {
806        let peer: SocketAddr = "127.0.0.1:5000".parse().expect("an address");
807        let tunnel = |end| ConnectionEnd::of_tunnel(&TunnelConnection { peer, end });
808        let listed = |end: Option<ConnectionEnd>| {
809            end.map(|end| (end.peer, end.kind.as_str(), end.code, end.detail))
810        };
811        assert_eq!(tunnel(Ok(TunnelEnd::Closed)), None);
812        assert_eq!(
813            listed(tunnel(Ok(TunnelEnd::ClosedUnproven))),
814            Some((peer, "unproven", None, String::new()))
815        );
816        assert_eq!(
817            listed(tunnel(Ok(TunnelEnd::Truncated { code: Some(1006) }))),
818            Some((peer, "truncated", Some(1006), String::new()))
819        );
820        assert_eq!(
821            listed(tunnel(Ok(TunnelEnd::Refused {
822                code: 4502,
823                reason: "nothing is listening".into(),
824            }))),
825            Some((peer, "refused", Some(4502), "nothing is listening".into()))
826        );
827        let failed = listed(tunnel(Err(Error::new(ErrorKind::Unexpected, "a reset"))));
828        assert!(
829            matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
830            "{failed:?}"
831        );
832
833        let forward = |event| ConnectionEnd::of_forward(&ForwardNotice::Event { peer, event });
834        assert_eq!(
835            listed(forward(ForwardEvent::Refused {
836                status: 403,
837                explanation: "out of scope".into(),
838            })),
839            Some((peer, "refused", Some(403), "out of scope".into()))
840        );
841        assert_eq!(
842            forward(ForwardEvent::Forwarded {
843                status: 200,
844                upgraded: false,
845            }),
846            None
847        );
848        let failed = listed(ConnectionEnd::of_forward(&ForwardNotice::Failed {
849            peer,
850            error: Error::new(ErrorKind::Unexpected, "a reset"),
851        }));
852        assert!(
853            matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
854            "{failed:?}"
855        );
856    }
857
858    /// A connection's task that panics reads as ended abnormally, and one the caller cuts
859    /// reads as cut: the two are the caller's to tell apart.
860    #[tokio::test]
861    async fn a_panicked_connection_and_a_cut_one_are_told_apart() {
862        let listener = TcpListener::bind("127.0.0.1:0").await.expect("a free port");
863        let local = listener.local_addr().expect("bound");
864        let (cut, cut_rx) = oneshot::channel::<()>();
865        tokio::spawn(async move {
866            let _first = tokio::net::TcpStream::connect(local).await;
867            let _second = tokio::net::TcpStream::connect(local).await;
868            tokio::time::sleep(Duration::from_millis(200)).await;
869            let _ = cut.send(());
870            std::future::pending::<()>().await;
871        });
872        let mut accepted = 0;
873        let mut ended = Vec::new();
874        let serving = serve(
875            &listener,
876            ServeLimits {
877                max_connections: Some(2),
878            },
879            std::future::pending(),
880            signalled(cut_rx),
881            |stream, _| {
882                accepted += 1;
883                let panics = accepted == 1;
884                async move {
885                    let _held = stream;
886                    assert!(!panics, "the first connection's task panics");
887                    std::future::pending::<Result<(), Error>>().await
888                }
889            },
890            |_, result| ended.push(result.expect_err("neither ends clean").to_string()),
891        );
892        let (stopped, served) = tokio::time::timeout(Duration::from_secs(10), serving)
893            .await
894            .expect("the cut ends the drain");
895        assert_eq!((stopped, served), (StopReason::Limit, 2));
896        assert_eq!(ended.len(), 2, "{ended:?}");
897        let panicked = ended.iter().filter(|end| end.contains("ended abnormally"));
898        let cut = ended.iter().filter(|end| end.contains("was cut"));
899        assert_eq!((panicked.count(), cut.count()), (1, 1), "{ended:?}");
900    }
901}