1pub mod complete;
33pub mod exec;
34pub mod files;
35pub mod http;
36pub mod keepalive;
37pub mod proxy;
38pub mod split;
39pub mod sse;
40
41use std::sync::Arc;
42use std::time::Duration;
43
44pub use complete::{
45 ClientDeadline, CompletionOptions, CompletionPlan, DEFAULT_CLIENT_GRACE, KillAnswer,
46 NO_TIMEOUT_CEILING, OutputFlow, OutputSink,
47};
48pub use exec::{
49 DEFAULT_EXEC_WAIT, EndReason, ExecHandle, ExecResult, StreamEnd, StreamOptions,
50 TIMED_OUT_EXIT_CODE, exec_id_at,
51};
52pub use http::{ChunkSource, HttpBackend, HttpRequest, HttpResponse, OpenStream, SharedBackend};
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, WS_AUTH_SUBPROTOCOL_PREFIX,
57 WS_PORT_SUBPROTOCOL_PREFIX, WS_SUBPROTOCOL,
58};
59pub use split::{GapLog, GapPolicy, OutputGap, Split, SplitItem};
60pub use sse::{ExecEvent, Frame, SseParser};
61
62use crate::error::{Error, ErrorKind, WireKind};
63
64pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
66
67pub const DEFAULT_BOOTSTRAP_TIMEOUT: Duration = Duration::from_secs(120);
71
72#[deprecated(
74 note = "renamed DEFAULT_BOOTSTRAP_TIMEOUT: `sandbox::DEFAULT_READY_TIMEOUT` shared this name \
75 for a different wait, the one for RUNNING"
76)]
77pub const DEFAULT_READY_TIMEOUT: Duration = DEFAULT_BOOTSTRAP_TIMEOUT;
78
79const READY_POLL_INTERVAL: Duration = Duration::from_secs(2);
81
82pub struct Transport {
87 backend: http::SharedBackend,
88 agent_token: String,
89 proxy: Option<Arc<ProxyAuth>>,
94 timeout: Duration,
95}
96
97impl Transport {
98 pub fn proxy(&self) -> Option<&Arc<ProxyAuth>> {
100 self.proxy.as_ref()
101 }
102
103 async fn headers(&self, token: Option<Option<&str>>) -> Result<Vec<(String, String)>, Error> {
112 let mut headers = match &self.proxy {
115 Some(proxy) => proxy.headers().await?,
116 None => Vec::new(),
117 };
118 let bearer = match token {
119 Some(explicit) => explicit,
120 None => Some(self.agent_token.as_str()),
121 };
122 if let Some(bearer) = bearer {
123 headers.push(("Authorization".to_string(), format!("Bearer {bearer}")));
124 }
125 Ok(headers)
126 }
127
128 async fn request(&self, request: HttpRequest) -> Result<HttpResponse, Error> {
130 let timeout = request.timeout.unwrap_or(self.timeout);
131 let method = request.method;
132 let path = request.path.clone();
133 tokio::time::timeout(timeout, self.request_inner(request))
134 .await
135 .map_err(|_| Error::wire(WireKind::Transport,
136 format!("{method} {path} exceeded its request timeout (including proxy authentication)")))?
137 }
138
139 async fn request_inner(&self, mut request: HttpRequest) -> Result<HttpResponse, Error> {
140 let mut headers = self.headers(request_token(&request)).await?;
149 request.headers.retain(|(name, _)| name != TOKEN_INTENT);
150 headers.append(&mut request.headers);
151 request.headers = headers;
152 if request.timeout.is_none() {
153 request.timeout = Some(self.timeout);
154 }
155 self.backend.send(request).await
156 }
157
158 async fn send(&self, request: HttpRequest) -> Result<HttpResponse, Error> {
160 let (method, path) = (request.method, request.path.clone());
161 let response = self.request(request).await?;
162 response.error_for_status(method, &path)?;
163 Ok(response)
164 }
165
166 async fn send_json<T: serde::de::DeserializeOwned>(
168 &self,
169 request: HttpRequest,
170 ) -> Result<T, Error> {
171 let (method, path) = (request.method, request.path.clone());
172 let response = self.send(request).await?;
173 serde_json::from_slice(&response.body).map_err(|err| {
174 Error::wire(
178 WireKind::ProtocolError,
179 format!(
180 "{method} {path} answered a body this client cannot read: {err}; \
181 check /v1/schema for the daemon's protocol version"
182 ),
183 )
184 .with_source(err)
185 })
186 }
187}
188
189const TOKEN_INTENT: &str = "x-microvms-core-token-intent";
196
197const TOKEN_INTENT_NONE: &str = "none";
199
200fn request_token(request: &HttpRequest) -> Option<Option<&str>> {
201 match request.header(TOKEN_INTENT) {
202 None => None,
203 Some(TOKEN_INTENT_NONE) => Some(None),
204 Some(explicit) => Some(Some(explicit)),
205 }
206}
207
208fn unauthenticated(mut request: HttpRequest) -> HttpRequest {
210 request
211 .headers
212 .push((TOKEN_INTENT.to_string(), TOKEN_INTENT_NONE.to_string()));
213 request
214}
215
216#[derive(Clone)]
221pub struct Session {
222 transport: Arc<Transport>,
223 endpoint: String,
224 port: u16,
225 egress_posture: crate::control::EgressPosture,
226 clock: Arc<dyn Clock>,
229}
230
231impl Session {
232 pub fn agent_token(&self) -> &str {
240 &self.transport.agent_token
241 }
242
243 pub fn builder(endpoint: impl Into<String>, agent_token: impl Into<String>) -> SessionBuilder {
249 SessionBuilder {
250 endpoint: endpoint.into(),
251 agent_token: agent_token.into(),
252 minter: None,
253 backend: None,
254 port: DEFAULT_AGENT_PORT,
255 timeout: DEFAULT_REQUEST_TIMEOUT,
256 proxy: None,
257 egress_posture: crate::control::EgressPosture::default(),
258 clock: None,
259 }
260 }
261
262 pub fn endpoint(&self) -> &str {
264 &self.endpoint
265 }
266
267 pub(crate) fn mint_exec_id(&self) -> String {
269 exec_id_at(self.clock.unix_now())
270 }
271
272 pub fn egress_posture(&self) -> crate::control::EgressPosture {
280 self.egress_posture
281 }
282
283 pub fn port(&self) -> u16 {
285 self.port
286 }
287
288 pub fn proxy_auth(&self) -> Option<&Arc<ProxyAuth>> {
290 self.transport.proxy()
291 }
292
293 pub async fn connect_headers(&self, port: u16) -> Result<Vec<(String, String)>, Error> {
315 match self.transport.proxy() {
316 Some(proxy) => {
317 tokio::time::timeout(self.transport.timeout, proxy.headers_for_port(port))
318 .await
319 .map_err(|_| {
320 Error::wire(
321 WireKind::Transport,
322 "proxy authentication request timed out",
323 )
324 })?
325 }
326 None => Ok(Vec::new()),
327 }
328 }
329
330 pub async fn connect_subprotocols(&self, port: u16) -> Result<Option<[String; 3]>, Error> {
342 match self.transport.proxy() {
343 Some(proxy) => Ok(Some(
344 tokio::time::timeout(self.transport.timeout, proxy.subprotocols(port))
345 .await
346 .map_err(|_| {
347 Error::wire(
348 WireKind::Transport,
349 "proxy authentication request timed out",
350 )
351 })??,
352 )),
353 None => Ok(None),
354 }
355 }
356
357 pub fn rebind(&mut self, endpoint: String) {
369 self.endpoint = endpoint;
370 if let Some(proxy) = self.transport.proxy() {
371 proxy.invalidate();
372 }
373 }
374
375 pub async fn health(&self) -> Result<protocol::health::Health, Error> {
384 self.transport
385 .send_json(unauthenticated(HttpRequest::new("GET", "/v1/health")))
386 .await
387 }
388
389 pub async fn keep_awake(
394 &self,
395 policy: &KeepAwake,
396 stop: impl std::future::Future<Output = ()>,
397 ) -> Result<KeepAwakeReport, Error> {
398 policy
399 .run(|| async { self.health().await.map(Some) }, stop)
400 .await
401 }
402
403 pub async fn wait_until_ready(
411 &self,
412 timeout: Duration,
413 ) -> Result<protocol::health::Health, Error> {
414 let deadline = deadline_after(timeout);
415 let mut last: Option<Error> = None;
416 loop {
417 match self.health().await {
418 Ok(health) if health.bootstrapped => return Ok(health),
419 Ok(_) => {}
420 Err(err) if err.retryable() => last = Some(err),
421 Err(err) => return Err(err),
422 }
423 let now = tokio::time::Instant::now();
424 if now >= deadline {
425 let detail = match last {
426 Some(err) => format!(" (last error: {err})"),
427 None => String::new(),
428 };
429 return Err(Error::new(
430 ErrorKind::Timeout,
431 format!(
432 "the daemon was not bootstrapped within {}s{detail}",
433 timeout.as_secs()
434 ),
435 ));
436 }
437 tokio::time::sleep(READY_POLL_INTERVAL.min(deadline - now)).await;
438 }
439 }
440
441 pub async fn run(&self, req: protocol::exec::StartRequest) -> Result<ExecHandle, Error> {
449 let requested = req.exec_id.clone();
450 let mut request = HttpRequest::new("POST", "/v1/exec/start");
451 request
452 .headers
453 .push(("content-type".into(), "application/json".into()));
454 request.body = serde_json::to_vec(&req).map_err(|err| {
455 Error::invalid_arg(format!("the start request will not serialize: {err}"))
456 })?;
457 let started: protocol::exec::StartResponse = self.transport.send_json(request).await?;
458 let exec_id = if started.exec_id.is_empty() {
463 requested
464 } else {
465 started.exec_id
466 };
467 Ok(ExecHandle::new(Arc::clone(&self.transport), exec_id))
468 }
469
470 pub fn exec(&self, exec_id: impl Into<String>) -> ExecHandle {
472 ExecHandle::new(Arc::clone(&self.transport), exec_id.into())
473 }
474
475 pub async fn run_sync(
477 &self,
478 req: protocol::exec::StartRequest,
479 timeout: Duration,
480 ) -> Result<ExecResult, Error> {
481 self.run(req).await?.wait_and_ack(timeout).await
482 }
483
484 pub async fn kill(&self, exec_id: &str) -> Result<bool, Error> {
486 self.exec(exec_id).kill().await
487 }
488
489 pub async fn procs(&self) -> Result<protocol::exec::ProcsResponse, Error> {
497 self.transport
498 .send_json(HttpRequest::new("GET", "/v1/procs"))
499 .await
500 }
501
502 pub async fn upload_file(
506 &self,
507 path: &str,
508 data: &[u8],
509 mode: Option<&str>,
510 ) -> Result<(), Error> {
511 files::upload_file(&self.transport, path, data, mode).await
512 }
513
514 pub async fn download_file(&self, path: &str) -> Result<Vec<u8>, Error> {
516 files::download_file(&self.transport, path).await
517 }
518
519 pub async fn download_file_lines(
524 &self,
525 path: &str,
526 start_line: Option<u64>,
527 end_line: Option<u64>,
528 ) -> Result<Vec<u8>, Error> {
529 files::download_file_lines(&self.transport, path, start_line, end_line).await
530 }
531
532 pub async fn file_exists(&self, path: &str) -> Result<bool, Error> {
534 files::file_exists(&self.transport, path).await
535 }
536
537 pub async fn upload_tar(&self, remote: &str, archive: &[u8]) -> Result<(), Error> {
539 files::upload_tar(&self.transport, remote, archive).await
540 }
541
542 pub async fn download_tar(&self, remote: &str) -> Result<Vec<u8>, Error> {
544 files::download_tar(&self.transport, remote).await
545 }
546}
547
548impl std::fmt::Debug for Session {
549 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
550 f.debug_struct("Session")
552 .field("endpoint", &self.endpoint)
553 .field("port", &self.port)
554 .field("proxy", &self.transport.proxy())
555 .finish()
556 }
557}
558
559pub struct SessionBuilder {
561 endpoint: String,
562 agent_token: String,
563 minter: Option<Arc<dyn TokenMinter>>,
564 backend: Option<http::SharedBackend>,
565 port: u16,
566 timeout: Duration,
567 proxy: Option<Arc<ProxyAuth>>,
568 egress_posture: crate::control::EgressPosture,
569 clock: Option<Arc<dyn Clock>>,
570}
571
572impl SessionBuilder {
573 #[must_use]
575 pub fn with_egress_posture(mut self, posture: crate::control::EgressPosture) -> Self {
576 self.egress_posture = posture;
577 self
578 }
579
580 #[must_use]
582 pub fn with_minter(mut self, minter: Arc<dyn TokenMinter>) -> Self {
583 self.minter = Some(minter);
584 self
585 }
586
587 #[must_use]
590 pub fn with_proxy_auth(mut self, proxy: Arc<ProxyAuth>) -> Self {
591 self.proxy = Some(proxy);
592 self
593 }
594
595 #[must_use]
597 pub fn with_backend(mut self, backend: http::SharedBackend) -> Self {
598 self.backend = Some(backend);
599 self
600 }
601
602 #[must_use]
604 pub fn with_port(mut self, port: u16) -> Self {
605 self.port = port;
606 self
607 }
608
609 #[must_use]
611 pub fn with_timeout(mut self, timeout: Duration) -> Self {
612 self.timeout = timeout;
613 self
614 }
615
616 #[must_use]
622 pub fn with_clock(mut self, clock: Arc<dyn Clock>) -> Self {
623 self.clock = Some(clock);
624 self
625 }
626
627 pub fn build_with(
634 mut self,
635 adapters: &dyn crate::adapters::Adapters,
636 ) -> Result<Session, Error> {
637 if self.backend.is_none() {
638 self.backend = Some(adapters.http_backend(&self.endpoint, self.timeout)?);
639 }
640 if self.clock.is_none() {
641 self.clock = Some(adapters.clock());
642 }
643 self.try_build()
644 }
645
646 pub fn try_build(self) -> Result<Session, Error> {
651 let (Some(backend), Some(clock)) = (self.backend, self.clock) else {
652 return Err(Error::new(
653 ErrorKind::Precondition,
654 "a session needs a backend and a clock: set both with with_backend and \
655 with_clock, or build through the prelude's SessionBuilder::build",
656 ));
657 };
658 let proxy = match (self.proxy, self.minter) {
659 (Some(proxy), _) => Some(proxy),
660 (None, Some(minter)) => Some(Arc::new(ProxyAuth::with_clock(
661 minter,
662 self.port,
663 Arc::clone(&clock),
664 ))),
665 (None, None) => None,
666 };
667 Ok(Session {
668 transport: Arc::new(Transport {
669 backend,
670 agent_token: self.agent_token,
671 proxy,
672 timeout: self.timeout,
673 }),
674 endpoint: self.endpoint,
675 port: self.port,
676 egress_posture: self.egress_posture,
677 clock,
678 })
679 }
680}
681
682const FAR_FUTURE: Duration = Duration::from_secs(946_080_000);
686
687pub(crate) fn deadline_after(after: Duration) -> tokio::time::Instant {
695 let now = tokio::time::Instant::now();
696 now.checked_add(after).unwrap_or_else(|| now + FAR_FUTURE)
697}
698
699#[cfg(any(test, feature = "test-support"))]
700#[cfg_attr(test, mutants::skip)] pub(crate) mod testing {
702 use std::collections::HashMap;
710 use std::sync::{Mutex, PoisonError};
711
712 use futures_util::future::BoxFuture;
713
714 use super::http::{ChunkSource, HttpBackend, HttpRequest, HttpResponse, OpenStream};
715 use super::*;
716
717 pub enum Reply {
719 Body(u16, Vec<u8>),
721 Chunks(u16, Vec<Vec<u8>>),
723 Cut(&'static str),
725 Stalled(Vec<Vec<u8>>),
728 }
729
730 impl Reply {
731 pub fn json(status: u16, body: impl serde::Serialize) -> Self {
732 Reply::Body(status, serde_json::to_vec(&body).expect("test body"))
733 }
734
735 pub fn ok(body: impl serde::Serialize) -> Self {
736 Reply::json(200, body)
737 }
738 }
739
740 #[derive(Default)]
741 pub struct Recorder {
742 seen: Mutex<Vec<HttpRequest>>,
743 replies: Mutex<std::collections::VecDeque<Reply>>,
744 }
745
746 impl Recorder {
747 pub fn with(replies: impl IntoIterator<Item = Reply>) -> Arc<Self> {
748 Arc::new(Self {
749 seen: Mutex::new(Vec::new()),
750 replies: Mutex::new(replies.into_iter().collect()),
751 })
752 }
753
754 pub fn requests(&self) -> Vec<HttpRequest> {
756 self.seen
757 .lock()
758 .unwrap_or_else(PoisonError::into_inner)
759 .clone()
760 }
761
762 pub fn last(&self) -> HttpRequest {
763 self.requests()
764 .pop()
765 .expect("the recorder saw no request at all")
766 }
767
768 fn record(&self, request: HttpRequest) -> Reply {
769 self.seen
770 .lock()
771 .unwrap_or_else(PoisonError::into_inner)
772 .push(request.clone());
773 self.replies
774 .lock()
775 .unwrap_or_else(PoisonError::into_inner)
776 .pop_front()
777 .unwrap_or_else(|| {
778 panic!(
779 "the recorder ran out of replies at {} {}",
780 request.method, request.path
781 )
782 })
783 }
784 }
785
786 struct Queued(std::collections::VecDeque<Vec<u8>>);
787
788 impl ChunkSource for Queued {
789 fn next_chunk(&mut self) -> BoxFuture<'_, Result<Option<Vec<u8>>, Error>> {
790 Box::pin(async move { Ok(self.0.pop_front()) })
791 }
792 }
793
794 struct Stalling(std::collections::VecDeque<Vec<u8>>);
796
797 impl ChunkSource for Stalling {
798 fn next_chunk(&mut self) -> BoxFuture<'_, Result<Option<Vec<u8>>, Error>> {
799 Box::pin(async move {
800 match self.0.pop_front() {
801 Some(chunk) => Ok(Some(chunk)),
802 None => std::future::pending().await,
803 }
804 })
805 }
806 }
807
808 impl HttpBackend for Recorder {
809 fn send(&self, request: HttpRequest) -> BoxFuture<'_, Result<HttpResponse, Error>> {
810 let reply = self.record(request);
811 Box::pin(async move {
812 match reply {
813 Reply::Body(status, body) => Ok(HttpResponse {
814 status,
815 headers: HashMap::new(),
816 body,
817 }),
818 Reply::Chunks(status, chunks) => Ok(HttpResponse {
819 status,
820 headers: HashMap::new(),
821 body: chunks.concat(),
822 }),
823 Reply::Stalled(chunks) => Ok(HttpResponse {
824 status: 200,
825 headers: HashMap::new(),
826 body: chunks.concat(),
827 }),
828 Reply::Cut(why) => Err(Error::wire(WireKind::Transport, why)),
829 }
830 })
831 }
832
833 fn open_stream(
834 &self,
835 request: HttpRequest,
836 _idle_timeout: Duration,
837 ) -> BoxFuture<'_, Result<OpenStream, Error>> {
838 let reply = self.record(request);
839 Box::pin(async move {
840 let (status, chunks) = match reply {
841 Reply::Chunks(status, chunks) => (status, chunks),
842 Reply::Body(status, body) => (status, vec![body]),
843 Reply::Cut(why) => return Err(Error::wire(WireKind::Transport, why)),
844 Reply::Stalled(chunks) => {
845 let head = HttpResponse {
846 status: 200,
847 headers: HashMap::new(),
848 body: Vec::new(),
849 };
850 let source: Box<dyn ChunkSource> =
851 Box::new(Stalling(chunks.into_iter().collect()));
852 return Ok((head, source));
853 }
854 };
855 let head = HttpResponse {
856 status,
857 headers: HashMap::new(),
858 body: if (200..300).contains(&status) {
861 Vec::new()
862 } else {
863 chunks.concat()
864 },
865 };
866 let source: Box<dyn ChunkSource> = if (200..300).contains(&status) {
867 Box::new(Queued(chunks.into_iter().collect()))
868 } else {
869 Box::new(Queued(std::collections::VecDeque::new()))
870 };
871 Ok((head, source))
872 })
873 }
874 }
875
876 pub fn session_with(
878 recorder: Arc<Recorder>,
879 ) -> (
880 Session,
881 Arc<ProxyAuth>,
882 Arc<crate::clock::testing::TestClock>,
883 ) {
884 let clock = Arc::new(crate::clock::testing::TestClock::default());
885 let auth = Arc::new(
886 ProxyAuth::with_refresh_after(
887 Arc::new(proxy::testing::CountingMinter::default()),
888 DEFAULT_AGENT_PORT,
889 DEFAULT_REFRESH_AFTER,
890 Arc::clone(&clock) as Arc<dyn Clock>,
891 )
892 .expect("the default interval is below the ceiling"),
893 );
894 let session = Session::builder("https://vm.example", "agent-token-abcdef")
895 .with_backend(recorder)
896 .with_proxy_auth(Arc::clone(&auth))
897 .with_clock(Arc::clone(&clock) as Arc<dyn Clock>)
898 .try_build()
899 .expect("builds");
900 (session, auth, clock)
901 }
902
903 pub fn health_body(bootstrapped: bool) -> serde_json::Value {
904 serde_json::json!({
905 "version": "0.1.0",
906 "bootstrapped": bootstrapped,
907 "disk": null,
908 "identity_degraded": false,
909 "identity_repaired": true,
910 })
911 }
912
913 #[derive(Default)]
921 pub struct HealthyDaemon {
922 polls: std::sync::atomic::AtomicUsize,
923 }
924
925 impl HealthyDaemon {
926 pub fn new() -> Arc<Self> {
927 Arc::new(Self::default())
928 }
929
930 pub fn polls(&self) -> usize {
932 self.polls.load(std::sync::atomic::Ordering::SeqCst)
933 }
934
935 fn refuse(request: &HttpRequest) -> Error {
936 Error::wire(
937 WireKind::Transport,
938 format!(
939 "{} {} reached a test daemon that answers only health",
940 request.method, request.path
941 ),
942 )
943 }
944 }
945
946 impl HttpBackend for HealthyDaemon {
947 fn send(&self, request: HttpRequest) -> BoxFuture<'_, Result<HttpResponse, Error>> {
948 let answer = if request.method == "GET" && request.path == "/v1/health" {
949 self.polls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
950 Ok(HttpResponse {
951 status: 200,
952 headers: HashMap::new(),
953 body: health_body(true).to_string().into_bytes(),
954 })
955 } else {
956 Err(Self::refuse(&request))
957 };
958 Box::pin(async move { answer })
959 }
960
961 fn open_stream(
962 &self,
963 request: HttpRequest,
964 _idle_timeout: Duration,
965 ) -> BoxFuture<'_, Result<OpenStream, Error>> {
966 let refused = Self::refuse(&request);
967 Box::pin(async move { Err(refused) })
968 }
969 }
970}
971
972#[cfg(test)]
973mod tests {
974 use super::testing::{Recorder, Reply, health_body, session_with};
975 use super::*;
976 use crate::session::proxy::testing::CountingMinter;
977
978 fn start_request(exec_id: &str) -> protocol::exec::StartRequest {
979 protocol::exec::StartRequest::new(exec_id, vec!["/bin/true".to_string()])
980 }
981
982 fn header(request: &HttpRequest, name: &str) -> Option<String> {
983 request.header(name).map(str::to_string)
984 }
985
986 fn header_of(headers: &[(String, String)], name: &str) -> Option<String> {
988 headers
989 .iter()
990 .find(|(key, _)| key.eq_ignore_ascii_case(name))
991 .map(|(_, value)| value.clone())
992 }
993
994 #[test]
1001 fn try_build_needs_a_backend_and_a_clock() {
1002 let clock = Arc::new(crate::clock::testing::TestClock::new()) as Arc<dyn Clock>;
1003 let no_backend = Session::builder("https://vm.example", "token")
1004 .with_clock(Arc::clone(&clock))
1005 .try_build()
1006 .expect_err("no backend");
1007 assert_eq!(no_backend.kind(), ErrorKind::Precondition);
1008 let no_clock = Session::builder("https://vm.example", "token")
1009 .with_backend(Recorder::with([]))
1010 .try_build()
1011 .expect_err("no clock");
1012 assert_eq!(no_clock.kind(), ErrorKind::Precondition);
1013 Session::builder("https://vm.example", "token")
1014 .with_backend(Recorder::with([]))
1015 .with_clock(clock)
1016 .try_build()
1017 .expect("both set");
1018 }
1019
1020 #[tokio::test]
1026 async fn build_with_keeps_the_callers_backend() {
1027 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1028 let session = Session::builder("http://127.0.0.1:9000", "token")
1029 .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1030 .build_with(&crate::testing::TestAdapters::new())
1031 .expect("builds");
1032 session.health().await.expect("the recorder answered");
1033 assert_eq!(recorder.requests().len(), 1);
1034 }
1035
1036 #[tokio::test(start_paused = true)]
1037 async fn request_timeout_also_bounds_a_blocked_proxy_mint() {
1038 struct BlockedMinter;
1039 impl TokenMinter for BlockedMinter {
1040 fn mint(&self) -> futures_util::future::BoxFuture<'_, Result<ProxyToken, Error>> {
1041 Box::pin(std::future::pending())
1042 }
1043 }
1044 let backend = Recorder::with([]);
1045 let session = Session::builder("http://127.0.0.1:9", "private")
1046 .with_backend(backend.clone())
1047 .with_minter(Arc::new(BlockedMinter))
1048 .with_timeout(Duration::from_secs(3))
1049 .build_with(&crate::testing::TestAdapters::new())
1050 .expect("session");
1051 let result = tokio::time::timeout(Duration::from_secs(4), session.health())
1052 .await
1053 .expect("the request budget must include a stalled proxy mint")
1054 .expect_err("mint never completed");
1055 assert_eq!(result.wire_kind(), Some(WireKind::Transport));
1056 assert!(result.retryable());
1057 let result = tokio::time::timeout(Duration::from_secs(4), session.connect_headers(9000))
1058 .await
1059 .expect("external transfers must bound authentication too")
1060 .expect_err("mint remains blocked");
1061 assert!(result.retryable());
1062 let result =
1063 tokio::time::timeout(Duration::from_secs(4), session.connect_subprotocols(9000))
1064 .await
1065 .expect("external sockets must bound authentication too")
1066 .expect_err("mint remains blocked");
1067 assert!(result.retryable());
1068 assert!(
1069 backend.requests().is_empty(),
1070 "no request could authenticate"
1071 );
1072 }
1073
1074 #[tokio::test]
1082 async fn a_callers_content_type_survives_the_auth_header_injection() {
1083 let recorder = Recorder::with([Reply::ok(
1084 serde_json::json!({"exec_id":"e1","phase":"running"}),
1085 )]);
1086 let (session, _, _) = session_with(Arc::clone(&recorder));
1087
1088 session.run(start_request("e1")).await.expect("start");
1089
1090 let seen = recorder.requests();
1091 assert_eq!(
1092 header(&seen[0], "content-type").as_deref(),
1093 Some("application/json"),
1094 "the exec start lost its content type to the auth header injection"
1095 );
1096 }
1097
1098 #[tokio::test]
1103 async fn the_session_hands_back_the_token_it_authenticates_with() {
1104 let recorder = Recorder::with([]);
1105 let (session, _, _) = session_with(Arc::clone(&recorder));
1106 assert_eq!(session.agent_token(), "agent-token-abcdef");
1107 }
1108
1109 #[tokio::test]
1116 async fn every_endpoint_request_carries_both_proxy_headers() {
1117 let recorder = Recorder::with([
1118 Reply::ok(health_body(true)),
1119 Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
1120 ]);
1121 let (session, _, _) = session_with(Arc::clone(&recorder));
1122
1123 session.health().await.expect("health succeeds");
1124 session.run(start_request("e1")).await.expect("start");
1125
1126 let seen = recorder.requests();
1127 assert_eq!(seen.len(), 2);
1128 for request in &seen {
1129 assert_eq!(
1130 header(request, PROXY_AUTH_HEADER).as_deref(),
1131 Some(CountingMinter::value(0).as_str()),
1132 "{} {} went out without the proxy auth header",
1133 request.method,
1134 request.path
1135 );
1136 assert_eq!(
1137 header(request, PROXY_PORT_HEADER).as_deref(),
1138 Some("9000"),
1139 "{} {} went out without the proxy port header, which the proxy \
1140 rejects as if the token were bad",
1141 request.method,
1142 request.path
1143 );
1144 }
1145 }
1146
1147 #[tokio::test]
1153 async fn health_sends_no_bearer_and_an_exec_route_does() {
1154 let recorder = Recorder::with([
1155 Reply::ok(health_body(true)),
1156 Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
1157 ]);
1158 let (session, _, _) = session_with(Arc::clone(&recorder));
1159
1160 session.health().await.expect("health");
1161 let seen = recorder.requests();
1162 assert_eq!(header(&seen[0], "authorization"), None);
1163 assert_eq!(
1164 header(&seen[0], "x-microvms-core-token-intent"),
1165 None,
1166 "the internal token-intent marker must not reach the wire"
1167 );
1168
1169 session.run(start_request("e1")).await.expect("start");
1170 assert_eq!(
1171 header(&recorder.last(), "authorization").as_deref(),
1172 Some("Bearer agent-token-abcdef")
1173 );
1174 }
1175
1176 #[tokio::test(start_paused = true)]
1179 async fn keep_awake_polls_health_through_the_proxy_until_the_vm_is_idle() {
1180 let busy = |busy: bool| {
1181 let mut body = health_body(true);
1182 body["busy"] = serde_json::json!(busy);
1183 Reply::ok(body)
1184 };
1185 let recorder = Recorder::with([busy(true), busy(true), busy(false)]);
1186 let (session, _, _) = session_with(Arc::clone(&recorder));
1187 let report = session
1188 .clone()
1189 .keep_awake(
1190 &KeepAwake::new(None).while_busy(true),
1191 std::future::pending(),
1192 )
1193 .await
1194 .expect("keepalive");
1195 assert_eq!((report.end, report.polls), (KeepAwakeEnd::Idle, 3));
1196 let seen = recorder.requests();
1197 assert_eq!(seen.len(), 3);
1198 for request in &seen {
1199 assert_eq!(
1200 (request.method, request.path.as_str()),
1201 ("GET", "/v1/health")
1202 );
1203 assert_eq!(header(request, "authorization"), None);
1204 assert!(header(request, PROXY_AUTH_HEADER).is_some());
1205 assert!(header(request, PROXY_PORT_HEADER).is_some());
1206 }
1207 }
1208
1209 #[tokio::test]
1212 async fn a_direct_session_sends_no_proxy_headers() {
1213 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1214 let session = Session::builder("http://127.0.0.1:9000", "token")
1215 .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1216 .build_with(&crate::testing::TestAdapters::new())
1217 .expect("builds");
1218 assert!(session.proxy_auth().is_none());
1219
1220 session.health().await.expect("health");
1221 let request = recorder.last();
1222 assert_eq!(header(&request, PROXY_AUTH_HEADER), None);
1223 assert_eq!(header(&request, PROXY_PORT_HEADER), None);
1224 }
1225
1226 #[tokio::test]
1229 async fn a_mint_failure_surfaces_as_a_retryable_request_error() {
1230 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1231 let auth = Arc::new(
1232 ProxyAuth::with_refresh_after(
1233 Arc::new(CountingMinter::failing(1)),
1234 DEFAULT_AGENT_PORT,
1235 DEFAULT_REFRESH_AFTER,
1236 Arc::new(proxy::testing::TestClock::default()),
1237 )
1238 .expect("interval accepted"),
1239 );
1240 let session = Session::builder("https://vm.example", "token")
1241 .with_backend(Arc::clone(&recorder) as http::SharedBackend)
1242 .with_proxy_auth(Arc::clone(&auth))
1243 .build_with(&crate::testing::TestAdapters::new())
1244 .expect("builds");
1245
1246 let err = session.health().await.expect_err("the mint fails");
1247 assert_eq!(err.kind(), ErrorKind::Retryable);
1248 assert_eq!(err.wire_kind(), Some(WireKind::AuthTokenMint));
1249 assert!(
1250 recorder.requests().is_empty(),
1251 "a request went out without a token"
1252 );
1253
1254 session.health().await.expect("the retry succeeds");
1255 assert_eq!(auth.mint_count(), 1);
1256 }
1257
1258 #[tokio::test]
1261 async fn a_daemon_status_arrives_as_its_typed_error() {
1262 for (status, expected) in [
1263 (401, WireKind::Unauthorized),
1264 (404, WireKind::NotFound),
1265 (409, WireKind::Conflict),
1266 (503, WireKind::NotBootstrapped),
1267 ] {
1268 let recorder = Recorder::with([Reply::Body(status, b"{\"error\":\"x\"}".to_vec())]);
1269 let (session, _, _) = session_with(recorder);
1270 let err = session.health().await.expect_err("a failing status");
1271 assert_eq!(err.wire_kind(), Some(expected), "status {status}");
1272 }
1273 }
1274
1275 #[tokio::test(start_paused = true)]
1278 async fn wait_until_ready_polls_through_not_bootstrapped_and_a_dropped_connection() {
1279 let recorder = Recorder::with([
1280 Reply::Cut("connection refused"),
1281 Reply::Body(503, b"not bootstrapped".to_vec()),
1282 Reply::ok(health_body(false)),
1283 Reply::ok(health_body(true)),
1284 ]);
1285 let (session, _, _) = session_with(Arc::clone(&recorder));
1286
1287 let health = session
1288 .wait_until_ready(DEFAULT_BOOTSTRAP_TIMEOUT)
1289 .await
1290 .expect("the daemon comes up");
1291 assert!(health.bootstrapped);
1292 assert_eq!(recorder.requests().len(), 4);
1293 }
1294
1295 #[tokio::test(start_paused = true)]
1298 async fn wait_until_ready_gives_up_at_once_on_a_credential_failure() {
1299 let recorder = Recorder::with([Reply::Body(401, b"wrong token".to_vec())]);
1300 let (session, _, _) = session_with(Arc::clone(&recorder));
1301
1302 let err = session
1303 .wait_until_ready(DEFAULT_BOOTSTRAP_TIMEOUT)
1304 .await
1305 .expect_err("a 401 is fatal");
1306 assert_eq!(err.kind(), ErrorKind::Credentials);
1307 assert_eq!(
1308 recorder.requests().len(),
1309 1,
1310 "a fatal error was retried, so a wrong token costs the full timeout"
1311 );
1312 }
1313
1314 #[tokio::test(start_paused = true)]
1317 async fn wait_until_ready_times_out_naming_the_last_retryable_error() {
1318 let recorder = Recorder::with([
1319 Reply::Cut("connection refused"),
1320 Reply::Cut("connection refused"),
1321 Reply::Cut("connection refused"),
1322 ]);
1323 let (session, _, _) = session_with(recorder);
1324
1325 let err = session
1326 .wait_until_ready(Duration::from_secs(4))
1327 .await
1328 .expect_err("the daemon never comes up");
1329 assert_eq!(err.kind(), ErrorKind::Timeout);
1330 assert!(
1331 err.to_string().contains("connection refused"),
1332 "the timeout must name what it kept hitting: {err}"
1333 );
1334 }
1335
1336 #[tokio::test]
1342 async fn rebind_invalidates_the_proxy_token_so_the_next_request_mints() {
1343 let recorder = Recorder::with([Reply::ok(health_body(true)), Reply::ok(health_body(true))]);
1344 let (mut session, auth, _) = session_with(Arc::clone(&recorder));
1345
1346 session.health().await.expect("health");
1347 assert_eq!(auth.mint_count(), 1);
1348
1349 session.rebind("https://vm.example".to_string());
1350 assert!(!auth.is_cached(), "rebind did not drop the token");
1351
1352 session.health().await.expect("health");
1353 assert_eq!(auth.mint_count(), 2);
1354 assert_eq!(
1355 header(&recorder.last(), PROXY_AUTH_HEADER).as_deref(),
1356 Some(CountingMinter::value(1).as_str()),
1357 "the request after a rebind carried the pre-suspend token"
1358 );
1359 }
1360
1361 #[tokio::test]
1364 async fn a_start_request_serializes_through_the_protocol_crate() {
1365 let recorder = Recorder::with([Reply::ok(
1366 serde_json::json!({"exec_id":"e-confirmed","phase":"running"}),
1367 )]);
1368 let (session, _, _) = session_with(Arc::clone(&recorder));
1369
1370 let handle = session.run(start_request("e-asked")).await.expect("start");
1371 assert_eq!(
1372 handle.exec_id(),
1373 "e-confirmed",
1374 "the handle addresses the id the client asked for rather than the one the \
1375 daemon confirmed, so every later call would 404 if they diverged"
1376 );
1377
1378 let request = recorder.last();
1379 assert_eq!(request.method, "POST");
1380 assert_eq!(request.path, "/v1/exec/start");
1381 let body: serde_json::Value = serde_json::from_slice(&request.body).expect("json body");
1382 assert_eq!(body["exec_id"], "e-asked");
1383 assert_eq!(body["command"], serde_json::json!(["/bin/true"]));
1384 }
1385
1386 #[tokio::test]
1389 async fn an_unreadable_response_body_names_the_schema_route() {
1390 let recorder = Recorder::with([Reply::Body(200, b"{\"unexpected\":true}".to_vec())]);
1391 let (session, _, _) = session_with(recorder);
1392 let err = session.health().await.expect_err("the body will not parse");
1393 assert_eq!(err.wire_kind(), Some(WireKind::ProtocolError));
1394 assert!(err.to_string().contains("/v1/schema"), "{err}");
1395 }
1396
1397 #[tokio::test]
1418 async fn a_session_answers_connect_credentials_for_a_port_off_its_own_token_cache() {
1419 let recorder = Recorder::with([Reply::ok(health_body(true))]);
1420 let (session, auth, _) = session_with(Arc::clone(&recorder));
1421
1422 session.health().await.expect("health warms the cache");
1423 assert_eq!(auth.mint_count(), 1);
1424
1425 let headers = session
1426 .connect_headers(8080)
1427 .await
1428 .expect("the cached token answers");
1429 let offered = session
1430 .connect_subprotocols(8080)
1431 .await
1432 .expect("the cached token answers")
1433 .expect("a proxied session offers subprotocols");
1434
1435 assert_eq!(
1436 header_of(&headers, PROXY_PORT_HEADER).as_deref(),
1437 Some("8080"),
1438 "the caller asked for 8080 and would have targeted the session's port"
1439 );
1440 assert_eq!(
1442 header_of(&headers, PROXY_AUTH_HEADER).as_deref(),
1443 Some(CountingMinter::value(1).as_str()),
1444 "the header pair carried the agent-port-only token, which the proxy refuses for \
1445 8080 with 403 'Access to port denied'"
1446 );
1447 assert_eq!(offered[0], proxy::WS_SUBPROTOCOL);
1448 assert_eq!(
1449 offered[1],
1450 format!(
1451 "{}{}",
1452 proxy::WS_AUTH_SUBPROTOCOL_PREFIX,
1453 CountingMinter::value(1)
1454 ),
1455 "the handshake carried a different token than the header path, so they are two \
1456 caches rather than one"
1457 );
1458 assert_eq!(offered[2], "lambda-microvms.port.8080");
1459 assert_eq!(
1460 auth.mint_count(),
1461 2,
1462 "one mint covered both the agent port and 8080, which the control plane does not \
1463 do — or the second accessor minted a third, which would be two caches"
1464 );
1465 assert!(
1466 recorder.requests().len() == 1,
1467 "building credentials sent a request to the VM: {:?}",
1468 recorder.requests()
1469 );
1470 }
1471
1472 #[tokio::test]
1483 async fn a_direct_session_offers_no_subprotocols_and_empty_headers() {
1484 let recorder = Recorder::with([]);
1485 let session = Session::builder("http://127.0.0.1:9000", "token")
1486 .with_backend(recorder as http::SharedBackend)
1487 .build_with(&crate::testing::TestAdapters::new())
1488 .expect("builds");
1489
1490 assert!(
1491 session
1492 .connect_headers(9000)
1493 .await
1494 .expect("no proxy is not a failure")
1495 .is_empty(),
1496 "a direct session's requests carry no proxy headers, so neither should this"
1497 );
1498 assert!(
1499 session
1500 .connect_subprotocols(9000)
1501 .await
1502 .expect("no proxy is not a failure")
1503 .is_none(),
1504 "a handshake with the base subprotocol and no token is refused for a reason that \
1505 names neither"
1506 );
1507 }
1508
1509 #[tokio::test]
1520 async fn connect_credentials_re_mint_when_the_refresh_window_has_rolled_over() {
1521 let recorder = Recorder::with([]);
1522 let (session, auth, clock) = session_with(recorder);
1523
1524 let first = session
1525 .connect_subprotocols(9000)
1526 .await
1527 .expect("mints")
1528 .expect("proxied");
1529 assert!(first[1].ends_with(&CountingMinter::value(0)));
1530
1531 clock.advance(DEFAULT_REFRESH_AFTER + Duration::from_secs(1));
1532 let second = session
1533 .connect_subprotocols(9000)
1534 .await
1535 .expect("re-mints")
1536 .expect("proxied");
1537
1538 assert_eq!(auth.mint_count(), 2);
1539 assert!(
1540 second[1].ends_with(&CountingMinter::value(1)),
1541 "the handshake would carry a token past its refresh window: {}",
1542 second[1]
1543 );
1544 }
1545
1546 #[tokio::test]
1556 async fn procs_is_a_bearer_get_of_the_procs_route_with_a_typed_body() {
1557 let recorder = Recorder::with([Reply::ok(serde_json::json!({
1558 "procs": [{
1559 "exec_id": "e1",
1560 "pgid": 4242,
1561 "started_at": 1_756_500_000u64,
1562 "child_exited": true,
1563 "reap": false,
1564 "pids": [4243],
1565 }]
1566 }))]);
1567 let (session, _, _) = session_with(Arc::clone(&recorder));
1568
1569 let listed = session.procs().await.expect("procs");
1570 let request = recorder.last();
1571 assert_eq!(request.method, "GET");
1572 assert_eq!(request.path, "/v1/procs");
1573 assert_eq!(
1574 header(&request, "authorization").as_deref(),
1575 Some("Bearer agent-token-abcdef"),
1576 "the procs route is bearer-guarded; without the token it is a 503 or a 401"
1577 );
1578 assert_eq!(listed.procs.len(), 1);
1579 assert_eq!(listed.procs[0].exec_id, "e1");
1580 assert_eq!(listed.procs[0].pgid, Some(4242));
1581 assert!(listed.procs[0].child_exited);
1582 assert_eq!(listed.procs[0].pids, [4243]);
1583 }
1584
1585 #[tokio::test]
1593 async fn a_start_request_carries_the_reap_flag_on_the_wire() {
1594 let recorder = Recorder::with([Reply::ok(
1595 serde_json::json!({"exec_id":"e1","phase":"running"}),
1596 )]);
1597 let (session, _, _) = session_with(Arc::clone(&recorder));
1598
1599 let mut request = start_request("e1");
1600 request.reap_group_on_exit = true;
1601 session.run(request).await.expect("start");
1602
1603 let body: serde_json::Value =
1604 serde_json::from_slice(&recorder.last().body).expect("a JSON start body");
1605 assert_eq!(
1606 body["reap_group_on_exit"],
1607 serde_json::Value::Bool(true),
1608 "the daemon defaults a missing flag to false, so a dropped key is a silent \
1609 behaviour change: {body}"
1610 );
1611 }
1612
1613 #[test]
1615 fn a_session_debug_does_not_print_the_agent_token() {
1616 let recorder = Recorder::with([]);
1617 let session = Session::builder("https://vm.example", "super-secret-token")
1618 .with_backend(recorder as http::SharedBackend)
1619 .build_with(&crate::testing::TestAdapters::new())
1620 .expect("builds");
1621 let rendered = format!("{session:?}");
1622 assert!(rendered.contains("vm.example"), "{rendered}");
1623 assert!(
1624 !rendered.contains("super-secret"),
1625 "the agent token reached a Debug string: {rendered}"
1626 );
1627 }
1628
1629 #[tokio::test(start_paused = true)]
1644 async fn a_ready_wait_longer_than_the_clock_can_hold_still_polls() {
1645 let recorder =
1646 Recorder::with([Reply::ok(health_body(false)), Reply::ok(health_body(true))]);
1647 let (session, _, _) = session_with(Arc::clone(&recorder));
1648
1649 let health = session
1650 .wait_until_ready(Duration::MAX)
1651 .await
1652 .expect("a far deadline waits for the bootstrap");
1653 assert!(health.bootstrapped);
1654 assert_eq!(
1655 recorder.requests().len(),
1656 2,
1657 "the wait polled past the first not-bootstrapped answer"
1658 );
1659 }
1660}
1661
1662#[cfg(test)]
1663mod sse_fuzz;