1pub mod exec;
33pub mod files;
34pub mod forward;
35pub mod http;
36pub mod keepalive;
37pub mod proxy;
38pub mod shell;
39pub mod sse;
40pub mod tunnel;
41
42use std::sync::Arc;
43use std::time::Duration;
44
45pub use exec::{EndReason, ExecHandle, ExecResult, StreamEnd, StreamOptions, mint_exec_id};
46pub use forward::{
47 DEFAULT_EXCHANGE_TIMEOUT, ForwardClient, ForwardEvent, ForwardSpec, forwards_request_header,
48 refusal_explanation, upstream_url,
49};
50pub use http::{
51 ChunkSource, HttpBackend, HttpRequest, HttpResponse, OpenStream, ReqwestBackend, SharedBackend,
52};
53pub use keepalive::{KeepAwake, KeepAwakeEnd, KeepAwakeReport, KeepAwakeTask, RunningGate};
54pub use proxy::{
55 Clock, DEFAULT_AGENT_PORT, DEFAULT_REFRESH_AFTER, MAX_TOKEN_LIFETIME, PROXY_AUTH_HEADER,
56 PROXY_PORT_HEADER, ProxyAuth, ProxyToken, TokenMinter, TokioClock, WS_AUTH_SUBPROTOCOL_PREFIX,
57 WS_PORT_SUBPROTOCOL_PREFIX, WS_SUBPROTOCOL,
58};
59pub use sse::{ExecEvent, Frame, SseParser};
60pub use tunnel::{TUNNEL_CHUNK_BYTES, TunnelEnd, explain_close, relay_connection, tunnel_url};
61
62use crate::error::{Error, ErrorKind, WireKind};
63
64pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
66
67pub const DEFAULT_READY_TIMEOUT: Duration = Duration::from_secs(120);
69
70const READY_POLL_INTERVAL: Duration = Duration::from_secs(2);
72
73pub struct Transport {
78 backend: http::SharedBackend,
79 agent_token: String,
80 proxy: Option<Arc<ProxyAuth>>,
85 timeout: Duration,
86}
87
88impl Transport {
89 pub fn proxy(&self) -> Option<&Arc<ProxyAuth>> {
91 self.proxy.as_ref()
92 }
93
94 async fn headers(&self, token: Option<Option<&str>>) -> Result<Vec<(String, String)>, Error> {
103 let mut headers = match &self.proxy {
106 Some(proxy) => proxy.headers().await?,
107 None => Vec::new(),
108 };
109 let bearer = match token {
110 Some(explicit) => explicit,
111 None => Some(self.agent_token.as_str()),
112 };
113 if let Some(bearer) = bearer {
114 headers.push(("Authorization".to_string(), format!("Bearer {bearer}")));
115 }
116 Ok(headers)
117 }
118
119 async fn request(&self, request: HttpRequest) -> Result<HttpResponse, Error> {
121 let timeout = request.timeout.unwrap_or(self.timeout);
122 let method = request.method;
123 let path = request.path.clone();
124 tokio::time::timeout(timeout, self.request_inner(request))
125 .await
126 .map_err(|_| Error::wire(WireKind::Transport,
127 format!("{method} {path} exceeded its request timeout (including proxy authentication)")))?
128 }
129
130 async fn request_inner(&self, mut request: HttpRequest) -> Result<HttpResponse, Error> {
131 let mut headers = self.headers(request_token(&request)).await?;
140 request.headers.retain(|(name, _)| name != TOKEN_INTENT);
141 headers.append(&mut request.headers);
142 request.headers = headers;
143 if request.timeout.is_none() {
144 request.timeout = Some(self.timeout);
145 }
146 self.backend.send(request).await
147 }
148
149 async fn send(&self, request: HttpRequest) -> Result<HttpResponse, Error> {
151 let (method, path) = (request.method, request.path.clone());
152 let response = self.request(request).await?;
153 response.error_for_status(method, &path)?;
154 Ok(response)
155 }
156
157 async fn send_json<T: serde::de::DeserializeOwned>(
159 &self,
160 request: HttpRequest,
161 ) -> Result<T, Error> {
162 let (method, path) = (request.method, request.path.clone());
163 let response = self.send(request).await?;
164 serde_json::from_slice(&response.body).map_err(|err| {
165 Error::wire(
169 WireKind::ProtocolError,
170 format!(
171 "{method} {path} answered a body this client cannot read: {err}; \
172 check /v1/schema for the daemon's protocol version"
173 ),
174 )
175 .with_source(err)
176 })
177 }
178}
179
180const TOKEN_INTENT: &str = "x-microvms-core-token-intent";
187
188const TOKEN_INTENT_NONE: &str = "none";
190
191fn request_token(request: &HttpRequest) -> Option<Option<&str>> {
192 match request.header(TOKEN_INTENT) {
193 None => None,
194 Some(TOKEN_INTENT_NONE) => Some(None),
195 Some(explicit) => Some(Some(explicit)),
196 }
197}
198
199fn unauthenticated(mut request: HttpRequest) -> HttpRequest {
201 request
202 .headers
203 .push((TOKEN_INTENT.to_string(), TOKEN_INTENT_NONE.to_string()));
204 request
205}
206
207#[derive(Clone)]
212pub struct Session {
213 transport: Arc<Transport>,
214 endpoint: String,
215 port: u16,
216}
217
218impl Session {
219 pub fn agent_token(&self) -> &str {
227 &self.transport.agent_token
228 }
229
230 pub async fn connect(
236 endpoint: impl Into<String>,
237 agent_token: impl Into<String>,
238 minter: Arc<dyn TokenMinter>,
239 ) -> Result<Session, Error> {
240 Self::builder(endpoint, agent_token)
241 .with_minter(minter)
242 .build()
243 }
244
245 pub async fn attach(
252 region: crate::Region,
253 microvm_id: impl Into<String>,
254 endpoint: impl Into<String>,
255 agent_token: impl Into<String>,
256 port: Option<u16>,
257 request_timeout: Option<Duration>,
258 ) -> Result<Self, Error> {
259 if request_timeout.is_some_and(|timeout| timeout.is_zero()) {
260 return Err(Error::invalid_arg("request timeout must be positive"));
261 }
262 let mut control = crate::control::ControlPlane::new(region).await?;
263 if let Some(port) = port {
264 control = control.with_port(port)?;
265 }
266 let port = control.port();
267 let minter = Arc::new(crate::sandbox::ControlPlaneMinter {
268 control: Arc::new(control),
269 microvm_id: microvm_id.into(),
270 });
271 let mut builder = Self::builder(endpoint, agent_token)
272 .with_minter(minter)
273 .with_port(port);
274 if let Some(timeout) = request_timeout {
275 builder = builder.with_timeout(timeout);
276 }
277 builder.build()
278 }
279
280 pub fn direct(
286 endpoint: impl Into<String>,
287 agent_token: impl Into<String>,
288 ) -> Result<Session, Error> {
289 Self::builder(endpoint, agent_token).build()
290 }
291
292 pub fn builder(endpoint: impl Into<String>, agent_token: impl Into<String>) -> SessionBuilder {
294 SessionBuilder {
295 endpoint: endpoint.into(),
296 agent_token: agent_token.into(),
297 minter: None,
298 backend: None,
299 port: DEFAULT_AGENT_PORT,
300 timeout: DEFAULT_REQUEST_TIMEOUT,
301 proxy: None,
302 }
303 }
304
305 pub fn endpoint(&self) -> &str {
307 &self.endpoint
308 }
309
310 pub fn port(&self) -> u16 {
312 self.port
313 }
314
315 pub fn proxy_auth(&self) -> Option<&Arc<ProxyAuth>> {
317 self.transport.proxy()
318 }
319
320 pub async fn connect_headers(&self, port: u16) -> Result<Vec<(String, String)>, Error> {
342 match self.transport.proxy() {
343 Some(proxy) => {
344 tokio::time::timeout(self.transport.timeout, proxy.headers_for_port(port))
345 .await
346 .map_err(|_| {
347 Error::wire(
348 WireKind::Transport,
349 "proxy authentication request timed out",
350 )
351 })?
352 }
353 None => Ok(Vec::new()),
354 }
355 }
356
357 pub async fn connect_subprotocols(&self, port: u16) -> Result<Option<[String; 3]>, Error> {
369 match self.transport.proxy() {
370 Some(proxy) => Ok(Some(
371 tokio::time::timeout(self.transport.timeout, proxy.subprotocols(port))
372 .await
373 .map_err(|_| {
374 Error::wire(
375 WireKind::Transport,
376 "proxy authentication request timed out",
377 )
378 })??,
379 )),
380 None => Ok(None),
381 }
382 }
383
384 pub fn rebind(&mut self, endpoint: String) {
396 self.endpoint = endpoint;
397 if let Some(proxy) = self.transport.proxy() {
398 proxy.invalidate();
399 }
400 }
401
402 pub async fn health(&self) -> Result<protocol::health::Health, Error> {
411 self.transport
412 .send_json(unauthenticated(HttpRequest::new("GET", "/v1/health")))
413 .await
414 }
415
416 pub async fn keep_awake(
421 &self,
422 policy: &KeepAwake,
423 stop: impl std::future::Future<Output = ()>,
424 ) -> Result<KeepAwakeReport, Error> {
425 policy
426 .run(|| async { self.health().await.map(Some) }, stop)
427 .await
428 }
429
430 pub async fn wait_until_ready(
438 &self,
439 timeout: Duration,
440 ) -> Result<protocol::health::Health, Error> {
441 let deadline = tokio::time::Instant::now() + timeout;
442 let mut last: Option<Error> = None;
443 loop {
444 match self.health().await {
445 Ok(health) if health.bootstrapped => return Ok(health),
446 Ok(_) => {}
447 Err(err) if err.retryable() => last = Some(err),
448 Err(err) => return Err(err),
449 }
450 let now = tokio::time::Instant::now();
451 if now >= deadline {
452 let detail = match last {
453 Some(err) => format!(" (last error: {err})"),
454 None => String::new(),
455 };
456 return Err(Error::new(
457 ErrorKind::Timeout,
458 format!(
459 "the daemon was not bootstrapped within {}s{detail}",
460 timeout.as_secs()
461 ),
462 ));
463 }
464 tokio::time::sleep(READY_POLL_INTERVAL.min(deadline - now)).await;
465 }
466 }
467
468 pub async fn run(&self, req: protocol::exec::StartRequest) -> Result<ExecHandle, Error> {
476 let requested = req.exec_id.clone();
477 let mut request = HttpRequest::new("POST", "/v1/exec/start");
478 request
479 .headers
480 .push(("content-type".into(), "application/json".into()));
481 request.body = serde_json::to_vec(&req).map_err(|err| {
482 Error::invalid_arg(format!("the start request will not serialize: {err}"))
483 })?;
484 let started: protocol::exec::StartResponse = self.transport.send_json(request).await?;
485 let exec_id = if started.exec_id.is_empty() {
490 requested
491 } else {
492 started.exec_id
493 };
494 Ok(ExecHandle::new(Arc::clone(&self.transport), exec_id))
495 }
496
497 pub fn exec(&self, exec_id: impl Into<String>) -> ExecHandle {
499 ExecHandle::new(Arc::clone(&self.transport), exec_id.into())
500 }
501
502 pub async fn run_sync(
504 &self,
505 req: protocol::exec::StartRequest,
506 timeout: Duration,
507 ) -> Result<ExecResult, Error> {
508 self.run(req).await?.wait_and_ack(timeout).await
509 }
510
511 pub async fn kill(&self, exec_id: &str) -> Result<bool, Error> {
513 self.exec(exec_id).kill().await
514 }
515
516 pub async fn procs(&self) -> Result<protocol::exec::ProcsResponse, Error> {
524 self.transport
525 .send_json(HttpRequest::new("GET", "/v1/procs"))
526 .await
527 }
528
529 pub async fn upload_file(
533 &self,
534 path: &str,
535 data: &[u8],
536 mode: Option<&str>,
537 ) -> Result<(), Error> {
538 files::upload_file(&self.transport, path, data, mode).await
539 }
540
541 pub async fn download_file(&self, path: &str) -> Result<Vec<u8>, Error> {
543 files::download_file(&self.transport, path).await
544 }
545
546 pub async fn file_exists(&self, path: &str) -> Result<bool, Error> {
548 files::file_exists(&self.transport, path).await
549 }
550
551 pub async fn upload_tar(&self, remote: &str, archive: &[u8]) -> Result<(), Error> {
553 files::upload_tar(&self.transport, remote, archive).await
554 }
555
556 pub async fn download_tar(&self, remote: &str) -> Result<Vec<u8>, Error> {
558 files::download_tar(&self.transport, remote).await
559 }
560}
561
562impl std::fmt::Debug for Session {
563 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
564 f.debug_struct("Session")
566 .field("endpoint", &self.endpoint)
567 .field("port", &self.port)
568 .field("proxy", &self.transport.proxy())
569 .finish()
570 }
571}
572
573pub struct SessionBuilder {
575 endpoint: String,
576 agent_token: String,
577 minter: Option<Arc<dyn TokenMinter>>,
578 backend: Option<http::SharedBackend>,
579 port: u16,
580 timeout: Duration,
581 proxy: Option<Arc<ProxyAuth>>,
582}
583
584impl SessionBuilder {
585 #[must_use]
587 pub fn with_minter(mut self, minter: Arc<dyn TokenMinter>) -> Self {
588 self.minter = Some(minter);
589 self
590 }
591
592 #[must_use]
595 pub fn with_proxy_auth(mut self, proxy: Arc<ProxyAuth>) -> Self {
596 self.proxy = Some(proxy);
597 self
598 }
599
600 #[must_use]
602 pub fn with_backend(mut self, backend: http::SharedBackend) -> Self {
603 self.backend = Some(backend);
604 self
605 }
606
607 #[must_use]
609 pub fn with_port(mut self, port: u16) -> Self {
610 self.port = port;
611 self
612 }
613
614 #[must_use]
616 pub fn with_timeout(mut self, timeout: Duration) -> Self {
617 self.timeout = timeout;
618 self
619 }
620
621 pub fn build(self) -> Result<Session, Error> {
623 let backend = match self.backend {
624 Some(backend) => backend,
625 None => Arc::new(ReqwestBackend::new(&self.endpoint, self.timeout)?),
626 };
627 let proxy = match (self.proxy, self.minter) {
628 (Some(proxy), _) => Some(proxy),
629 (None, Some(minter)) => Some(Arc::new(ProxyAuth::new(minter, self.port))),
630 (None, None) => None,
631 };
632 Ok(Session {
633 transport: Arc::new(Transport {
634 backend,
635 agent_token: self.agent_token,
636 proxy,
637 timeout: self.timeout,
638 }),
639 endpoint: self.endpoint,
640 port: self.port,
641 })
642 }
643}
644
645#[cfg(test)]
646pub(crate) mod testing {
647 use std::collections::HashMap;
655 use std::sync::{Mutex, PoisonError};
656
657 use futures_util::future::BoxFuture;
658
659 use super::http::{ChunkSource, HttpBackend, HttpRequest, HttpResponse, OpenStream};
660 use super::*;
661
662 pub(crate) enum Reply {
664 Body(u16, Vec<u8>),
666 Chunks(u16, Vec<Vec<u8>>),
668 Cut(&'static str),
670 }
671
672 impl Reply {
673 pub(crate) fn json(status: u16, body: impl serde::Serialize) -> Self {
674 Reply::Body(status, serde_json::to_vec(&body).expect("test body"))
675 }
676
677 pub(crate) fn ok(body: impl serde::Serialize) -> Self {
678 Reply::json(200, body)
679 }
680 }
681
682 #[derive(Default)]
683 pub(crate) struct Recorder {
684 seen: Mutex<Vec<HttpRequest>>,
685 replies: Mutex<std::collections::VecDeque<Reply>>,
686 }
687
688 impl Recorder {
689 pub(crate) fn with(replies: impl IntoIterator<Item = Reply>) -> Arc<Self> {
690 Arc::new(Self {
691 seen: Mutex::new(Vec::new()),
692 replies: Mutex::new(replies.into_iter().collect()),
693 })
694 }
695
696 pub(crate) fn requests(&self) -> Vec<HttpRequest> {
698 self.seen
699 .lock()
700 .unwrap_or_else(PoisonError::into_inner)
701 .clone()
702 }
703
704 pub(crate) fn last(&self) -> HttpRequest {
705 self.requests()
706 .pop()
707 .expect("the recorder saw no request at all")
708 }
709
710 fn record(&self, request: HttpRequest) -> Reply {
711 self.seen
712 .lock()
713 .unwrap_or_else(PoisonError::into_inner)
714 .push(request.clone());
715 self.replies
716 .lock()
717 .unwrap_or_else(PoisonError::into_inner)
718 .pop_front()
719 .unwrap_or_else(|| {
720 panic!(
721 "the recorder ran out of replies at {} {}",
722 request.method, request.path
723 )
724 })
725 }
726 }
727
728 struct Queued(std::collections::VecDeque<Vec<u8>>);
729
730 impl ChunkSource for Queued {
731 fn next_chunk(&mut self) -> BoxFuture<'_, Result<Option<Vec<u8>>, Error>> {
732 Box::pin(async move { Ok(self.0.pop_front()) })
733 }
734 }
735
736 impl HttpBackend for Recorder {
737 fn send(&self, request: HttpRequest) -> BoxFuture<'_, Result<HttpResponse, Error>> {
738 let reply = self.record(request);
739 Box::pin(async move {
740 match reply {
741 Reply::Body(status, body) => Ok(HttpResponse {
742 status,
743 headers: HashMap::new(),
744 body,
745 }),
746 Reply::Chunks(status, chunks) => Ok(HttpResponse {
747 status,
748 headers: HashMap::new(),
749 body: chunks.concat(),
750 }),
751 Reply::Cut(why) => Err(Error::wire(WireKind::Transport, why)),
752 }
753 })
754 }
755
756 fn open_stream(
757 &self,
758 request: HttpRequest,
759 _idle_timeout: Duration,
760 ) -> BoxFuture<'_, Result<OpenStream, Error>> {
761 let reply = self.record(request);
762 Box::pin(async move {
763 let (status, chunks) = match reply {
764 Reply::Chunks(status, chunks) => (status, chunks),
765 Reply::Body(status, body) => (status, vec![body]),
766 Reply::Cut(why) => return Err(Error::wire(WireKind::Transport, why)),
767 };
768 let head = HttpResponse {
769 status,
770 headers: HashMap::new(),
771 body: if (200..300).contains(&status) {
774 Vec::new()
775 } else {
776 chunks.concat()
777 },
778 };
779 let source: Box<dyn ChunkSource> = if (200..300).contains(&status) {
780 Box::new(Queued(chunks.into_iter().collect()))
781 } else {
782 Box::new(Queued(std::collections::VecDeque::new()))
783 };
784 Ok((head, source))
785 })
786 }
787 }
788
789 pub(crate) fn session_with(
791 recorder: Arc<Recorder>,
792 ) -> (Session, Arc<ProxyAuth>, Arc<proxy::testing::ManualClock>) {
793 let clock = Arc::new(proxy::testing::ManualClock::default());
794 let auth = Arc::new(
795 ProxyAuth::with_refresh_after(
796 Arc::new(proxy::testing::CountingMinter::default()),
797 DEFAULT_AGENT_PORT,
798 DEFAULT_REFRESH_AFTER,
799 Arc::clone(&clock) as Arc<dyn Clock>,
800 )
801 .expect("the default interval is below the ceiling"),
802 );
803 let session = Session::builder("https://vm.example", "agent-token-abcdef")
804 .with_backend(recorder)
805 .with_proxy_auth(Arc::clone(&auth))
806 .build()
807 .expect("builds");
808 (session, auth, clock)
809 }
810
811 pub(crate) fn health_body(bootstrapped: bool) -> serde_json::Value {
812 serde_json::json!({
813 "version": "0.1.0",
814 "bootstrapped": bootstrapped,
815 "disk": null,
816 "identity_degraded": false,
817 "identity_repaired": true,
818 })
819 }
820}
821
822#[cfg(test)]
823mod tests {
824 use super::testing::{Recorder, Reply, health_body, session_with};
825 use super::*;
826 use crate::session::proxy::testing::CountingMinter;
827
828 fn start_request(exec_id: &str) -> protocol::exec::StartRequest {
829 protocol::exec::StartRequest {
830 exec_id: exec_id.to_string(),
831 command: vec!["/bin/true".to_string()],
832 shell: false,
833 cwd: None,
834 env: Default::default(),
835 user: None,
836 group: None,
837 timeout_sec: None,
838 stdin: false,
839 reap_group_on_exit: false,
840 }
841 }
842
843 fn header(request: &HttpRequest, name: &str) -> Option<String> {
844 request.header(name).map(str::to_string)
845 }
846
847 fn header_of(headers: &[(String, String)], name: &str) -> Option<String> {
849 headers
850 .iter()
851 .find(|(key, _)| key.eq_ignore_ascii_case(name))
852 .map(|(_, value)| value.clone())
853 }
854
855 #[tokio::test(start_paused = true)]
856 async fn request_timeout_also_bounds_a_blocked_proxy_mint() {
857 struct BlockedMinter;
858 impl TokenMinter for BlockedMinter {
859 fn mint(&self) -> futures_util::future::BoxFuture<'_, Result<ProxyToken, Error>> {
860 Box::pin(std::future::pending())
861 }
862 }
863 let backend = Recorder::with([]);
864 let session = Session::builder("http://127.0.0.1:9", "private")
865 .with_backend(backend.clone())
866 .with_minter(Arc::new(BlockedMinter))
867 .with_timeout(Duration::from_secs(3))
868 .build()
869 .expect("session");
870 let result = tokio::time::timeout(Duration::from_secs(4), session.health())
871 .await
872 .expect("the request budget must include a stalled proxy mint")
873 .expect_err("mint never completed");
874 assert_eq!(result.wire_kind(), Some(WireKind::Transport));
875 assert!(result.retryable());
876 let result = tokio::time::timeout(Duration::from_secs(4), session.connect_headers(9000))
877 .await
878 .expect("external transfers must bound authentication too")
879 .expect_err("mint remains blocked");
880 assert!(result.retryable());
881 let result =
882 tokio::time::timeout(Duration::from_secs(4), session.connect_subprotocols(9000))
883 .await
884 .expect("external sockets must bound authentication too")
885 .expect_err("mint remains blocked");
886 assert!(result.retryable());
887 assert!(
888 backend.requests().is_empty(),
889 "no request could authenticate"
890 );
891 }
892
893 #[tokio::test]
901 async fn a_callers_content_type_survives_the_auth_header_injection() {
902 let recorder = Recorder::with([Reply::ok(
903 serde_json::json!({"exec_id":"e1","phase":"running"}),
904 )]);
905 let (session, _, _) = session_with(Arc::clone(&recorder));
906
907 session.run(start_request("e1")).await.expect("start");
908
909 let seen = recorder.requests();
910 assert_eq!(
911 header(&seen[0], "content-type").as_deref(),
912 Some("application/json"),
913 "the exec start lost its content type to the auth header injection"
914 );
915 }
916
917 #[tokio::test]
922 async fn the_session_hands_back_the_token_it_authenticates_with() {
923 let recorder = Recorder::with([]);
924 let (session, _, _) = session_with(Arc::clone(&recorder));
925 assert_eq!(session.agent_token(), "agent-token-abcdef");
926 }
927
928 #[tokio::test]
935 async fn every_endpoint_request_carries_both_proxy_headers() {
936 let recorder = Recorder::with([
937 Reply::ok(health_body(true)),
938 Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
939 ]);
940 let (session, _, _) = session_with(Arc::clone(&recorder));
941
942 session.health().await.expect("health succeeds");
943 session.run(start_request("e1")).await.expect("start");
944
945 let seen = recorder.requests();
946 assert_eq!(seen.len(), 2);
947 for request in &seen {
948 assert_eq!(
949 header(request, PROXY_AUTH_HEADER).as_deref(),
950 Some(CountingMinter::value(0).as_str()),
951 "{} {} went out without the proxy auth header",
952 request.method,
953 request.path
954 );
955 assert_eq!(
956 header(request, PROXY_PORT_HEADER).as_deref(),
957 Some("9000"),
958 "{} {} went out without the proxy port header, which the proxy \
959 rejects as if the token were bad",
960 request.method,
961 request.path
962 );
963 }
964 }
965
966 #[tokio::test]
972 async fn health_sends_no_bearer_and_an_exec_route_does() {
973 let recorder = Recorder::with([
974 Reply::ok(health_body(true)),
975 Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
976 ]);
977 let (session, _, _) = session_with(Arc::clone(&recorder));
978
979 session.health().await.expect("health");
980 let seen = recorder.requests();
981 assert_eq!(header(&seen[0], "authorization"), None);
982 assert_eq!(
983 header(&seen[0], "x-microvms-core-token-intent"),
984 None,
985 "the internal token-intent marker must not reach the wire"
986 );
987
988 session.run(start_request("e1")).await.expect("start");
989 assert_eq!(
990 header(&recorder.last(), "authorization").as_deref(),
991 Some("Bearer agent-token-abcdef")
992 );
993 }
994
995 #[tokio::test(start_paused = true)]
998 async fn keep_awake_polls_health_through_the_proxy_until_the_vm_is_idle() {
999 let busy = |busy: bool| {
1000 let mut body = health_body(true);
1001 body["busy"] = serde_json::json!(busy);
1002 Reply::ok(body)
1003 };
1004 let recorder = Recorder::with([busy(true), busy(true), busy(false)]);
1005 let (session, _, _) = session_with(Arc::clone(&recorder));
1006 let report = session
1007 .clone()
1008 .keep_awake(
1009 &KeepAwake::new(None).while_busy(true),
1010 std::future::pending(),
1011 )
1012 .await
1013 .expect("keepalive");
1014 assert_eq!((report.end, report.polls), (KeepAwakeEnd::Idle, 3));
1015 let seen = recorder.requests();
1016 assert_eq!(seen.len(), 3);
1017 for request in &seen {
1018 assert_eq!(
1019 (request.method, request.path.as_str()),
1020 ("GET", "/v1/health")
1021 );
1022 assert_eq!(header(request, "authorization"), None);
1023 assert!(header(request, PROXY_AUTH_HEADER).is_some());
1024 assert!(header(request, PROXY_PORT_HEADER).is_some());
1025 }
1026 }
1027
1028 #[tokio::test]
1031 async fn a_direct_session_sends_no_proxy_headers() {
1032 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1033 let session = Session::builder("http://127.0.0.1:9000", "token")
1034 .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1035 .build()
1036 .expect("builds");
1037 assert!(session.proxy_auth().is_none());
1038
1039 session.health().await.expect("health");
1040 let request = recorder.last();
1041 assert_eq!(header(&request, PROXY_AUTH_HEADER), None);
1042 assert_eq!(header(&request, PROXY_PORT_HEADER), None);
1043 }
1044
1045 #[tokio::test]
1048 async fn a_mint_failure_surfaces_as_a_retryable_request_error() {
1049 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1050 let auth = Arc::new(
1051 ProxyAuth::with_refresh_after(
1052 Arc::new(CountingMinter::failing(1)),
1053 DEFAULT_AGENT_PORT,
1054 DEFAULT_REFRESH_AFTER,
1055 Arc::new(proxy::testing::ManualClock::default()),
1056 )
1057 .expect("interval accepted"),
1058 );
1059 let session = Session::builder("https://vm.example", "token")
1060 .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1061 .with_proxy_auth(Arc::clone(&auth))
1062 .build()
1063 .expect("builds");
1064
1065 let err = session.health().await.expect_err("the mint fails");
1066 assert_eq!(err.kind(), ErrorKind::Retryable);
1067 assert_eq!(err.wire_kind(), Some(WireKind::AuthTokenMint));
1068 assert!(
1069 recorder.requests().is_empty(),
1070 "a request went out without a token"
1071 );
1072
1073 session.health().await.expect("the retry succeeds");
1074 assert_eq!(auth.mint_count(), 1);
1075 }
1076
1077 #[tokio::test]
1080 async fn a_daemon_status_arrives_as_its_typed_error() {
1081 for (status, expected) in [
1082 (401, WireKind::Unauthorized),
1083 (404, WireKind::NotFound),
1084 (409, WireKind::Conflict),
1085 (503, WireKind::NotBootstrapped),
1086 ] {
1087 let recorder = Recorder::with([Reply::Body(status, b"{\"error\":\"x\"}".to_vec())]);
1088 let (session, _, _) = session_with(recorder);
1089 let err = session.health().await.expect_err("a failing status");
1090 assert_eq!(err.wire_kind(), Some(expected), "status {status}");
1091 }
1092 }
1093
1094 #[tokio::test(start_paused = true)]
1097 async fn wait_until_ready_polls_through_not_bootstrapped_and_a_dropped_connection() {
1098 let recorder = Recorder::with([
1099 Reply::Cut("connection refused"),
1100 Reply::Body(503, b"not bootstrapped".to_vec()),
1101 Reply::ok(health_body(false)),
1102 Reply::ok(health_body(true)),
1103 ]);
1104 let (session, _, _) = session_with(Arc::clone(&recorder));
1105
1106 let health = session
1107 .wait_until_ready(DEFAULT_READY_TIMEOUT)
1108 .await
1109 .expect("the daemon comes up");
1110 assert!(health.bootstrapped);
1111 assert_eq!(recorder.requests().len(), 4);
1112 }
1113
1114 #[tokio::test(start_paused = true)]
1117 async fn wait_until_ready_gives_up_at_once_on_a_credential_failure() {
1118 let recorder = Recorder::with([Reply::Body(401, b"wrong token".to_vec())]);
1119 let (session, _, _) = session_with(Arc::clone(&recorder));
1120
1121 let err = session
1122 .wait_until_ready(DEFAULT_READY_TIMEOUT)
1123 .await
1124 .expect_err("a 401 is fatal");
1125 assert_eq!(err.kind(), ErrorKind::Credentials);
1126 assert_eq!(
1127 recorder.requests().len(),
1128 1,
1129 "a fatal error was retried, so a wrong token costs the full timeout"
1130 );
1131 }
1132
1133 #[tokio::test(start_paused = true)]
1136 async fn wait_until_ready_times_out_naming_the_last_retryable_error() {
1137 let recorder = Recorder::with([
1138 Reply::Cut("connection refused"),
1139 Reply::Cut("connection refused"),
1140 Reply::Cut("connection refused"),
1141 ]);
1142 let (session, _, _) = session_with(recorder);
1143
1144 let err = session
1145 .wait_until_ready(Duration::from_secs(4))
1146 .await
1147 .expect_err("the daemon never comes up");
1148 assert_eq!(err.kind(), ErrorKind::Timeout);
1149 assert!(
1150 err.to_string().contains("connection refused"),
1151 "the timeout must name what it kept hitting: {err}"
1152 );
1153 }
1154
1155 #[tokio::test]
1161 async fn rebind_invalidates_the_proxy_token_so_the_next_request_mints() {
1162 let recorder = Recorder::with([Reply::ok(health_body(true)), Reply::ok(health_body(true))]);
1163 let (mut session, auth, _) = session_with(Arc::clone(&recorder));
1164
1165 session.health().await.expect("health");
1166 assert_eq!(auth.mint_count(), 1);
1167
1168 session.rebind("https://vm.example".to_string());
1169 assert!(!auth.is_cached(), "rebind did not drop the token");
1170
1171 session.health().await.expect("health");
1172 assert_eq!(auth.mint_count(), 2);
1173 assert_eq!(
1174 header(&recorder.last(), PROXY_AUTH_HEADER).as_deref(),
1175 Some(CountingMinter::value(1).as_str()),
1176 "the request after a rebind carried the pre-suspend token"
1177 );
1178 }
1179
1180 #[tokio::test]
1183 async fn a_start_request_serializes_through_the_protocol_crate() {
1184 let recorder = Recorder::with([Reply::ok(
1185 serde_json::json!({"exec_id":"e-confirmed","phase":"running"}),
1186 )]);
1187 let (session, _, _) = session_with(Arc::clone(&recorder));
1188
1189 let handle = session.run(start_request("e-asked")).await.expect("start");
1190 assert_eq!(
1191 handle.exec_id(),
1192 "e-confirmed",
1193 "the handle addresses the id the client asked for rather than the one the \
1194 daemon confirmed, so every later call would 404 if they diverged"
1195 );
1196
1197 let request = recorder.last();
1198 assert_eq!(request.method, "POST");
1199 assert_eq!(request.path, "/v1/exec/start");
1200 let body: serde_json::Value = serde_json::from_slice(&request.body).expect("json body");
1201 assert_eq!(body["exec_id"], "e-asked");
1202 assert_eq!(body["command"], serde_json::json!(["/bin/true"]));
1203 }
1204
1205 #[tokio::test]
1208 async fn an_unreadable_response_body_names_the_schema_route() {
1209 let recorder = Recorder::with([Reply::Body(200, b"{\"unexpected\":true}".to_vec())]);
1210 let (session, _, _) = session_with(recorder);
1211 let err = session.health().await.expect_err("the body will not parse");
1212 assert_eq!(err.wire_kind(), Some(WireKind::ProtocolError));
1213 assert!(err.to_string().contains("/v1/schema"), "{err}");
1214 }
1215
1216 #[tokio::test]
1237 async fn a_session_answers_connect_credentials_for_a_port_off_its_own_token_cache() {
1238 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1239 let (session, auth, _) = session_with(Arc::clone(&recorder));
1240
1241 session.health().await.expect("health warms the cache");
1242 assert_eq!(auth.mint_count(), 1);
1243
1244 let headers = session
1245 .connect_headers(8080)
1246 .await
1247 .expect("the cached token answers");
1248 let offered = session
1249 .connect_subprotocols(8080)
1250 .await
1251 .expect("the cached token answers")
1252 .expect("a proxied session offers subprotocols");
1253
1254 assert_eq!(
1255 header_of(&headers, PROXY_PORT_HEADER).as_deref(),
1256 Some("8080"),
1257 "the caller asked for 8080 and would have targeted the session's port"
1258 );
1259 assert_eq!(
1261 header_of(&headers, PROXY_AUTH_HEADER).as_deref(),
1262 Some(CountingMinter::value(1).as_str()),
1263 "the header pair carried the agent-port-only token, which the proxy refuses for \
1264 8080 with 403 'Access to port denied'"
1265 );
1266 assert_eq!(offered[0], proxy::WS_SUBPROTOCOL);
1267 assert_eq!(
1268 offered[1],
1269 format!(
1270 "{}{}",
1271 proxy::WS_AUTH_SUBPROTOCOL_PREFIX,
1272 CountingMinter::value(1)
1273 ),
1274 "the handshake carried a different token than the header path, so they are two \
1275 caches rather than one"
1276 );
1277 assert_eq!(offered[2], "lambda-microvms.port.8080");
1278 assert_eq!(
1279 auth.mint_count(),
1280 2,
1281 "one mint covered both the agent port and 8080, which the control plane does not \
1282 do — or the second accessor minted a third, which would be two caches"
1283 );
1284 assert!(
1285 recorder.requests().len() == 1,
1286 "building credentials sent a request to the VM: {:?}",
1287 recorder.requests()
1288 );
1289 }
1290
1291 #[tokio::test]
1302 async fn a_direct_session_offers_no_subprotocols_and_empty_headers() {
1303 let recorder = Recorder::with([]);
1304 let session = Session::builder("http://127.0.0.1:9000", "token")
1305 .with_backend(recorder as http::SharedBackend)
1306 .build()
1307 .expect("builds");
1308
1309 assert!(
1310 session
1311 .connect_headers(9000)
1312 .await
1313 .expect("no proxy is not a failure")
1314 .is_empty(),
1315 "a direct session's requests carry no proxy headers, so neither should this"
1316 );
1317 assert!(
1318 session
1319 .connect_subprotocols(9000)
1320 .await
1321 .expect("no proxy is not a failure")
1322 .is_none(),
1323 "a handshake with the base subprotocol and no token is refused for a reason that \
1324 names neither"
1325 );
1326 }
1327
1328 #[tokio::test]
1339 async fn connect_credentials_re_mint_when_the_refresh_window_has_rolled_over() {
1340 let recorder = Recorder::with([]);
1341 let (session, auth, clock) = session_with(recorder);
1342
1343 let first = session
1344 .connect_subprotocols(9000)
1345 .await
1346 .expect("mints")
1347 .expect("proxied");
1348 assert!(first[1].ends_with(&CountingMinter::value(0)));
1349
1350 clock.advance(DEFAULT_REFRESH_AFTER + Duration::from_secs(1));
1351 let second = session
1352 .connect_subprotocols(9000)
1353 .await
1354 .expect("re-mints")
1355 .expect("proxied");
1356
1357 assert_eq!(auth.mint_count(), 2);
1358 assert!(
1359 second[1].ends_with(&CountingMinter::value(1)),
1360 "the handshake would carry a token past its refresh window: {}",
1361 second[1]
1362 );
1363 }
1364
1365 #[tokio::test]
1375 async fn procs_is_a_bearer_get_of_the_procs_route_with_a_typed_body() {
1376 let recorder = Recorder::with([Reply::ok(serde_json::json!({
1377 "procs": [{
1378 "exec_id": "e1",
1379 "pgid": 4242,
1380 "started_at": 1_756_500_000u64,
1381 "child_exited": true,
1382 "reap": false,
1383 "pids": [4243],
1384 }]
1385 }))]);
1386 let (session, _, _) = session_with(Arc::clone(&recorder));
1387
1388 let listed = session.procs().await.expect("procs");
1389 let request = recorder.last();
1390 assert_eq!(request.method, "GET");
1391 assert_eq!(request.path, "/v1/procs");
1392 assert_eq!(
1393 header(&request, "authorization").as_deref(),
1394 Some("Bearer agent-token-abcdef"),
1395 "the procs route is bearer-guarded; without the token it is a 503 or a 401"
1396 );
1397 assert_eq!(listed.procs.len(), 1);
1398 assert_eq!(listed.procs[0].exec_id, "e1");
1399 assert_eq!(listed.procs[0].pgid, Some(4242));
1400 assert!(listed.procs[0].child_exited);
1401 assert_eq!(listed.procs[0].pids, [4243]);
1402 }
1403
1404 #[tokio::test]
1412 async fn a_start_request_carries_the_reap_flag_on_the_wire() {
1413 let recorder = Recorder::with([Reply::ok(
1414 serde_json::json!({"exec_id":"e1","phase":"running"}),
1415 )]);
1416 let (session, _, _) = session_with(Arc::clone(&recorder));
1417
1418 let mut request = start_request("e1");
1419 request.reap_group_on_exit = true;
1420 session.run(request).await.expect("start");
1421
1422 let body: serde_json::Value =
1423 serde_json::from_slice(&recorder.last().body).expect("a JSON start body");
1424 assert_eq!(
1425 body["reap_group_on_exit"],
1426 serde_json::Value::Bool(true),
1427 "the daemon defaults a missing flag to false, so a dropped key is a silent \
1428 behaviour change: {body}"
1429 );
1430 }
1431
1432 #[test]
1434 fn a_session_debug_does_not_print_the_agent_token() {
1435 let recorder = Recorder::with([]);
1436 let session = Session::builder("https://vm.example", "super-secret-token")
1437 .with_backend(recorder as http::SharedBackend)
1438 .build()
1439 .expect("builds");
1440 let rendered = format!("{session:?}");
1441 assert!(rendered.contains("vm.example"), "{rendered}");
1442 assert!(
1443 !rendered.contains("super-secret"),
1444 "the agent token reached a Debug string: {rendered}"
1445 );
1446 }
1447}