use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use actix_web::http::header::{HeaderMap, CONTENT_TYPE};
use actix_web::HttpResponse;
use bytes::{Bytes, BytesMut};
use serde::Deserialize;
use serde_json::Value;
use thiserror::Error;
use tokio::sync::{mpsc, oneshot};
use tokio_stream::{Stream, StreamExt};
use webui::streaming::{ChunkPool, StreamingWriter};
use webui::{BoundaryId, BoundaryMode, HandlerError, Protocol, RenderOptions, StreamingResponse};
use super::create_handler;
use crate::commands::common::Plugin;
pub(super) const MEDIA_TYPE: &str = "application/x-webui-stream";
pub(super) const ACCEPT: &str = "application/x-webui-stream, application/json";
const VERSION: u8 = 1;
const MAX_RECORD_BYTES: usize = 2_000_000;
const MAX_INITIAL_SHELL_BYTES: usize = 4_000_000;
const INITIAL_RECORD_CAPACITY: usize = 4 * 1024;
pub(super) struct StateDefaults {
token_css: Option<HashMap<String, String>>,
base_path: String,
route_params: HashMap<String, String>,
}
impl StateDefaults {
pub(super) fn new(
token_css: Option<HashMap<String, String>>,
base_path: String,
route_params: HashMap<String, String>,
) -> Self {
Self {
token_css,
base_path,
route_params,
}
}
pub(super) fn apply(&self, state: &mut Value) {
if let Some(token_css) = &self.token_css {
webui_tokens::inject_token_css(state, token_css);
}
let Value::Object(map) = state else {
return;
};
map.insert("basePath".to_owned(), Value::String(self.base_path.clone()));
for (key, value) in &self.route_params {
map.insert(key.clone(), Value::String(value.clone()));
}
}
fn apply_record(&self, state: &mut Value, record: usize) -> Result<(), ApiStreamError> {
if !state.is_object() {
return Err(ApiStreamError::StateMustBeObject { record });
}
self.apply(state);
Ok(())
}
}
pub(super) struct RenderConfig {
pub(super) protocol: Arc<Protocol>,
pub(super) entry: String,
pub(super) route_path: String,
pub(super) plugin: Option<Plugin>,
pub(super) body_inject: Option<Arc<str>>,
pub(super) chunk_pool: Arc<ChunkPool>,
}
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase", deny_unknown_fields)]
enum ApiStreamCommand {
Shell {
version: u8,
state: Value,
},
Boundary {
name: String,
#[serde(default)]
mode: ApiBoundaryMode,
state: Option<Value>,
},
Update {
name: String,
state: Value,
},
Finish {
state: Option<Value>,
},
}
impl ApiStreamCommand {
fn prepare(&mut self, defaults: &StateDefaults, record: usize) -> Result<(), ApiStreamError> {
match self {
Self::Shell { state, .. } => defaults.apply_record(state, record),
Self::Boundary {
state: Some(state), ..
}
| Self::Finish { state: Some(state) } => defaults.apply_record(state, record),
Self::Update { state, .. } if !state.is_object() => {
Err(ApiStreamError::StateMustBeObject { record })
}
Self::Boundary { state: None, .. }
| Self::Finish { state: None }
| Self::Update { .. } => Ok(()),
}
}
fn shell_version(&self) -> Option<u8> {
match self {
Self::Shell { version, .. } => Some(*version),
_ => None,
}
}
fn is_finish(&self) -> bool {
matches!(self, Self::Finish { .. })
}
}
#[derive(Clone, Copy, Debug, Default, Deserialize)]
#[serde(rename_all = "lowercase")]
enum ApiBoundaryMode {
#[default]
Final,
Updatable,
}
impl From<ApiBoundaryMode> for BoundaryMode {
fn from(value: ApiBoundaryMode) -> Self {
match value {
ApiBoundaryMode::Final => Self::Final,
ApiBoundaryMode::Updatable => Self::Updatable,
}
}
}
#[derive(Debug, Error)]
enum ApiStreamError {
#[error(
"WebUI API stream record {record} exceeds the {MAX_RECORD_BYTES}-byte limit; reduce the state payload"
)]
RecordTooLarge { record: usize },
#[error("WebUI API stream record {record} is not valid JSON: {source}")]
InvalidJson {
record: usize,
#[source]
source: serde_json::Error,
},
#[error("WebUI API stream record {record} must contain an object-valued state")]
StateMustBeObject { record: usize },
#[error("WebUI API stream record {record} must be the initial shell record")]
MissingShell { record: usize },
#[error("WebUI API stream record {record} repeats the shell record")]
DuplicateShell { record: usize },
#[error(
"WebUI API stream record {record} uses unsupported version {version}; send version {VERSION}"
)]
UnsupportedVersion { record: usize, version: u8 },
#[error("WebUI API stream ended before a finish record; send {{\"type\":\"finish\"}}")]
MissingFinish,
#[error("WebUI API stream transport failed: {0}")]
Backend(String),
#[error("WebUI API stream renderer stopped before accepting record {record}")]
RendererStopped { record: usize },
}
#[derive(Debug, Error)]
enum ShellInitializationError {
#[error("{0}")]
Render(String),
#[error("WebUI streaming renderer stopped during initialization")]
RendererStopped,
#[error(
"WebUI initial streaming shell exceeds the 4,000,000-byte precommit limit; reduce the initial state or move content behind a boundary"
)]
TooLarge,
}
struct RecordDecoder {
bytes: BytesMut,
next_record: usize,
search_from: usize,
}
impl RecordDecoder {
fn new() -> Self {
Self {
bytes: BytesMut::with_capacity(INITIAL_RECORD_CAPACITY),
next_record: 0,
search_from: 0,
}
}
fn push(&mut self, chunk: &[u8]) {
self.bytes.extend_from_slice(chunk);
}
fn next(&mut self) -> Result<Option<Bytes>, ApiStreamError> {
let Some(relative_end) = memchr::memchr(b'\n', &self.bytes[self.search_from..]) else {
self.search_from = self.bytes.len();
self.ensure_pending_limit()?;
return Ok(None);
};
let end = self.search_from + relative_end;
self.ensure_line_limit(end)?;
let mut line = self.bytes.split_to(end + 1);
line.truncate(end);
self.next_record += 1;
self.search_from = 0;
Ok(Some(line.freeze()))
}
fn finish(&mut self) -> Result<Option<Bytes>, ApiStreamError> {
if self.bytes.is_empty() {
return Ok(None);
}
self.ensure_pending_limit()?;
self.next_record += 1;
Ok(Some(self.bytes.split().freeze()))
}
fn ensure_pending_limit(&self) -> Result<(), ApiStreamError> {
self.ensure_line_limit(self.bytes.len())
}
fn ensure_line_limit(&self, length: usize) -> Result<(), ApiStreamError> {
if length > MAX_RECORD_BYTES {
return Err(record_too_large(self.next_record));
}
Ok(())
}
}
#[cold]
#[inline(never)]
fn record_too_large(record: usize) -> ApiStreamError {
ApiStreamError::RecordTooLarge { record }
}
struct CommandIngest {
sender: mpsc::Sender<ApiStreamCommand>,
state_defaults: StateDefaults,
record: usize,
shell_seen: bool,
}
impl CommandIngest {
async fn dispatch(&mut self, bytes: &[u8]) -> Result<bool, ApiStreamError> {
let record = self.record;
self.record += 1;
let mut command = serde_json::from_slice::<ApiStreamCommand>(bytes)
.map_err(|source| ApiStreamError::InvalidJson { record, source })?;
self.validate_order(&command, record)?;
command.prepare(&self.state_defaults, record)?;
let finished = command.is_finish();
self.sender
.send(command)
.await
.map_err(|_| ApiStreamError::RendererStopped { record })?;
Ok(finished)
}
fn validate_order(
&mut self,
command: &ApiStreamCommand,
record: usize,
) -> Result<(), ApiStreamError> {
if let Some(version) = command.shell_version() {
if self.shell_seen {
return Err(duplicate_shell(record));
}
if version != VERSION {
return Err(unsupported_version(record, version));
}
self.shell_seen = true;
return Ok(());
}
if !self.shell_seen {
return Err(missing_shell(record));
}
Ok(())
}
}
#[cold]
#[inline(never)]
fn duplicate_shell(record: usize) -> ApiStreamError {
ApiStreamError::DuplicateShell { record }
}
#[cold]
#[inline(never)]
fn missing_shell(record: usize) -> ApiStreamError {
ApiStreamError::MissingShell { record }
}
#[cold]
#[inline(never)]
fn unsupported_version(record: usize, version: u8) -> ApiStreamError {
ApiStreamError::UnsupportedVersion { record, version }
}
pub(super) fn is_stream(headers: &HeaderMap) -> bool {
headers
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.is_some_and(|value| value.trim().eq_ignore_ascii_case(MEDIA_TYPE))
}
struct BrowserHtmlStream {
initial: std::vec::IntoIter<Bytes>,
live: tokio_stream::wrappers::ReceiverStream<Bytes>,
cancel: Option<oneshot::Sender<()>>,
}
impl BrowserHtmlStream {
fn new(initial: Vec<Bytes>, live: mpsc::Receiver<Bytes>, cancel: oneshot::Sender<()>) -> Self {
Self {
initial: initial.into_iter(),
live: tokio_stream::wrappers::ReceiverStream::new(live),
cancel: Some(cancel),
}
}
fn cancel(&mut self) {
if let Some(cancel) = self.cancel.take() {
let _ = cancel.send(());
}
}
}
impl Stream for BrowserHtmlStream {
type Item = Result<Bytes, actix_web::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
if let Some(chunk) = this.initial.next() {
return Poll::Ready(Some(Ok(chunk)));
}
match Pin::new(&mut this.live).poll_next(cx) {
Poll::Ready(Some(chunk)) => Poll::Ready(Some(Ok(chunk))),
Poll::Ready(None) => {
this.cancel();
Poll::Ready(None)
}
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for BrowserHtmlStream {
fn drop(&mut self) {
self.cancel();
}
}
pub(super) async fn render<S>(
mut backend: S,
config: RenderConfig,
state_defaults: StateDefaults,
) -> HttpResponse
where
S: Stream<Item = Result<Bytes, String>> + Unpin + 'static,
{
let (command_tx, command_rx) = mpsc::channel(1);
let (html_tx, html_rx) = mpsc::channel(StreamingWriter::DEFAULT_CHANNEL_CAPACITY);
let (ready_tx, ready_rx) = oneshot::channel();
let render_route_path = config.route_path.clone();
let ingest_route_path = config.route_path.clone();
let initialization_route_path = config.route_path.clone();
let mut decoder = RecordDecoder::new();
let mut ingest = CommandIngest {
sender: command_tx,
state_defaults,
record: 0,
shell_seen: false,
};
if let Err(error) = ingest_first(&mut backend, &mut decoder, &mut ingest).await {
log::error!("streaming API command stream failed for {ingest_route_path}: {error}");
return HttpResponse::BadGateway()
.content_type("text/plain; charset=utf-8")
.body(error.to_string());
}
actix_web::rt::task::spawn_blocking(move || {
let mut writer = StreamingWriter::new_pooled(html_tx, Arc::clone(&config.chunk_pool))
.with_flush_timeout(Duration::from_secs(30));
if let Err(error) = run_renderer(config, command_rx, &mut writer, ready_tx) {
log::error!("streaming API render failed for {render_route_path}: {error}");
}
});
let (initial_html, html_rx) = match stage_initial_shell(ready_rx, html_rx).await {
Ok(output) => output,
Err(error) => {
log::error!(
"streaming API initialization failed for {initialization_route_path}: {error}"
);
return HttpResponse::InternalServerError()
.content_type("text/plain; charset=utf-8")
.body(error.to_string());
}
};
let (cancel_tx, cancel_rx) = oneshot::channel();
actix_web::rt::spawn(async move {
tokio::select! {
result = ingest_remaining(backend, decoder, ingest) => {
if let Err(error) = result {
log::error!(
"streaming API command stream failed for {ingest_route_path}: {error}"
);
}
}
_ = cancel_rx => {}
}
});
let stream = BrowserHtmlStream::new(initial_html, html_rx, cancel_tx);
HttpResponse::Ok()
.content_type("text/html; charset=utf-8")
.insert_header(("Cache-Control", "no-store"))
.streaming(stream)
}
async fn stage_initial_shell(
mut ready: oneshot::Receiver<Result<(), String>>,
mut html: mpsc::Receiver<Bytes>,
) -> Result<(Vec<Bytes>, mpsc::Receiver<Bytes>), ShellInitializationError> {
let mut initial = Vec::with_capacity(StreamingWriter::DEFAULT_CHANNEL_CAPACITY);
let mut total = 0usize;
loop {
tokio::select! {
result = &mut ready => {
validate_renderer_ready(result)?;
while let Ok(chunk) = html.try_recv() {
stage_initial_chunk(&mut initial, &mut total, chunk)?;
}
return Ok((initial, html));
}
chunk = html.recv() => {
let Some(chunk) = chunk else {
validate_renderer_ready(ready.await)?;
return Ok((initial, html));
};
stage_initial_chunk(&mut initial, &mut total, chunk)?;
}
}
}
}
fn validate_renderer_ready(
result: Result<Result<(), String>, oneshot::error::RecvError>,
) -> Result<(), ShellInitializationError> {
match result {
Ok(Ok(())) => Ok(()),
Ok(Err(error)) => Err(ShellInitializationError::Render(error)),
Err(_) => Err(ShellInitializationError::RendererStopped),
}
}
fn stage_initial_chunk(
initial: &mut Vec<Bytes>,
total: &mut usize,
chunk: Bytes,
) -> Result<(), ShellInitializationError> {
let next_total = total.saturating_add(chunk.len());
if next_total > MAX_INITIAL_SHELL_BYTES {
return Err(ShellInitializationError::TooLarge);
}
*total = next_total;
initial.push(chunk);
Ok(())
}
fn run_renderer(
config: RenderConfig,
mut commands: mpsc::Receiver<ApiStreamCommand>,
writer: &mut StreamingWriter,
ready: oneshot::Sender<Result<(), String>>,
) -> Result<(), HandlerError> {
let handler = create_handler(config.plugin);
let options = RenderOptions::new(&config.entry, &config.route_path);
let options = match config.body_inject.as_deref() {
Some(body) => options.with_body_inject(body),
None => options,
};
let mut response = match handler.stream_response(&config.protocol, &options, writer) {
Ok(response) => response,
Err(error) => {
let _ = ready.send(Err(error.to_string()));
return Err(error);
}
};
let mut boundaries = HashMap::with_capacity(response.boundary_count());
let base_state = match commands.blocking_recv() {
Some(ApiStreamCommand::Shell { state, .. }) => {
if let Err(error) = response.write_shell(&state) {
let _ = ready.send(Err(error.to_string()));
return Err(error);
}
Some(state)
}
Some(_) => {
let error = HandlerError::Invariant(
"streaming API renderer did not receive the validated shell first".to_owned(),
);
let _ = ready.send(Err(error.to_string()));
return Err(error);
}
None => {
let error = HandlerError::Writer(
"streaming API command producer stopped before shell rendering".to_owned(),
);
let _ = ready.send(Err(error.to_string()));
return Err(error);
}
};
let _ = ready.send(Ok(()));
while let Some(command) = commands.blocking_recv() {
match command {
ApiStreamCommand::Shell { .. } => {
return Err(HandlerError::Invariant(
"streaming API renderer received a duplicate shell".to_owned(),
));
}
ApiStreamCommand::Boundary { name, mode, state } => {
let boundary = resolve_boundary(&response, &mut boundaries, name)?;
let state = state.as_ref().or(base_state.as_ref()).ok_or_else(|| {
HandlerError::Invariant(
"streaming API boundary has no shell or explicit state".to_owned(),
)
})?;
response.write_boundary(boundary, state, mode.into())?;
}
ApiStreamCommand::Update { name, state } => {
let boundary = resolve_boundary(&response, &mut boundaries, name)?;
response.update(boundary, &state)?;
}
ApiStreamCommand::Finish { state } => {
let state = state.as_ref().or(base_state.as_ref()).ok_or_else(|| {
HandlerError::Invariant(
"streaming API finish has no shell or explicit state".to_owned(),
)
})?;
return response.finish(state);
}
}
}
Err(HandlerError::Writer(
"streaming API command producer stopped before finish".to_owned(),
))
}
fn resolve_boundary(
response: &StreamingResponse<'_, StreamingWriter>,
boundaries: &mut HashMap<String, BoundaryId>,
name: String,
) -> Result<BoundaryId, HandlerError> {
if let Some(boundary) = boundaries.get(&name) {
return Ok(*boundary);
}
let boundary = response.boundary(&name)?;
boundaries.insert(name, boundary);
Ok(boundary)
}
async fn ingest_first<S>(
backend: &mut S,
decoder: &mut RecordDecoder,
ingest: &mut CommandIngest,
) -> Result<(), ApiStreamError>
where
S: Stream<Item = Result<Bytes, String>> + Unpin,
{
loop {
if let Some(record) = decoder.next()? {
ingest.dispatch(&record).await?;
return Ok(());
}
let Some(chunk) = backend.next().await else {
if let Some(record) = decoder.finish()? {
ingest.dispatch(&record).await?;
return Ok(());
}
return Err(missing_shell(0));
};
decoder.push(&chunk.map_err(ApiStreamError::Backend)?);
}
}
async fn ingest_remaining<S>(
mut backend: S,
mut decoder: RecordDecoder,
mut ingest: CommandIngest,
) -> Result<(), ApiStreamError>
where
S: Stream<Item = Result<Bytes, String>> + Unpin,
{
while let Some(record) = decoder.next()? {
if ingest.dispatch(&record).await? {
return Ok(());
}
}
while let Some(chunk) = backend.next().await {
let chunk = chunk.map_err(ApiStreamError::Backend)?;
decoder.push(&chunk);
while let Some(record) = decoder.next()? {
if ingest.dispatch(&record).await? {
return Ok(());
}
}
}
if let Some(record) = decoder.finish()? {
if ingest.dispatch(&record).await? {
return Ok(());
}
}
Err(ApiStreamError::MissingFinish)
}
#[cfg(test)]
async fn ingest<S>(
mut backend: S,
sender: mpsc::Sender<ApiStreamCommand>,
state_defaults: StateDefaults,
) -> Result<(), ApiStreamError>
where
S: Stream<Item = Result<Bytes, String>> + Unpin,
{
let mut decoder = RecordDecoder::new();
let mut ingest = CommandIngest {
sender,
state_defaults,
record: 0,
shell_seen: false,
};
ingest_first(&mut backend, &mut decoder, &mut ingest).await?;
ingest_remaining(backend, decoder, ingest).await
}
#[cfg(test)]
mod tests {
use super::*;
use actix_web::body::to_bytes;
use actix_web::http::header::HeaderValue;
use webui_protocol::{FragmentList, WebUIFragment, WebUIProtocol};
fn defaults() -> StateDefaults {
StateDefaults::new(None, "/".to_owned(), HashMap::new())
}
fn shell_render_failure_config() -> RenderConfig {
fn structural(value: &str) -> WebUIFragment {
let mut signal = String::with_capacity(value.len() + 9);
signal.push_str("}}}webui:");
signal.push_str(value);
WebUIFragment::signal(signal, true)
}
let fragments = vec![
WebUIFragment::raw("<html><head>"),
structural("head_start"),
structural("head_end"),
WebUIFragment::raw("</head><body>"),
structural("body_start"),
structural("streaming_root:outside-boundary"),
structural("boundary_start:0"),
WebUIFragment::raw("<main>ready</main>"),
structural("boundary_end:0"),
structural("body_end"),
WebUIFragment::raw("</body></html>"),
];
let document = WebUIProtocol::new(HashMap::from([(
"index.html".to_owned(),
FragmentList { fragments },
)]));
render_config(document)
}
fn valid_streaming_config(shell_fragments: Vec<WebUIFragment>) -> RenderConfig {
fn structural(value: &str) -> WebUIFragment {
let mut signal = String::with_capacity(value.len() + 9);
signal.push_str("}}}webui:");
signal.push_str(value);
WebUIFragment::signal(signal, true)
}
let mut fragments = Vec::with_capacity(shell_fragments.len() + 11);
fragments.extend([
WebUIFragment::raw("<html><head>"),
structural("head_start"),
structural("head_end"),
WebUIFragment::raw("</head><body>"),
structural("body_start"),
]);
fragments.extend(shell_fragments);
fragments.extend([
structural("boundary_start:0"),
WebUIFragment::raw("<main>ready</main>"),
structural("boundary_end:0"),
structural("body_end"),
WebUIFragment::raw("</body></html>"),
]);
let mut document = WebUIProtocol::new(HashMap::from([(
"index.html".to_owned(),
FragmentList { fragments },
)]));
document.streaming_boundaries.insert(
"index.html".to_owned(),
webui_protocol::StreamingBoundaryList {
names: vec!["content".to_owned()],
},
);
render_config(document)
}
fn render_config(document: WebUIProtocol) -> RenderConfig {
RenderConfig {
protocol: Arc::new(Protocol::new(document)),
entry: "index.html".to_owned(),
route_path: "/".to_owned(),
plugin: None,
body_inject: None,
chunk_pool: Arc::new(ChunkPool::new(
StreamingWriter::DEFAULT_CHANNEL_CAPACITY,
StreamingWriter::CHUNK_TARGET + 1024,
)),
}
}
struct StalledBackend {
shell: Option<Bytes>,
dropped: Option<oneshot::Sender<()>>,
}
impl Stream for StalledBackend {
type Item = Result<Bytes, String>;
fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match self.shell.take() {
Some(shell) => Poll::Ready(Some(Ok(shell))),
None => Poll::Pending,
}
}
}
impl Drop for StalledBackend {
fn drop(&mut self) {
if let Some(dropped) = self.dropped.take() {
let _ = dropped.send(());
}
}
}
fn next_record(decoder: &mut RecordDecoder) -> Option<Bytes> {
decoder
.next()
.unwrap_or_else(|error| panic!("decode failed: {error}"))
}
#[test]
fn decoder_reassembles_split_records() {
let mut decoder = RecordDecoder::new();
decoder.push(br#"{"type":"shell","#);
assert!(next_record(&mut decoder).is_none());
decoder.push(
br#""state":{}}
{"type":"finish"}
"#,
);
assert_eq!(
next_record(&mut decoder).unwrap_or_else(|| panic!("missing shell record")),
br#"{"type":"shell","state":{}}"#[..]
);
assert_eq!(
next_record(&mut decoder).unwrap_or_else(|| panic!("missing finish record")),
br#"{"type":"finish"}"#[..]
);
assert!(next_record(&mut decoder).is_none());
assert!(decoder
.finish()
.unwrap_or_else(|error| panic!("finish failed: {error}"))
.is_none());
}
#[tokio::test]
async fn command_stream_requires_shell_and_finish() {
let chunks = tokio_stream::iter([Ok(Bytes::from_static(
br#"{"type":"shell","version":1,"state":{}}
{"type":"boundary","name":"ready"}
{"type":"finish"}
"#,
))]);
let (sender, mut receiver) = mpsc::channel(4);
ingest(chunks, sender, defaults())
.await
.unwrap_or_else(|error| panic!("ingest failed: {error}"));
assert!(matches!(
receiver.recv().await,
Some(ApiStreamCommand::Shell {
version: VERSION,
..
})
));
assert!(matches!(
receiver.recv().await,
Some(ApiStreamCommand::Boundary { .. })
));
assert!(matches!(
receiver.recv().await,
Some(ApiStreamCommand::Finish { .. })
));
}
#[tokio::test]
async fn command_before_shell_is_rejected() {
let chunks = tokio_stream::iter([Ok(Bytes::from_static(
br#"{"type":"boundary","name":"ready"}
"#,
))]);
let (sender, _receiver) = mpsc::channel(1);
let error = ingest(chunks, sender, defaults())
.await
.err()
.unwrap_or_else(|| panic!("command before shell was accepted"));
assert!(matches!(error, ApiStreamError::MissingShell { record: 0 }));
}
#[tokio::test]
async fn unsupported_stream_version_is_rejected() {
let chunks = tokio_stream::iter([Ok(Bytes::from_static(
br#"{"type":"shell","version":2,"state":{}}
"#,
))]);
let (sender, _receiver) = mpsc::channel(1);
let error = ingest(chunks, sender, defaults())
.await
.err()
.unwrap_or_else(|| panic!("unsupported version was accepted"));
assert!(matches!(
error,
ApiStreamError::UnsupportedVersion {
record: 0,
version: 2
}
));
}
#[tokio::test]
async fn empty_stream_is_rejected_before_rendering() {
let chunks = tokio_stream::empty();
let (sender, _receiver) = mpsc::channel(1);
let error = ingest(chunks, sender, defaults())
.await
.err()
.unwrap_or_else(|| panic!("empty stream was accepted"));
assert!(matches!(error, ApiStreamError::MissingShell { record: 0 }));
}
#[actix_web::test]
async fn shell_render_failure_is_reported_before_http_success() {
let chunks = tokio_stream::iter([Ok::<Bytes, String>(Bytes::from_static(
br#"{"type":"shell","version":1,"state":{}}
{"type":"finish"}
"#,
))]);
let response = render(chunks, shell_render_failure_config(), defaults()).await;
assert_eq!(
response.status(),
actix_web::http::StatusCode::INTERNAL_SERVER_ERROR
);
}
#[actix_web::test]
async fn initial_shell_larger_than_the_live_channel_is_staged_without_deadlock() {
let shell_bytes =
(StreamingWriter::DEFAULT_CHANNEL_CAPACITY + 1) * StreamingWriter::CHUNK_TARGET;
let shell_fragments = (0..=StreamingWriter::DEFAULT_CHANNEL_CAPACITY)
.map(|_| WebUIFragment::raw("~".repeat(StreamingWriter::CHUNK_TARGET)))
.collect();
let chunks = tokio_stream::iter([Ok::<Bytes, String>(Bytes::from_static(
br#"{"type":"shell","version":1,"state":{}}
{"type":"boundary","name":"content"}
{"type":"finish"}
"#,
))]);
let response = tokio::time::timeout(
Duration::from_secs(2),
render(chunks, valid_streaming_config(shell_fragments), defaults()),
)
.await
.unwrap_or_else(|_| panic!("initial shell staging deadlocked"));
assert_eq!(response.status(), actix_web::http::StatusCode::OK);
let body = tokio::time::timeout(Duration::from_secs(2), to_bytes(response.into_body()))
.await
.unwrap_or_else(|_| panic!("staged response did not finish"))
.unwrap_or_else(|error| panic!("failed to read staged response: {error}"));
assert_eq!(memchr::memchr_iter(b'~', &body).count(), shell_bytes);
let shell_position = body
.iter()
.position(|byte| *byte == b'~')
.unwrap_or_else(|| panic!("staged shell bytes are missing"));
let boundary_position = body
.windows(b"<main>ready</main>".len())
.position(|window| window == b"<main>ready</main>")
.unwrap_or_else(|| panic!("boundary output is missing"));
assert!(shell_position < boundary_position);
}
#[actix_web::test]
async fn initial_shell_over_precommit_limit_is_rejected() {
let chunks = tokio_stream::iter([Ok::<Bytes, String>(Bytes::from_static(
br#"{"type":"shell","version":1,"state":{}}
"#,
))]);
let shell = WebUIFragment::raw("x".repeat(MAX_INITIAL_SHELL_BYTES + 1));
let response = tokio::time::timeout(
Duration::from_secs(2),
render(chunks, valid_streaming_config(vec![shell]), defaults()),
)
.await
.unwrap_or_else(|_| panic!("oversized initial shell timed out"));
assert_eq!(
response.status(),
actix_web::http::StatusCode::INTERNAL_SERVER_ERROR
);
let body = to_bytes(response.into_body())
.await
.unwrap_or_else(|error| panic!("failed to read error response: {error}"));
assert!(String::from_utf8_lossy(&body).contains("4,000,000-byte precommit limit"));
}
#[actix_web::test]
async fn dropping_browser_response_cancels_stalled_backend() {
let (dropped_tx, dropped_rx) = oneshot::channel();
let backend = StalledBackend {
shell: Some(Bytes::from_static(
br#"{"type":"shell","version":1,"state":{}}
"#,
)),
dropped: Some(dropped_tx),
};
let response = render(backend, valid_streaming_config(Vec::new()), defaults()).await;
assert_eq!(response.status(), actix_web::http::StatusCode::OK);
drop(response);
tokio::time::timeout(Duration::from_secs(1), dropped_rx)
.await
.unwrap_or_else(|_| panic!("stalled backend was not cancelled"))
.unwrap_or_else(|_| panic!("backend drop signal was lost"));
}
#[test]
fn stream_content_type_accepts_parameters_case_insensitively() {
let mut headers = HeaderMap::new();
headers.insert(
CONTENT_TYPE,
HeaderValue::from_static("Application/X-WebUI-Stream; charset=utf-8"),
);
assert!(is_stream(&headers));
}
#[test]
fn decoder_rejects_an_oversized_record() {
let mut decoder = RecordDecoder::new();
decoder.push(&vec![b'x'; MAX_RECORD_BYTES + 1]);
assert!(matches!(
decoder.next(),
Err(ApiStreamError::RecordTooLarge { record: 0 })
));
}
}