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