use std::{
collections::BTreeMap,
io,
net::SocketAddr,
pin::Pin,
task::{Context, Poll},
time::{Duration, Instant},
};
use bytes::Bytes;
use tokio::{
io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader},
sync::{mpsc, oneshot},
};
use ts_control::SshRecorderFailureAction;
use ts_http_util::{Client, Method, Request, ResponseExt, StatusCode};
const CAST_VERSION: u32 = 2;
const PER_DIAL_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(5);
const HTTP2_PROBE_TIMEOUT: Duration = Duration::from_secs(10);
const ALL_DIAL_ATTEMPTS_TIMEOUT: Duration = Duration::from_secs(30);
pub const UPLOAD_ACK_WINDOW: Duration = Duration::from_secs(30);
const EXPECT_CONTINUE_TIMEOUT: Duration = PER_DIAL_ATTEMPT_TIMEOUT;
const MAX_RESPONSE_HEAD: usize = 8 * 1024;
const MAX_ACK_BUFFER: usize = 64 * 1024;
const CAST_QUEUE_DEPTH: usize = 64;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordingAttempt {
pub recorder: SocketAddr,
pub failure_message: String,
}
#[derive(Debug, thiserror::Error)]
pub enum RecorderError {
#[error("recording: no recorders configured")]
NoRecorders,
#[error("{0}")]
AllFailed(String),
#[error("recording: timed out connecting to recorders")]
DialBudgetElapsed,
#[error("{0}")]
Recorder(String),
#[error("recording: {0}")]
Io(#[from] io::Error),
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize)]
pub struct CastHeader {
pub version: u32,
pub width: u16,
pub height: u16,
pub timestamp: i64,
#[serde(skip_serializing_if = "String::is_empty")]
pub command: String,
#[serde(rename = "srcNode")]
pub src_node: String,
#[serde(rename = "srcNodeID")]
pub src_node_id: String,
#[serde(rename = "srcNodeTags", skip_serializing_if = "Vec::is_empty")]
pub src_node_tags: Vec<String>,
#[serde(rename = "srcNodeUserID", skip_serializing_if = "is_zero")]
pub src_node_user_id: i64,
#[serde(rename = "srcNodeUser", skip_serializing_if = "String::is_empty")]
pub src_node_user: String,
pub env: BTreeMap<String, String>,
#[serde(rename = "sshUser")]
pub ssh_user: String,
#[serde(rename = "localUser")]
pub local_user: String,
#[serde(rename = "connectionID")]
pub connection_id: String,
}
fn is_zero(v: &i64) -> bool {
*v == 0
}
impl CastHeader {
pub fn new(timestamp_unix: i64, term: &str) -> Self {
let term = if term.is_empty() {
"xterm-256color"
} else {
term
};
Self {
version: CAST_VERSION,
timestamp: timestamp_unix,
env: BTreeMap::from([("TERM".to_string(), term.to_string())]),
..Default::default()
}
}
pub fn to_line(&self) -> Result<Vec<u8>, serde_json::Error> {
let mut line = serde_json::to_vec(self)?;
line.push(b'\n');
Ok(line)
}
}
pub fn cast_output_line(elapsed: Duration, data: &[u8]) -> Vec<u8> {
let frame = (elapsed.as_secs_f64(), "o", String::from_utf8_lossy(data));
let mut line = serde_json::to_vec(&frame).unwrap_or_default();
line.push(b'\n');
line
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StartFailure {
FailOpen,
Reject(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UploadFailure {
FailOpen,
Terminate(String),
}
pub fn start_failure_action(on_failure: Option<&SshRecorderFailureAction>) -> StartFailure {
match on_failure {
Some(f) if !f.reject_session_with_message.is_empty() => {
StartFailure::Reject(f.reject_session_with_message.clone())
}
_ => StartFailure::FailOpen,
}
}
pub fn upload_failure_action(on_failure: Option<&SshRecorderFailureAction>) -> UploadFailure {
match on_failure {
Some(f) if !f.terminate_session_with_message.is_empty() => {
UploadFailure::Terminate(f.terminate_session_with_message.clone())
}
_ => UploadFailure::FailOpen,
}
}
pub struct TailnetDialer(std::sync::Arc<crate::Device>);
impl TailnetDialer {
pub fn new(dev: std::sync::Arc<crate::Device>) -> Self {
Self(dev)
}
}
impl RecorderDialer for TailnetDialer {
type Io = crate::netstack::TcpStream;
async fn dial(&self, addr: SocketAddr) -> io::Result<Self::Io> {
self.0.tcp_connect(addr).await.map_err(io::Error::other)
}
}
pub trait RecorderDialer: Send + Sync {
type Io: AsyncRead + AsyncWrite + Unpin + Send + 'static;
fn dial(&self, addr: SocketAddr) -> impl Future<Output = io::Result<Self::Io>> + Send;
}
struct CastBody {
rx: Option<mpsc::Receiver<Bytes>>,
}
impl CastBody {
fn empty() -> Self {
Self { rx: None }
}
fn channel(rx: mpsc::Receiver<Bytes>) -> Self {
Self { rx: Some(rx) }
}
}
impl hyper::body::Body for CastBody {
type Data = Bytes;
type Error = io::Error;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<hyper::body::Frame<Bytes>, io::Error>>> {
match self.get_mut().rx.as_mut() {
None => Poll::Ready(None),
Some(rx) => rx
.poll_recv(cx)
.map(|frame| frame.map(|b| Ok(hyper::body::Frame::data(b)))),
}
}
}
#[derive(Debug, Default, serde::Deserialize)]
struct V2ResponseFrame {
#[serde(default)]
#[allow(
dead_code,
reason = "the ack's arrival is the signal; its value is advisory"
)]
ack: i64,
#[serde(default)]
error: String,
}
struct RecorderUpload {
recorder: SocketAddr,
body: mpsc::Sender<Bytes>,
done: oneshot::Receiver<Result<(), String>>,
}
#[derive(Debug, thiserror::Error)]
#[error("{message}")]
pub struct RecordingRejected {
pub message: String,
#[source]
pub cause: RecorderError,
}
#[derive(Debug)]
pub struct SessionRecording {
start: Instant,
terminate_message: Option<String>,
body: mpsc::Sender<Bytes>,
stopped: bool,
terminate: Option<oneshot::Receiver<String>>,
_alive: oneshot::Sender<()>,
recorder: SocketAddr,
}
impl SessionRecording {
pub async fn start<D: RecorderDialer>(
recorders: &[SocketAddr],
on_failure: Option<&SshRecorderFailureAction>,
header: &CastHeader,
dialer: &D,
) -> Result<Option<Self>, RecordingRejected> {
let (result, attempts) = connect_to_recorder(recorders, dialer).await;
let upload = match result {
Ok(upload) => upload,
Err(e) => {
notify_unsupported(on_failure, &attempts);
return match start_failure_action(on_failure) {
StartFailure::Reject(message) => {
tracing::warn!(error = %e, "recording: error starting recording (rejecting session)");
Err(RecordingRejected { message, cause: e })
}
StartFailure::FailOpen => {
tracing::warn!(error = %e, "recording: error starting recording (failing open)");
Ok(None)
}
};
}
};
let line = header.to_line().map_err(|e| RecordingRejected {
message: "can't start new recording".to_string(),
cause: RecorderError::Recorder(format!("recording: encoding cast header: {e}")),
})?;
if upload.body.send(Bytes::from(line)).await.is_err() {
return Err(RecordingRejected {
message: "can't start new recording".to_string(),
cause: RecorderError::Recorder(
"recording: recorder closed the upload before the cast header".to_string(),
),
});
}
let (terminate_tx, terminate_rx) = oneshot::channel();
let (alive_tx, mut alive_rx) = oneshot::channel::<()>();
let action = upload_failure_action(on_failure);
let terminate_message = match &action {
UploadFailure::Terminate(message) => Some(message.clone()),
UploadFailure::FailOpen => None,
};
let recorder = upload.recorder;
let done = upload.done;
tokio::spawn(async move {
let err = match done.await {
Ok(Ok(())) => {
if matches!(
alive_rx.try_recv(),
Err(oneshot::error::TryRecvError::Closed)
) {
tracing::debug!(%recorder, "recording: finished uploading recording");
return;
}
"recording upload ended before the SSH session".to_string()
}
Ok(Err(e)) => e,
Err(_) => return,
};
match action {
UploadFailure::Terminate(message) => {
tracing::warn!(%recorder, error = %err, "recording: error uploading recording (closing session)");
if terminate_tx.send(message).is_err() {
tracing::debug!(%recorder, "recording: session ended before it could be terminated");
}
}
UploadFailure::FailOpen => {
tracing::warn!(%recorder, error = %err, "recording: error uploading recording (failing open)");
}
}
});
Ok(Some(Self {
start: Instant::now(),
terminate_message,
body: upload.body,
stopped: false,
terminate: Some(terminate_rx),
_alive: alive_tx,
recorder,
}))
}
pub fn recorder(&self) -> SocketAddr {
self.recorder
}
pub fn take_terminate(&mut self) -> Option<oneshot::Receiver<String>> {
self.terminate.take()
}
pub async fn record_output(&mut self, data: &[u8]) -> Result<(), String> {
if self.stopped {
return Ok(());
}
let line = cast_output_line(self.start.elapsed(), data);
if self.body.send(Bytes::from(line)).await.is_err() {
if let Some(message) = &self.terminate_message {
return Err(message.clone());
}
tracing::warn!(
recorder = %self.recorder,
"recording: recorder upload closed; continuing unrecorded (failing open)"
);
self.stopped = true;
}
Ok(())
}
}
fn notify_unsupported(
on_failure: Option<&SshRecorderFailureAction>,
attempts: &[RecordingAttempt],
) {
let Some(url) = on_failure
.map(|f| f.notify_url.as_str())
.filter(|u| !u.is_empty())
else {
return;
};
tracing::warn!(
notify_url = %url,
attempts = attempts.len(),
"recording: onRecordingFailure.notifyURL is set but this server has no control channel to \
notify; recording failure is reported here only"
);
}
async fn connect_to_recorder<D: RecorderDialer>(
recorders: &[SocketAddr],
dialer: &D,
) -> (Result<RecorderUpload, RecorderError>, Vec<RecordingAttempt>) {
if recorders.is_empty() {
return (Err(RecorderError::NoRecorders), Vec::new());
}
let deadline = Instant::now() + ALL_DIAL_ATTEMPTS_TIMEOUT;
let mut attempts = Vec::with_capacity(recorders.len());
let mut failures = Vec::new();
for &addr in recorders {
let Some(budget) = deadline.checked_duration_since(Instant::now()) else {
attempts.push(RecordingAttempt {
recorder: addr,
failure_message: RecorderError::DialBudgetElapsed.to_string(),
});
failures.push(RecorderError::DialBudgetElapsed.to_string());
break;
};
match tokio::time::timeout(budget, connect_one(addr, dialer)).await {
Ok(Ok(upload)) => {
attempts.push(RecordingAttempt {
recorder: addr,
failure_message: String::new(),
});
return (Ok(upload), attempts);
}
Ok(Err(e)) => {
let msg = format!("recording: error starting recording on {addr}: {e}");
attempts.push(RecordingAttempt {
recorder: addr,
failure_message: msg.clone(),
});
failures.push(msg);
}
Err(_) => {
let msg = format!("recording: error starting recording on {addr}: timed out");
attempts.push(RecordingAttempt {
recorder: addr,
failure_message: msg.clone(),
});
failures.push(msg);
}
}
}
(Err(RecorderError::AllFailed(failures.join("; "))), attempts)
}
async fn connect_one<D: RecorderDialer>(
addr: SocketAddr,
dialer: &D,
) -> Result<RecorderUpload, RecorderError> {
let io = dial(addr, dialer).await?;
let v2 = match ts_http_util::http2::connect::<CastBody>(io).await {
Ok(client) => supports_v2(&client, addr).await.then_some(client),
Err(e) => {
tracing::debug!(%addr, error = %e, "recording: h2c handshake failed; trying V1");
None
}
};
match v2 {
Some(client) => connect_v2(client, addr).await,
None => connect_v1(dial(addr, dialer).await?, addr).await,
}
}
async fn dial<D: RecorderDialer>(addr: SocketAddr, dialer: &D) -> Result<D::Io, RecorderError> {
match tokio::time::timeout(PER_DIAL_ATTEMPT_TIMEOUT, dialer.dial(addr)).await {
Ok(io) => Ok(io?),
Err(_) => Err(RecorderError::Recorder(format!("dialing {addr} timed out"))),
}
}
async fn supports_v2(client: &ts_http_util::Http2<CastBody>, addr: SocketAddr) -> bool {
let req = match Request::builder()
.method(Method::HEAD)
.uri(format!("http://{addr}/v2/record"))
.body(CastBody::empty())
{
Ok(req) => req,
Err(e) => {
tracing::debug!(%addr, error = %e, "recording: building V2 probe");
return false;
}
};
match tokio::time::timeout(HTTP2_PROBE_TIMEOUT, client.send(req)).await {
Ok(Ok(resp)) => {
resp.status() == StatusCode::OK && resp.version() >= hyper::http::Version::HTTP_2
}
Ok(Err(e)) => {
tracing::debug!(%addr, error = %e, "recording: V2 probe failed; falling back to V1");
false
}
Err(_) => {
tracing::debug!(%addr, "recording: V2 probe timed out; falling back to V1");
false
}
}
}
async fn connect_v2(
client: ts_http_util::Http2<CastBody>,
addr: SocketAddr,
) -> Result<RecorderUpload, RecorderError> {
let (body_tx, body_rx) = mpsc::channel(CAST_QUEUE_DEPTH);
let req = Request::builder()
.method(Method::POST)
.uri(format!("http://{addr}/v2/record"))
.body(CastBody::channel(body_rx))
.map_err(|e| RecorderError::Recorder(format!("building V2 request: {e}")))?;
let resp = client
.send(req)
.await
.map_err(|e| RecorderError::Recorder(format!("V2 upload: {e}")))?;
if resp.status() != StatusCode::OK {
return Err(RecorderError::Recorder(format!(
"recording: unexpected status: {}",
resp.status()
)));
}
let (done_tx, done_rx) = oneshot::channel();
let mut acks = resp.into_read();
tokio::spawn(async move {
let _client = client;
if done_tx.send(read_acks(&mut acks).await).is_err() {
tracing::debug!(%addr, "recording: session ended before the upload result was read");
}
});
Ok(RecorderUpload {
recorder: addr,
body: body_tx,
done: done_rx,
})
}
async fn read_acks<R: AsyncRead + Unpin>(acks: &mut R) -> Result<(), String> {
let mut buf = Vec::new();
let mut chunk = [0u8; 4096];
loop {
let n = match tokio::time::timeout(UPLOAD_ACK_WINDOW, acks.read(&mut chunk)).await {
Ok(Ok(0)) => return Ok(()),
Ok(Ok(n)) => n,
Ok(Err(e)) => return Err(format!("recording: unexpected error receiving acks: {e}")),
Err(_) => {
return Err(format!(
"did not receive ack frames from the recorder in {}s",
UPLOAD_ACK_WINDOW.as_secs()
));
}
};
buf.extend_from_slice(&chunk[..n]);
match take_ack_frames(&mut buf) {
Ok(frames) => {
for frame in frames {
if !frame.error.is_empty() {
return Err(format!(
"recording: received error from the recorder: {:?}",
frame.error
));
}
}
}
Err(e) => return Err(e),
}
if buf.len() > MAX_ACK_BUFFER {
return Err("recording: recorder sent an oversized ack frame".to_string());
}
}
}
fn take_ack_frames(buf: &mut Vec<u8>) -> Result<Vec<V2ResponseFrame>, String> {
let mut frames = Vec::new();
let consumed = {
let mut stream = serde_json::Deserializer::from_slice(buf).into_iter::<V2ResponseFrame>();
loop {
match stream.next() {
Some(Ok(frame)) => frames.push(frame),
Some(Err(e)) if e.is_eof() => break stream.byte_offset(),
Some(Err(e)) => {
return Err(format!("recording: unexpected error receiving acks: {e}"));
}
None => break stream.byte_offset(),
}
}
};
buf.drain(..consumed);
Ok(frames)
}
async fn connect_v1<Io: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
io: Io,
addr: SocketAddr,
) -> Result<RecorderUpload, RecorderError> {
let (read, mut write) = tokio::io::split(io);
let mut read = BufReader::new(read);
write.write_all(v1_request_head(addr).as_bytes()).await?;
write.flush().await?;
let head = match tokio::time::timeout(EXPECT_CONTINUE_TIMEOUT, read_head(&mut read)).await {
Ok(head) => head?,
Err(_) => {
return Err(RecorderError::Recorder(
"recording: recorder did not answer Expect: 100-continue".to_string(),
));
}
};
match parse_status(&head) {
Some(100) => {}
Some(status) => {
return Err(RecorderError::Recorder(format!(
"recording: unexpected status: {status}"
)));
}
None => {
return Err(RecorderError::Recorder(
"recording: unparseable response from recorder".to_string(),
));
}
}
let (body_tx, mut body_rx) = mpsc::channel::<Bytes>(CAST_QUEUE_DEPTH);
let (done_tx, done_rx) = oneshot::channel();
tokio::spawn(async move {
if done_tx
.send(pump_v1(&mut body_rx, &mut read, &mut write).await)
.is_err()
{
tracing::debug!(%addr, "recording: session ended before the upload result was read");
}
});
Ok(RecorderUpload {
recorder: addr,
body: body_tx,
done: done_rx,
})
}
async fn pump_v1<R, W>(
body: &mut mpsc::Receiver<Bytes>,
read: &mut BufReader<R>,
write: &mut W,
) -> Result<(), String>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
while let Some(chunk) = body.recv().await {
if chunk.is_empty() {
continue;
}
let framed = format!("{:x}\r\n", chunk.len());
write
.write_all(framed.as_bytes())
.await
.map_err(|e| format!("recording: upload write: {e}"))?;
write
.write_all(&chunk)
.await
.map_err(|e| format!("recording: upload write: {e}"))?;
write
.write_all(b"\r\n")
.await
.map_err(|e| format!("recording: upload write: {e}"))?;
write
.flush()
.await
.map_err(|e| format!("recording: upload flush: {e}"))?;
}
write
.write_all(b"0\r\n\r\n")
.await
.map_err(|e| format!("recording: upload close: {e}"))?;
write
.flush()
.await
.map_err(|e| format!("recording: upload close: {e}"))?;
let head = read_head(read)
.await
.map_err(|e| format!("recording: reading final response: {e}"))?;
match parse_status(&head) {
Some(200) => Ok(()),
Some(status) => Err(format!("recording: unexpected status: {status}")),
None => Err("recording: unparseable response from recorder".to_string()),
}
}
fn v1_request_head(addr: SocketAddr) -> String {
format!(
"POST /record HTTP/1.1\r\n\
Host: {addr}\r\n\
User-Agent: tailscale-rs/{version}\r\n\
Transfer-Encoding: chunked\r\n\
Expect: 100-continue\r\n\
\r\n",
version = env!("CARGO_PKG_VERSION"),
)
}
async fn read_head<R: AsyncRead + Unpin>(read: &mut BufReader<R>) -> io::Result<String> {
let mut head = String::new();
loop {
let mut line = String::new();
let n = read.read_line(&mut line).await?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"recorder closed the connection before answering",
));
}
head.push_str(&line);
if line == "\r\n" || line == "\n" {
return Ok(head);
}
if head.len() > MAX_RESPONSE_HEAD {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"recorder response head exceeds the maximum size",
));
}
}
}
fn parse_status(head: &str) -> Option<u16> {
let line = head.lines().next()?;
let mut parts = line.split_whitespace();
let version = parts.next()?;
if !version.starts_with("HTTP/") {
return None;
}
parts.next()?.parse().ok()
}
#[cfg(all(test, feature = "ssh"))]
mod tests {
use std::{
collections::VecDeque,
sync::Mutex as StdMutex,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::io::DuplexStream;
use super::*;
fn recorder_addr() -> SocketAddr {
"192.0.2.10:8080".parse().unwrap()
}
struct ScriptedDialer(StdMutex<VecDeque<DuplexStream>>);
impl ScriptedDialer {
fn new(conns: impl IntoIterator<Item = DuplexStream>) -> Self {
Self(StdMutex::new(conns.into_iter().collect()))
}
}
impl RecorderDialer for ScriptedDialer {
type Io = DuplexStream;
async fn dial(&self, _addr: SocketAddr) -> io::Result<Self::Io> {
self.0
.lock()
.expect("scripted dialer lock")
.pop_front()
.ok_or_else(|| io::Error::other("scripted dialer is out of connections"))
}
}
struct DeadDialer;
impl RecorderDialer for DeadDialer {
type Io = DuplexStream;
async fn dial(&self, _addr: SocketAddr) -> io::Result<Self::Io> {
Err(io::Error::other("no route to recorder"))
}
}
fn dead_connection() -> DuplexStream {
let (near, far) = tokio::io::duplex(64);
drop(far);
near
}
#[test]
fn cast_header_line_carries_the_go_field_names() {
let mut header = CastHeader::new(1_700_000_000, "screen-256color");
header.ssh_user = "alice".to_string();
header.local_user = "ubuntu".to_string();
header.src_node = "laptop.tail-scale.ts.net".to_string();
header.src_node_id = "nodeid-abc".to_string();
header.connection_id = "ssh-conn-20231114T221320-0011223344".to_string();
header.src_node_user_id = 42;
let line = header.to_line().expect("header must encode");
assert_eq!(line.last(), Some(&b'\n'), "the header is one cast line");
let v: serde_json::Value = serde_json::from_slice(&line).expect("header must be JSON");
assert_eq!(v["version"], 2);
assert_eq!(v["timestamp"], 1_700_000_000_i64);
assert_eq!(v["env"]["TERM"], "screen-256color");
assert_eq!(v["sshUser"], "alice");
assert_eq!(v["localUser"], "ubuntu");
assert_eq!(v["srcNode"], "laptop.tail-scale.ts.net");
assert_eq!(v["srcNodeID"], "nodeid-abc");
assert_eq!(v["srcNodeUserID"], 42);
assert_eq!(v["connectionID"], "ssh-conn-20231114T221320-0011223344");
assert!(v.get("command").is_none(), "empty command must be omitted");
assert!(v.get("srcNodeTags").is_none(), "no tags must be omitted");
assert!(v.get("srcNodeUser").is_none(), "no login must be omitted");
assert_eq!(v["width"], 0);
assert_eq!(v["height"], 0);
}
#[test]
fn cast_header_defaults_an_empty_term() {
let header = CastHeader::new(0, "");
assert_eq!(
header.env.get("TERM").map(String::as_str),
Some("xterm-256color")
);
}
#[test]
fn cast_header_tags_are_omitted_when_absent_and_present_when_set() {
let mut header = CastHeader::new(0, "vt100");
header.src_node_tags = vec!["tag:prod".to_string()];
let v: serde_json::Value =
serde_json::from_slice(&header.to_line().expect("encodes")).expect("JSON");
assert_eq!(v["srcNodeTags"][0], "tag:prod");
assert!(v.get("srcNodeUserID").is_none());
}
#[test]
fn cast_output_line_is_a_castv2_output_frame() {
let line = cast_output_line(Duration::from_millis(1500), b"hi there");
assert_eq!(line.last(), Some(&b'\n'));
let v: serde_json::Value = serde_json::from_slice(&line).expect("frame must be JSON");
assert_eq!(v[0], 1.5);
assert_eq!(v[1], "o", "only output is recorded, never input");
assert_eq!(v[2], "hi there");
}
#[test]
fn cast_output_line_survives_invalid_utf8() {
let line = cast_output_line(Duration::ZERO, &[0xff, 0xfe, b'!']);
let v: serde_json::Value = serde_json::from_slice(&line).expect("frame must be JSON");
assert_eq!(v[2], "\u{fffd}\u{fffd}!");
}
#[test]
fn start_failure_is_fail_open_unless_reject_message_is_set() {
assert_eq!(start_failure_action(None), StartFailure::FailOpen);
assert_eq!(
start_failure_action(Some(&SshRecorderFailureAction::default())),
StartFailure::FailOpen,
"an empty action must not be read as fail-closed"
);
assert_eq!(
start_failure_action(Some(&SshRecorderFailureAction {
terminate_session_with_message: "gone".to_string(),
..Default::default()
})),
StartFailure::FailOpen,
"terminate-on-upload-failure says nothing about starting"
);
assert_eq!(
start_failure_action(Some(&SshRecorderFailureAction {
reject_session_with_message: "no recorder, no shell".to_string(),
..Default::default()
})),
StartFailure::Reject("no recorder, no shell".to_string()),
);
}
#[test]
fn upload_failure_is_fail_open_unless_terminate_message_is_set() {
assert_eq!(upload_failure_action(None), UploadFailure::FailOpen);
assert_eq!(
upload_failure_action(Some(&SshRecorderFailureAction {
reject_session_with_message: "no recorder, no shell".to_string(),
..Default::default()
})),
UploadFailure::FailOpen,
"reject-at-start says nothing about a session already running"
);
assert_eq!(
upload_failure_action(Some(&SshRecorderFailureAction {
terminate_session_with_message: "recording lost".to_string(),
..Default::default()
})),
UploadFailure::Terminate("recording lost".to_string()),
);
}
#[tokio::test]
async fn unreachable_recorder_fails_open_by_default() {
let header = CastHeader::new(0, "xterm");
let rec = SessionRecording::start(&[recorder_addr()], None, &header, &DeadDialer)
.await
.expect("the default policy must not refuse the session");
assert!(rec.is_none(), "the session runs, just without a recording");
}
#[tokio::test]
async fn unreachable_recorder_is_fail_closed_when_the_policy_says_so() {
let on_failure = SshRecorderFailureAction {
reject_session_with_message: "this session must be recorded".to_string(),
..Default::default()
};
let header = CastHeader::new(0, "xterm");
let err =
SessionRecording::start(&[recorder_addr()], Some(&on_failure), &header, &DeadDialer)
.await
.expect_err("a fail-closed policy must refuse the session");
assert_eq!(err.message, "this session must be recorded");
}
#[tokio::test]
async fn every_recorder_is_attempted_in_order() {
let first: SocketAddr = "192.0.2.10:8080".parse().unwrap();
let second: SocketAddr = "198.51.100.20:9000".parse().unwrap();
let (result, attempts) = connect_to_recorder(&[first, second], &DeadDialer).await;
assert!(result.is_err());
assert_eq!(
attempts.iter().map(|a| a.recorder).collect::<Vec<_>>(),
vec![first, second],
);
assert!(attempts.iter().all(|a| !a.failure_message.is_empty()));
}
#[tokio::test]
async fn no_recorders_is_an_error_not_a_silent_success() {
let (result, attempts) = connect_to_recorder(&[], &DeadDialer).await;
assert!(matches!(result, Err(RecorderError::NoRecorders)));
assert!(attempts.is_empty());
}
#[test]
fn parse_status_reads_the_status_line() {
assert_eq!(parse_status("HTTP/1.1 100 Continue\r\n\r\n"), Some(100));
assert_eq!(parse_status("HTTP/1.1 200 OK\r\nX: y\r\n\r\n"), Some(200));
assert_eq!(parse_status("HTTP/1.0 404 Not Found\r\n\r\n"), Some(404));
assert_eq!(parse_status("hello\r\n"), None);
assert_eq!(parse_status(""), None);
assert_eq!(parse_status("HTTP/1.1 nope\r\n"), None);
}
#[test]
fn ack_frames_are_decoded_and_partials_are_kept() {
let mut buf = br#"{"ack":1}{"ack":2}{"ac"#.to_vec();
let frames = take_ack_frames(&mut buf).expect("two whole frames decode");
assert_eq!(frames.len(), 2);
assert_eq!(
buf,
br#"{"ac"#.to_vec(),
"the partial frame waits for more bytes"
);
let mut buf = br#"{"error":"disk full"}"#.to_vec();
let frames = take_ack_frames(&mut buf).expect("an error frame is still a frame");
assert_eq!(frames[0].error, "disk full");
let mut buf = b"not json at all".to_vec();
assert!(take_ack_frames(&mut buf).is_err());
}
#[tokio::test]
async fn response_head_is_bounded() {
let (near, mut far) = tokio::io::duplex(64 * 1024);
tokio::spawn(async move {
let junk = format!("X-Pad: {}\r\n", "a".repeat(1024));
for _ in 0..32 {
if far.write_all(junk.as_bytes()).await.is_err() {
return;
}
}
});
let mut read = BufReader::new(near);
let err = read_head(&mut read)
.await
.expect_err("an endless head must fail");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn v1_request_head_announces_expect_continue() {
let head = v1_request_head(recorder_addr());
assert!(head.starts_with("POST /record HTTP/1.1\r\n"));
assert!(head.contains("Host: 192.0.2.10:8080\r\n"));
assert!(head.contains("Expect: 100-continue\r\n"));
assert!(head.contains("Transfer-Encoding: chunked\r\n"));
assert!(head.ends_with("\r\n\r\n"));
}
async fn read_request_head<R: AsyncRead + Unpin>(read: &mut BufReader<R>) -> String {
let mut head = String::new();
loop {
let mut line = String::new();
let n = read.read_line(&mut line).await.expect("head line");
assert_ne!(n, 0, "connection closed mid-head");
head.push_str(&line);
if line == "\r\n" {
return head;
}
}
}
async fn read_chunked_body<R: AsyncRead + Unpin>(read: &mut BufReader<R>) -> Vec<u8> {
let mut body = Vec::new();
loop {
let mut size_line = String::new();
if read.read_line(&mut size_line).await.expect("chunk size") == 0 {
return body;
}
let size = usize::from_str_radix(size_line.trim(), 16).expect("chunk size is hex");
if size == 0 {
let mut end = String::new();
drop(read.read_line(&mut end).await);
return body;
}
let mut chunk = vec![0u8; size];
read.read_exact(&mut chunk).await.expect("chunk data");
let mut crlf = [0u8; 2];
read.read_exact(&mut crlf).await.expect("chunk CRLF");
body.extend_from_slice(&chunk);
}
}
async fn fake_v1_recorder(io: DuplexStream) -> Vec<u8> {
let (read, mut write) = tokio::io::split(io);
let mut read = BufReader::new(read);
let head = read_request_head(&mut read).await;
assert!(
head.starts_with("POST /record HTTP/1.1\r\n"),
"head was {head:?}"
);
assert!(
head.contains("Expect: 100-continue\r\n"),
"head was {head:?}"
);
write
.write_all(b"HTTP/1.1 100 Continue\r\n\r\n")
.await
.expect("100-continue");
write.flush().await.expect("flush");
let cast = read_chunked_body(&mut read).await;
write
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")
.await
.expect("final response");
write.flush().await.expect("flush");
cast
}
#[tokio::test]
async fn v1_recorder_receives_the_whole_cast() {
let (client_io, server_io) = tokio::io::duplex(64 * 1024);
let dialer = ScriptedDialer::new([dead_connection(), client_io]);
let recorder = tokio::spawn(fake_v1_recorder(server_io));
let mut header = CastHeader::new(1_700_000_000, "xterm-256color");
header.local_user = "ubuntu".to_string();
let mut rec = SessionRecording::start(&[recorder_addr()], None, &header, &dialer)
.await
.expect("the recorder accepts the recording")
.expect("recording must be live");
assert_eq!(rec.recorder(), recorder_addr());
rec.record_output(b"$ whoami\r\n").await.expect("recorded");
rec.record_output(b"ubuntu\r\n").await.expect("recorded");
drop(rec);
let cast = tokio::time::timeout(Duration::from_secs(10), recorder)
.await
.expect("recorder must finish")
.expect("recorder task");
let cast = String::from_utf8(cast).expect("the cast is UTF-8 JSON lines");
let mut lines = cast.lines();
let head: serde_json::Value =
serde_json::from_str(lines.next().expect("header line")).expect("header JSON");
assert_eq!(head["version"], 2);
assert_eq!(head["localUser"], "ubuntu");
let first: serde_json::Value =
serde_json::from_str(lines.next().expect("first frame")).expect("frame JSON");
assert_eq!(first[1], "o");
assert_eq!(first[2], "$ whoami\r\n");
let second: serde_json::Value =
serde_json::from_str(lines.next().expect("second frame")).expect("frame JSON");
assert_eq!(second[2], "ubuntu\r\n");
assert!(lines.next().is_none(), "nothing beyond what was recorded");
}
#[tokio::test]
async fn v1_recorder_that_refuses_is_not_treated_as_connected() {
let (client_io, server_io) = tokio::io::duplex(64 * 1024);
tokio::spawn(async move {
let (read, mut write) = tokio::io::split(server_io);
let mut read = BufReader::new(read);
read_request_head(&mut read).await;
drop(
write
.write_all(b"HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\n\r\n")
.await,
);
drop(write.flush().await);
});
let dialer = ScriptedDialer::new([dead_connection(), client_io]);
let on_failure = SshRecorderFailureAction {
reject_session_with_message: "this session must be recorded".to_string(),
..Default::default()
};
let header = CastHeader::new(0, "xterm");
let err = SessionRecording::start(&[recorder_addr()], Some(&on_failure), &header, &dialer)
.await
.expect_err("a refusing recorder must not look like a live recording");
assert_eq!(err.message, "this session must be recorded");
assert!(
err.cause.to_string().contains("403"),
"the refusal reason must be preserved: {}",
err.cause
);
}
fn spawn_recorder_that_hangs_up(server_io: DuplexStream) {
tokio::spawn(async move {
let (read, mut write) = tokio::io::split(server_io);
let mut read = BufReader::new(read);
read_request_head(&mut read).await;
drop(write.write_all(b"HTTP/1.1 100 Continue\r\n\r\n").await);
drop(write.flush().await);
});
}
async fn record_until_it_notices(rec: &mut SessionRecording) -> Result<(), String> {
tokio::time::timeout(Duration::from_secs(10), async {
loop {
rec.record_output(b"x").await?;
tokio::task::yield_now().await;
if rec.stopped {
return Ok(());
}
}
})
.await
.expect("the broken upload must be noticed")
}
#[tokio::test]
async fn a_broken_upload_is_fail_closed_with_the_policy_message() {
let (client_io, server_io) = tokio::io::duplex(64 * 1024);
spawn_recorder_that_hangs_up(server_io);
let dialer = ScriptedDialer::new([dead_connection(), client_io]);
let on_failure = SshRecorderFailureAction {
terminate_session_with_message: "recording lost; ending session".to_string(),
..Default::default()
};
let header = CastHeader::new(0, "xterm");
let mut rec =
SessionRecording::start(&[recorder_addr()], Some(&on_failure), &header, &dialer)
.await
.expect("the recorder accepted the recording")
.expect("recording must be live");
assert_eq!(
record_until_it_notices(&mut rec).await,
Err("recording lost; ending session".to_string()),
);
}
#[tokio::test]
async fn a_broken_upload_is_fail_open_by_default() {
let (client_io, server_io) = tokio::io::duplex(64 * 1024);
spawn_recorder_that_hangs_up(server_io);
let dialer = ScriptedDialer::new([dead_connection(), client_io]);
let header = CastHeader::new(0, "xterm");
let mut rec = SessionRecording::start(&[recorder_addr()], None, &header, &dialer)
.await
.expect("the recorder accepted the recording")
.expect("recording must be live");
assert_eq!(record_until_it_notices(&mut rec).await, Ok(()));
assert!(rec.stopped, "no further cast lines are attempted");
rec.record_output(b"still alive").await.expect("fail-open");
}
async fn collect_incoming(mut body: hyper::body::Incoming) -> Vec<u8> {
use hyper::body::Body as _;
let mut out = Vec::new();
while let Some(frame) = std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await
{
match frame {
Ok(frame) => {
if let Some(data) = frame.data_ref() {
out.extend_from_slice(data);
}
}
Err(_) => break,
}
}
out
}
async fn fake_v2_recorder(io: DuplexStream, cast_tx: oneshot::Sender<Vec<u8>>) {
let cast_tx = std::sync::Arc::new(StdMutex::new(Some(cast_tx)));
let service = hyper::service::service_fn(move |req: Request<hyper::body::Incoming>| {
let cast_tx = cast_tx.clone();
async move {
let response = |status: StatusCode, body: CastBody| {
ts_http_util::Response::builder()
.status(status)
.body(body)
.expect("response builds")
};
match (req.method().clone(), req.uri().path()) {
(Method::HEAD, "/v2/record") => {
Ok::<_, io::Error>(response(StatusCode::OK, CastBody::empty()))
}
(Method::POST, "/v2/record") => {
let (ack_tx, ack_rx) = mpsc::channel(4);
tokio::spawn(async move {
drop(ack_tx.send(Bytes::from_static(br#"{"ack":0}"#)).await);
let cast = collect_incoming(req.into_body()).await;
drop(
ack_tx
.send(Bytes::from(format!(r#"{{"ack":{}}}"#, cast.len())))
.await,
);
if let Some(tx) = cast_tx.lock().expect("cast lock").take() {
drop(tx.send(cast));
}
});
Ok(response(StatusCode::OK, CastBody::channel(ack_rx)))
}
_ => Ok(response(StatusCode::NOT_FOUND, CastBody::empty())),
}
}
});
drop(
hyper::server::conn::http2::Builder::new(hyper_util::rt::TokioExecutor::new())
.serve_connection(hyper_util::rt::TokioIo::new(io), service)
.await,
);
}
#[tokio::test]
async fn v2_recorder_receives_the_whole_cast() {
let (client_io, server_io) = tokio::io::duplex(64 * 1024);
let (cast_tx, cast_rx) = oneshot::channel();
tokio::spawn(fake_v2_recorder(server_io, cast_tx));
let dialer = ScriptedDialer::new([client_io]);
let mut header = CastHeader::new(1_700_000_000, "xterm-256color");
header.ssh_user = "alice".to_string();
let mut rec = SessionRecording::start(&[recorder_addr()], None, &header, &dialer)
.await
.expect("the recorder accepts the recording")
.expect("recording must be live");
rec.record_output(b"hello\r\n").await.expect("recorded");
drop(rec);
let cast = tokio::time::timeout(Duration::from_secs(10), cast_rx)
.await
.expect("recorder must finish")
.expect("recorder sends the cast");
let cast = String::from_utf8(cast).expect("the cast is UTF-8 JSON lines");
let mut lines = cast.lines();
let head: serde_json::Value =
serde_json::from_str(lines.next().expect("header line")).expect("header JSON");
assert_eq!(head["sshUser"], "alice");
let frame: serde_json::Value =
serde_json::from_str(lines.next().expect("frame")).expect("frame JSON");
assert_eq!(frame[1], "o");
assert_eq!(frame[2], "hello\r\n");
}
#[test]
fn cast_header_timestamp_is_unix_seconds() {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("clock after the epoch")
.as_secs() as i64;
let header = CastHeader::new(now, "xterm");
assert!(header.timestamp > 1_600_000_000, "{}", header.timestamp);
}
}