Skip to main content

microvms_core/session/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2//! The in-VM client: one MicroVM's control API, with the proxy auth handled for you.
3//!
4//! Talks to `agentd` through the endpoint proxy — the minted proxy token read from the
5//! auth-token header map and both proxy headers on every request (TRAP-7), minting
6//! inside the retry path below the sixty-minute ceiling (TRAP-9), and the byte-offset
7//! cursor that makes an interrupted output stream resumable.
8//!
9//! # A session holds no state worth keeping
10//!
11//! Every exec record, every file, and the bootstrap token live in the VM. So a
12//! [`Session`] rebuilt from an endpoint and an agent token reattaches to everything a
13//! previous process was doing, and [`ExecHandle`] rebuilt from an exec id addresses the
14//! same server-side exec. That is a property of the protocol rather than of this type,
15//! and it is why the exec id is caller-minted.
16//!
17//! # The wire types are the daemon's
18//!
19//! Every request body and response shape comes from the `protocol` crate, which the
20//! daemon also compiles against (ARCH-2). There is deliberately no mirror here: a field
21//! renamed on the daemon side breaks this build, which is the earliest and cheapest
22//! place a protocol change can fail.
23//!
24//! # Layout
25//!
26//! [`proxy`] is the two-header auth and the mint schedule. [`http`] is the transport
27//! seam. [`sse`] is the frame parser and the typed events. [`exec`] is the handle and
28//! the cursor-driven stream. [`files`] is file and tar transfer. [`forward`] is the
29//! port-forwarder's listener and its 403-vs-502 diagnostic. [`tunnel`] is the WebSocket
30//! client that carries raw TCP to the daemon's relay.
31
32pub mod exec;
33pub mod files;
34pub mod forward;
35pub mod http;
36pub mod keepalive;
37pub mod proxy;
38pub mod shell;
39pub mod sse;
40pub mod tunnel;
41
42use std::sync::Arc;
43use std::time::Duration;
44
45pub use exec::{EndReason, ExecHandle, ExecResult, StreamEnd, StreamOptions, mint_exec_id};
46pub use forward::{
47    DEFAULT_EXCHANGE_TIMEOUT, ForwardClient, ForwardEvent, ForwardSpec, forwards_request_header,
48    refusal_explanation, upstream_url,
49};
50pub use http::{
51    ChunkSource, HttpBackend, HttpRequest, HttpResponse, OpenStream, ReqwestBackend, SharedBackend,
52};
53pub use keepalive::{KeepAwake, KeepAwakeEnd, KeepAwakeReport, KeepAwakeTask, RunningGate};
54pub use proxy::{
55    Clock, DEFAULT_AGENT_PORT, DEFAULT_REFRESH_AFTER, MAX_TOKEN_LIFETIME, PROXY_AUTH_HEADER,
56    PROXY_PORT_HEADER, ProxyAuth, ProxyToken, TokenMinter, TokioClock, WS_AUTH_SUBPROTOCOL_PREFIX,
57    WS_PORT_SUBPROTOCOL_PREFIX, WS_SUBPROTOCOL,
58};
59pub use sse::{ExecEvent, Frame, SseParser};
60pub use tunnel::{TUNNEL_CHUNK_BYTES, TunnelEnd, explain_close, relay_connection, tunnel_url};
61
62use crate::error::{Error, ErrorKind, WireKind};
63
64/// How long one non-streaming request may take.
65pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
66
67/// How long to wait for a daemon to report bootstrapped.
68pub const DEFAULT_READY_TIMEOUT: Duration = Duration::from_secs(120);
69
70/// How often to re-poll health while waiting for bootstrap.
71const READY_POLL_INTERVAL: Duration = Duration::from_secs(2);
72
73/// The transport: a backend, the agent token, and the proxy auth every request needs.
74///
75/// Separate from [`Session`] because [`ExecHandle`] needs it and holding a whole session
76/// would make the two mutually recursive.
77pub struct Transport {
78    backend: http::SharedBackend,
79    agent_token: String,
80    /// `None` means send no proxy headers at all, which is the shape for talking to a
81    /// daemon directly — a local binary, a test server, or a VM reached over a tunnel.
82    /// Requiring a control-plane client for that case would make this library
83    /// untestable without AWS.
84    proxy: Option<Arc<ProxyAuth>>,
85    timeout: Duration,
86}
87
88impl Transport {
89    /// The proxy auth, when this transport has one.
90    pub fn proxy(&self) -> Option<&Arc<ProxyAuth>> {
91        self.proxy.as_ref()
92    }
93
94    /// Every header this request needs: both proxy headers, then the bearer.
95    ///
96    /// `token` is three-way rather than two: `Some(Some(..))` overrides the session's
97    /// token (which is how a 401 is provoked), `Some(None)` sends no `Authorization`
98    /// header at all (which is how the health and hook routes are exercised), and
99    /// `None` means the session's own token. A two-valued parameter could not express
100    /// the middle case, and `None`-means-no-header would make the common case the
101    /// dangerous one.
102    async fn headers(&self, token: Option<Option<&str>>) -> Result<Vec<(String, String)>, Error> {
103        // The mint sits here, inside the path every request takes, which is what makes
104        // it happen at all (TRAP-9). A failure is retryable and propagates as one.
105        let mut headers = match &self.proxy {
106            Some(proxy) => proxy.headers().await?,
107            None => Vec::new(),
108        };
109        let bearer = match token {
110            Some(explicit) => explicit,
111            None => Some(self.agent_token.as_str()),
112        };
113        if let Some(bearer) = bearer {
114            headers.push(("Authorization".to_string(), format!("Bearer {bearer}")));
115        }
116        Ok(headers)
117    }
118
119    /// Sends one request and returns the response whatever its status.
120    async fn request(&self, request: HttpRequest) -> Result<HttpResponse, Error> {
121        let timeout = request.timeout.unwrap_or(self.timeout);
122        let method = request.method;
123        let path = request.path.clone();
124        tokio::time::timeout(timeout, self.request_inner(request))
125            .await
126            .map_err(|_| Error::wire(WireKind::Transport,
127                format!("{method} {path} exceeded its request timeout (including proxy authentication)")))?
128    }
129
130    async fn request_inner(&self, mut request: HttpRequest) -> Result<HttpResponse, Error> {
131        // Prepend, never replace: the caller's headers carry the content type
132        // (exec start is application/json, file upload is octet-stream), and
133        // replacing the vec silently stripped them — the daemon answered 400
134        // "body is not a valid start request" on the first live run while every
135        // fake-backed test stayed green, because the fakes parse bodies without
136        // reading content-type. The replacement was also what stripped the
137        // token-intent marker, so keeping the caller's headers means dropping
138        // that one explicitly — it is not a real header and must not leave.
139        let mut headers = self.headers(request_token(&request)).await?;
140        request.headers.retain(|(name, _)| name != TOKEN_INTENT);
141        headers.append(&mut request.headers);
142        request.headers = headers;
143        if request.timeout.is_none() {
144            request.timeout = Some(self.timeout);
145        }
146        self.backend.send(request).await
147    }
148
149    /// [`Self::request`], failing with the typed error for any non-2xx.
150    async fn send(&self, request: HttpRequest) -> Result<HttpResponse, Error> {
151        let (method, path) = (request.method, request.path.clone());
152        let response = self.request(request).await?;
153        response.error_for_status(method, &path)?;
154        Ok(response)
155    }
156
157    /// [`Self::send`], deserializing the body.
158    async fn send_json<T: serde::de::DeserializeOwned>(
159        &self,
160        request: HttpRequest,
161    ) -> Result<T, Error> {
162        let (method, path) = (request.method, request.path.clone());
163        let response = self.send(request).await?;
164        serde_json::from_slice(&response.body).map_err(|err| {
165            // A body the client cannot read is the daemon and the client disagreeing
166            // about a shape, which is what the shared protocol crate exists to prevent —
167            // so reaching here means a version mismatch rather than a bad request.
168            Error::wire(
169                WireKind::ProtocolError,
170                format!(
171                    "{method} {path} answered a body this client cannot read: {err}; \
172                     check /v1/schema for the daemon's protocol version"
173                ),
174            )
175            .with_source(err)
176        })
177    }
178}
179
180/// A marker header naming which bearer a request wants.
181///
182/// Carried on the request rather than passed alongside it so [`Transport::request`] has
183/// one parameter instead of two that can be transposed. Stripped before the request
184/// leaves — the name is not a real header and the daemon would ignore it, but sending
185/// it would leak the choice onto the wire.
186const TOKEN_INTENT: &str = "x-microvms-core-token-intent";
187
188/// The value meaning "no Authorization header at all".
189const TOKEN_INTENT_NONE: &str = "none";
190
191fn request_token(request: &HttpRequest) -> Option<Option<&str>> {
192    match request.header(TOKEN_INTENT) {
193        None => None,
194        Some(TOKEN_INTENT_NONE) => Some(None),
195        Some(explicit) => Some(Some(explicit)),
196    }
197}
198
199/// Marks a request as unauthenticated.
200fn unauthenticated(mut request: HttpRequest) -> HttpRequest {
201    request
202        .headers
203        .push((TOKEN_INTENT.to_string(), TOKEN_INTENT_NONE.to_string()));
204    request
205}
206
207/// The control API of one running MicroVM.
208///
209/// `Clone` shares the transport, including the cached proxy token, so a clone moved into a
210/// background task (a keepalive) sees the same refresh a resume triggers.
211#[derive(Clone)]
212pub struct Session {
213    transport: Arc<Transport>,
214    endpoint: String,
215    port: u16,
216}
217
218impl Session {
219    /// The bearer this session authenticates with.
220    ///
221    /// Public because a caller who launched with a minted token needs to read it back
222    /// to reattach later (the Python oracle's `Session.agent_token` property, and what
223    /// the CLI's run envelope publishes for `microvm exec --agent-token`). The value is
224    /// still kept out of `Debug` — readable on purpose is different from printed by
225    /// accident.
226    pub fn agent_token(&self) -> &str {
227        &self.transport.agent_token
228    }
229
230    /// A session against `endpoint`, minting proxy tokens through `minter`.
231    ///
232    /// Does not talk to the VM. Constructing a session is free and re-doable, and a
233    /// constructor that probed would make "do I have a session" mean "is the VM up",
234    /// which are different questions with different answers during a launch.
235    pub async fn connect(
236        endpoint: impl Into<String>,
237        agent_token: impl Into<String>,
238        minter: Arc<dyn TokenMinter>,
239    ) -> Result<Session, Error> {
240        Self::builder(endpoint, agent_token)
241            .with_minter(minter)
242            .build()
243    }
244
245    /// Reattach to an existing AWS VM using a privately retained agent token.
246    ///
247    /// Resolves AWS credentials but does not probe the guest or bootstrap it again.
248    /// The returned session is independent of any sandbox lock: use a separate attach
249    /// with a short request timeout for supervisor keepalives during long transfers.
250    /// Keep `agent_token` in private encrypted storage, never in public job status.
251    pub async fn attach(
252        region: crate::Region,
253        microvm_id: impl Into<String>,
254        endpoint: impl Into<String>,
255        agent_token: impl Into<String>,
256        port: Option<u16>,
257        request_timeout: Option<Duration>,
258    ) -> Result<Self, Error> {
259        if request_timeout.is_some_and(|timeout| timeout.is_zero()) {
260            return Err(Error::invalid_arg("request timeout must be positive"));
261        }
262        let mut control = crate::control::ControlPlane::new(region).await?;
263        if let Some(port) = port {
264            control = control.with_port(port)?;
265        }
266        let port = control.port();
267        let minter = Arc::new(crate::sandbox::ControlPlaneMinter {
268            control: Arc::new(control),
269            microvm_id: microvm_id.into(),
270        });
271        let mut builder = Self::builder(endpoint, agent_token)
272            .with_minter(minter)
273            .with_port(port);
274        if let Some(timeout) = request_timeout {
275            builder = builder.with_timeout(timeout);
276        }
277        builder.build()
278    }
279
280    /// A session with no proxy auth, for a daemon reached directly.
281    ///
282    /// The conformance path and every local-binary test go through here. See
283    /// [`Transport::proxy`] on why this is a supported shape rather than a test-only
284    /// escape hatch.
285    pub fn direct(
286        endpoint: impl Into<String>,
287        agent_token: impl Into<String>,
288    ) -> Result<Session, Error> {
289        Self::builder(endpoint, agent_token).build()
290    }
291
292    /// A builder, for the cases that need a port, a timeout, or a custom backend.
293    pub fn builder(endpoint: impl Into<String>, agent_token: impl Into<String>) -> SessionBuilder {
294        SessionBuilder {
295            endpoint: endpoint.into(),
296            agent_token: agent_token.into(),
297            minter: None,
298            backend: None,
299            port: DEFAULT_AGENT_PORT,
300            timeout: DEFAULT_REQUEST_TIMEOUT,
301            proxy: None,
302        }
303    }
304
305    /// The endpoint this session addresses.
306    pub fn endpoint(&self) -> &str {
307        &self.endpoint
308    }
309
310    /// The port the proxy token is scoped to.
311    pub fn port(&self) -> u16 {
312        self.port
313    }
314
315    /// The proxy auth, for a caller that needs the mint count or an invalidation.
316    pub fn proxy_auth(&self) -> Option<&Arc<ProxyAuth>> {
317        self.transport.proxy()
318    }
319
320    // -- credentials for a connection this session does not open ------------
321    //
322    // Both of the next two exist for one caller shape: something that opens its own socket
323    // to a port on this VM — a WebSocket to an in-VM bridge, an HTTP client the caller
324    // already owns — and needs the auth that would otherwise only be reachable by going
325    // through this session's own request path. They mint through [`ProxyAuth`]'s existing
326    // cache, so a consumer never reimplements TRAP-9's schedule, and neither stores what it
327    // returns.
328
329    /// Both proxy headers for `port`, minted through this session's token cache.
330    ///
331    /// The `{url, headers}` half of a port-endpoint resolution: pair these with
332    /// [`Self::endpoint`] and a caller can open its own connection with no knowledge of how
333    /// the token is minted or when it refreshes.
334    ///
335    /// **Empty for a direct session**, which is the correct answer rather than a missing one:
336    /// a daemon reached directly takes no proxy headers, so an empty set is exactly what its
337    /// requests carry today (see [`Transport::proxy`]).
338    ///
339    /// The auth value is a bearer credential and this returns it. Nothing in this crate logs
340    /// or stores the result.
341    pub async fn connect_headers(&self, port: u16) -> Result<Vec<(String, String)>, Error> {
342        match self.transport.proxy() {
343            Some(proxy) => {
344                tokio::time::timeout(self.transport.timeout, proxy.headers_for_port(port))
345                    .await
346                    .map_err(|_| {
347                        Error::wire(
348                            WireKind::Transport,
349                            "proxy authentication request timed out",
350                        )
351                    })?
352            }
353            None => Ok(Vec::new()),
354        }
355    }
356
357    /// The three WebSocket subprotocols for `port`, minted through this session's token
358    /// cache.
359    ///
360    /// `None` for a direct session, and that is deliberately not an empty list: the
361    /// subprotocol form exists only because a request through the *endpoint proxy* cannot set
362    /// headers, so there is nothing for a directly-reached daemon to offer, and a caller
363    /// handed `["lambda-microvms"]` with no token would open a handshake that is refused for a
364    /// reason naming neither.
365    ///
366    /// See [`ProxyAuth::subprotocols`] for the wire format and for the platform's stripping
367    /// behaviour. The middle value carries the credential.
368    pub async fn connect_subprotocols(&self, port: u16) -> Result<Option<[String; 3]>, Error> {
369        match self.transport.proxy() {
370            Some(proxy) => Ok(Some(
371                tokio::time::timeout(self.transport.timeout, proxy.subprotocols(port))
372                    .await
373                    .map_err(|_| {
374                        Error::wire(
375                            WireKind::Transport,
376                            "proxy authentication request timed out",
377                        )
378                    })??,
379            )),
380            None => Ok(None),
381        }
382    }
383
384    /// Points the session at a new endpoint and drops the cached proxy token (STATE-8).
385    ///
386    /// The measured behaviour is that the endpoint URL does not change across
387    /// suspend/resume, so this is usually a no-op on the URL — but the token drop is
388    /// not: a token minted against the pre-suspend instance may no longer validate, and
389    /// that rejection reads exactly like a dead daemon.
390    ///
391    /// Takes `&mut self` because it changes where the session points, which is not
392    /// something a shared reference should be able to do mid-request. The invalidation
393    /// itself is `&self` on [`ProxyAuth`], so a resume path holding only the auth can
394    /// still drop the token.
395    pub fn rebind(&mut self, endpoint: String) {
396        self.endpoint = endpoint;
397        if let Some(proxy) = self.transport.proxy() {
398            proxy.invalidate();
399        }
400    }
401
402    // -- health ------------------------------------------------------------
403
404    /// Unauthenticated liveness. `bootstrapped` is the useful field.
405    ///
406    /// Reachable through the endpoint at all implies bootstrapped in practice — the
407    /// platform forwards no external traffic until the run hook returns 200 — but a
408    /// caller inside the VM, or one talking to the daemon directly, can observe the
409    /// pre-bootstrap state.
410    pub async fn health(&self) -> Result<protocol::health::Health, Error> {
411        self.transport
412            .send_json(unauthenticated(HttpRequest::new("GET", "/v1/health")))
413            .await
414    }
415
416    /// Keeps the VM awake by polling health until `stop` resolves or the policy ends it.
417    ///
418    /// See [`keepalive`]: idleness is inbound traffic only, so this is the supported way to
419    /// hold a VM up while an exec works without client traffic.
420    pub async fn keep_awake(
421        &self,
422        policy: &KeepAwake,
423        stop: impl std::future::Future<Output = ()>,
424    ) -> Result<KeepAwakeReport, Error> {
425        policy
426            .run(|| async { self.health().await.map(Some) }, stop)
427            .await
428    }
429
430    /// Polls health until the daemon reports bootstrapped.
431    ///
432    /// Connection errors are expected here rather than exceptional: a VM that has just
433    /// reached RUNNING commonly refuses a connection or two before the proxy path is
434    /// wired up, and a mint failure is retryable by construction. A *fatal* error ends
435    /// the wait immediately — retrying a 401 until the deadline is one of the two
436    /// mistakes the retryable split exists to prevent.
437    pub async fn wait_until_ready(
438        &self,
439        timeout: Duration,
440    ) -> Result<protocol::health::Health, Error> {
441        let deadline = tokio::time::Instant::now() + timeout;
442        let mut last: Option<Error> = None;
443        loop {
444            match self.health().await {
445                Ok(health) if health.bootstrapped => return Ok(health),
446                Ok(_) => {}
447                Err(err) if err.retryable() => last = Some(err),
448                Err(err) => return Err(err),
449            }
450            let now = tokio::time::Instant::now();
451            if now >= deadline {
452                let detail = match last {
453                    Some(err) => format!(" (last error: {err})"),
454                    None => String::new(),
455                };
456                return Err(Error::new(
457                    ErrorKind::Timeout,
458                    format!(
459                        "the daemon was not bootstrapped within {}s{detail}",
460                        timeout.as_secs()
461                    ),
462                ));
463            }
464            tokio::time::sleep(READY_POLL_INTERVAL.min(deadline - now)).await;
465        }
466    }
467
468    // -- exec --------------------------------------------------------------
469
470    /// Starts a command and returns its handle. Does not wait.
471    ///
472    /// The `exec_id` on the request is the idempotency key. The daemon returns success
473    /// for a known id without spawning a second child, so a caller whose retry must be
474    /// safe across its own restart supplies a stable one.
475    pub async fn run(&self, req: protocol::exec::StartRequest) -> Result<ExecHandle, Error> {
476        let requested = req.exec_id.clone();
477        let mut request = HttpRequest::new("POST", "/v1/exec/start");
478        request
479            .headers
480            .push(("content-type".into(), "application/json".into()));
481        request.body = serde_json::to_vec(&req).map_err(|err| {
482            Error::invalid_arg(format!("the start request will not serialize: {err}"))
483        })?;
484        let started: protocol::exec::StartResponse = self.transport.send_json(request).await?;
485        // The daemon's id wins if it answered with one. It should always be the
486        // requested one — the request carries the key — but a handle built from what
487        // the client *asked for* rather than what the daemon *confirmed* would address
488        // nothing if that ever diverged, and every later call would 404.
489        let exec_id = if started.exec_id.is_empty() {
490            requested
491        } else {
492            started.exec_id
493        };
494        Ok(ExecHandle::new(Arc::clone(&self.transport), exec_id))
495    }
496
497    /// A handle for an exec started earlier, possibly by another process.
498    pub fn exec(&self, exec_id: impl Into<String>) -> ExecHandle {
499        ExecHandle::new(Arc::clone(&self.transport), exec_id.into())
500    }
501
502    /// Start, wait, ack. The one-shot shape, for when output is all you want.
503    pub async fn run_sync(
504        &self,
505        req: protocol::exec::StartRequest,
506        timeout: Duration,
507    ) -> Result<ExecResult, Error> {
508        self.run(req).await?.wait_and_ack(timeout).await
509    }
510
511    /// Signals an exec's whole process group. Returns whether anything was signalled.
512    pub async fn kill(&self, exec_id: &str) -> Result<bool, Error> {
513        self.exec(exec_id).kill().await
514    }
515
516    /// Process accounting: every registered exec with its group's live pids.
517    ///
518    /// The daemon reads `/proc` itself, so this works against a guest with no `ps`.
519    /// `child_exited: true` beside a non-empty `pids` is a command that finished while
520    /// something it backgrounded did not; the `exec_id` beside it is what [`Self::kill`]
521    /// takes. Bearer-guarded, unlike [`Self::health`]: the answer names every pid in the
522    /// guest, and a caller without the token has no business with it.
523    pub async fn procs(&self) -> Result<protocol::exec::ProcsResponse, Error> {
524        self.transport
525            .send_json(HttpRequest::new("GET", "/v1/procs"))
526            .await
527    }
528
529    // -- file transfer -----------------------------------------------------
530
531    /// Writes one file, creating parents. `mode` is octal as a string.
532    pub async fn upload_file(
533        &self,
534        path: &str,
535        data: &[u8],
536        mode: Option<&str>,
537    ) -> Result<(), Error> {
538        files::upload_file(&self.transport, path, data, mode).await
539    }
540
541    /// Reads one file.
542    pub async fn download_file(&self, path: &str) -> Result<Vec<u8>, Error> {
543        files::download_file(&self.transport, path).await
544    }
545
546    /// Whether a path exists, distinguishing absence from every other refusal.
547    pub async fn file_exists(&self, path: &str) -> Result<bool, Error> {
548        files::file_exists(&self.transport, path).await
549    }
550
551    /// Extracts pre-built tar bytes under `remote`.
552    pub async fn upload_tar(&self, remote: &str, archive: &[u8]) -> Result<(), Error> {
553        files::upload_tar(&self.transport, remote, archive).await
554    }
555
556    /// The raw tar bytes of a remote tree.
557    pub async fn download_tar(&self, remote: &str) -> Result<Vec<u8>, Error> {
558        files::download_tar(&self.transport, remote).await
559    }
560}
561
562impl std::fmt::Debug for Session {
563    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
564        // The agent token is a credential and stays out.
565        f.debug_struct("Session")
566            .field("endpoint", &self.endpoint)
567            .field("port", &self.port)
568            .field("proxy", &self.transport.proxy())
569            .finish()
570    }
571}
572
573/// Builds a [`Session`]. See [`Session::builder`].
574pub struct SessionBuilder {
575    endpoint: String,
576    agent_token: String,
577    minter: Option<Arc<dyn TokenMinter>>,
578    backend: Option<http::SharedBackend>,
579    port: u16,
580    timeout: Duration,
581    proxy: Option<Arc<ProxyAuth>>,
582}
583
584impl SessionBuilder {
585    /// Mints proxy tokens through `minter`, with the default refresh schedule.
586    #[must_use]
587    pub fn with_minter(mut self, minter: Arc<dyn TokenMinter>) -> Self {
588        self.minter = Some(minter);
589        self
590    }
591
592    /// Uses an already-built [`ProxyAuth`], for a caller that needs a custom refresh
593    /// interval or clock. Wins over [`Self::with_minter`].
594    #[must_use]
595    pub fn with_proxy_auth(mut self, proxy: Arc<ProxyAuth>) -> Self {
596        self.proxy = Some(proxy);
597        self
598    }
599
600    /// Replaces the HTTP backend. The seam a simulated transport plugs into.
601    #[must_use]
602    pub fn with_backend(mut self, backend: http::SharedBackend) -> Self {
603        self.backend = Some(backend);
604        self
605    }
606
607    /// The port the proxy token is scoped to and the port header names.
608    #[must_use]
609    pub fn with_port(mut self, port: u16) -> Self {
610        self.port = port;
611        self
612    }
613
614    /// How long one non-streaming request may take.
615    #[must_use]
616    pub fn with_timeout(mut self, timeout: Duration) -> Self {
617        self.timeout = timeout;
618        self
619    }
620
621    /// Builds the session.
622    pub fn build(self) -> Result<Session, Error> {
623        let backend = match self.backend {
624            Some(backend) => backend,
625            None => Arc::new(ReqwestBackend::new(&self.endpoint, self.timeout)?),
626        };
627        let proxy = match (self.proxy, self.minter) {
628            (Some(proxy), _) => Some(proxy),
629            (None, Some(minter)) => Some(Arc::new(ProxyAuth::new(minter, self.port))),
630            (None, None) => None,
631        };
632        Ok(Session {
633            transport: Arc::new(Transport {
634                backend,
635                agent_token: self.agent_token,
636                proxy,
637                timeout: self.timeout,
638            }),
639            endpoint: self.endpoint,
640            port: self.port,
641        })
642    }
643}
644
645#[cfg(test)]
646pub(crate) mod testing {
647    //! A recording backend the session tests share.
648    //!
649    //! It is a *recorder* rather than an assertion sink: it keeps every request head
650    //! and answers from a queue the test loaded, so an assertion is written against
651    //! what actually went out. A fake that asserted internally would be a fake whose
652    //! expectations are invisible at the call site.
653
654    use std::collections::HashMap;
655    use std::sync::{Mutex, PoisonError};
656
657    use futures_util::future::BoxFuture;
658
659    use super::http::{ChunkSource, HttpBackend, HttpRequest, HttpResponse, OpenStream};
660    use super::*;
661
662    /// What a recorder should answer next.
663    pub(crate) enum Reply {
664        /// A status and a body.
665        Body(u16, Vec<u8>),
666        /// A stream: the head status, then these chunks in order.
667        Chunks(u16, Vec<Vec<u8>>),
668        /// A transport failure, i.e. the request never produced a status.
669        Cut(&'static str),
670    }
671
672    impl Reply {
673        pub(crate) fn json(status: u16, body: impl serde::Serialize) -> Self {
674            Reply::Body(status, serde_json::to_vec(&body).expect("test body"))
675        }
676
677        pub(crate) fn ok(body: impl serde::Serialize) -> Self {
678            Reply::json(200, body)
679        }
680    }
681
682    #[derive(Default)]
683    pub(crate) struct Recorder {
684        seen: Mutex<Vec<HttpRequest>>,
685        replies: Mutex<std::collections::VecDeque<Reply>>,
686    }
687
688    impl Recorder {
689        pub(crate) fn with(replies: impl IntoIterator<Item = Reply>) -> Arc<Self> {
690            Arc::new(Self {
691                seen: Mutex::new(Vec::new()),
692                replies: Mutex::new(replies.into_iter().collect()),
693            })
694        }
695
696        /// Every request the recorder saw, in order.
697        pub(crate) fn requests(&self) -> Vec<HttpRequest> {
698            self.seen
699                .lock()
700                .unwrap_or_else(PoisonError::into_inner)
701                .clone()
702        }
703
704        pub(crate) fn last(&self) -> HttpRequest {
705            self.requests()
706                .pop()
707                .expect("the recorder saw no request at all")
708        }
709
710        fn record(&self, request: HttpRequest) -> Reply {
711            self.seen
712                .lock()
713                .unwrap_or_else(PoisonError::into_inner)
714                .push(request.clone());
715            self.replies
716                .lock()
717                .unwrap_or_else(PoisonError::into_inner)
718                .pop_front()
719                .unwrap_or_else(|| {
720                    panic!(
721                        "the recorder ran out of replies at {} {}",
722                        request.method, request.path
723                    )
724                })
725        }
726    }
727
728    struct Queued(std::collections::VecDeque<Vec<u8>>);
729
730    impl ChunkSource for Queued {
731        fn next_chunk(&mut self) -> BoxFuture<'_, Result<Option<Vec<u8>>, Error>> {
732            Box::pin(async move { Ok(self.0.pop_front()) })
733        }
734    }
735
736    impl HttpBackend for Recorder {
737        fn send(&self, request: HttpRequest) -> BoxFuture<'_, Result<HttpResponse, Error>> {
738            let reply = self.record(request);
739            Box::pin(async move {
740                match reply {
741                    Reply::Body(status, body) => Ok(HttpResponse {
742                        status,
743                        headers: HashMap::new(),
744                        body,
745                    }),
746                    Reply::Chunks(status, chunks) => Ok(HttpResponse {
747                        status,
748                        headers: HashMap::new(),
749                        body: chunks.concat(),
750                    }),
751                    Reply::Cut(why) => Err(Error::wire(WireKind::Transport, why)),
752                }
753            })
754        }
755
756        fn open_stream(
757            &self,
758            request: HttpRequest,
759            _idle_timeout: Duration,
760        ) -> BoxFuture<'_, Result<OpenStream, Error>> {
761            let reply = self.record(request);
762            Box::pin(async move {
763                let (status, chunks) = match reply {
764                    Reply::Chunks(status, chunks) => (status, chunks),
765                    Reply::Body(status, body) => (status, vec![body]),
766                    Reply::Cut(why) => return Err(Error::wire(WireKind::Transport, why)),
767                };
768                let head = HttpResponse {
769                    status,
770                    headers: HashMap::new(),
771                    // A failing head carries its body, as the real backend does, so the
772                    // typed error keeps the daemon's detail string.
773                    body: if (200..300).contains(&status) {
774                        Vec::new()
775                    } else {
776                        chunks.concat()
777                    },
778                };
779                let source: Box<dyn ChunkSource> = if (200..300).contains(&status) {
780                    Box::new(Queued(chunks.into_iter().collect()))
781                } else {
782                    Box::new(Queued(std::collections::VecDeque::new()))
783                };
784                Ok((head, source))
785            })
786        }
787    }
788
789    /// A session over a recorder, with proxy auth wired to a counting minter.
790    pub(crate) fn session_with(
791        recorder: Arc<Recorder>,
792    ) -> (Session, Arc<ProxyAuth>, Arc<proxy::testing::ManualClock>) {
793        let clock = Arc::new(proxy::testing::ManualClock::default());
794        let auth = Arc::new(
795            ProxyAuth::with_refresh_after(
796                Arc::new(proxy::testing::CountingMinter::default()),
797                DEFAULT_AGENT_PORT,
798                DEFAULT_REFRESH_AFTER,
799                Arc::clone(&clock) as Arc<dyn Clock>,
800            )
801            .expect("the default interval is below the ceiling"),
802        );
803        let session = Session::builder("https://vm.example", "agent-token-abcdef")
804            .with_backend(recorder)
805            .with_proxy_auth(Arc::clone(&auth))
806            .build()
807            .expect("builds");
808        (session, auth, clock)
809    }
810
811    pub(crate) fn health_body(bootstrapped: bool) -> serde_json::Value {
812        serde_json::json!({
813            "version": "0.1.0",
814            "bootstrapped": bootstrapped,
815            "disk": null,
816            "identity_degraded": false,
817            "identity_repaired": true,
818        })
819    }
820}
821
822#[cfg(test)]
823mod tests {
824    use super::testing::{Recorder, Reply, health_body, session_with};
825    use super::*;
826    use crate::session::proxy::testing::CountingMinter;
827
828    fn start_request(exec_id: &str) -> protocol::exec::StartRequest {
829        protocol::exec::StartRequest {
830            exec_id: exec_id.to_string(),
831            command: vec!["/bin/true".to_string()],
832            shell: false,
833            cwd: None,
834            env: Default::default(),
835            user: None,
836            group: None,
837            timeout_sec: None,
838            stdin: false,
839            reap_group_on_exit: false,
840        }
841    }
842
843    fn header(request: &HttpRequest, name: &str) -> Option<String> {
844        request.header(name).map(str::to_string)
845    }
846
847    /// One value out of a header pair list, matched as HTTP matches names.
848    fn header_of(headers: &[(String, String)], name: &str) -> Option<String> {
849        headers
850            .iter()
851            .find(|(key, _)| key.eq_ignore_ascii_case(name))
852            .map(|(_, value)| value.clone())
853    }
854
855    #[tokio::test(start_paused = true)]
856    async fn request_timeout_also_bounds_a_blocked_proxy_mint() {
857        struct BlockedMinter;
858        impl TokenMinter for BlockedMinter {
859            fn mint(&self) -> futures_util::future::BoxFuture<'_, Result<ProxyToken, Error>> {
860                Box::pin(std::future::pending())
861            }
862        }
863        let backend = Recorder::with([]);
864        let session = Session::builder("http://127.0.0.1:9", "private")
865            .with_backend(backend.clone())
866            .with_minter(Arc::new(BlockedMinter))
867            .with_timeout(Duration::from_secs(3))
868            .build()
869            .expect("session");
870        let result = tokio::time::timeout(Duration::from_secs(4), session.health())
871            .await
872            .expect("the request budget must include a stalled proxy mint")
873            .expect_err("mint never completed");
874        assert_eq!(result.wire_kind(), Some(WireKind::Transport));
875        assert!(result.retryable());
876        let result = tokio::time::timeout(Duration::from_secs(4), session.connect_headers(9000))
877            .await
878            .expect("external transfers must bound authentication too")
879            .expect_err("mint remains blocked");
880        assert!(result.retryable());
881        let result =
882            tokio::time::timeout(Duration::from_secs(4), session.connect_subprotocols(9000))
883                .await
884                .expect("external sockets must bound authentication too")
885                .expect_err("mint remains blocked");
886        assert!(result.retryable());
887        assert!(
888            backend.requests().is_empty(),
889            "no request could authenticate"
890        );
891    }
892
893    /// The live-run regression: `Transport::request` REPLACED the header vec with
894    /// the auth headers, stripping the content type the caller set — the daemon
895    /// answered 400 "body is not a valid start request" against real axum while
896    /// every fake stayed green (fakes parse bodies without reading content-type).
897    ///
898    /// **Guard proof.** Change the prepend back to `request.headers = auth_headers`
899    /// and this test is red on the content-type assertion.
900    #[tokio::test]
901    async fn a_callers_content_type_survives_the_auth_header_injection() {
902        let recorder = Recorder::with([Reply::ok(
903            serde_json::json!({"exec_id":"e1","phase":"running"}),
904        )]);
905        let (session, _, _) = session_with(Arc::clone(&recorder));
906
907        session.run(start_request("e1")).await.expect("start");
908
909        let seen = recorder.requests();
910        assert_eq!(
911            header(&seen[0], "content-type").as_deref(),
912            Some("application/json"),
913            "the exec start lost its content type to the auth header injection"
914        );
915    }
916
917    /// The read-back the CLI's run envelope depends on: a caller who launched with a
918    /// minted token can only reattach if the session hands the token back. The first
919    /// live run published `agentToken: null` because nothing read it — the daemon then
920    /// saw a bootstrap replay from the literal string "None".
921    #[tokio::test]
922    async fn the_session_hands_back_the_token_it_authenticates_with() {
923        let recorder = Recorder::with([]);
924        let (session, _, _) = session_with(Arc::clone(&recorder));
925        assert_eq!(session.agent_token(), "agent-token-abcdef");
926    }
927
928    /// TRAP-7's send side, asserted on a recorded request rather than on a return
929    /// value: both proxy headers went out.
930    ///
931    /// **Guard proof.** Delete the `headers.push((PROXY_PORT_HEADER, ..))` branch in
932    /// `ProxyAuth::headers_from` and this test is red on the port assertion while every
933    /// other test in the crate stays green.
934    #[tokio::test]
935    async fn every_endpoint_request_carries_both_proxy_headers() {
936        let recorder = Recorder::with([
937            Reply::ok(health_body(true)),
938            Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
939        ]);
940        let (session, _, _) = session_with(Arc::clone(&recorder));
941
942        session.health().await.expect("health succeeds");
943        session.run(start_request("e1")).await.expect("start");
944
945        let seen = recorder.requests();
946        assert_eq!(seen.len(), 2);
947        for request in &seen {
948            assert_eq!(
949                header(request, PROXY_AUTH_HEADER).as_deref(),
950                Some(CountingMinter::value(0).as_str()),
951                "{} {} went out without the proxy auth header",
952                request.method,
953                request.path
954            );
955            assert_eq!(
956                header(request, PROXY_PORT_HEADER).as_deref(),
957                Some("9000"),
958                "{} {} went out without the proxy port header, which the proxy \
959                 rejects as if the token were bad",
960                request.method,
961                request.path
962            );
963        }
964    }
965
966    /// Health is unauthenticated and everything else is not.
967    ///
968    /// Both halves matter: health carrying a bearer would still work, but it would make
969    /// the route untestable as the unauthenticated route it is, and an exec route
970    /// *without* one would 401.
971    #[tokio::test]
972    async fn health_sends_no_bearer_and_an_exec_route_does() {
973        let recorder = Recorder::with([
974            Reply::ok(health_body(true)),
975            Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
976        ]);
977        let (session, _, _) = session_with(Arc::clone(&recorder));
978
979        session.health().await.expect("health");
980        let seen = recorder.requests();
981        assert_eq!(header(&seen[0], "authorization"), None);
982        assert_eq!(
983            header(&seen[0], "x-microvms-core-token-intent"),
984            None,
985            "the internal token-intent marker must not reach the wire"
986        );
987
988        session.run(start_request("e1")).await.expect("start");
989        assert_eq!(
990            header(&recorder.last(), "authorization").as_deref(),
991            Some("Bearer agent-token-abcdef")
992        );
993    }
994
995    /// A keepalive poll is exactly the health route, unauthenticated but proxy-authorized:
996    /// the proxy headers are what make it inbound traffic the idle meter counts.
997    #[tokio::test(start_paused = true)]
998    async fn keep_awake_polls_health_through_the_proxy_until_the_vm_is_idle() {
999        let busy = |busy: bool| {
1000            let mut body = health_body(true);
1001            body["busy"] = serde_json::json!(busy);
1002            Reply::ok(body)
1003        };
1004        let recorder = Recorder::with([busy(true), busy(true), busy(false)]);
1005        let (session, _, _) = session_with(Arc::clone(&recorder));
1006        let report = session
1007            .clone()
1008            .keep_awake(
1009                &KeepAwake::new(None).while_busy(true),
1010                std::future::pending(),
1011            )
1012            .await
1013            .expect("keepalive");
1014        assert_eq!((report.end, report.polls), (KeepAwakeEnd::Idle, 3));
1015        let seen = recorder.requests();
1016        assert_eq!(seen.len(), 3);
1017        for request in &seen {
1018            assert_eq!(
1019                (request.method, request.path.as_str()),
1020                ("GET", "/v1/health")
1021            );
1022            assert_eq!(header(request, "authorization"), None);
1023            assert!(header(request, PROXY_AUTH_HEADER).is_some());
1024            assert!(header(request, PROXY_PORT_HEADER).is_some());
1025        }
1026    }
1027
1028    /// A session with no minter sends no proxy headers at all, which is the shape for
1029    /// a daemon reached directly.
1030    #[tokio::test]
1031    async fn a_direct_session_sends_no_proxy_headers() {
1032        let recorder = Recorder::with([Reply::ok(health_body(true))]);
1033        let session = Session::builder("http://127.0.0.1:9000", "token")
1034            .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1035            .build()
1036            .expect("builds");
1037        assert!(session.proxy_auth().is_none());
1038
1039        session.health().await.expect("health");
1040        let request = recorder.last();
1041        assert_eq!(header(&request, PROXY_AUTH_HEADER), None);
1042        assert_eq!(header(&request, PROXY_PORT_HEADER), None);
1043    }
1044
1045    /// A mint failure surfaces as a retryable error from the request, not from
1046    /// construction — the whole point of minting in the request path.
1047    #[tokio::test]
1048    async fn a_mint_failure_surfaces_as_a_retryable_request_error() {
1049        let recorder = Recorder::with([Reply::ok(health_body(true))]);
1050        let auth = Arc::new(
1051            ProxyAuth::with_refresh_after(
1052                Arc::new(CountingMinter::failing(1)),
1053                DEFAULT_AGENT_PORT,
1054                DEFAULT_REFRESH_AFTER,
1055                Arc::new(proxy::testing::ManualClock::default()),
1056            )
1057            .expect("interval accepted"),
1058        );
1059        let session = Session::builder("https://vm.example", "token")
1060            .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1061            .with_proxy_auth(Arc::clone(&auth))
1062            .build()
1063            .expect("builds");
1064
1065        let err = session.health().await.expect_err("the mint fails");
1066        assert_eq!(err.kind(), ErrorKind::Retryable);
1067        assert_eq!(err.wire_kind(), Some(WireKind::AuthTokenMint));
1068        assert!(
1069            recorder.requests().is_empty(),
1070            "a request went out without a token"
1071        );
1072
1073        session.health().await.expect("the retry succeeds");
1074        assert_eq!(auth.mint_count(), 1);
1075    }
1076
1077    /// Every status the daemon chooses arrives as its own wire kind, from the session's
1078    /// own surface rather than from the transport in isolation.
1079    #[tokio::test]
1080    async fn a_daemon_status_arrives_as_its_typed_error() {
1081        for (status, expected) in [
1082            (401, WireKind::Unauthorized),
1083            (404, WireKind::NotFound),
1084            (409, WireKind::Conflict),
1085            (503, WireKind::NotBootstrapped),
1086        ] {
1087            let recorder = Recorder::with([Reply::Body(status, b"{\"error\":\"x\"}".to_vec())]);
1088            let (session, _, _) = session_with(recorder);
1089            let err = session.health().await.expect_err("a failing status");
1090            assert_eq!(err.wire_kind(), Some(expected), "status {status}");
1091        }
1092    }
1093
1094    /// `wait_until_ready` returns as soon as the daemon reports bootstrapped, and
1095    /// tolerates the pre-bootstrap answers before it.
1096    #[tokio::test(start_paused = true)]
1097    async fn wait_until_ready_polls_through_not_bootstrapped_and_a_dropped_connection() {
1098        let recorder = Recorder::with([
1099            Reply::Cut("connection refused"),
1100            Reply::Body(503, b"not bootstrapped".to_vec()),
1101            Reply::ok(health_body(false)),
1102            Reply::ok(health_body(true)),
1103        ]);
1104        let (session, _, _) = session_with(Arc::clone(&recorder));
1105
1106        let health = session
1107            .wait_until_ready(DEFAULT_READY_TIMEOUT)
1108            .await
1109            .expect("the daemon comes up");
1110        assert!(health.bootstrapped);
1111        assert_eq!(recorder.requests().len(), 4);
1112    }
1113
1114    /// A fatal error ends the wait immediately: retrying a 401 until the deadline is
1115    /// the mistake the retryable split exists to prevent.
1116    #[tokio::test(start_paused = true)]
1117    async fn wait_until_ready_gives_up_at_once_on_a_credential_failure() {
1118        let recorder = Recorder::with([Reply::Body(401, b"wrong token".to_vec())]);
1119        let (session, _, _) = session_with(Arc::clone(&recorder));
1120
1121        let err = session
1122            .wait_until_ready(DEFAULT_READY_TIMEOUT)
1123            .await
1124            .expect_err("a 401 is fatal");
1125        assert_eq!(err.kind(), ErrorKind::Credentials);
1126        assert_eq!(
1127            recorder.requests().len(),
1128            1,
1129            "a fatal error was retried, so a wrong token costs the full timeout"
1130        );
1131    }
1132
1133    /// The timeout is a client-side deadline and names the last retryable error, so a
1134    /// caller can tell "never answered" from "answered, not ready".
1135    #[tokio::test(start_paused = true)]
1136    async fn wait_until_ready_times_out_naming_the_last_retryable_error() {
1137        let recorder = Recorder::with([
1138            Reply::Cut("connection refused"),
1139            Reply::Cut("connection refused"),
1140            Reply::Cut("connection refused"),
1141        ]);
1142        let (session, _, _) = session_with(recorder);
1143
1144        let err = session
1145            .wait_until_ready(Duration::from_secs(4))
1146            .await
1147            .expect_err("the daemon never comes up");
1148        assert_eq!(err.kind(), ErrorKind::Timeout);
1149        assert!(
1150            err.to_string().contains("connection refused"),
1151            "the timeout must name what it kept hitting: {err}"
1152        );
1153    }
1154
1155    /// `rebind` drops the cached token, so the next request mints (STATE-8).
1156    ///
1157    /// The mint count is the assertion because a resumed VM keeps its endpoint URL:
1158    /// nothing about the *request* changes, and the only observable is that a fresh
1159    /// token was fetched.
1160    #[tokio::test]
1161    async fn rebind_invalidates_the_proxy_token_so_the_next_request_mints() {
1162        let recorder = Recorder::with([Reply::ok(health_body(true)), Reply::ok(health_body(true))]);
1163        let (mut session, auth, _) = session_with(Arc::clone(&recorder));
1164
1165        session.health().await.expect("health");
1166        assert_eq!(auth.mint_count(), 1);
1167
1168        session.rebind("https://vm.example".to_string());
1169        assert!(!auth.is_cached(), "rebind did not drop the token");
1170
1171        session.health().await.expect("health");
1172        assert_eq!(auth.mint_count(), 2);
1173        assert_eq!(
1174            header(&recorder.last(), PROXY_AUTH_HEADER).as_deref(),
1175            Some(CountingMinter::value(1).as_str()),
1176            "the request after a rebind carried the pre-suspend token"
1177        );
1178    }
1179
1180    /// A start request's body is the protocol crate's serialization, and the handle
1181    /// addresses the id the daemon confirmed.
1182    #[tokio::test]
1183    async fn a_start_request_serializes_through_the_protocol_crate() {
1184        let recorder = Recorder::with([Reply::ok(
1185            serde_json::json!({"exec_id":"e-confirmed","phase":"running"}),
1186        )]);
1187        let (session, _, _) = session_with(Arc::clone(&recorder));
1188
1189        let handle = session.run(start_request("e-asked")).await.expect("start");
1190        assert_eq!(
1191            handle.exec_id(),
1192            "e-confirmed",
1193            "the handle addresses the id the client asked for rather than the one the \
1194             daemon confirmed, so every later call would 404 if they diverged"
1195        );
1196
1197        let request = recorder.last();
1198        assert_eq!(request.method, "POST");
1199        assert_eq!(request.path, "/v1/exec/start");
1200        let body: serde_json::Value = serde_json::from_slice(&request.body).expect("json body");
1201        assert_eq!(body["exec_id"], "e-asked");
1202        assert_eq!(body["command"], serde_json::json!(["/bin/true"]));
1203    }
1204
1205    /// A body this client cannot read is reported as a protocol failure naming the
1206    /// schema route, not as a transport error.
1207    #[tokio::test]
1208    async fn an_unreadable_response_body_names_the_schema_route() {
1209        let recorder = Recorder::with([Reply::Body(200, b"{\"unexpected\":true}".to_vec())]);
1210        let (session, _, _) = session_with(recorder);
1211        let err = session.health().await.expect_err("the body will not parse");
1212        assert_eq!(err.wire_kind(), Some(WireKind::ProtocolError));
1213        assert!(err.to_string().contains("/v1/schema"), "{err}");
1214    }
1215
1216    // ── credentials for a caller's own connection ─────────────────────────────
1217
1218    /// **A session hands out connect credentials for a port it was not built with, off one
1219    /// token cache, minting exactly once more to cover that port.**
1220    ///
1221    /// The property that makes an external `getPortEndpoint` a few lines rather than a
1222    /// reimplementation: one health request warms the cache, the first workload-port call
1223    /// widens the token, and the two accessors then share it.
1224    ///
1225    /// This test previously required `mint_count() == 1` after the workload-port calls —
1226    /// "resolving a port endpoint burned a second control-plane call". That was the defect
1227    /// written down as a requirement: a token scoped to the agent port does not authorize 8080,
1228    /// so the credentials it praised are refused with 403 `Access to port denied` on the header
1229    /// path and close code 1006 on the WebSocket path. Measured 2026-08-15 against real AWS.
1230    /// The one-cache claim survives in the form that is true — **both** accessors answer off
1231    /// the same widened token, which is still what stops TRAP-9 becoming two schedules.
1232    ///
1233    /// **Guard proof.** Give `connect_subprotocols` its own `ProxyAuth::new(..)` — the shape a
1234    /// consumer that could not reach this session's auth would be forced into — and the mint
1235    /// count reads 3 rather than 2, and the two accessors disagree on the auth value. Verified.
1236    #[tokio::test]
1237    async fn a_session_answers_connect_credentials_for_a_port_off_its_own_token_cache() {
1238        let recorder = Recorder::with([Reply::ok(health_body(true))]);
1239        let (session, auth, _) = session_with(Arc::clone(&recorder));
1240
1241        session.health().await.expect("health warms the cache");
1242        assert_eq!(auth.mint_count(), 1);
1243
1244        let headers = session
1245            .connect_headers(8080)
1246            .await
1247            .expect("the cached token answers");
1248        let offered = session
1249            .connect_subprotocols(8080)
1250            .await
1251            .expect("the cached token answers")
1252            .expect("a proxied session offers subprotocols");
1253
1254        assert_eq!(
1255            header_of(&headers, PROXY_PORT_HEADER).as_deref(),
1256            Some("8080"),
1257            "the caller asked for 8080 and would have targeted the session's port"
1258        );
1259        // The *widened* token, which is the second one minted: the first covers 9000 only.
1260        assert_eq!(
1261            header_of(&headers, PROXY_AUTH_HEADER).as_deref(),
1262            Some(CountingMinter::value(1).as_str()),
1263            "the header pair carried the agent-port-only token, which the proxy refuses for \
1264             8080 with 403 'Access to port denied'"
1265        );
1266        assert_eq!(offered[0], proxy::WS_SUBPROTOCOL);
1267        assert_eq!(
1268            offered[1],
1269            format!(
1270                "{}{}",
1271                proxy::WS_AUTH_SUBPROTOCOL_PREFIX,
1272                CountingMinter::value(1)
1273            ),
1274            "the handshake carried a different token than the header path, so they are two \
1275             caches rather than one"
1276        );
1277        assert_eq!(offered[2], "lambda-microvms.port.8080");
1278        assert_eq!(
1279            auth.mint_count(),
1280            2,
1281            "one mint covered both the agent port and 8080, which the control plane does not \
1282             do — or the second accessor minted a third, which would be two caches"
1283        );
1284        assert!(
1285            recorder.requests().len() == 1,
1286            "building credentials sent a request to the VM: {:?}",
1287            recorder.requests()
1288        );
1289    }
1290
1291    /// **A direct session offers no subprotocols at all, and that is not an empty list.**
1292    ///
1293    /// The asymmetry between the two accessors, asserted because it is the one thing a caller
1294    /// could reasonably get backwards. Empty headers are what a direct session's requests
1295    /// really carry, so `[]` there is a true answer. A subprotocol list *without* the token
1296    /// would be a false one: the handshake would be offered, refused, and the refusal would
1297    /// name neither the token nor the port.
1298    ///
1299    /// **Guard proof.** Return `Some([WS_SUBPROTOCOL.into(), String::new(), ..])` for the
1300    /// direct case and the `is_none` assertion is red. Verified.
1301    #[tokio::test]
1302    async fn a_direct_session_offers_no_subprotocols_and_empty_headers() {
1303        let recorder = Recorder::with([]);
1304        let session = Session::builder("http://127.0.0.1:9000", "token")
1305            .with_backend(recorder as http::SharedBackend)
1306            .build()
1307            .expect("builds");
1308
1309        assert!(
1310            session
1311                .connect_headers(9000)
1312                .await
1313                .expect("no proxy is not a failure")
1314                .is_empty(),
1315            "a direct session's requests carry no proxy headers, so neither should this"
1316        );
1317        assert!(
1318            session
1319                .connect_subprotocols(9000)
1320                .await
1321                .expect("no proxy is not a failure")
1322                .is_none(),
1323            "a handshake with the base subprotocol and no token is refused for a reason that \
1324             names neither"
1325        );
1326    }
1327
1328    /// **A connect credential built after the refresh window carries the fresh token.**
1329    ///
1330    /// TRAP-9 through this door. A WebSocket outlives a request, so a connection opened at
1331    /// minute fifty must hold a token minted at minute fifty — and the only observable is the
1332    /// value, because a stale token and a fresh one produce identical successful handshakes
1333    /// until one is rejected.
1334    ///
1335    /// **Guard proof.** Snapshot the subprotocols at `build()` and answer from the snapshot —
1336    /// the shape a caller who resolved a port endpoint once would get — and the second value
1337    /// still reads `jwe-0`. Verified.
1338    #[tokio::test]
1339    async fn connect_credentials_re_mint_when_the_refresh_window_has_rolled_over() {
1340        let recorder = Recorder::with([]);
1341        let (session, auth, clock) = session_with(recorder);
1342
1343        let first = session
1344            .connect_subprotocols(9000)
1345            .await
1346            .expect("mints")
1347            .expect("proxied");
1348        assert!(first[1].ends_with(&CountingMinter::value(0)));
1349
1350        clock.advance(DEFAULT_REFRESH_AFTER + Duration::from_secs(1));
1351        let second = session
1352            .connect_subprotocols(9000)
1353            .await
1354            .expect("re-mints")
1355            .expect("proxied");
1356
1357        assert_eq!(auth.mint_count(), 2);
1358        assert!(
1359            second[1].ends_with(&CountingMinter::value(1)),
1360            "the handshake would carry a token past its refresh window: {}",
1361            second[1]
1362        );
1363    }
1364
1365    /// **`procs` is a bearer GET of `/v1/procs`, and its body is the typed account.**
1366    ///
1367    /// Mirrors `health` in shape and differs in the one thing that matters: the token goes
1368    /// out. The route is bearer-guarded because it names every pid in the guest, and a
1369    /// client that sent it unauthenticated would read the daemon's 503-before-bootstrap as
1370    /// a broken VM.
1371    ///
1372    /// **Guard proof.** Wrap the request in `unauthenticated(..)` the way `health` does and
1373    /// the authorization assertion is red; point it at `/v1/health` and the path one is.
1374    #[tokio::test]
1375    async fn procs_is_a_bearer_get_of_the_procs_route_with_a_typed_body() {
1376        let recorder = Recorder::with([Reply::ok(serde_json::json!({
1377            "procs": [{
1378                "exec_id": "e1",
1379                "pgid": 4242,
1380                "started_at": 1_756_500_000u64,
1381                "child_exited": true,
1382                "reap": false,
1383                "pids": [4243],
1384            }]
1385        }))]);
1386        let (session, _, _) = session_with(Arc::clone(&recorder));
1387
1388        let listed = session.procs().await.expect("procs");
1389        let request = recorder.last();
1390        assert_eq!(request.method, "GET");
1391        assert_eq!(request.path, "/v1/procs");
1392        assert_eq!(
1393            header(&request, "authorization").as_deref(),
1394            Some("Bearer agent-token-abcdef"),
1395            "the procs route is bearer-guarded; without the token it is a 503 or a 401"
1396        );
1397        assert_eq!(listed.procs.len(), 1);
1398        assert_eq!(listed.procs[0].exec_id, "e1");
1399        assert_eq!(listed.procs[0].pgid, Some(4242));
1400        assert!(listed.procs[0].child_exited);
1401        assert_eq!(listed.procs[0].pids, [4243]);
1402    }
1403
1404    /// **The reap flag rides the start body under its wire name.**
1405    ///
1406    /// The field is `#[serde(default)]` on the daemon side, so a client that dropped it
1407    /// would still be accepted — and the exec would silently keep today's leave-it-running
1408    /// behaviour. Asserted on the recorded body, not on the request struct.
1409    ///
1410    /// **Guard proof.** Add `#[serde(skip)]` to the field and this is red on the missing key.
1411    #[tokio::test]
1412    async fn a_start_request_carries_the_reap_flag_on_the_wire() {
1413        let recorder = Recorder::with([Reply::ok(
1414            serde_json::json!({"exec_id":"e1","phase":"running"}),
1415        )]);
1416        let (session, _, _) = session_with(Arc::clone(&recorder));
1417
1418        let mut request = start_request("e1");
1419        request.reap_group_on_exit = true;
1420        session.run(request).await.expect("start");
1421
1422        let body: serde_json::Value =
1423            serde_json::from_slice(&recorder.last().body).expect("a JSON start body");
1424        assert_eq!(
1425            body["reap_group_on_exit"],
1426            serde_json::Value::Bool(true),
1427            "the daemon defaults a missing flag to false, so a dropped key is a silent \
1428             behaviour change: {body}"
1429        );
1430    }
1431
1432    /// A session's `Debug` does not print the agent token.
1433    #[test]
1434    fn a_session_debug_does_not_print_the_agent_token() {
1435        let recorder = Recorder::with([]);
1436        let session = Session::builder("https://vm.example", "super-secret-token")
1437            .with_backend(recorder as http::SharedBackend)
1438            .build()
1439            .expect("builds");
1440        let rendered = format!("{session:?}");
1441        assert!(rendered.contains("vm.example"), "{rendered}");
1442        assert!(
1443            !rendered.contains("super-secret"),
1444            "the agent token reached a Debug string: {rendered}"
1445        );
1446    }
1447}