use std::sync::Arc;
use futures_util::{SinkExt as _, StreamExt as _};
use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _};
use tokio_tungstenite::tungstenite::client::IntoClientRequest as _;
use tokio_tungstenite::tungstenite::{Message, protocol::CloseFrame};
use microvms_app::error::{Error, ErrorKind};
use microvms_app::session::proxy::ProxyAuth;
pub const TUNNEL_CHUNK_BYTES: usize = 64 * 1024;
const TUNNEL_PATH: &str = "/v1/tcp";
pub const VERIFIED_CHUNK_BYTES: usize = 32 * 1024;
const NOISE_MAX_MESSAGE_BYTES: usize = 65535;
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum TunnelEnd {
Closed,
Refused { code: u16, reason: String },
Truncated { code: Option<u16> },
ClosedUnproven,
}
pub fn tunnel_url(endpoint: &str, guest_port: u16) -> String {
tunnel_url_with_identity(endpoint, guest_port, false)
}
pub fn tunnel_url_with_identity(endpoint: &str, guest_port: u16, identity: bool) -> String {
let host = endpoint
.trim_end_matches('/')
.trim_start_matches("https://")
.trim_start_matches("http://")
.trim_start_matches("wss://")
.trim_start_matches("ws://");
let scheme = if endpoint.starts_with("http://") || endpoint.starts_with("ws://") {
"ws"
} else {
"wss"
};
let flag = if identity { "&identity=true" } else { "" };
format!("{scheme}://{host}{TUNNEL_PATH}?port={guest_port}{flag}")
}
pub fn explain_close(code: u16, guest_port: u16) -> Option<String> {
if let Some(explanation) = protocol::tunnel::close::explanation(code, guest_port) {
return Some(explanation);
}
if code == 1006 {
return Some(format!(
"the connection closed abnormally (1006) with no reason, which is what every \
endpoint-proxy WebSocket failure looks like: a refused handshake, a token whose \
scope does not cover port {guest_port}, and a dropped TCP connection are all \
this code. Retry the same port over HTTPS to tell a scope mistake (403) from a \
dead server (502)."
));
}
None
}
pub async fn relay_connection<S>(
local: S,
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: &Arc<ProxyAuth>,
) -> Result<TunnelEnd, Error>
where
S: AsyncRead + AsyncWrite + Unpin,
{
relay_connection_inner(local, endpoint, guest_port, agent_token, Some(auth), None).await
}
pub async fn relay_connection_verified<S>(
local: S,
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: &Arc<ProxyAuth>,
identity: µvms_app::identity::TunnelIdentity,
) -> Result<TunnelEnd, Error>
where
S: AsyncRead + AsyncWrite + Unpin,
{
relay_connection_inner(
local,
endpoint,
guest_port,
agent_token,
Some(auth),
Some(identity),
)
.await
}
pub async fn verify_identity(
endpoint: &str,
agent_token: &str,
auth: &Arc<ProxyAuth>,
identity: µvms_app::identity::TunnelIdentity,
) -> Result<TunnelEnd, Error> {
let guest_port = auth.port();
match open(
endpoint,
guest_port,
agent_token,
Some(auth),
Some(identity),
)
.await?
{
Opened::Refused(end) => Ok(end),
Opened::Ready(ready) => {
let Ready {
mut socket,
mut noise,
} = *ready;
let mut scratch = vec![0_u8; NOISE_MAX_MESSAGE_BYTES];
end_first(&mut socket, &mut noise, &mut scratch).await;
Ok(TunnelEnd::Closed)
}
}
}
enum Opened {
Ready(Box<Ready>),
Refused(TunnelEnd),
}
struct Ready {
socket: TunnelSocket,
noise: Option<Verified>,
}
pub(crate) struct Verified {
transport: snow::TransportState,
pub(crate) daemon_proves_end: bool,
}
type TunnelSocket =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
async fn open(
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: Option<&Arc<ProxyAuth>>,
identity: Option<µvms_app::identity::TunnelIdentity>,
) -> Result<Opened, Error> {
let offered = match auth {
Some(auth) => Some(auth.subprotocols(auth.port()).await?),
None => None,
};
let url = tunnel_url_with_identity(endpoint, guest_port, identity.is_some());
let mut request = url.as_str().into_client_request().map_err(|err| {
Error::new(
ErrorKind::InvalidArg,
format!("{url} is not a usable WebSocket URL: {err}"),
)
})?;
{
let headers = request.headers_mut();
if let Some(offered) = offered {
headers.insert(
"sec-websocket-protocol",
offered.join(", ").parse().map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("the minted subprotocols are not a legal header value: {err}"),
)
})?,
);
}
headers.insert(
"authorization",
format!("Bearer {agent_token}").parse().map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("the agent token is not a legal header value: {err}"),
)
})?,
);
}
let (mut socket, _response) =
tokio_tungstenite::connect_async(request)
.await
.map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!(
"the tunnel handshake to port {guest_port} failed: {err}. Every \
endpoint-proxy WebSocket failure is close code 1006 with no reason, so \
if this names nothing, retry the port over HTTPS to tell a scope \
mistake (403) from a dead server (502)."
),
)
})?;
let noise = match identity {
None => None,
Some(identity) => match initiate(&mut socket, identity, guest_port).await? {
Initiated::Transport(verified) => Some(verified),
Initiated::Refused(end) => return Ok(Opened::Refused(end)),
},
};
Ok(Opened::Ready(Box::new(Ready { socket, noise })))
}
pub(crate) async fn relay_connection_inner<S>(
mut local: S,
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: Option<&Arc<ProxyAuth>>,
identity: Option<µvms_app::identity::TunnelIdentity>,
) -> Result<TunnelEnd, Error>
where
S: AsyncRead + AsyncWrite + Unpin,
{
let Ready {
mut socket,
mut noise,
} = match open(endpoint, guest_port, agent_token, auth, identity).await? {
Opened::Ready(ready) => *ready,
Opened::Refused(end) => return Ok(end),
};
let chunk = if noise.is_some() {
VERIFIED_CHUNK_BYTES
} else {
TUNNEL_CHUNK_BYTES
};
let mut buffer = vec![0_u8; chunk];
let mut scratch = vec![0_u8; NOISE_MAX_MESSAGE_BYTES];
loop {
tokio::select! {
read = local.read(&mut buffer) => match read {
Ok(0) => {
end_first(&mut socket, &mut noise, &mut scratch).await;
return Ok(TunnelEnd::Closed);
}
Ok(count) => {
let frame = match seal(&mut noise, &buffer[..count], &mut scratch)? {
Some(sealed) => sealed,
None => buffer[..count].to_vec(),
};
if socket.send(Message::Binary(frame.into())).await.is_err() {
return Ok(unproven_end(noise.as_ref(), None));
}
}
Err(err) => {
return Err(Error::new(
ErrorKind::Unexpected,
format!("reading the local connection failed: {err}"),
));
}
},
frame = socket.next() => match frame {
Some(Ok(Message::Binary(bytes))) => {
let plain: &[u8] = match noise.as_mut() {
None => &bytes,
Some(verified) => match open_frame(verified, &bytes, &mut scratch)? {
Frame::Chunk(count) => &scratch[..count],
Frame::End => {
let _ = local.flush().await;
let _ = socket.send(Message::Close(None)).await;
return Ok(TunnelEnd::Closed);
}
},
};
if local.write_all(plain).await.is_err() {
end_first(&mut socket, &mut noise, &mut scratch).await;
return Ok(TunnelEnd::Closed);
}
}
Some(Ok(Message::Text(text))) => {
if noise.is_some() {
return Err(Error::new(
ErrorKind::Unexpected,
"a verified tunnel received a text frame, which cannot be a Noise \
message".to_string(),
));
}
if local.write_all(text.as_bytes()).await.is_err() {
return Ok(TunnelEnd::Closed);
}
}
Some(Ok(Message::Close(frame))) => {
let _ = local.flush().await;
return Ok(match classify_close(frame.as_ref()) {
TunnelEnd::Closed => {
unproven_end(noise.as_ref(), frame.as_ref().map(|f| f.code.into()))
}
other => other,
});
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => return Ok(unproven_end(noise.as_ref(), None)),
},
}
}
}
fn unproven_end(noise: Option<&Verified>, code: Option<u16>) -> TunnelEnd {
match noise {
None => TunnelEnd::Closed,
Some(verified) if verified.daemon_proves_end => TunnelEnd::Truncated { code },
Some(_) => TunnelEnd::ClosedUnproven,
}
}
async fn end_first(socket: &mut TunnelSocket, noise: &mut Option<Verified>, scratch: &mut [u8]) {
if let Ok(Some(end)) = seal(noise, &[], scratch) {
let _ = socket.send(Message::Binary(end.into())).await;
}
let _ = socket.send(Message::Close(None)).await;
}
fn seal(
noise: &mut Option<Verified>,
plain: &[u8],
scratch: &mut [u8],
) -> Result<Option<Vec<u8>>, Error> {
let Some(verified) = noise.as_mut() else {
return Ok(None);
};
let count = verified
.transport
.write_message(plain, scratch)
.map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("encrypting a tunnel chunk failed: {err}"),
)
})?;
Ok(Some(scratch[..count].to_vec()))
}
enum Initiated {
Transport(Verified),
Refused(TunnelEnd),
}
async fn initiate(
socket: &mut TunnelSocket,
identity: µvms_app::identity::TunnelIdentity,
guest_port: u16,
) -> Result<Initiated, Error> {
let mut initiator = crate::identity::initiator(identity)?;
let mut scratch = vec![0_u8; NOISE_MAX_MESSAGE_BYTES];
let written = initiator
.write_message(&protocol::identity::HANDSHAKE_PAYLOAD, &mut scratch)
.map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("writing the identity handshake failed: {err}"),
)
})?;
if socket
.send(Message::Binary(scratch[..written].to_vec().into()))
.await
.is_err()
{
return Err(Error::new(
ErrorKind::Unexpected,
"the tunnel closed before the identity handshake could be sent".to_string(),
));
}
loop {
match socket.next().await {
Some(Ok(Message::Binary(reply))) => {
return read_reply(initiator, &reply, &mut scratch, guest_port)
.map(Initiated::Transport);
}
Some(Ok(Message::Close(frame))) => {
return Ok(Initiated::Refused(classify_close(frame.as_ref())));
}
Some(Ok(_)) => continue,
Some(Err(err)) => {
return Err(Error::new(
ErrorKind::Unexpected,
format!("the tunnel failed mid-handshake: {err}"),
));
}
None => {
return Err(Error::new(
ErrorKind::Unexpected,
"the tunnel closed mid-handshake with no close frame".to_string(),
));
}
}
}
}
pub(crate) fn read_reply(
mut initiator: snow::HandshakeState,
reply: &[u8],
scratch: &mut [u8],
guest_port: u16,
) -> Result<Verified, Error> {
let offer = initiator.read_message(reply, scratch).map_err(|_| {
Error::new(
ErrorKind::Unexpected,
format!(
"the identity handshake reply did not verify against the pinned key for this \
VM. The far end holds a different seed than the one this record was created \
with: a record copied from another VM, or a VM relaunched with a fresh seed, \
would both do this. (guest port {guest_port})"
),
)
})?;
let daemon_proves_end = protocol::identity::offers_end_of_stream(&scratch[..offer]);
let transport = initiator.into_transport_mode().map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("entering transport mode failed: {err}"),
)
})?;
Ok(Verified {
transport,
daemon_proves_end,
})
}
#[derive(Debug, Eq, PartialEq)]
pub(crate) enum Frame {
Chunk(usize),
End,
}
pub(crate) fn open_frame(
verified: &mut Verified,
frame: &[u8],
scratch: &mut [u8],
) -> Result<Frame, Error> {
let count = verified
.transport
.read_message(frame, scratch)
.map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("a tunnel frame did not authenticate: {err}"),
)
})?;
Ok(if protocol::identity::is_end_of_stream(&scratch[..count]) {
Frame::End
} else {
Frame::Chunk(count)
})
}
fn classify_close(frame: Option<&CloseFrame>) -> TunnelEnd {
let Some(frame) = frame else {
return TunnelEnd::Closed;
};
let code: u16 = frame.code.into();
if code == protocol::tunnel::close::NORMAL {
return TunnelEnd::Closed;
}
TunnelEnd::Refused {
code,
reason: frame.reason.to_string(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_bare_endpoint_host_becomes_wss() {
assert_eq!(
tunnel_url("vm-abc.example.aws", 8080),
"wss://vm-abc.example.aws/v1/tcp?port=8080"
);
assert_eq!(
tunnel_url("https://vm-abc.example.aws/", 5432),
"wss://vm-abc.example.aws/v1/tcp?port=5432"
);
}
#[test]
fn an_explicit_plaintext_endpoint_stays_plaintext() {
assert_eq!(
tunnel_url("http://127.0.0.1:9000", 8080),
"ws://127.0.0.1:9000/v1/tcp?port=8080"
);
assert_eq!(
tunnel_url("ws://127.0.0.1:9000", 1),
"ws://127.0.0.1:9000/v1/tcp?port=1"
);
}
#[test]
fn the_daemons_close_codes_are_explained_and_the_platforms_ambiguity_is_named() {
let dead = explain_close(protocol::tunnel::close::NO_LISTENER, 5432)
.expect("4502 explains itself");
assert!(dead.contains("5432"), "{dead}");
assert!(dead.contains("not an auth problem"), "{dead}");
let opaque = explain_close(1006, 8080).expect("1006 is worth naming");
assert!(
opaque.contains("scope does not cover port 8080"),
"{opaque}"
);
assert!(opaque.contains("403"), "{opaque}");
assert!(opaque.contains("502"), "{opaque}");
assert!(explain_close(4999, 8080).is_none());
}
#[test]
fn a_clean_close_is_not_a_refusal() {
assert_eq!(classify_close(None), TunnelEnd::Closed);
let normal = CloseFrame {
code: protocol::tunnel::close::NORMAL.into(),
reason: "done".into(),
};
assert_eq!(classify_close(Some(&normal)), TunnelEnd::Closed);
}
#[test]
fn a_daemon_refusal_carries_its_code_so_a_caller_can_branch() {
let refused = CloseFrame {
code: protocol::tunnel::close::NO_LISTENER.into(),
reason: "nothing is listening on 127.0.0.1:5432 in the guest".into(),
};
match classify_close(Some(&refused)) {
TunnelEnd::Refused { code, reason } => {
assert_eq!(code, protocol::tunnel::close::NO_LISTENER);
assert!(reason.contains("5432"), "{reason}");
}
other => panic!("expected a refusal, got {other:?}"),
}
}
#[test]
fn the_offered_marker_is_the_daemons_echoed_marker() {
assert_eq!(
microvms_app::session::proxy::WS_SUBPROTOCOL,
protocol::tunnel::WS_MARKER_SUBPROTOCOL,
"the marker must stay aliased to the protocol crate's, not copied from it"
);
}
const _: () = {
assert!(
TUNNEL_CHUNK_BYTES == 64 * 1024,
"the client chunk must match the daemon's RELAY_CHUNK_BYTES"
);
};
}