use std::collections::{BTreeMap, VecDeque};
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::time::{Duration, Instant};
use futures::future::BoxFuture;
use futures::Stream;
use nostr_sdk::prelude::EventId;
use crate::core::types::JsonRpcNotification;
use super::constants::MAX_OPEN_STREAM_PING_NONCE_BYTES;
use super::errors::OpenStreamError;
use super::frame::OpenStreamFrame;
pub type PublishFrame =
Arc<dyn Fn(JsonRpcNotification) -> BoxFuture<'static, crate::Result<EventId>> + Send + Sync>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FrameOutcome {
None,
SendPong(String),
Closed,
Aborted(Option<String>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum KeepaliveAction {
None,
SendPing(String),
Abort(String),
}
pub struct OpenStreamSessionOptions {
pub progress_token: String,
pub max_buffered_chunks: usize,
pub max_buffered_bytes: u64,
pub idle_timeout_ms: u64,
pub probe_timeout_ms: u64,
pub close_grace_period_ms: u64,
pub publish_frame: Option<PublishFrame>,
}
struct SessionState {
progress_token: String,
max_buffered_chunks: usize,
max_buffered_bytes: u64,
idle_timeout: Duration,
probe_timeout: Duration,
close_grace_period: Duration,
queue: VecDeque<String>,
buffered: BTreeMap<u64, String>,
buffered_bytes: u64,
queued_bytes: u64,
next_expected_chunk: u64,
last_progress: i64,
started: bool,
active: bool,
closed_remotely: bool,
expected_last_chunk_index: Option<u64>,
terminal: Option<OpenStreamError>,
waker: Option<Waker>,
closed_wakers: Vec<Waker>,
last_activity: Instant,
pending_probe_nonce: Option<String>,
control_nonce: u64,
probe_deadline: Option<Instant>,
close_grace_deadline: Option<Instant>,
outbound_progress: u64,
}
impl SessionState {
fn assert_active(&self) -> Result<(), OpenStreamError> {
if !self.active {
return Err(OpenStreamError::Sequence(format!(
"Received frame for inactive stream {}",
self.progress_token
)));
}
Ok(())
}
fn assert_started(&self) -> Result<(), OpenStreamError> {
if !self.started {
return Err(OpenStreamError::Sequence(format!(
"Received non-start frame before start for {}",
self.progress_token
)));
}
Ok(())
}
fn assert_progress(&mut self, progress: i64) -> Result<(), OpenStreamError> {
if progress <= self.last_progress {
return Err(OpenStreamError::Sequence(format!(
"Non-increasing progress for stream {}",
self.progress_token
)));
}
self.last_progress = progress;
Ok(())
}
fn buffer_chunk(&mut self, chunk_index: u64, data: String) -> Result<(), OpenStreamError> {
if chunk_index < self.next_expected_chunk {
return Err(OpenStreamError::Sequence(format!(
"Stale chunkIndex {chunk_index} for {}",
self.progress_token
)));
}
if self.buffered.contains_key(&chunk_index) {
return Err(OpenStreamError::Sequence(format!(
"Duplicate chunkIndex {chunk_index} for {}",
self.progress_token
)));
}
let chunk_bytes = data.len() as u64;
if self.buffered.len() + self.queue.len() >= self.max_buffered_chunks {
return Err(OpenStreamError::Sequence(format!(
"Buffered chunk limit exceeded for stream {}",
self.progress_token
)));
}
if self.buffered_bytes + self.queued_bytes + chunk_bytes > self.max_buffered_bytes {
return Err(OpenStreamError::Sequence(format!(
"Buffered byte limit exceeded for stream {}",
self.progress_token
)));
}
self.buffered.insert(chunk_index, data);
self.buffered_bytes += chunk_bytes;
Ok(())
}
fn flush_contiguous_chunks(&mut self) -> Result<bool, OpenStreamError> {
while let Some(data) = self.buffered.remove(&self.next_expected_chunk) {
self.buffered_bytes = self.buffered_bytes.saturating_sub(data.len() as u64);
self.emit(data);
self.next_expected_chunk += 1;
}
if self.closed_remotely {
return self.maybe_finish_gracefully();
}
Ok(false)
}
fn maybe_finish_gracefully(&mut self) -> Result<bool, OpenStreamError> {
if !self.closed_remotely || !self.buffered.is_empty() {
return Ok(false);
}
if let Some(expected) = self.expected_last_chunk_index {
if self.next_expected_chunk != expected + 1 {
return Err(OpenStreamError::Sequence(format!(
"Incomplete stream for {}: expected chunks through {expected}",
self.progress_token
)));
}
}
self.finalize(None);
Ok(true)
}
fn emit(&mut self, data: String) {
self.queued_bytes += data.len() as u64;
self.queue.push_back(data);
self.wake_stream();
}
fn handle_pong(&mut self, nonce: &str) {
if self.pending_probe_nonce.as_deref() == Some(nonce) {
self.pending_probe_nonce = None;
self.probe_deadline = None;
}
}
fn next_control_nonce(&mut self) -> String {
self.control_nonce += 1;
format!("{}:{}", self.progress_token, self.control_nonce)
}
fn tick(&mut self, now: Instant) -> KeepaliveAction {
if !self.active {
return KeepaliveAction::None;
}
if let (Some(deadline), true) = (self.probe_deadline, self.pending_probe_nonce.is_some()) {
if now >= deadline {
let token = self.progress_token.clone();
self.finalize(Some(OpenStreamError::abort(
token,
Some("Probe timeout".to_string()),
)));
return KeepaliveAction::Abort("Probe timeout".to_string());
}
}
if let Some(deadline) = self.close_grace_deadline {
if self.closed_remotely && !self.buffered.is_empty() && now >= deadline {
let token = self.progress_token.clone();
self.finalize(Some(OpenStreamError::abort(
token,
Some("Close grace period expired".to_string()),
)));
return KeepaliveAction::Abort("Close grace period expired".to_string());
}
}
if !self.closed_remotely
&& self.pending_probe_nonce.is_none()
&& now.saturating_duration_since(self.last_activity) >= self.idle_timeout
{
let nonce = self.next_control_nonce();
self.pending_probe_nonce = Some(nonce.clone());
self.probe_deadline = Some(now + self.probe_timeout);
return KeepaliveAction::SendPing(nonce);
}
KeepaliveAction::None
}
fn finalize(&mut self, error: Option<OpenStreamError>) {
if !self.active {
return;
}
self.active = false;
self.terminal = error;
self.pending_probe_nonce = None;
self.probe_deadline = None;
self.close_grace_deadline = None;
self.wake_stream();
for waker in self.closed_wakers.drain(..) {
waker.wake();
}
}
fn wake_stream(&mut self) {
if let Some(waker) = self.waker.take() {
waker.wake();
}
}
}
#[derive(Clone)]
pub struct OpenStreamSession {
progress_token: String,
state: Arc<Mutex<SessionState>>,
publish_frame: Option<PublishFrame>,
}
impl std::fmt::Debug for OpenStreamSession {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OpenStreamSession")
.field("progress_token", &self.progress_token)
.finish_non_exhaustive()
}
}
impl OpenStreamSession {
pub fn new(options: OpenStreamSessionOptions) -> Self {
let publish_frame = options.publish_frame.clone();
let progress_token = options.progress_token.clone();
let state = SessionState {
progress_token: options.progress_token,
max_buffered_chunks: options.max_buffered_chunks,
max_buffered_bytes: options.max_buffered_bytes,
idle_timeout: Duration::from_millis(options.idle_timeout_ms),
probe_timeout: Duration::from_millis(options.probe_timeout_ms),
close_grace_period: Duration::from_millis(options.close_grace_period_ms),
queue: VecDeque::new(),
buffered: BTreeMap::new(),
buffered_bytes: 0,
queued_bytes: 0,
next_expected_chunk: 0,
last_progress: -1,
started: false,
active: true,
closed_remotely: false,
expected_last_chunk_index: None,
terminal: None,
waker: None,
closed_wakers: Vec::new(),
last_activity: Instant::now(),
pending_probe_nonce: None,
control_nonce: 0,
probe_deadline: None,
close_grace_deadline: None,
outbound_progress: 0,
};
Self {
progress_token,
state: Arc::new(Mutex::new(state)),
publish_frame,
}
}
pub fn progress_token(&self) -> &str {
&self.progress_token
}
pub fn is_active(&self) -> bool {
self.state.lock().unwrap().active
}
pub fn has_started(&self) -> bool {
self.state.lock().unwrap().started
}
pub fn process_frame(
&self,
now: Instant,
progress: i64,
frame: OpenStreamFrame,
) -> Result<FrameOutcome, OpenStreamError> {
let mut s = self.state.lock().unwrap();
s.assert_active()?;
s.assert_progress(progress)?;
s.last_activity = now;
match frame {
OpenStreamFrame::Start { .. } => {
if s.started {
return Err(OpenStreamError::Sequence(format!(
"Duplicate start frame for stream {}",
s.progress_token
)));
}
s.started = true;
Ok(FrameOutcome::None)
}
OpenStreamFrame::Accept => Ok(FrameOutcome::None),
OpenStreamFrame::Ping { nonce } => {
s.assert_started()?;
if nonce.len() > MAX_OPEN_STREAM_PING_NONCE_BYTES {
return Err(OpenStreamError::Sequence(format!(
"ping nonce exceeds {MAX_OPEN_STREAM_PING_NONCE_BYTES}-byte cap (got {} bytes)",
nonce.len()
)));
}
Ok(FrameOutcome::SendPong(nonce))
}
OpenStreamFrame::Pong { nonce } => {
s.assert_started()?;
s.handle_pong(&nonce);
Ok(FrameOutcome::None)
}
OpenStreamFrame::Chunk { chunk_index, data } => {
s.assert_started()?;
s.buffer_chunk(chunk_index, data)?;
let finished = s.flush_contiguous_chunks()?;
Ok(if finished {
FrameOutcome::Closed
} else {
FrameOutcome::None
})
}
OpenStreamFrame::Close { last_chunk_index } => {
s.assert_started()?;
s.closed_remotely = true;
s.expected_last_chunk_index = last_chunk_index;
let finished = s.flush_contiguous_chunks()?;
if finished {
Ok(FrameOutcome::Closed)
} else {
s.close_grace_deadline = Some(now + s.close_grace_period);
Ok(FrameOutcome::None)
}
}
OpenStreamFrame::Abort { reason } => {
let token = s.progress_token.clone();
s.finalize(Some(OpenStreamError::abort(token, reason.clone())));
Ok(FrameOutcome::Aborted(reason))
}
}
}
pub fn tick(&self, now: Instant) -> KeepaliveAction {
self.state.lock().unwrap().tick(now)
}
pub fn fail(&self, error: OpenStreamError) {
self.state.lock().unwrap().finalize(Some(error));
}
pub fn dispose(&self) {
self.state.lock().unwrap().finalize(None);
}
pub async fn abort(&self, reason: Option<String>) {
let publish = {
let mut s = self.state.lock().unwrap();
if !s.active {
return;
}
let token = s.progress_token.clone();
s.finalize(Some(OpenStreamError::abort(token.clone(), reason.clone())));
self.publish_frame.as_ref().map(|publisher| {
s.outbound_progress += 1;
(publisher.clone(), token, s.outbound_progress)
})
};
if let Some((publisher, token, progress)) = publish {
if let Ok(notification) = (OpenStreamFrame::Abort {
reason: reason.clone(),
})
.into_progress_notification(&token, progress, None)
{
let _ = publisher(notification).await;
}
}
}
pub fn closed(&self) -> Closed {
Closed {
state: self.state.clone(),
}
}
#[cfg(test)]
pub(crate) fn shares_state_with(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.state, &other.state)
}
}
impl Stream for OpenStreamSession {
type Item = Result<String, OpenStreamError>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut s = self.state.lock().unwrap();
if let Some(item) = s.queue.pop_front() {
s.queued_bytes = s.queued_bytes.saturating_sub(item.len() as u64);
return Poll::Ready(Some(Ok(item)));
}
if !s.active {
return match s.terminal.take() {
Some(error) => Poll::Ready(Some(Err(error))),
None => Poll::Ready(None),
};
}
s.waker = Some(cx.waker().clone());
Poll::Pending
}
}
pub struct Closed {
state: Arc<Mutex<SessionState>>,
}
impl Future for Closed {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut s = self.state.lock().unwrap();
if !s.active {
Poll::Ready(())
} else {
s.closed_wakers.push(cx.waker().clone());
Poll::Pending
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::{poll, StreamExt};
fn make_session(token: &str, max_chunks: usize, max_bytes: u64) -> OpenStreamSession {
OpenStreamSession::new(OpenStreamSessionOptions {
progress_token: token.to_string(),
max_buffered_chunks: max_chunks,
max_buffered_bytes: max_bytes,
idle_timeout_ms: 30_000,
probe_timeout_ms: 20_000,
close_grace_period_ms: 5_000,
publish_frame: None,
})
}
fn make_session_timers(token: &str, idle: u64, probe: u64, grace: u64) -> OpenStreamSession {
OpenStreamSession::new(OpenStreamSessionOptions {
progress_token: token.to_string(),
max_buffered_chunks: 8,
max_buffered_bytes: 1024,
idle_timeout_ms: idle,
probe_timeout_ms: probe,
close_grace_period_ms: grace,
publish_frame: None,
})
}
fn start() -> OpenStreamFrame {
OpenStreamFrame::Start { content_type: None }
}
fn chunk(index: u64, data: &str) -> OpenStreamFrame {
OpenStreamFrame::Chunk {
chunk_index: index,
data: data.to_string(),
}
}
fn close(last: Option<u64>) -> OpenStreamFrame {
OpenStreamFrame::Close {
last_chunk_index: last,
}
}
async fn drain_ok(session: &mut OpenStreamSession) -> Vec<String> {
let mut out = Vec::new();
while let Some(item) = session.next().await {
out.push(item.expect("graceful stream must not yield an error"));
}
out
}
#[tokio::test]
async fn yields_ordered_chunks_and_finishes_after_close() {
let now = Instant::now();
let mut s = make_session("token-1", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "hello")).unwrap();
s.process_frame(now, 3, chunk(1, " world")).unwrap();
assert_eq!(
s.process_frame(now, 4, close(None)).unwrap(),
FrameOutcome::Closed
);
assert_eq!(drain_ok(&mut s).await, vec!["hello", " world"]);
s.closed().await;
}
#[tokio::test]
async fn buffers_out_of_order_chunks_until_contiguous() {
let now = Instant::now();
let mut s = make_session("token-2", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(1, "world")).unwrap();
s.process_frame(now, 3, chunk(0, "hello ")).unwrap();
s.process_frame(now, 4, close(None)).unwrap();
assert_eq!(drain_ok(&mut s).await, vec!["hello ", "world"]);
}
#[tokio::test]
async fn fails_when_progress_does_not_increase() {
let now = Instant::now();
let s = make_session("token-3", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
let err = s.process_frame(now, 1, chunk(0, "repeat")).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn rejects_a_duplicate_start_frame() {
let now = Instant::now();
let s = make_session("token-dup-start", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
let err = s.process_frame(now, 2, start()).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn abort_frame_terminates_stream_with_error_on_next_poll() {
let now = Instant::now();
let mut s = make_session("token-4", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
assert_eq!(
s.process_frame(
now,
2,
OpenStreamFrame::Abort {
reason: Some("boom".to_string())
}
)
.unwrap(),
FrameOutcome::Aborted(Some("boom".to_string()))
);
match s.next().await {
Some(Err(OpenStreamError::Abort { reason, .. })) => {
assert_eq!(reason.as_deref(), Some("boom"));
}
other => panic!("expected abort error, got {other:?}"),
}
assert!(s.next().await.is_none());
s.closed().await;
}
#[tokio::test]
async fn fails_when_buffered_chunk_count_exceeds_limit() {
let now = Instant::now();
let s = make_session("token-buffer-count", 1, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(1, "late")).unwrap();
let err = s.process_frame(now, 3, chunk(2, "later")).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn fails_when_buffered_byte_count_exceeds_limit_utf8() {
let now = Instant::now();
let s = make_session("token-buffer-bytes", 4, 4);
s.process_frame(now, 1, start()).unwrap();
let err = s.process_frame(now, 2, chunk(1, "héllo")).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn fail_terminates_stream_and_closed_resolves() {
let now = Instant::now();
let mut s = make_session("token-explicit-fail", 4, 16);
s.process_frame(now, 1, start()).unwrap();
s.fail(OpenStreamError::Sequence("synthetic failure".to_string()));
match s.next().await {
Some(Err(OpenStreamError::Sequence(msg))) => assert!(msg.contains("synthetic")),
other => panic!("expected sequence error, got {other:?}"),
}
s.closed().await;
}
#[tokio::test]
async fn counts_unread_queued_chunks_against_the_byte_limit() {
let now = Instant::now();
let s = make_session("token-queued-bytes", 4, 5);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "abc")).unwrap();
let err = s.process_frame(now, 3, chunk(1, "def")).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn releases_queued_byte_budget_after_consume() {
let now = Instant::now();
let mut s = make_session("token-queued-release", 4, 6);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "abc")).unwrap();
s.process_frame(now, 3, chunk(1, "def")).unwrap();
assert!(matches!(
s.process_frame(now, 4, chunk(2, "g")).unwrap_err(),
OpenStreamError::Sequence(_)
));
assert_eq!(s.next().await.unwrap().unwrap(), "abc");
s.process_frame(now, 5, chunk(2, "g")).unwrap();
}
#[tokio::test]
async fn rejects_frames_after_close() {
let now = Instant::now();
let s = make_session("token-post-close", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
assert_eq!(
s.process_frame(now, 2, close(None)).unwrap(),
FrameOutcome::Closed
);
let err = s.process_frame(now, 3, chunk(0, "late")).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn rejects_frames_after_abort() {
let now = Instant::now();
let s = make_session("token-post-abort", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(
now,
2,
OpenStreamFrame::Abort {
reason: Some("boom".to_string()),
},
)
.unwrap();
assert!(matches!(
s.process_frame(now, 3, chunk(0, "late")).unwrap_err(),
OpenStreamError::Sequence(_)
));
assert!(matches!(
s.process_frame(now, 4, close(None)).unwrap_err(),
OpenStreamError::Sequence(_)
));
}
#[tokio::test]
async fn rejects_stale_and_duplicate_chunk_indexes() {
let now = Instant::now();
let s = make_session("token-stale", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "hello")).unwrap();
assert!(matches!(
s.process_frame(now, 3, chunk(0, "late-duplicate"))
.unwrap_err(),
OpenStreamError::Sequence(_)
));
let s2 = make_session("token-dup", 8, 1024);
s2.process_frame(now, 1, start()).unwrap();
s2.process_frame(now, 2, chunk(2, "a")).unwrap();
assert!(matches!(
s2.process_frame(now, 3, chunk(2, "b")).unwrap_err(),
OpenStreamError::Sequence(_)
));
}
#[tokio::test]
async fn requires_all_chunks_through_close_last_chunk_index() {
let now = Instant::now();
let s = make_session("token-last-chunk-index", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "hello")).unwrap();
let err = s.process_frame(now, 3, close(Some(1))).unwrap_err();
assert!(matches!(err, OpenStreamError::Sequence(_)));
}
#[tokio::test]
async fn allows_graceful_close_when_last_chunk_index_matches() {
let now = Instant::now();
let mut s = make_session("token-last-chunk-complete", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "hello")).unwrap();
assert_eq!(
s.process_frame(now, 3, close(Some(0))).unwrap(),
FrameOutcome::Closed
);
assert_eq!(drain_ok(&mut s).await, vec!["hello"]);
s.closed().await;
}
#[tokio::test]
async fn rejects_malformed_close_last_chunk_index_and_stays_active() {
let now = Instant::now();
let mut s = make_session("token-malformed-last", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(0, "hello")).unwrap();
assert!(matches!(
s.process_frame(now, 3, close(Some(5))).unwrap_err(),
OpenStreamError::Sequence(_)
));
assert!(s.is_active());
s.abort(Some("cleanup".to_string())).await;
assert_eq!(s.next().await.unwrap().unwrap(), "hello");
match s.next().await {
Some(Err(OpenStreamError::Abort { reason, .. })) => {
assert_eq!(reason.as_deref(), Some("cleanup"));
}
other => panic!("expected abort error, got {other:?}"),
}
s.closed().await;
}
#[tokio::test]
async fn rejects_negative_chunk_index_and_stays_active() {
use crate::transport::open_stream::OpenStreamRegistry;
use serde_json::json;
let negative = json!({
"type": "open-stream",
"frameType": "chunk",
"chunkIndex": -1,
"data": "oops",
});
assert_eq!(OpenStreamFrame::from_cvm_value(&negative), None);
let fractional = json!({
"type": "open-stream",
"frameType": "chunk",
"chunkIndex": 1.5,
"data": "oops",
});
assert_eq!(OpenStreamFrame::from_cvm_value(&fractional), None);
let now = Instant::now();
let mut registry = OpenStreamRegistry::new();
registry
.process_frame(
now,
&start()
.into_progress_notification("tok-neg", 1, None)
.unwrap(),
)
.await
.unwrap();
assert_eq!(registry.size(), 1);
let bad = JsonRpcNotification {
jsonrpc: "2.0".to_string(),
method: "notifications/progress".to_string(),
params: Some(json!({
"progressToken": "tok-neg",
"progress": 2,
"cvm": negative,
})),
};
assert!(matches!(
registry.process_frame(now, &bad).await.unwrap_err(),
OpenStreamError::Sequence(_)
));
assert_eq!(registry.size(), 1);
assert!(registry.get_session("tok-neg").unwrap().is_active());
}
#[tokio::test]
async fn zero_chunk_close_succeeds() {
let now = Instant::now();
let mut s = make_session("token-zero-chunk", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
assert_eq!(
s.process_frame(now, 2, close(None)).unwrap(),
FrameOutcome::Closed
);
assert!(s.next().await.is_none());
s.closed().await;
}
#[tokio::test]
async fn chunk_filling_last_gap_after_close_yields_closed() {
let now = Instant::now();
let mut s = make_session("token-fill-gap", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
s.process_frame(now, 2, chunk(1, "world")).unwrap();
assert_eq!(
s.process_frame(now, 3, close(Some(1))).unwrap(),
FrameOutcome::None
);
assert_eq!(
s.process_frame(now, 4, chunk(0, "hello ")).unwrap(),
FrameOutcome::Closed
);
assert_eq!(drain_ok(&mut s).await, vec!["hello ", "world"]);
}
#[tokio::test]
async fn parked_reader_is_woken_by_an_arriving_chunk() {
let now = Instant::now();
let mut s = make_session("token-park", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
assert!(poll!(s.next()).is_pending());
s.process_frame(now, 2, chunk(0, "x")).unwrap();
match poll!(s.next()) {
Poll::Ready(Some(Ok(value))) => assert_eq!(value, "x"),
other => panic!("expected the queued chunk, got {other:?}"),
}
}
#[tokio::test]
async fn consumer_abort_publishes_an_abort_frame_then_finalizes() {
use nostr_sdk::prelude::EventId;
let now = Instant::now();
let published: Arc<Mutex<Vec<JsonRpcNotification>>> = Arc::new(Mutex::new(Vec::new()));
let recorder = published.clone();
let publish: PublishFrame = Arc::new(move |frame: JsonRpcNotification| {
let recorder = recorder.clone();
Box::pin(async move {
recorder.lock().unwrap().push(frame);
Ok(EventId::all_zeros())
})
});
let mut s = OpenStreamSession::new(OpenStreamSessionOptions {
progress_token: "token-consumer-abort".to_string(),
max_buffered_chunks: 8,
max_buffered_bytes: 1024,
idle_timeout_ms: 30_000,
probe_timeout_ms: 20_000,
close_grace_period_ms: 5_000,
publish_frame: Some(publish),
});
s.process_frame(now, 1, start()).unwrap();
s.abort(Some("user cancelled".to_string())).await;
{
let frames = published.lock().unwrap();
assert_eq!(frames.len(), 1);
let cvm = &frames[0].params.as_ref().unwrap()["cvm"];
assert_eq!(cvm["frameType"], "abort");
assert_eq!(cvm["reason"], "user cancelled");
}
match s.next().await {
Some(Err(OpenStreamError::Abort { reason, .. })) => {
assert_eq!(reason.as_deref(), Some("user cancelled"));
}
other => panic!("expected abort error, got {other:?}"),
}
}
#[tokio::test]
async fn ping_frame_requests_a_pong() {
let now = Instant::now();
let s = make_session("token-ping", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
assert_eq!(
s.process_frame(
now,
2,
OpenStreamFrame::Ping {
nonce: "nonce-1".to_string()
}
)
.unwrap(),
FrameOutcome::SendPong("nonce-1".to_string())
);
}
#[tokio::test]
async fn rejects_ping_nonce_exceeding_64_bytes() {
let now = Instant::now();
let s = make_session("token-ping-cap", 8, 1024);
s.process_frame(now, 1, start()).unwrap();
let at_cap = "x".repeat(MAX_OPEN_STREAM_PING_NONCE_BYTES);
assert_eq!(at_cap.len(), 64);
assert_eq!(
s.process_frame(
now,
2,
OpenStreamFrame::Ping {
nonce: at_cap.clone()
}
)
.unwrap(),
FrameOutcome::SendPong(at_cap)
);
let over_cap = "x".repeat(MAX_OPEN_STREAM_PING_NONCE_BYTES + 1);
assert!(matches!(
s.process_frame(now, 3, OpenStreamFrame::Ping { nonce: over_cap }),
Err(OpenStreamError::Sequence(_))
));
assert_eq!(
s.process_frame(
now,
4,
OpenStreamFrame::Ping {
nonce: "ok".to_string()
}
)
.unwrap(),
FrameOutcome::SendPong("ok".to_string())
);
}
#[tokio::test]
async fn idle_sends_ping_then_probe_timeout_aborts() {
let t0 = Instant::now();
let mut s = make_session_timers("token-probe", 10, 10, 100);
s.process_frame(t0, 1, start()).unwrap();
assert_eq!(s.tick(t0 + Duration::from_millis(5)), KeepaliveAction::None);
assert_eq!(
s.tick(t0 + Duration::from_millis(10)),
KeepaliveAction::SendPing("token-probe:1".to_string())
);
assert_eq!(
s.tick(t0 + Duration::from_millis(15)),
KeepaliveAction::None
);
assert_eq!(
s.tick(t0 + Duration::from_millis(20)),
KeepaliveAction::Abort("Probe timeout".to_string())
);
match s.next().await {
Some(Err(OpenStreamError::Abort { reason, .. })) => {
assert_eq!(reason.as_deref(), Some("Probe timeout"));
}
other => panic!("expected probe-timeout abort, got {other:?}"),
}
}
#[tokio::test]
async fn matching_pong_clears_probe_and_stream_survives() {
let t0 = Instant::now();
let s = make_session_timers("token-probe-ok", 10, 20, 100);
s.process_frame(t0, 1, start()).unwrap();
let nonce = match s.tick(t0 + Duration::from_millis(10)) {
KeepaliveAction::SendPing(nonce) => nonce,
other => panic!("expected a ping, got {other:?}"),
};
s.process_frame(
t0 + Duration::from_millis(15),
2,
OpenStreamFrame::Pong { nonce },
)
.unwrap();
assert!(!matches!(
s.tick(t0 + Duration::from_millis(30)),
KeepaliveAction::Abort(_)
));
assert!(s.is_active());
}
#[tokio::test]
async fn unexpected_pong_is_ignored_for_liveness() {
let t0 = Instant::now();
let mut s = make_session_timers("token-bad-pong", 10, 10, 100);
s.process_frame(t0, 1, start()).unwrap();
assert!(matches!(
s.tick(t0 + Duration::from_millis(10)),
KeepaliveAction::SendPing(_)
));
s.process_frame(
t0 + Duration::from_millis(12),
2,
OpenStreamFrame::Pong {
nonce: "unexpected".to_string(),
},
)
.unwrap();
assert_eq!(
s.tick(t0 + Duration::from_millis(20)),
KeepaliveAction::Abort("Probe timeout".to_string())
);
match s.next().await {
Some(Err(OpenStreamError::Abort { .. })) => {}
other => panic!("expected abort, got {other:?}"),
}
}
#[tokio::test]
async fn close_grace_deadline_with_missing_chunks_aborts() {
let t0 = Instant::now();
let mut s = make_session_timers("token-grace", 100, 100, 10);
s.process_frame(t0, 1, start()).unwrap();
s.process_frame(t0, 2, chunk(1, "late")).unwrap();
assert_eq!(
s.process_frame(t0, 3, close(None)).unwrap(),
FrameOutcome::None
);
assert_eq!(s.tick(t0 + Duration::from_millis(5)), KeepaliveAction::None);
assert_eq!(
s.tick(t0 + Duration::from_millis(10)),
KeepaliveAction::Abort("Close grace period expired".to_string())
);
match s.next().await {
Some(Err(OpenStreamError::Abort { reason, .. })) => {
assert_eq!(reason.as_deref(), Some("Close grace period expired"));
}
other => panic!("expected grace abort, got {other:?}"),
}
}
#[tokio::test]
async fn close_with_missing_tail_and_nothing_buffered_fails_immediately_without_grace() {
let t0 = Instant::now();
let s = make_session_timers("token-no-grace", 100, 100, 10);
s.process_frame(t0, 1, start()).unwrap();
s.process_frame(t0, 2, chunk(0, "hello")).unwrap();
assert!(matches!(
s.process_frame(t0, 3, close(Some(1))).unwrap_err(),
OpenStreamError::Sequence(_)
));
assert_eq!(s.tick(t0 + Duration::from_secs(60)), KeepaliveAction::None);
}
}