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 crate::error::{Error, ErrorKind};
use crate::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 },
}
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, auth, None).await
}
pub async fn relay_connection_verified<S>(
local: S,
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: &Arc<ProxyAuth>,
identity: &crate::identity::TunnelIdentity,
) -> Result<TunnelEnd, Error>
where
S: AsyncRead + AsyncWrite + Unpin,
{
relay_connection_inner(
local,
endpoint,
guest_port,
agent_token,
auth,
Some(identity),
)
.await
}
pub async fn verify_identity(
endpoint: &str,
agent_token: &str,
auth: &Arc<ProxyAuth>,
identity: &crate::identity::TunnelIdentity,
) -> Result<TunnelEnd, Error> {
let guest_port = auth.port();
match open(endpoint, guest_port, agent_token, auth, Some(identity)).await? {
Opened::Refused(end) => Ok(end),
Opened::Ready(mut ready) => {
let _ = ready.socket.send(Message::Close(None)).await;
Ok(TunnelEnd::Closed)
}
}
}
enum Opened {
Ready(Box<Ready>),
Refused(TunnelEnd),
}
struct Ready {
socket: TunnelSocket,
noise: Option<snow::TransportState>,
}
type TunnelSocket =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
async fn open(
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: &Arc<ProxyAuth>,
identity: Option<&crate::identity::TunnelIdentity>,
) -> Result<Opened, Error> {
let offered = auth.subprotocols(auth.port()).await?;
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();
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(transport) => Some(transport),
Initiated::Refused(end) => return Ok(Opened::Refused(end)),
},
};
Ok(Opened::Ready(Box::new(Ready { socket, noise })))
}
async fn relay_connection_inner<S>(
mut local: S,
endpoint: &str,
guest_port: u16,
agent_token: &str,
auth: &Arc<ProxyAuth>,
identity: Option<&crate::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) => {
let _ = socket.send(Message::Close(None)).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(TunnelEnd::Closed);
}
}
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(transport) => {
let count = transport.read_message(&bytes, &mut scratch).map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("a tunnel frame did not authenticate: {err}"),
)
})?;
&scratch[..count]
}
};
if local.write_all(plain).await.is_err() {
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(classify_close(frame.as_ref()));
}
Some(Ok(_)) => continue,
Some(Err(_)) => return Ok(TunnelEnd::Closed),
None => return Ok(TunnelEnd::Closed),
},
}
}
}
fn seal(
noise: &mut Option<snow::TransportState>,
plain: &[u8],
scratch: &mut [u8],
) -> Result<Option<Vec<u8>>, Error> {
let Some(transport) = noise.as_mut() else {
return Ok(None);
};
let count = 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(snow::TransportState),
Refused(TunnelEnd),
}
async fn initiate(
socket: &mut TunnelSocket,
identity: &crate::identity::TunnelIdentity,
guest_port: u16,
) -> Result<Initiated, Error> {
let mut initiator = identity.initiator()?;
let mut scratch = vec![0_u8; NOISE_MAX_MESSAGE_BYTES];
let written = initiator.write_message(&[], &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))) => {
initiator.read_message(&reply, &mut 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 transport = initiator.into_transport_mode().map_err(|err| {
Error::new(
ErrorKind::Unexpected,
format!("entering transport mode failed: {err}"),
)
})?;
return Ok(Initiated::Transport(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(),
));
}
}
}
}
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!(
crate::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"
);
};
}