1pub 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
71pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
73
74pub const DEFAULT_READY_TIMEOUT: Duration = Duration::from_secs(120);
76
77const READY_POLL_INTERVAL: Duration = Duration::from_secs(2);
79
80pub struct Transport {
85 backend: http::SharedBackend,
86 agent_token: String,
87 proxy: Option<Arc<ProxyAuth>>,
92 timeout: Duration,
93}
94
95impl Transport {
96 pub fn proxy(&self) -> Option<&Arc<ProxyAuth>> {
98 self.proxy.as_ref()
99 }
100
101 async fn headers(&self, token: Option<Option<&str>>) -> Result<Vec<(String, String)>, Error> {
110 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 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 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 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 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 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
187const TOKEN_INTENT: &str = "x-microvms-core-token-intent";
194
195const 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
206fn 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#[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 pub fn agent_token(&self) -> &str {
235 &self.transport.agent_token
236 }
237
238 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 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 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 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 pub fn endpoint(&self) -> &str {
316 &self.endpoint
317 }
318
319 pub fn egress_posture(&self) -> crate::control::EgressPosture {
327 self.egress_posture
328 }
329
330 pub fn port(&self) -> u16 {
332 self.port
333 }
334
335 pub fn proxy_auth(&self) -> Option<&Arc<ProxyAuth>> {
337 self.transport.proxy()
338 }
339
340 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 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 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 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 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 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 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 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 pub fn exec(&self, exec_id: impl Into<String>) -> ExecHandle {
519 ExecHandle::new(Arc::clone(&self.transport), exec_id.into())
520 }
521
522 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 pub async fn kill(&self, exec_id: &str) -> Result<bool, Error> {
533 self.exec(exec_id).kill().await
534 }
535
536 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 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 pub async fn download_file(&self, path: &str) -> Result<Vec<u8>, Error> {
563 files::download_file(&self.transport, path).await
564 }
565
566 pub async fn file_exists(&self, path: &str) -> Result<bool, Error> {
568 files::file_exists(&self.transport, path).await
569 }
570
571 pub async fn upload_tar(&self, remote: &str, archive: &[u8]) -> Result<(), Error> {
573 files::upload_tar(&self.transport, remote, archive).await
574 }
575
576 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 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
593pub 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 #[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 #[must_use]
615 pub fn with_minter(mut self, minter: Arc<dyn TokenMinter>) -> Self {
616 self.minter = Some(minter);
617 self
618 }
619
620 #[must_use]
623 pub fn with_proxy_auth(mut self, proxy: Arc<ProxyAuth>) -> Self {
624 self.proxy = Some(proxy);
625 self
626 }
627
628 #[must_use]
630 pub fn with_backend(mut self, backend: http::SharedBackend) -> Self {
631 self.backend = Some(backend);
632 self
633 }
634
635 #[must_use]
637 pub fn with_port(mut self, port: u16) -> Self {
638 self.port = port;
639 self
640 }
641
642 #[must_use]
644 pub fn with_timeout(mut self, timeout: Duration) -> Self {
645 self.timeout = timeout;
646 self
647 }
648
649 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 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 pub(crate) enum Reply {
693 Body(u16, Vec<u8>),
695 Chunks(u16, Vec<Vec<u8>>),
697 Cut(&'static str),
699 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 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 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 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 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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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}