use std::future::Future;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::{oneshot, watch};
use tokio::task::JoinSet;
use microvms_app::error::{Error, ErrorKind};
use microvms_app::identity::TunnelIdentity;
use microvms_app::session::Session;
use microvms_app::session::proxy::ProxyAuth;
use super::forward::{ForwardClient, ForwardEvent, ForwardSpec, serve_connection_via};
use super::tunnel::{TunnelEnd, relay_connection_inner};
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct ServeLimits {
pub max_connections: Option<u32>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum StopReason {
Requested,
Limit,
ListenerFailed(String),
}
impl StopReason {
pub fn was_requested(&self) -> bool {
matches!(self, StopReason::Requested)
}
}
#[derive(Clone)]
pub struct TunnelTarget {
pub endpoint: String,
pub agent_token: String,
pub auth: Option<Arc<ProxyAuth>>,
pub guest_port: u16,
pub identity: Option<TunnelIdentity>,
}
impl TunnelTarget {
pub fn for_session(
session: &Session,
guest_port: u16,
identity: Option<TunnelIdentity>,
) -> Self {
Self {
endpoint: session.endpoint().to_string(),
agent_token: session.agent_token().to_string(),
auth: session.proxy_auth().cloned(),
guest_port,
identity,
}
}
}
#[derive(Debug)]
pub struct TunnelConnection {
pub peer: SocketAddr,
pub end: Result<TunnelEnd, Error>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TunnelReport {
pub served: u32,
pub refused: u32,
pub truncated: u32,
pub unproven: u32,
pub proxy_token_mints: u64,
pub stopped: StopReason,
}
impl TunnelReport {
fn count(&mut self, end: &Result<TunnelEnd, Error>) {
match end {
Ok(TunnelEnd::Closed) => {}
Ok(TunnelEnd::ClosedUnproven) => self.unproven += 1,
Ok(TunnelEnd::Truncated { .. }) => self.truncated += 1,
Ok(TunnelEnd::Refused { .. }) | Err(_) => self.refused += 1,
}
}
}
pub async fn serve_tunnel<F>(
listener: TcpListener,
target: TunnelTarget,
limits: ServeLimits,
stop: impl Future<Output = ()>,
on_end: F,
) -> TunnelReport
where
F: FnMut(TunnelConnection),
{
tunnel_loop(
listener,
target,
limits,
stop,
std::future::pending(),
on_end,
)
.await
}
async fn tunnel_loop<F>(
listener: TcpListener,
target: TunnelTarget,
limits: ServeLimits,
stop: impl Future<Output = ()>,
cut: impl Future<Output = ()>,
mut on_end: F,
) -> TunnelReport
where
F: FnMut(TunnelConnection),
{
let target = Arc::new(target);
let mut report = TunnelReport {
served: 0,
refused: 0,
truncated: 0,
unproven: 0,
proxy_token_mints: 0,
stopped: StopReason::Requested,
};
let (stopped, served) = serve(
&listener,
limits,
stop,
cut,
|mut stream, _| {
let target = Arc::clone(&target);
async move {
let result = relay_connection_inner(
&mut stream,
&target.endpoint,
target.guest_port,
&target.agent_token,
target.auth.as_ref(),
target.identity.as_ref(),
)
.await;
if target.identity.is_some()
&& matches!(
result,
Err(_) | Ok(TunnelEnd::Truncated { .. } | TunnelEnd::Refused { .. })
)
{
stream.set_zero_linger().map_err(|error| {
Error::new(
ErrorKind::Unexpected,
format!("could not reset the local tunnel connection: {error}"),
)
})?;
}
result
}
},
|peer, end| {
report.count(&end);
on_end(TunnelConnection { peer, end });
},
)
.await;
report.served = served;
report.stopped = stopped;
report.proxy_token_mints = target.auth.as_ref().map_or(0, |auth| auth.mint_count());
report
}
#[derive(Debug)]
pub enum ForwardNotice {
Event {
peer: SocketAddr,
event: ForwardEvent,
},
Failed { peer: SocketAddr, error: Error },
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ForwardReport {
pub served: u32,
pub refused: u32,
pub upgrades: u32,
pub proxy_token_mints: u64,
pub stopped: StopReason,
}
impl ForwardReport {
fn count(&mut self, event: &ForwardEvent) {
match event {
ForwardEvent::Refused { .. } => self.refused += 1,
ForwardEvent::Forwarded { upgraded: true, .. } => self.upgrades += 1,
_ => {}
}
}
}
pub async fn serve_forward<F>(
listener: TcpListener,
spec: ForwardSpec,
auth: Option<Arc<ProxyAuth>>,
limits: ServeLimits,
stop: impl Future<Output = ()>,
on_notice: F,
) -> Result<ForwardReport, Error>
where
F: FnMut(ForwardNotice),
{
forward_loop(
listener,
spec,
auth,
limits,
stop,
std::future::pending(),
on_notice,
)
.await
}
async fn forward_loop<F>(
listener: TcpListener,
spec: ForwardSpec,
auth: Option<Arc<ProxyAuth>>,
limits: ServeLimits,
stop: impl Future<Output = ()>,
cut: impl Future<Output = ()>,
mut on_notice: F,
) -> Result<ForwardReport, Error>
where
F: FnMut(ForwardNotice),
{
let client = Arc::new(ForwardClient::new()?);
let mut report = ForwardReport {
served: 0,
refused: 0,
upgrades: 0,
proxy_token_mints: 0,
stopped: StopReason::Requested,
};
let (stopped, served) = serve(
&listener,
limits,
stop,
cut,
|stream, _| {
let (mut spec, auth, client) = (spec.clone(), auth.clone(), Arc::clone(&client));
async move {
spec.bind = stream.local_addr().map_err(|error| {
Error::new(
ErrorKind::Unexpected,
format!("reading the local forward address failed: {error}"),
)
})?;
let mut events = Vec::new();
let outcome =
serve_connection_via(stream, &spec, auth.as_ref(), &client, |event| {
events.push(event);
})
.await;
Ok((events, outcome))
}
},
|peer, done| {
let (events, failure) = match done {
Ok((events, outcome)) => (events, outcome.err()),
Err(error) => (Vec::new(), Some(error)),
};
for event in events {
report.count(&event);
on_notice(ForwardNotice::Event { peer, event });
}
if let Some(error) = failure {
on_notice(ForwardNotice::Failed { peer, error });
}
},
)
.await;
report.served = served;
report.stopped = stopped;
report.proxy_token_mints = auth.as_ref().map_or(0, |auth| auth.mint_count());
Ok(report)
}
pub const DEFAULT_BIND: SocketAddr = microvms_app::session::DEFAULT_BIND;
pub fn bind_address(bind: Option<&str>) -> Result<SocketAddr, Error> {
match bind {
None => Ok(DEFAULT_BIND),
Some(bind) => bind.parse().map_err(|err| {
Error::new(
ErrorKind::InvalidArg,
format!("{bind:?} is not a local address to listen on, such as 127.0.0.1:0: {err}"),
)
}),
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EndKind {
Refused,
Truncated,
Unproven,
Failed,
}
impl EndKind {
pub fn as_str(self) -> &'static str {
match self {
EndKind::Refused => "refused",
EndKind::Truncated => "truncated",
EndKind::Unproven => "unproven",
EndKind::Failed => "failed",
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ConnectionEnd {
pub peer: SocketAddr,
pub kind: EndKind,
pub code: Option<u16>,
pub detail: String,
}
impl ConnectionEnd {
pub fn of_tunnel(connection: &TunnelConnection) -> Option<Self> {
let (kind, code, detail) = match &connection.end {
Ok(TunnelEnd::Closed) => return None,
Ok(TunnelEnd::ClosedUnproven) => (EndKind::Unproven, None, String::new()),
Ok(TunnelEnd::Truncated { code }) => (EndKind::Truncated, *code, String::new()),
Ok(TunnelEnd::Refused { code, reason }) => {
(EndKind::Refused, Some(*code), reason.clone())
}
Err(error) => (EndKind::Failed, None, error.to_string()),
};
Some(Self {
peer: connection.peer,
kind,
code,
detail,
})
}
pub fn of_forward(notice: &ForwardNotice) -> Option<Self> {
match notice {
ForwardNotice::Event {
peer,
event:
ForwardEvent::Refused {
status,
explanation,
},
} => Some(Self {
peer: *peer,
kind: EndKind::Refused,
code: Some(*status),
detail: explanation.clone(),
}),
ForwardNotice::Event { .. } => None,
ForwardNotice::Failed { peer, error } => Some(Self {
peer: *peer,
kind: EndKind::Failed,
code: None,
detail: error.to_string(),
}),
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TunnelSummary {
pub report: TunnelReport,
pub ended: Vec<ConnectionEnd>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ForwardSummary {
pub report: ForwardReport,
pub ended: Vec<ConnectionEnd>,
}
type Outcome<S> = Result<S, (ErrorKind, String)>;
pub struct Serving<S> {
local: SocketAddr,
stop: Mutex<Option<oneshot::Sender<()>>>,
cut: Mutex<Option<oneshot::Sender<()>>>,
outcome: Arc<Mutex<Option<Outcome<S>>>>,
done: watch::Receiver<bool>,
}
impl<S: Clone + Send + 'static> Serving<S> {
fn spawn<Fut>(
local: SocketAddr,
serving: impl FnOnce(oneshot::Receiver<()>, oneshot::Receiver<()>) -> Fut,
) -> Self
where
Fut: Future<Output = Result<S, Error>> + Send + 'static,
{
let (stop_tx, stop_rx) = oneshot::channel();
let (cut_tx, cut_rx) = oneshot::channel();
let (done_tx, done_rx) = watch::channel(false);
let outcome = Arc::new(Mutex::new(None));
let written = Arc::clone(&outcome);
let serving = serving(stop_rx, cut_rx);
tokio::spawn(async move {
let result = serving
.await
.map_err(|error| (error.kind(), error.to_string()));
*written.lock().unwrap_or_else(PoisonError::into_inner) = Some(result);
done_tx.send_replace(true);
});
Self {
local,
stop: Mutex::new(Some(stop_tx)),
cut: Mutex::new(Some(cut_tx)),
outcome,
done: done_rx,
}
}
pub fn local_address(&self) -> SocketAddr {
self.local
}
pub fn is_running(&self) -> bool {
!*self.done.borrow()
}
pub fn request_stop(&self) {
signal(&self.stop);
}
pub async fn stop(&self, grace: Option<Duration>) -> Result<S, Error> {
self.request_stop();
if let Some(grace) = grace {
if let Ok(outcome) = tokio::time::timeout(grace, self.finished()).await {
return outcome;
}
signal(&self.cut);
}
self.finished().await
}
pub async fn finished(&self) -> Result<S, Error> {
let mut done = self.done.clone();
let _ = done.wait_for(|finished| *finished).await;
self.outcome
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone()
.expect("the loop records its outcome before signalling done")
.map_err(|(kind, message)| Error::new(kind, message))
}
}
fn signal(sender: &Mutex<Option<oneshot::Sender<()>>>) {
if let Some(sender) = sender.lock().unwrap_or_else(PoisonError::into_inner).take() {
let _ = sender.send(());
}
}
async fn signalled(signal: oneshot::Receiver<()>) {
let _ = signal.await;
}
pub async fn start_tunnel(
bind: SocketAddr,
target: TunnelTarget,
limits: ServeLimits,
) -> Result<Serving<TunnelSummary>, Error> {
let listener =
super::forward::bind(&ForwardSpec::new(bind, target.guest_port, &target.endpoint)).await?;
let local = local_address(&listener)?;
Ok(Serving::spawn(local, move |stop, cut| async move {
let mut ended = Vec::new();
let report = tunnel_loop(
listener,
target,
limits,
signalled(stop),
signalled(cut),
|connection| ended.extend(ConnectionEnd::of_tunnel(&connection)),
)
.await;
Ok(TunnelSummary { report, ended })
}))
}
pub async fn start_forward(
bind: SocketAddr,
endpoint: &str,
guest_port: u16,
auth: Option<Arc<ProxyAuth>>,
limits: ServeLimits,
) -> Result<Serving<ForwardSummary>, Error> {
let spec = ForwardSpec::new(bind, guest_port, endpoint);
let listener = super::forward::bind(&spec).await?;
let local = local_address(&listener)?;
Ok(Serving::spawn(local, move |stop, cut| async move {
let mut ended = Vec::new();
let report = forward_loop(
listener,
spec,
auth,
limits,
signalled(stop),
signalled(cut),
|notice| ended.extend(ConnectionEnd::of_forward(¬ice)),
)
.await?;
Ok(ForwardSummary { report, ended })
}))
}
fn local_address(listener: &TcpListener) -> Result<SocketAddr, Error> {
listener.local_addr().map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("the listener has no address: {err}"),
)
})
}
async fn serve<T, Fut>(
listener: &TcpListener,
limits: ServeLimits,
stop: impl Future<Output = ()>,
cut: impl Future<Output = ()>,
mut spawn: impl FnMut(tokio::net::TcpStream, SocketAddr) -> Fut,
mut ended: impl FnMut(SocketAddr, Result<T, Error>),
) -> (StopReason, u32)
where
T: Send + 'static,
Fut: Future<Output = Result<T, Error>> + Send + 'static,
{
let mut stop = std::pin::pin!(stop);
let mut tasks = JoinSet::new();
let mut peers = Peers::new();
type Peers = std::collections::HashMap<tokio::task::Id, SocketAddr>;
let mut finish =
|peers: &mut Peers,
joined: Result<(tokio::task::Id, Result<T, Error>), tokio::task::JoinError>| {
let (id, result) = match joined {
Ok((id, result)) => (id, result),
Err(error) if error.is_cancelled() => (
error.id(),
Err(Error::new(
ErrorKind::Unexpected,
"the connection was cut when its loop stopped",
)),
),
Err(error) => (
error.id(),
Err(Error::new(
ErrorKind::Unexpected,
format!("the connection's task ended abnormally: {error}"),
)),
),
};
let peer = peers
.remove(&id)
.unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
ended(peer, result);
};
let mut served = 0_u32;
let stopped = loop {
if limits.max_connections.is_some_and(|max| served >= max) {
break StopReason::Limit;
}
tokio::select! {
biased;
() = &mut stop => break StopReason::Requested,
Some(joined) = tasks.join_next_with_id(), if !tasks.is_empty() => finish(&mut peers, joined),
next = listener.accept() => match next {
Ok((stream, peer)) => {
served += 1;
let handle = tasks.spawn(spawn(stream, peer));
peers.insert(handle.id(), peer);
}
Err(error) => {
break StopReason::ListenerFailed(format!(
"the local listener stopped accepting: {error}"
));
}
},
}
};
let mut cut = std::pin::pin!(cut);
let mut cutting = false;
loop {
tokio::select! {
joined = tasks.join_next_with_id() => match joined {
Some(joined) => finish(&mut peers, joined),
None => break,
},
() = &mut cut, if !cutting => {
cutting = true;
tasks.abort_all();
}
}
}
(stopped, served)
}
#[cfg(test)]
mod tests {
use super::*;
fn tunnel_report() -> TunnelReport {
TunnelReport {
served: 0,
refused: 0,
truncated: 0,
unproven: 0,
proxy_token_mints: 0,
stopped: StopReason::Requested,
}
}
#[test]
fn each_tunnel_end_is_counted_once_in_its_column() {
let mut report = tunnel_report();
for end in [
Ok(TunnelEnd::Closed),
Ok(TunnelEnd::ClosedUnproven),
Ok(TunnelEnd::ClosedUnproven),
Ok(TunnelEnd::Truncated { code: None }),
Ok(TunnelEnd::Truncated { code: Some(1000) }),
Ok(TunnelEnd::Truncated { code: Some(1000) }),
Ok(TunnelEnd::Refused {
code: 4502,
reason: String::new(),
}),
Err(Error::new(ErrorKind::Unexpected, "a transport error")),
] {
report.count(&end);
}
assert_eq!(
(report.refused, report.truncated, report.unproven),
(2, 3, 2),
"{report:?}"
);
}
#[test]
fn only_the_callers_stop_was_requested() {
assert!(StopReason::Requested.was_requested());
assert!(!StopReason::Limit.was_requested());
assert!(!StopReason::ListenerFailed("gone".to_string()).was_requested());
}
#[test]
fn each_forward_event_is_counted_once_in_its_column() {
let mut report = ForwardReport {
served: 0,
refused: 0,
upgrades: 0,
proxy_token_mints: 0,
stopped: StopReason::Requested,
};
for event in [
ForwardEvent::Refused {
status: 502,
explanation: String::new(),
},
ForwardEvent::Forwarded {
status: 101,
upgraded: true,
},
ForwardEvent::Forwarded {
status: 101,
upgraded: true,
},
ForwardEvent::Forwarded {
status: 200,
upgraded: false,
},
ForwardEvent::ConnectionError {
detail: String::new(),
},
] {
report.count(&event);
}
assert_eq!((report.refused, report.upgrades), (1, 2), "{report:?}");
}
#[test]
fn each_end_but_a_clean_one_is_listed_with_its_kind() {
let peer: SocketAddr = "127.0.0.1:5000".parse().expect("an address");
let tunnel = |end| ConnectionEnd::of_tunnel(&TunnelConnection { peer, end });
let listed = |end: Option<ConnectionEnd>| {
end.map(|end| (end.peer, end.kind.as_str(), end.code, end.detail))
};
assert_eq!(tunnel(Ok(TunnelEnd::Closed)), None);
assert_eq!(
listed(tunnel(Ok(TunnelEnd::ClosedUnproven))),
Some((peer, "unproven", None, String::new()))
);
assert_eq!(
listed(tunnel(Ok(TunnelEnd::Truncated { code: Some(1006) }))),
Some((peer, "truncated", Some(1006), String::new()))
);
assert_eq!(
listed(tunnel(Ok(TunnelEnd::Refused {
code: 4502,
reason: "nothing is listening".into(),
}))),
Some((peer, "refused", Some(4502), "nothing is listening".into()))
);
let failed = listed(tunnel(Err(Error::new(ErrorKind::Unexpected, "a reset"))));
assert!(
matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
"{failed:?}"
);
let forward = |event| ConnectionEnd::of_forward(&ForwardNotice::Event { peer, event });
assert_eq!(
listed(forward(ForwardEvent::Refused {
status: 403,
explanation: "out of scope".into(),
})),
Some((peer, "refused", Some(403), "out of scope".into()))
);
assert_eq!(
forward(ForwardEvent::Forwarded {
status: 200,
upgraded: false,
}),
None
);
let failed = listed(ConnectionEnd::of_forward(&ForwardNotice::Failed {
peer,
error: Error::new(ErrorKind::Unexpected, "a reset"),
}));
assert!(
matches!(&failed, Some((_, "failed", None, detail)) if detail.contains("a reset")),
"{failed:?}"
);
}
#[tokio::test]
async fn a_panicked_connection_and_a_cut_one_are_told_apart() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("a free port");
let local = listener.local_addr().expect("bound");
let (cut, cut_rx) = oneshot::channel::<()>();
tokio::spawn(async move {
let _first = tokio::net::TcpStream::connect(local).await;
let _second = tokio::net::TcpStream::connect(local).await;
tokio::time::sleep(Duration::from_millis(200)).await;
let _ = cut.send(());
std::future::pending::<()>().await;
});
let mut accepted = 0;
let mut ended = Vec::new();
let serving = serve(
&listener,
ServeLimits {
max_connections: Some(2),
},
std::future::pending(),
signalled(cut_rx),
|stream, _| {
accepted += 1;
let panics = accepted == 1;
async move {
let _held = stream;
assert!(!panics, "the first connection's task panics");
std::future::pending::<Result<(), Error>>().await
}
},
|_, result| ended.push(result.expect_err("neither ends clean").to_string()),
);
let (stopped, served) = tokio::time::timeout(Duration::from_secs(10), serving)
.await
.expect("the cut ends the drain");
assert_eq!((stopped, served), (StopReason::Limit, 2));
assert_eq!(ended.len(), 2, "{ended:?}");
let panicked = ended.iter().filter(|end| end.contains("ended abnormally"));
let cut = ended.iter().filter(|end| end.contains("was cut"));
assert_eq!((panicked.count(), cut.count()), (1, 1), "{ended:?}");
}
}