326 lines
11 KiB
Rust
326 lines
11 KiB
Rust
use std::time::Duration;
|
|
|
|
use bytes::Bytes;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use http_body_util::{BodyExt, Empty};
|
|
use hyper::body::Incoming;
|
|
use hyper::client::conn::http1;
|
|
use hyper::header::{CONTENT_TYPE, HOST};
|
|
use hyper::{Request, StatusCode};
|
|
use hyper_util::rt::TokioIo;
|
|
use tokio::net::TcpStream;
|
|
use tokio::task::JoinHandle;
|
|
use tokio::time::{Instant, timeout, timeout_at};
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
use tokio_tungstenite::{WebSocketStream, client_async};
|
|
|
|
use crate::config::{Protocol, Target};
|
|
|
|
const INTERACTION_INTERVAL: Duration = Duration::from_millis(100);
|
|
const MAX_SSE_BUFFER: usize = 8_192;
|
|
const SSE_EVENT: &[u8] = b"event: live\nid: 1\ndata: ready\n\n";
|
|
const SSE_HEARTBEAT: &[u8] = b":\n\n";
|
|
|
|
pub struct Session {
|
|
connection: Connection,
|
|
cycle: usize,
|
|
index: usize,
|
|
timeout: Duration,
|
|
messages: u64,
|
|
}
|
|
|
|
enum Connection {
|
|
Tcp(TcpStream),
|
|
WebSocket(Box<WebSocketStream<TcpStream>>),
|
|
Sse(SseConnection),
|
|
}
|
|
|
|
struct SseConnection {
|
|
body: Incoming,
|
|
driver: ConnectionDriver,
|
|
pending: Vec<u8>,
|
|
}
|
|
|
|
struct ConnectionDriver(Option<JoinHandle<Result<(), hyper::Error>>>);
|
|
|
|
impl ConnectionDriver {
|
|
fn new(driver: JoinHandle<Result<(), hyper::Error>>) -> Self {
|
|
Self(Some(driver))
|
|
}
|
|
|
|
async fn close(&mut self) {
|
|
if let Some(driver) = self.0.take() {
|
|
driver.abort();
|
|
let _ = driver.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for ConnectionDriver {
|
|
fn drop(&mut self) {
|
|
if let Some(driver) = &self.0 {
|
|
driver.abort();
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Session {
|
|
pub async fn connect(
|
|
protocol: Protocol,
|
|
target: &Target,
|
|
cycle: usize,
|
|
index: usize,
|
|
operation_timeout: Duration,
|
|
) -> Result<Self, String> {
|
|
let stream = timeout(
|
|
operation_timeout,
|
|
TcpStream::connect(("127.0.0.1", target.port)),
|
|
)
|
|
.await
|
|
.map_err(|_| "connection timed out".to_owned())?
|
|
.map_err(|error| error.to_string())?;
|
|
stream
|
|
.set_nodelay(true)
|
|
.map_err(|error| error.to_string())?;
|
|
let connection = match protocol {
|
|
Protocol::Tcp => Connection::Tcp(stream),
|
|
Protocol::WebSocket => {
|
|
let (socket, response) =
|
|
timeout(operation_timeout, client_async(target.url.as_str(), stream))
|
|
.await
|
|
.map_err(|_| "WebSocket handshake timed out".to_owned())?
|
|
.map_err(|error| error.to_string())?;
|
|
if response.status() != StatusCode::SWITCHING_PROTOCOLS {
|
|
return Err(format!(
|
|
"WebSocket handshake returned {}",
|
|
response.status()
|
|
));
|
|
}
|
|
Connection::WebSocket(Box::new(socket))
|
|
}
|
|
Protocol::Sse => {
|
|
Connection::Sse(SseConnection::open(stream, target, operation_timeout).await?)
|
|
}
|
|
};
|
|
let mut session = Self {
|
|
connection,
|
|
cycle,
|
|
index,
|
|
timeout: operation_timeout,
|
|
messages: 0,
|
|
};
|
|
session.exchange(0).await?;
|
|
Ok(session)
|
|
}
|
|
|
|
pub async fn hold(&mut self, duration: Duration) -> Result<(), String> {
|
|
let deadline = Instant::now() + duration;
|
|
let mut sequence = 1_u64;
|
|
while Instant::now() < deadline {
|
|
if let Connection::Sse(connection) = &mut self.connection {
|
|
let heartbeats = connection.read_heartbeats(deadline).await?;
|
|
if heartbeats == 0 {
|
|
return Err("SSE hold ended without a heartbeat".into());
|
|
}
|
|
self.messages += heartbeats;
|
|
break;
|
|
}
|
|
self.exchange(sequence).await?;
|
|
sequence += 1;
|
|
let next = (Instant::now() + INTERACTION_INTERVAL).min(deadline);
|
|
tokio::time::sleep_until(next).await;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn close(mut self) -> Result<u64, String> {
|
|
match &mut self.connection {
|
|
Connection::Tcp(_) => {}
|
|
Connection::WebSocket(socket) => timeout(self.timeout, socket.as_mut().close(None))
|
|
.await
|
|
.map_err(|_| "WebSocket close timed out".to_owned())?
|
|
.map_err(|error| error.to_string())?,
|
|
Connection::Sse(connection) => connection.close().await,
|
|
}
|
|
Ok(self.messages)
|
|
}
|
|
|
|
async fn exchange(&mut self, sequence: u64) -> Result<(), String> {
|
|
match &mut self.connection {
|
|
Connection::Tcp(stream) => {
|
|
let payload = payload(self.cycle, self.index, sequence);
|
|
timeout(self.timeout, write_all(stream, &payload))
|
|
.await
|
|
.map_err(|_| "TCP write timed out".to_owned())??;
|
|
let mut received = vec![0; payload.len()];
|
|
timeout(self.timeout, read_exact(stream, &mut received))
|
|
.await
|
|
.map_err(|_| "TCP read timed out".to_owned())??;
|
|
if received != payload {
|
|
return Err("TCP echo body mismatch".into());
|
|
}
|
|
}
|
|
Connection::WebSocket(socket) => {
|
|
let payload = String::from_utf8(payload(self.cycle, self.index, sequence))
|
|
.expect("benchmark payload is ASCII");
|
|
timeout(
|
|
self.timeout,
|
|
socket.send(Message::Text(payload.clone().into())),
|
|
)
|
|
.await
|
|
.map_err(|_| "WebSocket write timed out".to_owned())?
|
|
.map_err(|error| error.to_string())?;
|
|
let message = timeout(self.timeout, socket.next())
|
|
.await
|
|
.map_err(|_| "WebSocket read timed out".to_owned())?
|
|
.ok_or("WebSocket ended before echo")?
|
|
.map_err(|error| error.to_string())?;
|
|
if !matches!(message, Message::Text(value) if value.as_str() == payload) {
|
|
return Err("WebSocket echo body mismatch".into());
|
|
}
|
|
}
|
|
Connection::Sse(connection) => connection.read_event(self.timeout).await?,
|
|
}
|
|
self.messages += 1;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl SseConnection {
|
|
async fn open(
|
|
stream: TcpStream,
|
|
target: &Target,
|
|
operation_timeout: Duration,
|
|
) -> Result<Self, String> {
|
|
let (mut sender, connection) =
|
|
timeout(operation_timeout, http1::handshake(TokioIo::new(stream)))
|
|
.await
|
|
.map_err(|_| "SSE HTTP handshake timed out".to_owned())?
|
|
.map_err(|error| error.to_string())?;
|
|
let driver = ConnectionDriver::new(tokio::spawn(connection));
|
|
let request = Request::builder()
|
|
.method("GET")
|
|
.uri(&target.path)
|
|
.header(HOST, format!("127.0.0.1:{}", target.port))
|
|
.body(Empty::<Bytes>::new())
|
|
.map_err(|error| error.to_string())?;
|
|
let response = timeout(operation_timeout, sender.send_request(request))
|
|
.await
|
|
.map_err(|_| "SSE request timed out".to_owned())?
|
|
.map_err(|error| error.to_string())?;
|
|
if response.status() != StatusCode::OK {
|
|
return Err(format!("SSE request returned {}", response.status()));
|
|
}
|
|
let content_type = response
|
|
.headers()
|
|
.get(CONTENT_TYPE)
|
|
.and_then(|value| value.to_str().ok())
|
|
.unwrap_or_default();
|
|
if content_type.split(';').next().map(str::trim) != Some("text/event-stream") {
|
|
return Err("SSE response has the wrong Content-Type".into());
|
|
}
|
|
Ok(Self {
|
|
body: response.into_body(),
|
|
driver,
|
|
pending: Vec::new(),
|
|
})
|
|
}
|
|
|
|
async fn read_event(&mut self, operation_timeout: Duration) -> Result<(), String> {
|
|
while self.pending.len() < SSE_EVENT.len() {
|
|
let frame = timeout(operation_timeout, self.body.frame())
|
|
.await
|
|
.map_err(|_| "SSE event timed out".to_owned())?
|
|
.ok_or("SSE body ended before the initial event")?
|
|
.map_err(|error| error.to_string())?;
|
|
if let Ok(bytes) = frame.into_data() {
|
|
self.extend(bytes.as_ref())?;
|
|
}
|
|
if self.pending.len() < SSE_EVENT.len() && !SSE_EVENT.starts_with(&self.pending) {
|
|
return Err("SSE initial event body mismatch".into());
|
|
}
|
|
}
|
|
if !self.pending.starts_with(SSE_EVENT) {
|
|
return Err("SSE initial event body mismatch".into());
|
|
}
|
|
self.pending.drain(..SSE_EVENT.len());
|
|
Ok(())
|
|
}
|
|
|
|
async fn read_heartbeats(&mut self, deadline: Instant) -> Result<u64, String> {
|
|
let mut messages = consume_heartbeats(&mut self.pending)?;
|
|
loop {
|
|
let frame = match timeout_at(deadline, self.body.frame()).await {
|
|
Err(_) => break,
|
|
Ok(None) => return Err("SSE body ended during hold".into()),
|
|
Ok(Some(Err(error))) => return Err(error.to_string()),
|
|
Ok(Some(Ok(frame))) => frame,
|
|
};
|
|
if let Ok(bytes) = frame.into_data() {
|
|
self.extend(bytes.as_ref())?;
|
|
messages += consume_heartbeats(&mut self.pending)?;
|
|
}
|
|
}
|
|
Ok(messages)
|
|
}
|
|
|
|
fn extend(&mut self, bytes: &[u8]) -> Result<(), String> {
|
|
if self.pending.len().saturating_add(bytes.len()) > MAX_SSE_BUFFER {
|
|
return Err("SSE validation buffer exceeded its limit".into());
|
|
}
|
|
self.pending.extend_from_slice(bytes);
|
|
Ok(())
|
|
}
|
|
|
|
async fn close(&mut self) {
|
|
self.driver.close().await;
|
|
}
|
|
}
|
|
|
|
fn consume_heartbeats(buffer: &mut Vec<u8>) -> Result<u64, String> {
|
|
let mut messages = 0;
|
|
while buffer.len() >= SSE_HEARTBEAT.len() {
|
|
if !buffer.starts_with(SSE_HEARTBEAT) {
|
|
return Err("SSE hold data was not a heartbeat".into());
|
|
}
|
|
buffer.drain(..SSE_HEARTBEAT.len());
|
|
messages += 1;
|
|
}
|
|
if !SSE_HEARTBEAT.starts_with(buffer) {
|
|
return Err("SSE heartbeat body mismatch".into());
|
|
}
|
|
Ok(messages)
|
|
}
|
|
|
|
fn payload(cycle: usize, index: usize, sequence: u64) -> Vec<u8> {
|
|
format!("live:{cycle}:{index}:{sequence}\n").into_bytes()
|
|
}
|
|
|
|
async fn write_all(stream: &TcpStream, bytes: &[u8]) -> Result<(), String> {
|
|
let mut offset = 0;
|
|
while offset < bytes.len() {
|
|
stream.writable().await.map_err(|error| error.to_string())?;
|
|
match stream.try_write(&bytes[offset..]) {
|
|
Ok(0) => return Err("TCP write returned zero bytes".into()),
|
|
Ok(length) => offset += length,
|
|
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
|
|
Err(error) => return Err(error.to_string()),
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn read_exact(stream: &TcpStream, bytes: &mut [u8]) -> Result<(), String> {
|
|
let mut offset = 0;
|
|
while offset < bytes.len() {
|
|
stream.readable().await.map_err(|error| error.to_string())?;
|
|
match stream.try_read(&mut bytes[offset..]) {
|
|
Ok(0) => return Err("TCP connection ended before the echo completed".into()),
|
|
Ok(length) => offset += length,
|
|
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
|
|
Err(error) => return Err(error.to_string()),
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|