1use std::collections::HashMap;
34
35use chia_bls::SecretKey;
36use chia_protocol::Bytes32;
37
38use crate::constants::MAX_CONCURRENT_STREAMS;
39use crate::envelope::{DigMessageEnvelope, InteractionShape, StreamFrame, StreamHeader};
40use crate::error::{MessageError, Result};
41use crate::replay::ReplayGuard;
42use crate::seal::{open_message, seal_message, SealParams};
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47enum Role {
48 Initiator,
49 Responder,
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54enum Handshake {
55 Opening,
56 Established,
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum StreamState {
62 Opening,
64 Open,
66 HalfClosedLocal,
68 HalfClosedRemote,
70 Closed,
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77enum Accepted {
78 Established,
80 Data,
82 Credit(u64),
84 RemoteClosed,
86 CloseAcked,
88 Reset,
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
94pub enum StreamEvent {
95 Opened,
98 Established,
100 Data(Vec<u8>),
102 CreditGranted(u64),
104 RemoteClosed,
106 CloseAcked,
108 PeerReset,
110}
111
112#[derive(Debug, PartialEq, Eq)]
121pub enum StreamAccept {
122 Event(StreamEvent),
124 Dropped {
130 cause: MessageError,
132 },
133 Reset {
139 frame: Box<DigMessageEnvelope>,
141 cause: MessageError,
143 },
144}
145
146#[derive(Debug)]
150pub struct StreamSession {
151 role: Role,
152 handshake: Handshake,
153 send_seq: u64,
155 recv_seq: u64,
158 send_credit: u64,
160 recv_window_remaining: u64,
162 local_closed: bool,
164 remote_closed: bool,
166 reset: bool,
168}
169
170impl StreamSession {
171 fn initiator(recv_window: u32) -> Self {
174 Self {
175 role: Role::Initiator,
176 handshake: Handshake::Opening,
177 send_seq: 0,
178 recv_seq: 0,
179 send_credit: 0,
180 recv_window_remaining: u64::from(recv_window),
181 local_closed: false,
182 remote_closed: false,
183 reset: false,
184 }
185 }
186
187 fn responder(granted_credit: u32) -> Self {
190 Self {
191 role: Role::Responder,
192 handshake: Handshake::Opening,
193 send_seq: 0,
194 recv_seq: 0,
195 send_credit: u64::from(granted_credit),
196 recv_window_remaining: 0,
197 local_closed: false,
198 remote_closed: false,
199 reset: false,
200 }
201 }
202
203 #[must_use]
205 pub fn state(&self) -> StreamState {
206 if self.reset || (self.local_closed && self.remote_closed) {
207 return StreamState::Closed;
208 }
209 if self.handshake == Handshake::Opening {
210 return StreamState::Opening;
211 }
212 match (self.local_closed, self.remote_closed) {
213 (true, false) => StreamState::HalfClosedLocal,
214 (false, true) => StreamState::HalfClosedRemote,
215 _ => StreamState::Open,
216 }
217 }
218
219 #[must_use]
221 pub fn is_closed(&self) -> bool {
222 self.state() == StreamState::Closed
223 }
224
225 #[must_use]
227 pub fn send_credit(&self) -> u64 {
228 self.send_credit
229 }
230
231 fn build_open_ack(&mut self, recv_window: u32) -> Result<StreamHeader> {
237 if self.role != Role::Responder || self.handshake != Handshake::Opening {
238 return Err(MessageError::StreamProtocol(
239 "OPEN_ACK only from a responder mid-handshake",
240 ));
241 }
242 self.handshake = Handshake::Established;
243 self.recv_window_remaining = u64::from(recv_window);
244 Ok(header(StreamFrame::OpenAck, 0, recv_window))
245 }
246
247 fn build_data(&mut self) -> Result<StreamHeader> {
253 if self.handshake != Handshake::Established {
254 return Err(MessageError::StreamProtocol(
255 "DATA before the stream is established",
256 ));
257 }
258 if self.local_closed {
259 return Err(MessageError::StreamProtocol("DATA after our own CLOSE"));
260 }
261 if self.send_credit == 0 {
262 return Err(MessageError::StreamProtocol(
263 "DATA exceeds the granted credit window",
264 ));
265 }
266 self.send_credit -= 1;
267 let seq = self.send_seq;
268 self.send_seq += 1;
269 Ok(header(StreamFrame::Data, seq, 0))
270 }
271
272 fn build_credit(&mut self, n: u32) -> Result<StreamHeader> {
278 if self.handshake != Handshake::Established {
279 return Err(MessageError::StreamProtocol(
280 "CREDIT before the stream is established",
281 ));
282 }
283 self.recv_window_remaining = self.recv_window_remaining.saturating_add(u64::from(n));
284 Ok(header(StreamFrame::Credit, 0, n))
285 }
286
287 fn build_close(&mut self) -> Result<StreamHeader> {
292 if self.local_closed {
293 return Err(MessageError::StreamProtocol("CLOSE after our own CLOSE"));
294 }
295 self.local_closed = true;
296 Ok(header(StreamFrame::Close, 0, 0))
297 }
298
299 fn build_reset(&mut self) -> StreamHeader {
301 self.reset = true;
302 header(StreamFrame::Reset, 0, 0)
303 }
304
305 fn on_recv(&mut self, frame: StreamFrame, hdr: StreamHeader) -> Result<Accepted> {
314 if self.reset {
315 return Err(MessageError::StreamProtocol("frame after RESET"));
316 }
317 match frame {
318 StreamFrame::Open => Err(MessageError::StreamProtocol(
319 "duplicate OPEN for a live stream",
320 )),
321 StreamFrame::OpenAck => {
322 if self.role != Role::Initiator || self.handshake != Handshake::Opening {
323 return Err(MessageError::StreamProtocol("unexpected OPEN_ACK"));
324 }
325 self.handshake = Handshake::Established;
326 self.send_credit = u64::from(hdr.window);
327 Ok(Accepted::Established)
328 }
329 StreamFrame::Data => {
330 if self.handshake != Handshake::Established {
331 return Err(MessageError::StreamProtocol(
332 "DATA before the stream is established",
333 ));
334 }
335 if self.remote_closed {
336 return Err(MessageError::StreamProtocol("DATA after the peer's CLOSE"));
337 }
338 if hdr.seq != self.recv_seq {
339 return Err(MessageError::StreamProtocol(
340 "out-of-order / gap / replayed DATA seq",
341 ));
342 }
343 if self.recv_window_remaining == 0 {
344 return Err(MessageError::StreamProtocol(
345 "DATA exceeds the credit window we granted",
346 ));
347 }
348 self.recv_window_remaining -= 1;
349 self.recv_seq += 1;
350 Ok(Accepted::Data)
351 }
352 StreamFrame::Credit => {
353 if self.handshake != Handshake::Established {
354 return Err(MessageError::StreamProtocol(
355 "CREDIT before the stream is established",
356 ));
357 }
358 self.send_credit = self.send_credit.saturating_add(u64::from(hdr.window));
359 Ok(Accepted::Credit(u64::from(hdr.window)))
360 }
361 StreamFrame::Close => {
362 if self.handshake != Handshake::Established {
363 return Err(MessageError::StreamProtocol(
364 "CLOSE before the stream is established",
365 ));
366 }
367 if self.remote_closed {
368 return Err(MessageError::StreamProtocol(
369 "duplicate CLOSE from the peer",
370 ));
371 }
372 self.remote_closed = true;
373 Ok(Accepted::RemoteClosed)
374 }
375 StreamFrame::CloseAck => {
376 if !self.local_closed {
377 return Err(MessageError::StreamProtocol("CLOSE_ACK without our CLOSE"));
378 }
379 Ok(Accepted::CloseAcked)
380 }
381 StreamFrame::Reset => {
382 self.reset = true;
383 Ok(Accepted::Reset)
384 }
385 }
386 }
387}
388
389fn header(frame: StreamFrame, seq: u64, window: u32) -> StreamHeader {
391 StreamHeader {
392 frame: frame.as_u8(),
393 seq,
394 window,
395 }
396}
397
398pub struct StreamEndpoint<'a> {
406 identity_sk: &'a SecretKey,
408 local_did: Bytes32,
410 local_epoch: u32,
412 peer_did: Bytes32,
414 peer_pub: &'a [u8; 48],
416 message_type: u32,
418 send_counter: u64,
420 sessions: HashMap<Bytes32, StreamSession>,
422 max_concurrent: usize,
424 guard: ReplayGuard,
426}
427
428impl<'a> StreamEndpoint<'a> {
429 #[must_use]
434 pub fn new(
435 identity_sk: &'a SecretKey,
436 local_did: Bytes32,
437 local_epoch: u32,
438 peer_did: Bytes32,
439 peer_pub: &'a [u8; 48],
440 message_type: u32,
441 ) -> Self {
442 Self {
443 identity_sk,
444 local_did,
445 local_epoch,
446 peer_did,
447 peer_pub,
448 message_type,
449 send_counter: 0,
450 sessions: HashMap::new(),
451 max_concurrent: MAX_CONCURRENT_STREAMS,
452 guard: ReplayGuard::new(),
453 }
454 }
455
456 #[must_use]
459 pub fn with_max_concurrent(mut self, max: usize) -> Self {
460 self.max_concurrent = max;
461 self
462 }
463
464 #[must_use]
466 pub fn stream_count(&self) -> usize {
467 self.sessions.len()
468 }
469
470 #[must_use]
472 pub fn session(&self, correlation_id: Bytes32) -> Option<&StreamSession> {
473 self.sessions.get(&correlation_id)
474 }
475
476 pub fn open(
483 &mut self,
484 correlation_id: Bytes32,
485 recv_window: u32,
486 now_ms: u64,
487 expires_at: u64,
488 ) -> Result<DigMessageEnvelope> {
489 if self.sessions.contains_key(&correlation_id) {
490 return Err(MessageError::StreamProtocol(
491 "correlation_id already in use",
492 ));
493 }
494 if self.sessions.len() >= self.max_concurrent {
495 return Err(MessageError::StreamLimit {
496 cap: self.max_concurrent,
497 });
498 }
499 self.sessions
500 .insert(correlation_id, StreamSession::initiator(recv_window));
501 let hdr = header(StreamFrame::Open, 0, recv_window);
502 self.seal_frame(correlation_id, hdr, &[], now_ms, expires_at)
503 }
504
505 pub fn open_ack(
512 &mut self,
513 correlation_id: Bytes32,
514 recv_window: u32,
515 now_ms: u64,
516 expires_at: u64,
517 ) -> Result<DigMessageEnvelope> {
518 let hdr = self
519 .session_mut(correlation_id)?
520 .build_open_ack(recv_window)?;
521 self.seal_frame(correlation_id, hdr, &[], now_ms, expires_at)
522 }
523
524 pub fn send_data(
531 &mut self,
532 correlation_id: Bytes32,
533 payload: &[u8],
534 now_ms: u64,
535 expires_at: u64,
536 ) -> Result<DigMessageEnvelope> {
537 let hdr = self.session_mut(correlation_id)?.build_data()?;
538 self.seal_frame(correlation_id, hdr, payload, now_ms, expires_at)
539 }
540
541 pub fn grant_credit(
546 &mut self,
547 correlation_id: Bytes32,
548 n: u32,
549 now_ms: u64,
550 expires_at: u64,
551 ) -> Result<DigMessageEnvelope> {
552 let hdr = self.session_mut(correlation_id)?.build_credit(n)?;
553 self.seal_frame(correlation_id, hdr, &[], now_ms, expires_at)
554 }
555
556 pub fn close(
562 &mut self,
563 correlation_id: Bytes32,
564 now_ms: u64,
565 expires_at: u64,
566 ) -> Result<DigMessageEnvelope> {
567 let hdr = self.session_mut(correlation_id)?.build_close()?;
568 let env = self.seal_frame(correlation_id, hdr, &[], now_ms, expires_at)?;
569 self.drop_if_closed(correlation_id);
570 Ok(env)
571 }
572
573 pub fn reset(
579 &mut self,
580 correlation_id: Bytes32,
581 now_ms: u64,
582 expires_at: u64,
583 ) -> Result<DigMessageEnvelope> {
584 let hdr = self.session_mut(correlation_id)?.build_reset();
585 let env = self.seal_frame(correlation_id, hdr, &[], now_ms, expires_at)?;
586 self.sessions.remove(&correlation_id);
587 Ok(env)
588 }
589
590 pub fn accept(
606 &mut self,
607 envelope: &DigMessageEnvelope,
608 resolve_sender_pub: impl Fn(Bytes32, u32) -> Option<[u8; 48]>,
609 now_ms: u64,
610 ) -> Result<StreamAccept> {
611 let correlation_id = envelope.correlation_id;
612
613 let opened = match open_message(
621 self.identity_sk,
622 envelope,
623 &resolve_sender_pub,
624 &mut self.guard,
625 now_ms,
626 ) {
627 Ok(opened) => opened,
628 Err(cause) => return Ok(StreamAccept::Dropped { cause }),
629 };
630
631 if opened.sender != self.peer_did {
639 return Ok(StreamAccept::Dropped {
640 cause: MessageError::StreamProtocol(
641 "authenticated sender is not this stream's peer",
642 ),
643 });
644 }
645
646 let Some(hdr) = envelope.stream else {
649 return Ok(StreamAccept::Dropped {
650 cause: MessageError::StreamProtocol("stream event on a non-stream envelope"),
651 });
652 };
653 let Some(frame) = StreamFrame::from_u8(hdr.frame) else {
654 return Ok(StreamAccept::Dropped {
655 cause: MessageError::StreamProtocol("unknown stream frame kind"),
656 });
657 };
658 if opened.shape != InteractionShape::StreamFrame {
659 return Ok(StreamAccept::Dropped {
660 cause: MessageError::StreamProtocol("stream frame with a non-stream shape"),
661 });
662 }
663
664 if frame == StreamFrame::Reset {
667 return Ok(if self.sessions.remove(&correlation_id).is_some() {
668 StreamAccept::Event(StreamEvent::PeerReset)
669 } else {
670 StreamAccept::Dropped {
671 cause: MessageError::StreamProtocol("RESET for an unknown stream"),
672 }
673 });
674 }
675
676 if frame == StreamFrame::Open {
678 return self.accept_open(correlation_id, hdr, now_ms);
679 }
680
681 if !self.sessions.contains_key(&correlation_id) {
686 return Ok(StreamAccept::Dropped {
687 cause: MessageError::StreamProtocol("frame for an unknown stream"),
688 });
689 }
690 let transition = self
691 .sessions
692 .get_mut(&correlation_id)
693 .expect("presence checked above")
694 .on_recv(frame, hdr);
695 match transition {
696 Ok(accepted) => {
697 let event = to_event(accepted, opened.payload);
698 self.drop_if_closed(correlation_id);
699 Ok(StreamAccept::Event(event))
700 }
701 Err(cause) => {
702 self.sessions.remove(&correlation_id);
703 self.reset_response(correlation_id, now_ms, cause)
704 }
705 }
706 }
707
708 fn accept_open(
712 &mut self,
713 correlation_id: Bytes32,
714 hdr: StreamHeader,
715 now_ms: u64,
716 ) -> Result<StreamAccept> {
717 if self.sessions.remove(&correlation_id).is_some() {
718 return self.reset_response(
724 correlation_id,
725 now_ms,
726 MessageError::StreamProtocol("duplicate OPEN for a live stream"),
727 );
728 }
729 if self.sessions.len() >= self.max_concurrent {
730 return self.reset_response(
731 correlation_id,
732 now_ms,
733 MessageError::StreamLimit {
734 cap: self.max_concurrent,
735 },
736 );
737 }
738 self.sessions
739 .insert(correlation_id, StreamSession::responder(hdr.window));
740 Ok(StreamAccept::Event(StreamEvent::Opened))
741 }
742
743 fn reset_response(
748 &mut self,
749 correlation_id: Bytes32,
750 now_ms: u64,
751 cause: MessageError,
752 ) -> Result<StreamAccept> {
753 let hdr = header(StreamFrame::Reset, 0, 0);
754 let frame = self.seal_frame(correlation_id, hdr, &[], now_ms, 0)?;
755 Ok(StreamAccept::Reset {
756 frame: Box::new(frame),
757 cause,
758 })
759 }
760
761 fn seal_frame(
765 &mut self,
766 correlation_id: Bytes32,
767 hdr: StreamHeader,
768 payload: &[u8],
769 now_ms: u64,
770 expires_at: u64,
771 ) -> Result<DigMessageEnvelope> {
772 let counter = self.send_counter;
773 self.send_counter += 1;
774 let params = SealParams {
775 sender_sk: self.identity_sk,
776 sender: self.local_did,
777 sender_epoch: self.local_epoch,
778 recipient: self.peer_did,
779 recipient_pub: self.peer_pub,
780 message_type: self.message_type,
781 shape: InteractionShape::StreamFrame,
782 correlation_id,
783 stream: Some(hdr),
784 counter,
785 timestamp_ms: now_ms,
786 expires_at,
787 payload,
788 };
789 seal_message(¶ms)
790 }
791
792 fn session_mut(&mut self, correlation_id: Bytes32) -> Result<&mut StreamSession> {
794 self.sessions
795 .get_mut(&correlation_id)
796 .ok_or(MessageError::StreamProtocol("unknown stream"))
797 }
798
799 fn drop_if_closed(&mut self, correlation_id: Bytes32) {
801 if self
802 .sessions
803 .get(&correlation_id)
804 .is_some_and(StreamSession::is_closed)
805 {
806 self.sessions.remove(&correlation_id);
807 }
808 }
809}
810
811fn to_event(accepted: Accepted, payload: Vec<u8>) -> StreamEvent {
813 match accepted {
814 Accepted::Established => StreamEvent::Established,
815 Accepted::Data => StreamEvent::Data(payload),
816 Accepted::Credit(n) => StreamEvent::CreditGranted(n),
817 Accepted::RemoteClosed => StreamEvent::RemoteClosed,
818 Accepted::CloseAcked => StreamEvent::CloseAcked,
819 Accepted::Reset => StreamEvent::PeerReset,
820 }
821}
822
823#[cfg(test)]
824mod tests {
825 use super::*;
826 use dig_identity::{derive_identity_sk, master_secret_key_from_seed, public_key_bytes};
827 use sha2::{Digest, Sha256};
828
829 const NOW: u64 = 1_700_000_000_000;
830
831 fn sk(label: &str) -> SecretKey {
832 let seed: [u8; 32] = Sha256::digest(label.as_bytes()).into();
833 derive_identity_sk(&master_secret_key_from_seed(&seed))
834 }
835
836 fn cid(tag: &str) -> Bytes32 {
837 Bytes32::new(Sha256::digest(tag.as_bytes()).into())
838 }
839
840 #[test]
843 fn initiator_handshake_then_data() {
844 let mut s = StreamSession::initiator(4);
845 assert_eq!(s.state(), StreamState::Opening);
846 assert!(s.build_data().is_err());
848 assert_eq!(
850 s.on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 2)),
851 Ok(Accepted::Established)
852 );
853 assert_eq!(s.state(), StreamState::Open);
854 assert_eq!(s.send_credit(), 2);
855 assert!(s.build_data().is_ok());
857 assert!(s.build_data().is_ok());
858 assert!(s.build_data().is_err());
859 }
860
861 #[test]
862 fn outbound_data_seq_increments() {
863 let mut s = StreamSession::initiator(0);
864 s.on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 3))
865 .unwrap();
866 assert_eq!(s.build_data().unwrap().seq, 0);
867 assert_eq!(s.build_data().unwrap().seq, 1);
868 assert_eq!(s.build_data().unwrap().seq, 2);
869 }
870
871 #[test]
872 fn responder_rejects_out_of_order_recv_seq() {
873 let mut s = StreamSession::responder(0);
875 s.build_open_ack(4).unwrap();
876 assert_eq!(
878 s.on_recv(StreamFrame::Data, header(StreamFrame::Data, 0, 0)),
879 Ok(Accepted::Data)
880 );
881 assert!(s
883 .on_recv(StreamFrame::Data, header(StreamFrame::Data, 2, 0))
884 .is_err());
885 assert!(s
887 .on_recv(StreamFrame::Data, header(StreamFrame::Data, 0, 0))
888 .is_err());
889 }
890
891 #[test]
892 fn recv_credit_window_bounds_inbound_data() {
893 let mut s = StreamSession::responder(0);
894 s.build_open_ack(1).unwrap(); assert_eq!(
896 s.on_recv(StreamFrame::Data, header(StreamFrame::Data, 0, 0)),
897 Ok(Accepted::Data)
898 );
899 assert!(s
901 .on_recv(StreamFrame::Data, header(StreamFrame::Data, 1, 0))
902 .is_err());
903 }
904
905 #[test]
906 fn credit_frame_relieves_send_backpressure() {
907 let mut s = StreamSession::initiator(0);
908 s.on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 1))
909 .unwrap();
910 s.build_data().unwrap();
911 assert!(s.build_data().is_err(), "credit exhausted");
912 assert_eq!(
914 s.on_recv(StreamFrame::Credit, header(StreamFrame::Credit, 0, 2)),
915 Ok(Accepted::Credit(2))
916 );
917 assert!(s.build_data().is_ok());
918 assert!(s.build_data().is_ok());
919 assert!(s.build_data().is_err());
920 }
921
922 #[test]
923 fn bidirectional_half_close() {
924 let mut s = StreamSession::initiator(4);
925 s.on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 4))
926 .unwrap();
927 s.build_close().unwrap();
929 assert_eq!(s.state(), StreamState::HalfClosedLocal);
930 assert!(s.build_data().is_err(), "no DATA after our CLOSE");
931 assert_eq!(
933 s.on_recv(StreamFrame::Data, header(StreamFrame::Data, 0, 0)),
934 Ok(Accepted::Data)
935 );
936 assert_eq!(
938 s.on_recv(StreamFrame::Close, header(StreamFrame::Close, 0, 0)),
939 Ok(Accepted::RemoteClosed)
940 );
941 assert_eq!(s.state(), StreamState::Closed);
942 assert!(s.is_closed());
943 }
944
945 #[test]
946 fn reset_aborts_from_any_state() {
947 let mut s = StreamSession::initiator(4);
948 assert_eq!(
949 s.on_recv(StreamFrame::Reset, header(StreamFrame::Reset, 0, 0)),
950 Ok(Accepted::Reset)
951 );
952 assert_eq!(s.state(), StreamState::Closed);
953 assert!(s
955 .on_recv(StreamFrame::Data, header(StreamFrame::Data, 0, 0))
956 .is_err());
957 }
958
959 struct Pair {
962 a_sk: SecretKey,
963 a_did: Bytes32,
964 a_pub: [u8; 48],
965 b_sk: SecretKey,
966 b_did: Bytes32,
967 b_pub: [u8; 48],
968 }
969 fn pair(tag: &str) -> Pair {
970 let a_sk = sk(&format!("{tag}/a"));
971 let b_sk = sk(&format!("{tag}/b"));
972 Pair {
973 a_pub: public_key_bytes(&a_sk),
974 b_pub: public_key_bytes(&b_sk),
975 a_did: cid(&format!("{tag}/a-did")),
976 b_did: cid(&format!("{tag}/b-did")),
977 a_sk,
978 b_sk,
979 }
980 }
981
982 const MT: u32 = 0x0000_0200; #[test]
985 fn full_stream_round_trip_open_data_close() {
986 let p = pair("rt");
987 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
988 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
989 let sender_is_bob = |_d: Bytes32, _e: u32| Some(p.b_pub);
990 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
991 let stream = cid("rt/stream");
992
993 let open = alice.open(stream, 4, NOW, 0).unwrap();
995 assert!(matches!(
996 bob.accept(&open, sender_is_alice, NOW).unwrap(),
997 StreamAccept::Event(StreamEvent::Opened)
998 ));
999
1000 let ack = bob.open_ack(stream, 4, NOW, 0).unwrap();
1002 assert!(matches!(
1003 alice.accept(&ack, sender_is_bob, NOW).unwrap(),
1004 StreamAccept::Event(StreamEvent::Established)
1005 ));
1006
1007 let d0 = alice.send_data(stream, b"hello ", NOW, 0).unwrap();
1009 let d1 = alice.send_data(stream, b"world", NOW, 0).unwrap();
1010 assert_eq!(
1011 bob.accept(&d0, sender_is_alice, NOW).unwrap(),
1012 StreamAccept::Event(StreamEvent::Data(b"hello ".to_vec()))
1013 );
1014 assert_eq!(
1015 bob.accept(&d1, sender_is_alice, NOW).unwrap(),
1016 StreamAccept::Event(StreamEvent::Data(b"world".to_vec()))
1017 );
1018
1019 let close = alice.close(stream, NOW, 0).unwrap();
1021 assert!(matches!(
1022 bob.accept(&close, sender_is_alice, NOW).unwrap(),
1023 StreamAccept::Event(StreamEvent::RemoteClosed)
1024 ));
1025 }
1026
1027 #[test]
1028 fn concurrent_stream_cap_rejects_the_nth_plus_one_open() {
1029 let p = pair("cap");
1030 let mut bob =
1032 StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT).with_max_concurrent(2);
1033 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1034 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1035
1036 for i in 0..2 {
1037 let open = alice.open(cid(&format!("cap/{i}")), 1, NOW, 0).unwrap();
1038 assert!(matches!(
1039 bob.accept(&open, sender_is_alice, NOW).unwrap(),
1040 StreamAccept::Event(StreamEvent::Opened)
1041 ));
1042 }
1043 assert_eq!(bob.stream_count(), 2);
1044
1045 let open3 = alice.open(cid("cap/3"), 1, NOW, 0).unwrap();
1047 match bob.accept(&open3, sender_is_alice, NOW).unwrap() {
1048 StreamAccept::Reset { cause, .. } => {
1049 assert!(matches!(cause, MessageError::StreamLimit { cap: 2 }));
1050 }
1051 other => panic!("expected a StreamLimit RESET, got {other:?}"),
1052 }
1053 assert_eq!(
1054 bob.stream_count(),
1055 2,
1056 "the rejected OPEN created no session"
1057 );
1058 }
1059
1060 #[test]
1061 fn failed_verify_frame_is_dropped_never_reset() {
1062 let p = pair("badverify");
1065 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1066 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1067 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1068 let stream = cid("badverify/s");
1069
1070 let mut open = alice.open(stream, 4, NOW, 0).unwrap();
1071 let last = open.sealed.ciphertext.len() - 1;
1072 open.sealed.ciphertext[last] ^= 0x01;
1073 assert_eq!(
1074 bob.accept(&open, sender_is_alice, NOW).unwrap(),
1075 StreamAccept::Dropped {
1076 cause: MessageError::OpenFailed
1077 }
1078 );
1079 assert_eq!(bob.stream_count(), 0);
1080 }
1081
1082 #[test]
1083 fn frame_for_unknown_stream_is_dropped_never_reset() {
1084 let p = pair("proto");
1085 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1086 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1087 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1088 let stream = cid("proto/s");
1089
1090 alice.open(stream, 4, NOW, 0).unwrap();
1092 alice
1093 .sessions
1094 .get_mut(&stream)
1095 .unwrap()
1096 .on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 4))
1097 .unwrap();
1098 let data = alice.send_data(stream, b"early", NOW, 0).unwrap();
1099 assert!(matches!(
1101 bob.accept(&data, sender_is_alice, NOW).unwrap(),
1102 StreamAccept::Dropped { .. }
1103 ));
1104 }
1105
1106 #[test]
1107 fn replayed_data_on_a_live_stream_is_dropped_and_stream_survives() {
1108 let p = pair("replay");
1111 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1112 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1113 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1114 let stream = cid("replay/s");
1115
1116 let open = alice.open(stream, 4, NOW, 0).unwrap();
1117 bob.accept(&open, sender_is_alice, NOW).unwrap();
1118 let ack = bob.open_ack(stream, 4, NOW, 0).unwrap();
1119 alice.accept(&ack, |_d, _e| Some(p.b_pub), NOW).unwrap();
1120 let data = alice.send_data(stream, b"once", NOW, 0).unwrap();
1121 assert_eq!(
1122 bob.accept(&data, sender_is_alice, NOW).unwrap(),
1123 StreamAccept::Event(StreamEvent::Data(b"once".to_vec()))
1124 );
1125 assert_eq!(
1127 bob.accept(&data, sender_is_alice, NOW).unwrap(),
1128 StreamAccept::Dropped {
1129 cause: MessageError::Replay
1130 }
1131 );
1132 assert_eq!(bob.stream_count(), 1);
1134 assert_eq!(bob.session(stream).unwrap().state(), StreamState::Open);
1135 let next = alice.send_data(stream, b"twice", NOW, 0).unwrap();
1137 assert_eq!(
1138 bob.accept(&next, sender_is_alice, NOW).unwrap(),
1139 StreamAccept::Event(StreamEvent::Data(b"twice".to_vec()))
1140 );
1141 }
1142
1143 #[test]
1144 fn inbound_reset_for_unknown_stream_does_not_beget_a_reset() {
1145 let p = pair("resetstorm");
1147 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1148 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1149 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1150
1151 let ghost = cid("resetstorm/ghost");
1153 alice.open(ghost, 1, NOW, 0).unwrap();
1154 let reset = alice.reset(ghost, NOW, 0).unwrap();
1155 assert!(matches!(
1156 bob.accept(&reset, sender_is_alice, NOW).unwrap(),
1157 StreamAccept::Dropped { .. }
1158 ));
1159 assert_eq!(bob.stream_count(), 0);
1160 }
1161
1162 #[test]
1163 fn protocol_violation_on_established_stream_still_resets() {
1164 let p = pair("legitreset");
1167 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1168 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1169 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1170 let sender_is_bob = |_d: Bytes32, _e: u32| Some(p.b_pub);
1171 let stream = cid("legitreset/s");
1172
1173 let open = alice.open(stream, 4, NOW, 0).unwrap();
1174 bob.accept(&open, sender_is_alice, NOW).unwrap();
1175 let ack = bob.open_ack(stream, 4, NOW, 0).unwrap();
1176 alice.accept(&ack, sender_is_bob, NOW).unwrap();
1177
1178 alice.sessions.get_mut(&stream).unwrap().send_seq = 1;
1180 let bad = alice.send_data(stream, b"gap", NOW, 0).unwrap();
1181 match bob.accept(&bad, sender_is_alice, NOW).unwrap() {
1182 StreamAccept::Reset { cause, .. } => {
1183 assert!(matches!(cause, MessageError::StreamProtocol(_)));
1184 }
1185 other => panic!("expected a RESET on the authenticated violation, got {other:?}"),
1186 }
1187 assert_eq!(bob.stream_count(), 0, "the violated session is dropped");
1188 }
1189
1190 #[test]
1191 fn authenticated_frame_from_a_non_peer_sender_is_hard_dropped() {
1192 let p = pair("senderbind");
1199 let eve_sk = sk("senderbind/eve");
1200 let eve_did = cid("senderbind/eve-did");
1201 let eve_pub = public_key_bytes(&eve_sk);
1202 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1203 let mut eve = StreamEndpoint::new(&eve_sk, eve_did, 0, p.b_did, &p.b_pub, MT);
1205 let permissive_resolver = move |d: Bytes32, _e: u32| {
1207 if d == eve_did {
1208 Some(eve_pub)
1209 } else {
1210 Some(p.a_pub)
1211 }
1212 };
1213 let stream = cid("senderbind/s");
1214
1215 let open = eve.open(stream, 4, NOW, 0).unwrap();
1216 assert_eq!(
1217 bob.accept(&open, permissive_resolver, NOW).unwrap(),
1218 StreamAccept::Dropped {
1219 cause: MessageError::StreamProtocol(
1220 "authenticated sender is not this stream's peer"
1221 )
1222 }
1223 );
1224 assert_eq!(
1225 bob.stream_count(),
1226 0,
1227 "no session opened for the wrong peer"
1228 );
1229 }
1230
1231 #[test]
1232 fn duplicate_open_for_a_live_stream_drops_the_local_session_too() {
1233 let p = pair("dupopen");
1237 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1238 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1239 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1240 let stream = cid("dupopen/s");
1241
1242 let open = alice.open(stream, 4, NOW, 0).unwrap();
1243 assert!(matches!(
1244 bob.accept(&open, sender_is_alice, NOW).unwrap(),
1245 StreamAccept::Event(StreamEvent::Opened)
1246 ));
1247 assert_eq!(bob.stream_count(), 1);
1248
1249 alice.sessions.remove(&stream);
1252 let dup_open = alice.open(stream, 4, NOW, 0).unwrap();
1253 match bob.accept(&dup_open, sender_is_alice, NOW).unwrap() {
1254 StreamAccept::Reset { cause, .. } => {
1255 assert!(matches!(cause, MessageError::StreamProtocol(_)));
1256 }
1257 other => panic!("expected a RESET on the duplicate OPEN, got {other:?}"),
1258 }
1259 assert_eq!(
1261 bob.stream_count(),
1262 0,
1263 "the stale local session must not survive the duplicate OPEN"
1264 );
1265 }
1266
1267 #[test]
1268 fn every_frame_uses_a_unique_ephemeral() {
1269 let p = pair("uniq");
1272 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1273 let stream = cid("uniq/s");
1274 let mut kems = Vec::new();
1275 kems.push(alice.open(stream, 100, NOW, 0).unwrap().sealed.kem_enc);
1276 alice
1277 .sessions
1278 .get_mut(&stream)
1279 .unwrap()
1280 .on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 100))
1281 .unwrap();
1282 for _ in 0..16 {
1283 kems.push(
1284 alice
1285 .send_data(stream, b"x", NOW, 0)
1286 .unwrap()
1287 .sealed
1288 .kem_enc,
1289 );
1290 }
1291 let unique: std::collections::HashSet<_> = kems.iter().collect();
1292 assert_eq!(
1293 unique.len(),
1294 kems.len(),
1295 "every frame must use a fresh ephemeral"
1296 );
1297 }
1298
1299 #[test]
1300 fn credit_grant_close_ack_and_peer_reset_round_trip() {
1301 let p = pair("credit");
1303 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1304 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1305 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1306 let sender_is_bob = |_d: Bytes32, _e: u32| Some(p.b_pub);
1307 let stream = cid("credit/s");
1308
1309 let open = alice.open(stream, 1, NOW, 0).unwrap();
1310 bob.accept(&open, sender_is_alice, NOW).unwrap();
1311 let ack = bob.open_ack(stream, 1, NOW, 0).unwrap();
1312 alice.accept(&ack, sender_is_bob, NOW).unwrap();
1313
1314 let credit = bob.grant_credit(stream, 3, NOW, 0).unwrap();
1316 assert_eq!(
1317 alice.accept(&credit, sender_is_bob, NOW).unwrap(),
1318 StreamAccept::Event(StreamEvent::CreditGranted(3))
1319 );
1320 assert_eq!(alice.session(stream).unwrap().send_credit(), 4);
1321
1322 let reset = bob.reset(stream, NOW, 0).unwrap();
1324 assert_eq!(
1325 alice.accept(&reset, sender_is_bob, NOW).unwrap(),
1326 StreamAccept::Event(StreamEvent::PeerReset)
1327 );
1328 assert_eq!(alice.stream_count(), 0);
1329 assert_eq!(
1330 bob.stream_count(),
1331 0,
1332 "reset drops the sender's session too"
1333 );
1334 }
1335
1336 #[test]
1337 fn close_ack_is_delivered() {
1338 let p = pair("closeack");
1339 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1340 let mut bob = StreamEndpoint::new(&p.b_sk, p.b_did, 0, p.a_did, &p.a_pub, MT);
1341 let sender_is_alice = |_d: Bytes32, _e: u32| Some(p.a_pub);
1342 let sender_is_bob = |_d: Bytes32, _e: u32| Some(p.b_pub);
1343 let stream = cid("closeack/s");
1344
1345 let open = alice.open(stream, 2, NOW, 0).unwrap();
1346 bob.accept(&open, sender_is_alice, NOW).unwrap();
1347 let ack = bob.open_ack(stream, 2, NOW, 0).unwrap();
1348 alice.accept(&ack, sender_is_bob, NOW).unwrap();
1349
1350 let close = alice.close(stream, NOW, 0).unwrap();
1352 bob.accept(&close, sender_is_alice, NOW).unwrap();
1353 let close_ack = bob
1355 .seal_frame(stream, header(StreamFrame::CloseAck, 0, 0), &[], NOW, 0)
1356 .unwrap();
1357 assert_eq!(
1358 alice.accept(&close_ack, sender_is_bob, NOW).unwrap(),
1359 StreamAccept::Event(StreamEvent::CloseAcked)
1360 );
1361 }
1362
1363 #[test]
1364 fn endpoint_send_side_error_branches() {
1365 let p = pair("errs");
1366 let mut alice = StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT);
1367 let stream = cid("errs/s");
1368
1369 assert!(matches!(
1371 alice.send_data(stream, b"x", NOW, 0),
1372 Err(MessageError::StreamProtocol(_))
1373 ));
1374
1375 alice.open(stream, 1, NOW, 0).unwrap();
1376 assert!(matches!(
1378 alice.open(stream, 1, NOW, 0),
1379 Err(MessageError::StreamProtocol(_))
1380 ));
1381
1382 let mut tight =
1384 StreamEndpoint::new(&p.a_sk, p.a_did, 0, p.b_did, &p.b_pub, MT).with_max_concurrent(1);
1385 tight.open(cid("errs/only"), 1, NOW, 0).unwrap();
1386 assert!(matches!(
1387 tight.open(cid("errs/2"), 1, NOW, 0),
1388 Err(MessageError::StreamLimit { cap: 1 })
1389 ));
1390 }
1391
1392 #[test]
1393 fn half_closed_remote_state_then_local_close() {
1394 let mut s = StreamSession::initiator(4);
1396 s.on_recv(StreamFrame::OpenAck, header(StreamFrame::OpenAck, 0, 4))
1397 .unwrap();
1398 s.on_recv(StreamFrame::Close, header(StreamFrame::Close, 0, 0))
1399 .unwrap();
1400 assert_eq!(s.state(), StreamState::HalfClosedRemote);
1401 assert!(
1402 s.build_data().is_ok(),
1403 "we may still send after the peer's CLOSE"
1404 );
1405 s.build_close().unwrap();
1406 assert_eq!(s.state(), StreamState::Closed);
1407 }
1408}