1use std::future::Future;
30use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
31use std::sync::{Arc, Mutex, PoisonError};
32use std::time::Duration;
33
34use tokio::net::TcpListener;
35use tokio::sync::{oneshot, watch};
36use tokio::task::JoinSet;
37
38use microvms_app::error::{Error, ErrorKind};
39use microvms_app::identity::TunnelIdentity;
40use microvms_app::session::Session;
41use microvms_app::session::proxy::ProxyAuth;
42
43use super::forward::{ForwardClient, ForwardEvent, ForwardSpec, serve_connection_via};
44use super::tunnel::{TunnelEnd, relay_connection_inner};
45
46#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
48pub struct ServeLimits {
49 pub max_connections: Option<u32>,
51}
52
53#[derive(Clone, Debug, Eq, PartialEq)]
55pub enum StopReason {
56 Requested,
58 Limit,
60 ListenerFailed(String),
63}
64
65impl StopReason {
66 pub fn was_requested(&self) -> bool {
69 matches!(self, StopReason::Requested)
70 }
71}
72
73#[derive(Clone)]
75pub struct TunnelTarget {
76 pub endpoint: String,
77 pub agent_token: String,
78 pub auth: Option<Arc<ProxyAuth>>,
80 pub guest_port: u16,
81 pub identity: Option<TunnelIdentity>,
83}
84
85impl TunnelTarget {
86 pub fn for_session(
88 session: &Session,
89 guest_port: u16,
90 identity: Option<TunnelIdentity>,
91 ) -> Self {
92 Self {
93 endpoint: session.endpoint().to_string(),
94 agent_token: session.agent_token().to_string(),
95 auth: session.proxy_auth().cloned(),
96 guest_port,
97 identity,
98 }
99 }
100}
101
102#[derive(Debug)]
104pub struct TunnelConnection {
105 pub peer: SocketAddr,
106 pub end: Result<TunnelEnd, Error>,
107}
108
109#[derive(Clone, Debug, Eq, PartialEq)]
111pub struct TunnelReport {
112 pub served: u32,
114 pub refused: u32,
116 pub truncated: u32,
118 pub unproven: u32,
120 pub proxy_token_mints: u64,
122 pub stopped: StopReason,
123}
124
125impl TunnelReport {
126 fn count(&mut self, end: &Result<TunnelEnd, Error>) {
128 match end {
129 Ok(TunnelEnd::Closed) => {}
130 Ok(TunnelEnd::ClosedUnproven) => self.unproven += 1,
131 Ok(TunnelEnd::Truncated { .. }) => self.truncated += 1,
132 Ok(TunnelEnd::Refused { .. }) | Err(_) => self.refused += 1,
133 }
134 }
135}
136
137pub async fn serve_tunnel<F>(
142 listener: TcpListener,
143 target: TunnelTarget,
144 limits: ServeLimits,
145 stop: impl Future<Output = ()>,
146 on_end: F,
147) -> TunnelReport
148where
149 F: FnMut(TunnelConnection),
150{
151 tunnel_loop(
152 listener,
153 target,
154 limits,
155 stop,
156 std::future::pending(),
157 on_end,
158 )
159 .await
160}
161
162async fn tunnel_loop<F>(
164 listener: TcpListener,
165 target: TunnelTarget,
166 limits: ServeLimits,
167 stop: impl Future<Output = ()>,
168 cut: impl Future<Output = ()>,
169 mut on_end: F,
170) -> TunnelReport
171where
172 F: FnMut(TunnelConnection),
173{
174 let target = Arc::new(target);
175 let mut report = TunnelReport {
176 served: 0,
177 refused: 0,
178 truncated: 0,
179 unproven: 0,
180 proxy_token_mints: 0,
181 stopped: StopReason::Requested,
182 };
183 let (stopped, served) = serve(
184 &listener,
185 limits,
186 stop,
187 cut,
188 |stream, _| {
189 let target = Arc::clone(&target);
190 async move {
191 relay_connection_inner(
192 stream,
193 &target.endpoint,
194 target.guest_port,
195 &target.agent_token,
196 target.auth.as_ref(),
197 target.identity.as_ref(),
198 )
199 .await
200 }
201 },
202 |peer, end| {
203 report.count(&end);
204 on_end(TunnelConnection { peer, end });
205 },
206 )
207 .await;
208 report.served = served;
209 report.stopped = stopped;
210 report.proxy_token_mints = target.auth.as_ref().map_or(0, |auth| auth.mint_count());
211 report
212}
213
214#[derive(Debug)]
216pub enum ForwardNotice {
217 Event {
219 peer: SocketAddr,
220 event: ForwardEvent,
221 },
222 Failed { peer: SocketAddr, error: Error },
224}
225
226#[derive(Clone, Debug, Eq, PartialEq)]
228pub struct ForwardReport {
229 pub served: u32,
231 pub refused: u32,
233 pub upgrades: u32,
235 pub proxy_token_mints: u64,
237 pub stopped: StopReason,
238}
239
240impl ForwardReport {
241 fn count(&mut self, event: &ForwardEvent) {
243 match event {
244 ForwardEvent::Refused { .. } => self.refused += 1,
245 ForwardEvent::Forwarded { upgraded: true, .. } => self.upgrades += 1,
246 _ => {}
247 }
248 }
249}
250
251pub async fn serve_forward<F>(
258 listener: TcpListener,
259 spec: ForwardSpec,
260 auth: Option<Arc<ProxyAuth>>,
261 limits: ServeLimits,
262 stop: impl Future<Output = ()>,
263 on_notice: F,
264) -> Result<ForwardReport, Error>
265where
266 F: FnMut(ForwardNotice),
267{
268 forward_loop(
269 listener,
270 spec,
271 auth,
272 limits,
273 stop,
274 std::future::pending(),
275 on_notice,
276 )
277 .await
278}
279
280async fn forward_loop<F>(
282 listener: TcpListener,
283 spec: ForwardSpec,
284 auth: Option<Arc<ProxyAuth>>,
285 limits: ServeLimits,
286 stop: impl Future<Output = ()>,
287 cut: impl Future<Output = ()>,
288 mut on_notice: F,
289) -> Result<ForwardReport, Error>
290where
291 F: FnMut(ForwardNotice),
292{
293 let client = Arc::new(ForwardClient::new()?);
296 let spec = Arc::new(spec);
297 let mut report = ForwardReport {
298 served: 0,
299 refused: 0,
300 upgrades: 0,
301 proxy_token_mints: 0,
302 stopped: StopReason::Requested,
303 };
304 let (stopped, served) = serve(
305 &listener,
306 limits,
307 stop,
308 cut,
309 |stream, _| {
310 let (spec, auth, client) = (Arc::clone(&spec), auth.clone(), Arc::clone(&client));
311 async move {
312 let mut events = Vec::new();
313 let outcome =
314 serve_connection_via(stream, &spec, auth.as_ref(), &client, |event| {
315 events.push(event);
316 })
317 .await;
318 outcome.map(|()| events)
319 }
320 },
321 |peer, done| {
322 let (events, failure) = match done {
323 Ok(events) => (events, None),
324 Err(error) => (Vec::new(), Some(error)),
325 };
326 for event in events {
327 report.count(&event);
328 on_notice(ForwardNotice::Event { peer, event });
329 }
330 if let Some(error) = failure {
331 on_notice(ForwardNotice::Failed { peer, error });
332 }
333 },
334 )
335 .await;
336 report.served = served;
337 report.stopped = stopped;
338 report.proxy_token_mints = auth.as_ref().map_or(0, |auth| auth.mint_count());
339 Ok(report)
340}
341
342pub const DEFAULT_BIND: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0));
347
348pub fn bind_address(bind: Option<&str>) -> Result<SocketAddr, Error> {
350 match bind {
351 None => Ok(DEFAULT_BIND),
352 Some(bind) => bind.parse().map_err(|err| {
353 Error::new(
354 ErrorKind::InvalidArg,
355 format!("{bind:?} is not a local address to listen on, such as 127.0.0.1:0: {err}"),
356 )
357 }),
358 }
359}
360
361#[derive(Clone, Copy, Debug, Eq, PartialEq)]
363pub enum EndKind {
364 Refused,
366 Truncated,
368 Unproven,
370 Failed,
372}
373
374impl EndKind {
375 pub fn as_str(self) -> &'static str {
377 match self {
378 EndKind::Refused => "refused",
379 EndKind::Truncated => "truncated",
380 EndKind::Unproven => "unproven",
381 EndKind::Failed => "failed",
382 }
383 }
384}
385
386#[derive(Clone, Debug, Eq, PartialEq)]
389pub struct ConnectionEnd {
390 pub peer: SocketAddr,
391 pub kind: EndKind,
392 pub code: Option<u16>,
394 pub detail: String,
396}
397
398impl ConnectionEnd {
399 pub fn of_tunnel(connection: &TunnelConnection) -> Option<Self> {
401 let (kind, code, detail) = match &connection.end {
402 Ok(TunnelEnd::Closed) => return None,
403 Ok(TunnelEnd::ClosedUnproven) => (EndKind::Unproven, None, String::new()),
404 Ok(TunnelEnd::Truncated { code }) => (EndKind::Truncated, *code, String::new()),
405 Ok(TunnelEnd::Refused { code, reason }) => {
406 (EndKind::Refused, Some(*code), reason.clone())
407 }
408 Err(error) => (EndKind::Failed, None, error.to_string()),
409 };
410 Some(Self {
411 peer: connection.peer,
412 kind,
413 code,
414 detail,
415 })
416 }
417
418 pub fn of_forward(notice: &ForwardNotice) -> Option<Self> {
420 match notice {
421 ForwardNotice::Event {
422 peer,
423 event:
424 ForwardEvent::Refused {
425 status,
426 explanation,
427 },
428 } => Some(Self {
429 peer: *peer,
430 kind: EndKind::Refused,
431 code: Some(*status),
432 detail: explanation.clone(),
433 }),
434 ForwardNotice::Event { .. } => None,
435 ForwardNotice::Failed { peer, error } => Some(Self {
436 peer: *peer,
437 kind: EndKind::Failed,
438 code: None,
439 detail: error.to_string(),
440 }),
441 }
442 }
443}
444
445#[derive(Clone, Debug, Eq, PartialEq)]
447pub struct TunnelSummary {
448 pub report: TunnelReport,
449 pub ended: Vec<ConnectionEnd>,
450}
451
452#[derive(Clone, Debug, Eq, PartialEq)]
455pub struct ForwardSummary {
456 pub report: ForwardReport,
457 pub ended: Vec<ConnectionEnd>,
458}
459
460type Outcome<S> = Result<S, (ErrorKind, String)>;
461
462pub struct Serving<S> {
467 local: SocketAddr,
468 stop: Mutex<Option<oneshot::Sender<()>>>,
469 cut: Mutex<Option<oneshot::Sender<()>>>,
470 outcome: Arc<Mutex<Option<Outcome<S>>>>,
471 done: watch::Receiver<bool>,
472}
473
474impl<S: Clone + Send + 'static> Serving<S> {
475 fn spawn<Fut>(
478 local: SocketAddr,
479 serving: impl FnOnce(oneshot::Receiver<()>, oneshot::Receiver<()>) -> Fut,
480 ) -> Self
481 where
482 Fut: Future<Output = Result<S, Error>> + Send + 'static,
483 {
484 let (stop_tx, stop_rx) = oneshot::channel();
485 let (cut_tx, cut_rx) = oneshot::channel();
486 let (done_tx, done_rx) = watch::channel(false);
487 let outcome = Arc::new(Mutex::new(None));
488 let written = Arc::clone(&outcome);
489 let serving = serving(stop_rx, cut_rx);
490 tokio::spawn(async move {
491 let result = serving
492 .await
493 .map_err(|error| (error.kind(), error.to_string()));
494 *written.lock().unwrap_or_else(PoisonError::into_inner) = Some(result);
495 done_tx.send_replace(true);
496 });
497 Self {
498 local,
499 stop: Mutex::new(Some(stop_tx)),
500 cut: Mutex::new(Some(cut_tx)),
501 outcome,
502 done: done_rx,
503 }
504 }
505
506 pub fn local_address(&self) -> SocketAddr {
508 self.local
509 }
510
511 pub fn is_running(&self) -> bool {
513 !*self.done.borrow()
514 }
515
516 pub fn request_stop(&self) {
519 signal(&self.stop);
520 }
521
522 pub async fn stop(&self, grace: Option<Duration>) -> Result<S, Error> {
528 self.request_stop();
529 if let Some(grace) = grace {
530 if let Ok(outcome) = tokio::time::timeout(grace, self.finished()).await {
531 return outcome;
532 }
533 signal(&self.cut);
534 }
535 self.finished().await
536 }
537
538 pub async fn finished(&self) -> Result<S, Error> {
540 let mut done = self.done.clone();
541 let _ = done.wait_for(|finished| *finished).await;
543 self.outcome
544 .lock()
545 .unwrap_or_else(PoisonError::into_inner)
546 .clone()
547 .expect("the loop records its outcome before signalling done")
548 .map_err(|(kind, message)| Error::new(kind, message))
549 }
550}
551
552fn signal(sender: &Mutex<Option<oneshot::Sender<()>>>) {
554 if let Some(sender) = sender.lock().unwrap_or_else(PoisonError::into_inner).take() {
555 let _ = sender.send(());
556 }
557}
558
559async fn signalled(signal: oneshot::Receiver<()>) {
561 let _ = signal.await;
562}
563
564pub async fn start_tunnel(
567 bind: SocketAddr,
568 target: TunnelTarget,
569 limits: ServeLimits,
570) -> Result<Serving<TunnelSummary>, Error> {
571 let listener =
572 super::forward::bind(&ForwardSpec::new(bind, target.guest_port, &target.endpoint)).await?;
573 let local = local_address(&listener)?;
574 Ok(Serving::spawn(local, move |stop, cut| async move {
575 let mut ended = Vec::new();
576 let report = tunnel_loop(
577 listener,
578 target,
579 limits,
580 signalled(stop),
581 signalled(cut),
582 |connection| ended.extend(ConnectionEnd::of_tunnel(&connection)),
583 )
584 .await;
585 Ok(TunnelSummary { report, ended })
586 }))
587}
588
589pub async fn start_forward(
592 bind: SocketAddr,
593 endpoint: &str,
594 guest_port: u16,
595 auth: Option<Arc<ProxyAuth>>,
596 limits: ServeLimits,
597) -> Result<Serving<ForwardSummary>, Error> {
598 let spec = ForwardSpec::new(bind, guest_port, endpoint);
599 let listener = super::forward::bind(&spec).await?;
600 let local = local_address(&listener)?;
601 Ok(Serving::spawn(local, move |stop, cut| async move {
602 let mut ended = Vec::new();
603 let report = forward_loop(
604 listener,
605 spec,
606 auth,
607 limits,
608 signalled(stop),
609 signalled(cut),
610 |notice| ended.extend(ConnectionEnd::of_forward(¬ice)),
611 )
612 .await?;
613 Ok(ForwardSummary { report, ended })
614 }))
615}
616
617fn local_address(listener: &TcpListener) -> Result<SocketAddr, Error> {
618 listener.local_addr().map_err(|err| {
619 Error::new(
620 ErrorKind::Unexpected,
621 format!("the listener has no address: {err}"),
622 )
623 })
624}
625
626async fn serve<T, Fut>(
633 listener: &TcpListener,
634 limits: ServeLimits,
635 stop: impl Future<Output = ()>,
636 cut: impl Future<Output = ()>,
637 mut spawn: impl FnMut(tokio::net::TcpStream, SocketAddr) -> Fut,
638 mut ended: impl FnMut(SocketAddr, Result<T, Error>),
639) -> (StopReason, u32)
640where
641 T: Send + 'static,
642 Fut: Future<Output = Result<T, Error>> + Send + 'static,
643{
644 let mut stop = std::pin::pin!(stop);
645 let mut tasks = JoinSet::new();
646 let mut peers = Peers::new();
648 type Peers = std::collections::HashMap<tokio::task::Id, SocketAddr>;
649 let mut finish =
650 |peers: &mut Peers,
651 joined: Result<(tokio::task::Id, Result<T, Error>), tokio::task::JoinError>| {
652 let (id, result) = match joined {
653 Ok((id, result)) => (id, result),
654 Err(error) if error.is_cancelled() => (
655 error.id(),
656 Err(Error::new(
657 ErrorKind::Unexpected,
658 "the connection was cut when its loop stopped",
659 )),
660 ),
661 Err(error) => (
662 error.id(),
663 Err(Error::new(
664 ErrorKind::Unexpected,
665 format!("the connection's task ended abnormally: {error}"),
666 )),
667 ),
668 };
669 let peer = peers
670 .remove(&id)
671 .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
672 ended(peer, result);
673 };
674 let mut served = 0_u32;
675 let stopped = loop {
676 if limits.max_connections.is_some_and(|max| served >= max) {
677 break StopReason::Limit;
678 }
679 tokio::select! {
680 biased;
683 () = &mut stop => break StopReason::Requested,
684 Some(joined) = tasks.join_next_with_id(), if !tasks.is_empty() => finish(&mut peers, joined),
685 next = listener.accept() => match next {
686 Ok((stream, peer)) => {
687 served += 1;
688 let handle = tasks.spawn(spawn(stream, peer));
689 peers.insert(handle.id(), peer);
690 }
691 Err(error) => {
692 break StopReason::ListenerFailed(format!(
693 "the local listener stopped accepting: {error}"
694 ));
695 }
696 },
697 }
698 };
699 let mut cut = std::pin::pin!(cut);
701 let mut cutting = false;
702 loop {
703 tokio::select! {
704 joined = tasks.join_next_with_id() => match joined {
705 Some(joined) => finish(&mut peers, joined),
706 None => break,
707 },
708 () = &mut cut, if !cutting => {
709 cutting = true;
710 tasks.abort_all();
711 }
712 }
713 }
714 (stopped, served)
715}
716
717#[cfg(test)]
718mod tests {
719 use super::*;
720
721 fn tunnel_report() -> TunnelReport {
722 TunnelReport {
723 served: 0,
724 refused: 0,
725 truncated: 0,
726 unproven: 0,
727 proxy_token_mints: 0,
728 stopped: StopReason::Requested,
729 }
730 }
731
732 #[test]
734 fn each_tunnel_end_is_counted_once_in_its_column() {
735 let mut report = tunnel_report();
736 for end in [
737 Ok(TunnelEnd::Closed),
738 Ok(TunnelEnd::ClosedUnproven),
739 Ok(TunnelEnd::ClosedUnproven),
740 Ok(TunnelEnd::Truncated { code: None }),
741 Ok(TunnelEnd::Truncated { code: Some(1000) }),
742 Ok(TunnelEnd::Truncated { code: Some(1000) }),
743 Ok(TunnelEnd::Refused {
744 code: 4502,
745 reason: String::new(),
746 }),
747 Err(Error::new(ErrorKind::Unexpected, "a transport error")),
748 ] {
749 report.count(&end);
750 }
751 assert_eq!(
752 (report.refused, report.truncated, report.unproven),
753 (2, 3, 2),
754 "{report:?}"
755 );
756 }
757
758 #[test]
760 fn only_the_callers_stop_was_requested() {
761 assert!(StopReason::Requested.was_requested());
762 assert!(!StopReason::Limit.was_requested());
763 assert!(!StopReason::ListenerFailed("gone".to_string()).was_requested());
764 }
765
766 #[test]
768 fn each_forward_event_is_counted_once_in_its_column() {
769 let mut report = ForwardReport {
770 served: 0,
771 refused: 0,
772 upgrades: 0,
773 proxy_token_mints: 0,
774 stopped: StopReason::Requested,
775 };
776 for event in [
777 ForwardEvent::Refused {
778 status: 502,
779 explanation: String::new(),
780 },
781 ForwardEvent::Forwarded {
782 status: 101,
783 upgraded: true,
784 },
785 ForwardEvent::Forwarded {
786 status: 101,
787 upgraded: true,
788 },
789 ForwardEvent::Forwarded {
790 status: 200,
791 upgraded: false,
792 },
793 ForwardEvent::ConnectionError {
794 detail: String::new(),
795 },
796 ] {
797 report.count(&event);
798 }
799 assert_eq!((report.refused, report.upgrades), (1, 2), "{report:?}");
800 }
801
802 #[test]
805 fn each_end_but_a_clean_one_is_listed_with_its_kind() {
806 let peer: SocketAddr = "127.0.0.1:5000".parse().expect("an address");
807 let tunnel = |end| ConnectionEnd::of_tunnel(&TunnelConnection { peer, end });
808 let listed = |end: Option<ConnectionEnd>| {
809 end.map(|end| (end.peer, end.kind.as_str(), end.code, end.detail))
810 };
811 assert_eq!(tunnel(Ok(TunnelEnd::Closed)), None);
812 assert_eq!(
813 listed(tunnel(Ok(TunnelEnd::ClosedUnproven))),
814 Some((peer, "unproven", None, String::new()))
815 );
816 assert_eq!(
817 listed(tunnel(Ok(TunnelEnd::Truncated { code: Some(1006) }))),
818 Some((peer, "truncated", Some(1006), String::new()))
819 );
820 assert_eq!(
821 listed(tunnel(Ok(TunnelEnd::Refused {
822 code: 4502,
823 reason: "nothing is listening".into(),
824 }))),
825 Some((peer, "refused", Some(4502), "nothing is listening".into()))
826 );
827 let failed = listed(tunnel(Err(Error::new(ErrorKind::Unexpected, "a reset"))));
828 assert!(
829 matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
830 "{failed:?}"
831 );
832
833 let forward = |event| ConnectionEnd::of_forward(&ForwardNotice::Event { peer, event });
834 assert_eq!(
835 listed(forward(ForwardEvent::Refused {
836 status: 403,
837 explanation: "out of scope".into(),
838 })),
839 Some((peer, "refused", Some(403), "out of scope".into()))
840 );
841 assert_eq!(
842 forward(ForwardEvent::Forwarded {
843 status: 200,
844 upgraded: false,
845 }),
846 None
847 );
848 let failed = listed(ConnectionEnd::of_forward(&ForwardNotice::Failed {
849 peer,
850 error: Error::new(ErrorKind::Unexpected, "a reset"),
851 }));
852 assert!(
853 matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
854 "{failed:?}"
855 );
856 }
857
858 #[tokio::test]
861 async fn a_panicked_connection_and_a_cut_one_are_told_apart() {
862 let listener = TcpListener::bind("127.0.0.1:0").await.expect("a free port");
863 let local = listener.local_addr().expect("bound");
864 let (cut, cut_rx) = oneshot::channel::<()>();
865 tokio::spawn(async move {
866 let _first = tokio::net::TcpStream::connect(local).await;
867 let _second = tokio::net::TcpStream::connect(local).await;
868 tokio::time::sleep(Duration::from_millis(200)).await;
869 let _ = cut.send(());
870 std::future::pending::<()>().await;
871 });
872 let mut accepted = 0;
873 let mut ended = Vec::new();
874 let serving = serve(
875 &listener,
876 ServeLimits {
877 max_connections: Some(2),
878 },
879 std::future::pending(),
880 signalled(cut_rx),
881 |stream, _| {
882 accepted += 1;
883 let panics = accepted == 1;
884 async move {
885 let _held = stream;
886 assert!(!panics, "the first connection's task panics");
887 std::future::pending::<Result<(), Error>>().await
888 }
889 },
890 |_, result| ended.push(result.expect_err("neither ends clean").to_string()),
891 );
892 let (stopped, served) = tokio::time::timeout(Duration::from_secs(10), serving)
893 .await
894 .expect("the cut ends the drain");
895 assert_eq!((stopped, served), (StopReason::Limit, 2));
896 assert_eq!(ended.len(), 2, "{ended:?}");
897 let panicked = ended.iter().filter(|end| end.contains("ended abnormally"));
898 let cut = ended.iter().filter(|end| end.contains("was cut"));
899 assert_eq!((panicked.count(), cut.count()), (1, 1), "{ended:?}");
900 }
901}