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::SocketAddr;
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        |mut stream, _| {
189            let target = Arc::clone(&target);
190            async move {
191                let result = relay_connection_inner(
192                    &mut 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                // A verified stream only ends cleanly when its authenticated end arrives.
201                // Keep the socket until the verdict so a short stream reaches the local
202                // application as a reset, after any bytes already delivered.
203                if target.identity.is_some()
204                    && matches!(
205                        result,
206                        Err(_) | Ok(TunnelEnd::Truncated { .. } | TunnelEnd::Refused { .. })
207                    )
208                {
209                    stream.set_zero_linger().map_err(|error| {
210                        Error::new(
211                            ErrorKind::Unexpected,
212                            format!("could not reset the local tunnel connection: {error}"),
213                        )
214                    })?;
215                }
216                result
217            }
218        },
219        |peer, end| {
220            report.count(&end);
221            on_end(TunnelConnection { peer, end });
222        },
223    )
224    .await;
225    report.served = served;
226    report.stopped = stopped;
227    report.proxy_token_mints = target.auth.as_ref().map_or(0, |auth| auth.mint_count());
228    report
229}
230
231/// Something one forwarded connection did, for a caller's warnings.
232#[derive(Debug)]
233pub enum ForwardNotice {
234    /// An event the forwarder reported for the connection from `peer`.
235    Event {
236        peer: SocketAddr,
237        event: ForwardEvent,
238    },
239    /// The connection from `peer` ended early, with why.
240    Failed { peer: SocketAddr, error: Error },
241}
242
243/// What a port-forward loop did.
244#[derive(Clone, Debug, Eq, PartialEq)]
245pub struct ForwardReport {
246    /// Connections accepted.
247    pub served: u32,
248    /// Exchanges the proxy refused ([`ForwardEvent::Refused`]).
249    pub refused: u32,
250    /// Exchanges that upgraded, a WebSocket among them.
251    pub upgrades: u32,
252    /// Proxy tokens the session minted by the time the loop ended.
253    pub proxy_token_mints: u64,
254    pub stopped: StopReason,
255}
256
257impl ForwardReport {
258    /// Counts one of a connection's events in its column.
259    fn count(&mut self, event: &ForwardEvent) {
260        match event {
261            ForwardEvent::Refused { .. } => self.refused += 1,
262            ForwardEvent::Forwarded { upgraded: true, .. } => self.upgrades += 1,
263            _ => {}
264        }
265    }
266}
267
268/// Serves `listener` as an HTTP and WebSocket forward to `spec`'s guest port until `stop`
269/// resolves, the limit is reached, or the listener fails, then waits for the open connections
270/// to end.
271///
272/// `auth` is the session's proxy credential, `None` for a direct session. `on_notice` hears
273/// about each connection's events and failure as it ends, on the loop's own task.
274pub async fn serve_forward<F>(
275    listener: TcpListener,
276    spec: ForwardSpec,
277    auth: Option<Arc<ProxyAuth>>,
278    limits: ServeLimits,
279    stop: impl Future<Output = ()>,
280    on_notice: F,
281) -> Result<ForwardReport, Error>
282where
283    F: FnMut(ForwardNotice),
284{
285    forward_loop(
286        listener,
287        spec,
288        auth,
289        limits,
290        stop,
291        std::future::pending(),
292        on_notice,
293    )
294    .await
295}
296
297/// [`serve_forward`], cutting the connections still open once `cut` resolves after the stop.
298async fn forward_loop<F>(
299    listener: TcpListener,
300    spec: ForwardSpec,
301    auth: Option<Arc<ProxyAuth>>,
302    limits: ServeLimits,
303    stop: impl Future<Output = ()>,
304    cut: impl Future<Output = ()>,
305    mut on_notice: F,
306) -> Result<ForwardReport, Error>
307where
308    F: FnMut(ForwardNotice),
309{
310    // One client for the whole forward: connection reuse to the endpoint is what keeps a
311    // page-load of thirty assets from paying thirty TLS handshakes.
312    let client = Arc::new(ForwardClient::new()?);
313    let mut report = ForwardReport {
314        served: 0,
315        refused: 0,
316        upgrades: 0,
317        proxy_token_mints: 0,
318        stopped: StopReason::Requested,
319    };
320    let (stopped, served) = serve(
321        &listener,
322        limits,
323        stop,
324        cut,
325        |stream, _| {
326            let (mut spec, auth, client) = (spec.clone(), auth.clone(), Arc::clone(&client));
327            async move {
328                // A wildcard listener has no single address. The accepted socket names
329                // the concrete local interface, including the port chosen for a :0 bind.
330                spec.bind = stream.local_addr().map_err(|error| {
331                    Error::new(
332                        ErrorKind::Unexpected,
333                        format!("reading the local forward address failed: {error}"),
334                    )
335                })?;
336                let mut events = Vec::new();
337                let outcome =
338                    serve_connection_via(stream, &spec, auth.as_ref(), &client, |event| {
339                        events.push(event);
340                    })
341                    .await;
342                Ok((events, outcome))
343            }
344        },
345        |peer, done| {
346            let (events, failure) = match done {
347                Ok((events, outcome)) => (events, outcome.err()),
348                Err(error) => (Vec::new(), Some(error)),
349            };
350            for event in events {
351                report.count(&event);
352                on_notice(ForwardNotice::Event { peer, event });
353            }
354            if let Some(error) = failure {
355                on_notice(ForwardNotice::Failed { peer, error });
356            }
357        },
358    )
359    .await;
360    report.served = served;
361    report.stopped = stopped;
362    report.proxy_token_mints = auth.as_ref().map_or(0, |auth| auth.mint_count());
363    Ok(report)
364}
365
366// ── the SDKs' handles ────────────────────────────────────────────────────────
367
368/// The shared listener default, retained at its existing public path.
369// Keep a definition here: an external re-export omits the published constant from this
370// crate's rustdoc JSON. The value still comes from the shared application-layer owner.
371pub const DEFAULT_BIND: SocketAddr = microvms_app::session::DEFAULT_BIND;
372
373/// A local address to listen on: `bind` when given, [`DEFAULT_BIND`] otherwise.
374pub fn bind_address(bind: Option<&str>) -> Result<SocketAddr, Error> {
375    match bind {
376        None => Ok(DEFAULT_BIND),
377        Some(bind) => bind.parse().map_err(|err| {
378            Error::new(
379                ErrorKind::InvalidArg,
380                format!("{bind:?} is not a local address to listen on, such as 127.0.0.1:0: {err}"),
381            )
382        }),
383    }
384}
385
386/// How a connection that didn't end clean ended, as a handle's report lists it.
387#[derive(Clone, Copy, Debug, Eq, PartialEq)]
388pub enum EndKind {
389    /// The daemon refused the tunnel, or the proxy refused the exchange.
390    Refused,
391    /// A verified tunnel ended without the daemon's end of stream (#342).
392    Truncated,
393    /// A verified tunnel ended into a daemon from before the end of stream (#342).
394    Unproven,
395    /// The connection failed with an error.
396    Failed,
397}
398
399impl EndKind {
400    /// The wire spelling the bindings report.
401    pub fn as_str(self) -> &'static str {
402        match self {
403            EndKind::Refused => "refused",
404            EndKind::Truncated => "truncated",
405            EndKind::Unproven => "unproven",
406            EndKind::Failed => "failed",
407        }
408    }
409}
410
411/// One connection that didn't end clean. A clean end is only counted, so a long-lived handle's
412/// report grows with its failures, not with its traffic.
413#[derive(Clone, Debug, Eq, PartialEq)]
414pub struct ConnectionEnd {
415    pub peer: SocketAddr,
416    pub kind: EndKind,
417    /// The close code or the HTTP status, when the end carried one.
418    pub code: Option<u16>,
419    /// The daemon's reason, the forwarder's explanation, or the error.
420    pub detail: String,
421}
422
423impl ConnectionEnd {
424    /// A tunnel connection's end, unless it was clean.
425    pub fn of_tunnel(connection: &TunnelConnection) -> Option<Self> {
426        let (kind, code, detail) = match &connection.end {
427            Ok(TunnelEnd::Closed) => return None,
428            Ok(TunnelEnd::ClosedUnproven) => (EndKind::Unproven, None, String::new()),
429            Ok(TunnelEnd::Truncated { code }) => (EndKind::Truncated, *code, String::new()),
430            Ok(TunnelEnd::Refused { code, reason }) => {
431                (EndKind::Refused, Some(*code), reason.clone())
432            }
433            Err(error) => (EndKind::Failed, None, error.to_string()),
434        };
435        Some(Self {
436            peer: connection.peer,
437            kind,
438            code,
439            detail,
440        })
441    }
442
443    /// A forwarded connection's refusal or failure; its other events aren't ends worth listing.
444    pub fn of_forward(notice: &ForwardNotice) -> Option<Self> {
445        match notice {
446            ForwardNotice::Event {
447                peer,
448                event:
449                    ForwardEvent::Refused {
450                        status,
451                        explanation,
452                    },
453            } => Some(Self {
454                peer: *peer,
455                kind: EndKind::Refused,
456                code: Some(*status),
457                detail: explanation.clone(),
458            }),
459            ForwardNotice::Event { .. } => None,
460            ForwardNotice::Failed { peer, error } => Some(Self {
461                peer: *peer,
462                kind: EndKind::Failed,
463                code: None,
464                detail: error.to_string(),
465            }),
466        }
467    }
468}
469
470/// What a tunnel handle did: the loop's report, and each connection that didn't end clean.
471#[derive(Clone, Debug, Eq, PartialEq)]
472pub struct TunnelSummary {
473    pub report: TunnelReport,
474    pub ended: Vec<ConnectionEnd>,
475}
476
477/// What a port-forward handle did: the loop's report, and each connection that didn't end
478/// clean.
479#[derive(Clone, Debug, Eq, PartialEq)]
480pub struct ForwardSummary {
481    pub report: ForwardReport,
482    pub ended: Vec<ConnectionEnd>,
483}
484
485type Outcome<S> = Result<S, (ErrorKind, String)>;
486
487/// A serving loop on a background task: the SDKs' tunnel and port-forward handles.
488///
489/// Dropping it stops the loop and cuts the connections still open, so no relay outlives the
490/// handle that owns it: a relay carries the VM's credentials to whoever connected.
491pub struct Serving<S> {
492    local: SocketAddr,
493    stop: Mutex<Option<oneshot::Sender<()>>>,
494    cut: Mutex<Option<oneshot::Sender<()>>>,
495    outcome: Arc<Mutex<Option<Outcome<S>>>>,
496    done: watch::Receiver<bool>,
497}
498
499impl<S: Clone + Send + 'static> Serving<S> {
500    /// Runs `serving` on the current tokio runtime, handing it the stop and the cut futures.
501    /// Panics outside a runtime, like `tokio::spawn`.
502    fn spawn<Fut>(
503        local: SocketAddr,
504        serving: impl FnOnce(oneshot::Receiver<()>, oneshot::Receiver<()>) -> Fut,
505    ) -> Self
506    where
507        Fut: Future<Output = Result<S, Error>> + Send + 'static,
508    {
509        let (stop_tx, stop_rx) = oneshot::channel();
510        let (cut_tx, cut_rx) = oneshot::channel();
511        let (done_tx, done_rx) = watch::channel(false);
512        let outcome = Arc::new(Mutex::new(None));
513        let written = Arc::clone(&outcome);
514        let serving = serving(stop_rx, cut_rx);
515        tokio::spawn(async move {
516            let result = serving
517                .await
518                .map_err(|error| (error.kind(), error.to_string()));
519            *written.lock().unwrap_or_else(PoisonError::into_inner) = Some(result);
520            done_tx.send_replace(true);
521        });
522        Self {
523            local,
524            stop: Mutex::new(Some(stop_tx)),
525            cut: Mutex::new(Some(cut_tx)),
526            outcome,
527            done: done_rx,
528        }
529    }
530
531    /// The address the loop listens on, with the port the OS picked for a port 0.
532    pub fn local_address(&self) -> SocketAddr {
533        self.local
534    }
535
536    /// Whether the loop is still serving, or waiting for its open connections to end.
537    pub fn is_running(&self) -> bool {
538        !*self.done.borrow()
539    }
540
541    /// Asks the loop to stop accepting. Idempotent; [`Self::finished`] waits for it, and for the
542    /// connections still open.
543    pub fn request_stop(&self) {
544        signal(&self.stop);
545    }
546
547    /// Stops accepting and waits for the connections still open to end, cutting the ones
548    /// still open after `grace` when it's given, and returns what the loop did.
549    ///
550    /// Without a grace, a client that keeps its connection open keeps this waiting: the
551    /// relay holds bytes that client may still be reading, so the loop won't cut it unasked.
552    pub async fn stop(&self, grace: Option<Duration>) -> Result<S, Error> {
553        self.request_stop();
554        if let Some(grace) = grace {
555            if let Ok(outcome) = tokio::time::timeout(grace, self.finished()).await {
556                return outcome;
557            }
558            signal(&self.cut);
559        }
560        self.finished().await
561    }
562
563    /// Waits for the loop to end and returns what it did. Callable any number of times.
564    pub async fn finished(&self) -> Result<S, Error> {
565        let mut done = self.done.clone();
566        // The sender drops only after writing the outcome, so an error here still means done.
567        let _ = done.wait_for(|finished| *finished).await;
568        self.outcome
569            .lock()
570            .unwrap_or_else(PoisonError::into_inner)
571            .clone()
572            .expect("the loop records its outcome before signalling done")
573            .map_err(|(kind, message)| Error::new(kind, message))
574    }
575}
576
577/// Sends a handle's stop or cut, once.
578fn signal(sender: &Mutex<Option<oneshot::Sender<()>>>) {
579    if let Some(sender) = sender.lock().unwrap_or_else(PoisonError::into_inner).take() {
580        let _ = sender.send(());
581    }
582}
583
584/// A handle's stop or cut future: the sender's `send`, or its drop.
585async fn signalled(signal: oneshot::Receiver<()>) {
586    let _ = signal.await;
587}
588
589/// Binds `bind` and starts a tunnel to `target` on the current tokio runtime: the SDKs'
590/// `Session.tunnel`.
591pub async fn start_tunnel(
592    bind: SocketAddr,
593    target: TunnelTarget,
594    limits: ServeLimits,
595) -> Result<Serving<TunnelSummary>, Error> {
596    let listener =
597        super::forward::bind(&ForwardSpec::new(bind, target.guest_port, &target.endpoint)).await?;
598    let local = local_address(&listener)?;
599    Ok(Serving::spawn(local, move |stop, cut| async move {
600        let mut ended = Vec::new();
601        let report = tunnel_loop(
602            listener,
603            target,
604            limits,
605            signalled(stop),
606            signalled(cut),
607            |connection| ended.extend(ConnectionEnd::of_tunnel(&connection)),
608        )
609        .await;
610        Ok(TunnelSummary { report, ended })
611    }))
612}
613
614/// Binds `bind` and starts a port-forward to `guest_port` behind `endpoint` on the current
615/// tokio runtime: the SDKs' `Session.port_forward`. `auth` is `None` for a direct session.
616pub async fn start_forward(
617    bind: SocketAddr,
618    endpoint: &str,
619    guest_port: u16,
620    auth: Option<Arc<ProxyAuth>>,
621    limits: ServeLimits,
622) -> Result<Serving<ForwardSummary>, Error> {
623    let spec = ForwardSpec::new(bind, guest_port, endpoint);
624    let listener = super::forward::bind(&spec).await?;
625    let local = local_address(&listener)?;
626    Ok(Serving::spawn(local, move |stop, cut| async move {
627        let mut ended = Vec::new();
628        let report = forward_loop(
629            listener,
630            spec,
631            auth,
632            limits,
633            signalled(stop),
634            signalled(cut),
635            |notice| ended.extend(ConnectionEnd::of_forward(&notice)),
636        )
637        .await?;
638        Ok(ForwardSummary { report, ended })
639    }))
640}
641
642fn local_address(listener: &TcpListener) -> Result<SocketAddr, Error> {
643    listener.local_addr().map_err(|err| {
644        Error::new(
645            ErrorKind::Unexpected,
646            format!("the listener has no address: {err}"),
647        )
648    })
649}
650
651/// The accept loop both serving loops share: accept until stopped, a task per connection, then
652/// wait for the open ones, cutting them if `cut` resolves first. Answers why it stopped and how
653/// many connections it accepted.
654///
655/// `spawn` makes a connection's task. `ended` takes each connection's result as its task ends,
656/// on this task; a task that panicked or was cut reads as an error.
657async fn serve<T, Fut>(
658    listener: &TcpListener,
659    limits: ServeLimits,
660    stop: impl Future<Output = ()>,
661    cut: impl Future<Output = ()>,
662    mut spawn: impl FnMut(tokio::net::TcpStream, SocketAddr) -> Fut,
663    mut ended: impl FnMut(SocketAddr, Result<T, Error>),
664) -> (StopReason, u32)
665where
666    T: Send + 'static,
667    Fut: Future<Output = Result<T, Error>> + Send + 'static,
668{
669    let mut stop = std::pin::pin!(stop);
670    let mut tasks = JoinSet::new();
671    // Peers by task, so a task that panics is still reported against its connection.
672    let mut peers = Peers::new();
673    type Peers = std::collections::HashMap<tokio::task::Id, SocketAddr>;
674    let mut finish =
675        |peers: &mut Peers,
676         joined: Result<(tokio::task::Id, Result<T, Error>), tokio::task::JoinError>| {
677            let (id, result) = match joined {
678                Ok((id, result)) => (id, result),
679                Err(error) if error.is_cancelled() => (
680                    error.id(),
681                    Err(Error::new(
682                        ErrorKind::Unexpected,
683                        "the connection was cut when its loop stopped",
684                    )),
685                ),
686                Err(error) => (
687                    error.id(),
688                    Err(Error::new(
689                        ErrorKind::Unexpected,
690                        format!("the connection's task ended abnormally: {error}"),
691                    )),
692                ),
693            };
694            let peer = peers
695                .remove(&id)
696                .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
697            ended(peer, result);
698        };
699    let mut served = 0_u32;
700    let stopped = loop {
701        if limits.max_connections.is_some_and(|max| served >= max) {
702            break StopReason::Limit;
703        }
704        tokio::select! {
705            // Biased so a stop wins over a connection that arrived in the same wakeup: the
706            // caller asked to stop, and serving one more first looks like the stop did nothing.
707            biased;
708            () = &mut stop => break StopReason::Requested,
709            Some(joined) = tasks.join_next_with_id(), if !tasks.is_empty() => finish(&mut peers, joined),
710            next = listener.accept() => match next {
711                Ok((stream, peer)) => {
712                    served += 1;
713                    let handle = tasks.spawn(spawn(stream, peer));
714                    peers.insert(handle.id(), peer);
715                }
716                Err(error) => {
717                    break StopReason::ListenerFailed(format!(
718                        "the local listener stopped accepting: {error}"
719                    ));
720                }
721            },
722        }
723    };
724    // Drained rather than abandoned (see the module docs), until the caller cuts what's left.
725    let mut cut = std::pin::pin!(cut);
726    let mut cutting = false;
727    loop {
728        tokio::select! {
729            joined = tasks.join_next_with_id() => match joined {
730                Some(joined) => finish(&mut peers, joined),
731                None => break,
732            },
733            () = &mut cut, if !cutting => {
734                cutting = true;
735                tasks.abort_all();
736            }
737        }
738    }
739    (stopped, served)
740}
741
742#[cfg(test)]
743mod tests {
744    use super::*;
745
746    fn tunnel_report() -> TunnelReport {
747        TunnelReport {
748            served: 0,
749            refused: 0,
750            truncated: 0,
751            unproven: 0,
752            proxy_token_mints: 0,
753            stopped: StopReason::Requested,
754        }
755    }
756
757    /// Each way a tunnel connection ends lands in its own column, once.
758    #[test]
759    fn each_tunnel_end_is_counted_once_in_its_column() {
760        let mut report = tunnel_report();
761        for end in [
762            Ok(TunnelEnd::Closed),
763            Ok(TunnelEnd::ClosedUnproven),
764            Ok(TunnelEnd::ClosedUnproven),
765            Ok(TunnelEnd::Truncated { code: None }),
766            Ok(TunnelEnd::Truncated { code: Some(1000) }),
767            Ok(TunnelEnd::Truncated { code: Some(1000) }),
768            Ok(TunnelEnd::Refused {
769                code: 4502,
770                reason: String::new(),
771            }),
772            Err(Error::new(ErrorKind::Unexpected, "a transport error")),
773        ] {
774            report.count(&end);
775        }
776        assert_eq!(
777            (report.refused, report.truncated, report.unproven),
778            (2, 3, 2),
779            "{report:?}"
780        );
781    }
782
783    /// Only the caller's own stop reads as requested.
784    #[test]
785    fn only_the_callers_stop_was_requested() {
786        assert!(StopReason::Requested.was_requested());
787        assert!(!StopReason::Limit.was_requested());
788        assert!(!StopReason::ListenerFailed("gone".to_string()).was_requested());
789    }
790
791    /// A refusal and an upgrade each land in their column; other events count nowhere.
792    #[test]
793    fn each_forward_event_is_counted_once_in_its_column() {
794        let mut report = ForwardReport {
795            served: 0,
796            refused: 0,
797            upgrades: 0,
798            proxy_token_mints: 0,
799            stopped: StopReason::Requested,
800        };
801        for event in [
802            ForwardEvent::Refused {
803                status: 502,
804                explanation: String::new(),
805            },
806            ForwardEvent::Forwarded {
807                status: 101,
808                upgraded: true,
809            },
810            ForwardEvent::Forwarded {
811                status: 101,
812                upgraded: true,
813            },
814            ForwardEvent::Forwarded {
815                status: 200,
816                upgraded: false,
817            },
818            ForwardEvent::ConnectionError {
819                detail: String::new(),
820            },
821        ] {
822            report.count(&event);
823        }
824        assert_eq!((report.refused, report.upgrades), (1, 2), "{report:?}");
825    }
826
827    /// Each end but a clean one is listed with its kind, its code and its reason; a forward's
828    /// events other than a refusal aren't ends.
829    #[test]
830    fn each_end_but_a_clean_one_is_listed_with_its_kind() {
831        let peer: SocketAddr = "127.0.0.1:5000".parse().expect("an address");
832        let tunnel = |end| ConnectionEnd::of_tunnel(&TunnelConnection { peer, end });
833        let listed = |end: Option<ConnectionEnd>| {
834            end.map(|end| (end.peer, end.kind.as_str(), end.code, end.detail))
835        };
836        assert_eq!(tunnel(Ok(TunnelEnd::Closed)), None);
837        assert_eq!(
838            listed(tunnel(Ok(TunnelEnd::ClosedUnproven))),
839            Some((peer, "unproven", None, String::new()))
840        );
841        assert_eq!(
842            listed(tunnel(Ok(TunnelEnd::Truncated { code: Some(1006) }))),
843            Some((peer, "truncated", Some(1006), String::new()))
844        );
845        assert_eq!(
846            listed(tunnel(Ok(TunnelEnd::Refused {
847                code: 4502,
848                reason: "nothing is listening".into(),
849            }))),
850            Some((peer, "refused", Some(4502), "nothing is listening".into()))
851        );
852        let failed = listed(tunnel(Err(Error::new(ErrorKind::Unexpected, "a reset"))));
853        assert!(
854            matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
855            "{failed:?}"
856        );
857
858        let forward = |event| ConnectionEnd::of_forward(&ForwardNotice::Event { peer, event });
859        assert_eq!(
860            listed(forward(ForwardEvent::Refused {
861                status: 403,
862                explanation: "out of scope".into(),
863            })),
864            Some((peer, "refused", Some(403), "out of scope".into()))
865        );
866        assert_eq!(
867            forward(ForwardEvent::Forwarded {
868                status: 200,
869                upgraded: false,
870            }),
871            None
872        );
873        let failed = listed(ConnectionEnd::of_forward(&ForwardNotice::Failed {
874            peer,
875            error: Error::new(ErrorKind::Unexpected, "a reset"),
876        }));
877        assert!(
878            matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
879            "{failed:?}"
880        );
881    }
882
883    /// A connection's task that panics reads as ended abnormally, and one the caller cuts
884    /// reads as cut: the two are the caller's to tell apart.
885    #[tokio::test]
886    async fn a_panicked_connection_and_a_cut_one_are_told_apart() {
887        let listener = TcpListener::bind("127.0.0.1:0").await.expect("a free port");
888        let local = listener.local_addr().expect("bound");
889        let (cut, cut_rx) = oneshot::channel::<()>();
890        tokio::spawn(async move {
891            let _first = tokio::net::TcpStream::connect(local).await;
892            let _second = tokio::net::TcpStream::connect(local).await;
893            tokio::time::sleep(Duration::from_millis(200)).await;
894            let _ = cut.send(());
895            std::future::pending::<()>().await;
896        });
897        let mut accepted = 0;
898        let mut ended = Vec::new();
899        let serving = serve(
900            &listener,
901            ServeLimits {
902                max_connections: Some(2),
903            },
904            std::future::pending(),
905            signalled(cut_rx),
906            |stream, _| {
907                accepted += 1;
908                let panics = accepted == 1;
909                async move {
910                    let _held = stream;
911                    assert!(!panics, "the first connection's task panics");
912                    std::future::pending::<Result<(), Error>>().await
913                }
914            },
915            |_, result| ended.push(result.expect_err("neither ends clean").to_string()),
916        );
917        let (stopped, served) = tokio::time::timeout(Duration::from_secs(10), serving)
918            .await
919            .expect("the cut ends the drain");
920        assert_eq!((stopped, served), (StopReason::Limit, 2));
921        assert_eq!(ended.len(), 2, "{ended:?}");
922        let panicked = ended.iter().filter(|end| end.contains("ended abnormally"));
923        let cut = ended.iter().filter(|end| end.contains("was cut"));
924        assert_eq!((panicked.count(), cut.count()), (1, 1), "{ended:?}");
925    }
926}