1use std::collections::BTreeMap;
8
9use serde::{Deserialize, Serialize};
10use serde_json::Value as JsonValue;
11use telltale_types::{LocalTypeR, ValType};
12
13use crate::buffer::{BoundedBuffer, BufferConfig, SignedBuffer, SignedValue};
14use crate::coroutine::Value;
15use crate::instr::Endpoint;
16use crate::verification::{
17 signValue, signing_key_for_endpoint, verifySignedValue, verifying_key_for_endpoint, AuthTree,
18 DefaultVerificationModel, Hash, HashTag, Signature, VerificationModel,
19};
20
21pub type SessionId = usize;
23
24pub type HandlerId = String;
26
27pub const DEFAULT_HANDLER_ID: &str = "default_handler";
29
30fn default_handler_id() -> HandlerId {
31 DEFAULT_HANDLER_ID.to_string()
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
36pub struct Edge {
37 pub sid: SessionId,
39 pub sender: String,
41 pub receiver: String,
43}
44
45impl Edge {
46 #[must_use]
48 pub fn new(sid: SessionId, sender: impl Into<String>, receiver: impl Into<String>) -> Self {
49 Self {
50 sid,
51 sender: sender.into(),
52 receiver: receiver.into(),
53 }
54 }
55}
56
57#[derive(Debug, Deserialize)]
58struct EdgeJson {
59 sid: Option<SessionId>,
60 sender: String,
61 receiver: String,
62}
63
64pub fn decode_edge_json(
70 value: &JsonValue,
71 session_hint: Option<SessionId>,
72) -> Result<Edge, String> {
73 let raw: EdgeJson =
74 serde_json::from_value(value.clone()).map_err(|e| format!("invalid edge json: {e}"))?;
75
76 let sid = raw
77 .sid
78 .or(session_hint)
79 .ok_or_else(|| "missing sid in edge json".to_string())?;
80 Ok(Edge::new(sid, raw.sender, raw.receiver))
81}
82
83#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
85pub enum SessionStatus {
86 Active,
88 Draining,
90 Closed,
92 Cancelled,
94 Faulted {
96 reason: String,
98 },
99}
100
101#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct TypeEntry {
104 pub current: LocalTypeR,
106 pub original: LocalTypeR,
108}
109
110#[derive(Debug, Serialize, Deserialize)]
115pub struct SessionState {
116 pub sid: SessionId,
118 pub roles: Vec<String>,
120 pub local_types: BTreeMap<Endpoint, TypeEntry>,
124 pub buffers: BTreeMap<Edge, SignedBuffer<Signature>>,
126 pub auth_leaves: BTreeMap<Edge, Vec<Hash>>,
128 #[serde(default)]
130 pub auth_trees: BTreeMap<Edge, AuthTree>,
131 pub auth_roots: BTreeMap<Edge, Hash>,
133 pub edge_handlers: BTreeMap<Edge, HandlerId>,
135 #[serde(default = "default_handler_id")]
137 pub default_handler: HandlerId,
138 pub edge_traces: BTreeMap<Edge, Vec<ValType>>,
140 pub status: SessionStatus,
142 pub epoch: usize,
144}
145
146impl SessionState {
147 fn update_auth_tree(&mut self, edge: &Edge, signed: &SignedValue<Signature>) {
148 let bytes = serde_json::to_vec(signed).unwrap_or_default();
149 let leaf = DefaultVerificationModel::hash(HashTag::MerkleLeaf, &bytes);
150 self.auth_leaves.entry(edge.clone()).or_default().push(leaf);
151 let tree = self
152 .auth_trees
153 .entry(edge.clone())
154 .or_insert_with(|| AuthTree::new(Vec::new()));
155 tree.append_leaf(leaf);
156 self.auth_roots.insert(edge.clone(), tree.root());
157 }
158
159 pub fn send_signed(
165 &mut self,
166 from: &str,
167 to: &str,
168 signed: &SignedValue<Signature>,
169 ) -> Result<crate::buffer::EnqueueResult, String> {
170 let edge = Edge::new(self.sid, from, to);
171 let buf = self
172 .buffers
173 .get_mut(&edge)
174 .ok_or_else(|| format!("no buffer for edge {from} → {to}"))?;
175 let result = buf.enqueue(signed.clone());
176 if matches!(result, crate::buffer::EnqueueResult::Ok) {
177 self.update_auth_tree(&edge, signed);
178 }
179 Ok(result)
180 }
181
182 pub fn send(
190 &mut self,
191 from: &str,
192 to: &str,
193 val: Value,
194 ) -> Result<crate::buffer::EnqueueResult, String> {
195 let signer = signing_key_for_endpoint(&Endpoint {
196 sid: self.sid,
197 role: from.to_string(),
198 });
199 let signature = signValue(&val, &signer);
200 self.send_signed(
201 from,
202 to,
203 &SignedValue {
204 payload: val,
205 signature,
206 sequence_no: 0,
207 },
208 )
209 }
210
211 pub fn send_with_sequence(
217 &mut self,
218 from: &str,
219 to: &str,
220 val: Value,
221 sequence_no: u64,
222 ) -> Result<crate::buffer::EnqueueResult, String> {
223 let signer = signing_key_for_endpoint(&Endpoint {
224 sid: self.sid,
225 role: from.to_string(),
226 });
227 let signature = signValue(&val, &signer);
228 self.send_signed(
229 from,
230 to,
231 &SignedValue {
232 payload: val,
233 signature,
234 sequence_no,
235 },
236 )
237 }
238
239 pub fn recv_signed(&mut self, from: &str, to: &str) -> Option<SignedValue<Signature>> {
241 let edge = Edge::new(self.sid, from, to);
242 self.buffers.get_mut(&edge).and_then(|buf| buf.dequeue())
243 }
244
245 pub fn recv_verified_signed(
251 &mut self,
252 from: &str,
253 to: &str,
254 ) -> Result<Option<SignedValue<Signature>>, String> {
255 let sender = Endpoint {
256 sid: self.sid,
257 role: from.to_string(),
258 };
259 let verifying = verifying_key_for_endpoint(&sender);
260 let signed = self.recv_signed(from, to);
261 let Some(signed) = signed else {
262 return Ok(None);
263 };
264 if !verifySignedValue(&signed.payload, &signed.signature, &verifying) {
265 return Err(format!(
266 "signature verification failed on edge {from} -> {to}"
267 ));
268 }
269 Ok(Some(signed))
270 }
271
272 pub fn recv_verified(&mut self, from: &str, to: &str) -> Result<Option<Value>, String> {
278 Ok(self
279 .recv_verified_signed(from, to)?
280 .map(|signed| signed.payload))
281 }
282
283 pub fn recv(&mut self, from: &str, to: &str) -> Option<Value> {
285 self.recv_verified(from, to).ok().flatten()
286 }
287
288 #[must_use]
290 pub fn has_message(&self, from: &str, to: &str) -> bool {
291 let edge = Edge::new(self.sid, from, to);
292 self.buffers.get(&edge).is_some_and(|buf| !buf.is_empty())
293 }
294}
295
296#[derive(Debug, Default, Serialize, Deserialize)]
301pub struct SessionStore {
302 sessions: BTreeMap<SessionId, SessionState>,
303 next_id: SessionId,
304}
305
306impl SessionStore {
307 #[must_use]
309 pub fn new() -> Self {
310 Self::default()
311 }
312
313 pub fn open_with_sid(
317 &mut self,
318 sid: SessionId,
319 roles: Vec<String>,
320 buffer_config: &BufferConfig,
321 initial_types: &BTreeMap<String, LocalTypeR>,
322 ) -> SessionId {
323 let mut local_types = BTreeMap::new();
325 for role in &roles {
326 if let Some(lt) = initial_types.get(role) {
327 let ep = Endpoint {
328 sid,
329 role: role.clone(),
330 };
331 local_types.insert(
332 ep,
333 TypeEntry {
334 current: unfold_mu(lt),
335 original: lt.clone(),
336 },
337 );
338 }
339 }
340
341 let mut buffers = BTreeMap::new();
343 for from in &roles {
344 for to in &roles {
345 if from != to {
346 let edge = Edge::new(sid, from.clone(), to.clone());
347 buffers.insert(edge, BoundedBuffer::new(buffer_config));
348 }
349 }
350 }
351
352 let state = SessionState {
353 sid,
354 roles,
355 local_types,
356 buffers,
357 auth_leaves: BTreeMap::new(),
358 auth_trees: BTreeMap::new(),
359 auth_roots: BTreeMap::new(),
360 edge_handlers: BTreeMap::new(),
361 default_handler: default_handler_id(),
362 edge_traces: BTreeMap::new(),
363 status: SessionStatus::Active,
364 epoch: 0,
365 };
366
367 self.sessions.insert(sid, state);
368 self.next_id = self.next_id.max(sid.saturating_add(1));
369 sid
370 }
371
372 pub fn open(
376 &mut self,
377 roles: Vec<String>,
378 buffer_config: &BufferConfig,
379 initial_types: &BTreeMap<String, LocalTypeR>,
380 ) -> SessionId {
381 let sid = self.next_id;
382 self.open_with_sid(sid, roles, buffer_config, initial_types)
383 }
384
385 #[must_use]
387 pub fn next_session_id(&self) -> SessionId {
388 self.next_id
389 }
390
391 #[must_use]
397 pub fn lookup_type(&self, ep: &Endpoint) -> Option<&LocalTypeR> {
398 self.sessions
399 .get(&ep.sid)?
400 .local_types
401 .get(ep)
402 .map(|e| &e.current)
403 }
404
405 pub fn update_type(&mut self, ep: &Endpoint, new_type: LocalTypeR) {
409 if let Some(session) = self.sessions.get_mut(&ep.sid) {
410 if let Some(entry) = session.local_types.get_mut(ep) {
411 entry.current = new_type;
412 }
413 }
414 }
415
416 pub fn update_original(&mut self, ep: &Endpoint, new_original: LocalTypeR) {
418 if let Some(session) = self.sessions.get_mut(&ep.sid) {
419 if let Some(entry) = session.local_types.get_mut(ep) {
420 entry.original = new_original;
421 }
422 }
423 }
424
425 #[must_use]
427 pub fn original_type(&self, ep: &Endpoint) -> Option<&LocalTypeR> {
428 self.sessions
429 .get(&ep.sid)?
430 .local_types
431 .get(ep)
432 .map(|e| &e.original)
433 }
434
435 pub fn remove_type(&mut self, ep: &Endpoint) {
437 if let Some(session) = self.sessions.get_mut(&ep.sid) {
438 session.local_types.remove(ep);
439 }
440 }
441
442 #[must_use]
446 pub fn get(&self, sid: SessionId) -> Option<&SessionState> {
447 self.sessions.get(&sid)
448 }
449
450 pub fn get_mut(&mut self, sid: SessionId) -> Option<&mut SessionState> {
452 self.sessions.get_mut(&sid)
453 }
454
455 pub fn iter(&self) -> impl Iterator<Item = &SessionState> {
457 self.sessions.values()
458 }
459
460 pub fn close(&mut self, sid: SessionId) -> Result<(), String> {
466 let session = self
467 .sessions
468 .get_mut(&sid)
469 .ok_or_else(|| format!("session {sid} not found"))?;
470
471 session.status = SessionStatus::Closed;
472 session.buffers.clear();
473 session.edge_traces.clear();
474 session.epoch = session.epoch.saturating_add(1);
475 Ok(())
476 }
477
478 #[must_use]
480 pub fn active_count(&self) -> usize {
481 self.sessions
482 .values()
483 .filter(|s| s.status == SessionStatus::Active)
484 .count()
485 }
486
487 #[must_use]
489 pub fn session_ids(&self) -> Vec<SessionId> {
490 self.sessions.keys().copied().collect()
491 }
492
493 #[must_use]
495 pub fn lookup_handler(&self, edge: &Edge) -> Option<&HandlerId> {
496 self.sessions.get(&edge.sid)?.edge_handlers.get(edge)
497 }
498
499 #[must_use]
501 pub fn default_handler_for_session(&self, sid: SessionId) -> Option<&HandlerId> {
502 Some(&self.sessions.get(&sid)?.default_handler)
503 }
504
505 pub fn set_default_handler_for_session(&mut self, sid: SessionId, handler: HandlerId) {
507 if let Some(session) = self.sessions.get_mut(&sid) {
508 session.default_handler = handler;
509 }
510 }
511
512 pub fn update_handler(&mut self, edge: &Edge, handler: HandlerId) {
514 if let Some(session) = self.sessions.get_mut(&edge.sid) {
515 session.edge_handlers.insert(edge.clone(), handler);
516 }
517 }
518
519 #[must_use]
521 pub fn lookup_trace(&self, edge: &Edge) -> Option<&[ValType]> {
522 self.sessions
523 .get(&edge.sid)?
524 .edge_traces
525 .get(edge)
526 .map(Vec::as_slice)
527 }
528
529 pub fn update_trace(&mut self, edge: &Edge, trace: Vec<ValType>) {
531 if let Some(session) = self.sessions.get_mut(&edge.sid) {
532 session.edge_traces.insert(edge.clone(), trace);
533 }
534 }
535}
536
537#[must_use]
543pub fn unfold_mu(lt: &LocalTypeR) -> LocalTypeR {
545 match lt {
546 LocalTypeR::Mu { body, .. } => unfold_mu(body),
547 other => other.clone(),
548 }
549}
550
551#[must_use]
556pub fn unfold_if_var(cont: &LocalTypeR, original: &LocalTypeR) -> LocalTypeR {
557 match cont {
558 LocalTypeR::Var(_) => unfold_mu(original),
559 LocalTypeR::Mu { .. } => unfold_mu(cont),
560 other => other.clone(),
561 }
562}
563
564#[must_use]
570pub fn unfold_if_var_with_scope(
571 cont: &LocalTypeR,
572 original: &LocalTypeR,
573) -> (LocalTypeR, Option<LocalTypeR>) {
574 match cont {
575 LocalTypeR::Var(_) => (unfold_mu(original), None),
576 LocalTypeR::Mu { .. } => (unfold_mu(cont), Some(cont.clone())),
577 other => (other.clone(), None),
578 }
579}
580
581#[cfg(test)]
582mod tests {
583 use super::*;
584 use serde_json::json;
585 use telltale_types::Label;
586
587 fn default_types() -> BTreeMap<String, LocalTypeR> {
588 let mut m = BTreeMap::new();
589 m.insert(
590 "A".to_string(),
591 LocalTypeR::mu(
592 "step",
593 LocalTypeR::Send {
594 partner: "B".into(),
595 branches: vec![(Label::new("msg"), None, LocalTypeR::var("step"))],
596 },
597 ),
598 );
599 m.insert(
600 "B".to_string(),
601 LocalTypeR::mu(
602 "step",
603 LocalTypeR::Recv {
604 partner: "A".into(),
605 branches: vec![(Label::new("msg"), None, LocalTypeR::var("step"))],
606 },
607 ),
608 );
609 m
610 }
611
612 #[test]
613 fn test_session_open_with_types() {
614 let mut store = SessionStore::new();
615 let types = default_types();
616 let sid = store.open(
617 vec!["A".into(), "B".into()],
618 &BufferConfig::default(),
619 &types,
620 );
621
622 let ep_a = Endpoint {
623 sid,
624 role: "A".into(),
625 };
626 let ep_b = Endpoint {
627 sid,
628 role: "B".into(),
629 };
630
631 assert!(matches!(
633 store.lookup_type(&ep_a),
634 Some(LocalTypeR::Send { .. })
635 ));
636 assert!(matches!(
637 store.lookup_type(&ep_b),
638 Some(LocalTypeR::Recv { .. })
639 ));
640 }
641
642 #[test]
643 fn test_type_advance_and_unfold() {
644 let mut store = SessionStore::new();
645 let types = default_types();
646 let sid = store.open(
647 vec!["A".into(), "B".into()],
648 &BufferConfig::default(),
649 &types,
650 );
651
652 let ep_a = Endpoint {
653 sid,
654 role: "A".into(),
655 };
656
657 let lt = store.lookup_type(&ep_a).unwrap().clone();
659 let (_, _vt, continuation) = match < {
660 LocalTypeR::Send { branches, .. } => branches.first().unwrap().clone(),
661 _ => panic!("expected Send"),
662 };
663
664 let original = store.original_type(&ep_a).unwrap();
666 let resolved = unfold_if_var(&continuation, original);
667 assert!(matches!(resolved, LocalTypeR::Send { .. }));
668
669 store.update_type(&ep_a, resolved);
671 assert!(matches!(
672 store.lookup_type(&ep_a),
673 Some(LocalTypeR::Send { .. })
674 ));
675 }
676
677 #[test]
678 fn test_session_send_recv() {
679 let mut store = SessionStore::new();
680 let sid = store.open(
681 vec!["A".into(), "B".into()],
682 &BufferConfig::default(),
683 &BTreeMap::new(),
684 );
685
686 let session = store.get_mut(sid).unwrap();
687 session.send("A", "B", Value::Nat(42)).unwrap();
688 assert!(session.has_message("A", "B"));
689 assert!(!session.has_message("B", "A"));
690
691 let val = session.recv("A", "B");
692 assert_eq!(val, Some(Value::Nat(42)));
693 }
694
695 #[test]
696 fn test_close_clears_buffers_and_traces_even_when_messages_pending() {
697 let mut store = SessionStore::new();
698 let sid = store.open(
699 vec!["A".into(), "B".into()],
700 &BufferConfig::default(),
701 &BTreeMap::new(),
702 );
703 let edge = Edge::new(sid, "A", "B");
704 store
705 .get_mut(sid)
706 .expect("session exists")
707 .send("A", "B", Value::Nat(7))
708 .expect("enqueue pending message");
709 store.update_trace(&edge, vec![ValType::Nat]);
710
711 store.close(sid).expect("close session");
712 let session = store.get(sid).expect("session exists after close");
713 assert_eq!(session.status, SessionStatus::Closed);
714 assert!(session.buffers.is_empty());
715 assert!(session.edge_traces.is_empty());
716 }
717
718 #[test]
719 fn test_namespace_isolation() {
720 let mut store = SessionStore::new();
721 let sid1 = store.open(
722 vec!["A".into(), "B".into()],
723 &BufferConfig::default(),
724 &BTreeMap::new(),
725 );
726 let sid2 = store.open(
727 vec!["A".into(), "B".into()],
728 &BufferConfig::default(),
729 &BTreeMap::new(),
730 );
731
732 assert_ne!(sid1, sid2);
733
734 store
735 .get_mut(sid1)
736 .unwrap()
737 .send("A", "B", Value::Nat(1))
738 .unwrap();
739 assert!(!store.get(sid2).unwrap().has_message("A", "B"));
740 }
741
742 #[test]
743 fn test_remove_type() {
744 let mut store = SessionStore::new();
745 let types = default_types();
746 let sid = store.open(
747 vec!["A".into(), "B".into()],
748 &BufferConfig::default(),
749 &types,
750 );
751
752 let ep_a = Endpoint {
753 sid,
754 role: "A".into(),
755 };
756 assert!(store.lookup_type(&ep_a).is_some());
757
758 store.remove_type(&ep_a);
759 assert!(store.lookup_type(&ep_a).is_none());
760 }
761
762 #[test]
763 fn test_cross_session_role_name_edge_collision_regression() {
764 let mut store = SessionStore::new();
765 let sid1 = store.open(
766 vec!["A".into(), "B".into()],
767 &BufferConfig::default(),
768 &BTreeMap::new(),
769 );
770 let sid2 = store.open(
771 vec!["A".into(), "B".into()],
772 &BufferConfig::default(),
773 &BTreeMap::new(),
774 );
775
776 let e1 = Edge::new(sid1, "A", "B");
777 let e2 = Edge::new(sid2, "A", "B");
778 assert_ne!(e1, e2, "edges from distinct sessions must not collide");
779 assert!(store
780 .get(sid1)
781 .expect("sid1 exists")
782 .buffers
783 .contains_key(&e1));
784 assert!(store
785 .get(sid2)
786 .expect("sid2 exists")
787 .buffers
788 .contains_key(&e2));
789 }
790
791 #[test]
792 fn test_edge_handler_and_trace_bindings() {
793 let mut store = SessionStore::new();
794 let sid = store.open(
795 vec!["A".into(), "B".into()],
796 &BufferConfig::default(),
797 &BTreeMap::new(),
798 );
799 let edge = Edge::new(sid, "A", "B");
800
801 assert!(store.lookup_handler(&edge).is_none());
802 store.update_handler(&edge, "handler/send".to_string());
803 assert_eq!(
804 store.lookup_handler(&edge).map(String::as_str),
805 Some("handler/send")
806 );
807
808 assert!(store.lookup_trace(&edge).is_none());
809 store.update_trace(&edge, vec![ValType::Nat]);
810 assert_eq!(store.lookup_trace(&edge), Some([ValType::Nat].as_slice()));
811 }
812
813 #[test]
814 fn test_decode_edge_json_requires_sid_sender_receiver() {
815 let sid_qualified = json!({
816 "sid": 7,
817 "sender": "A",
818 "receiver": "B"
819 });
820 let e = decode_edge_json(&sid_qualified, None).expect("decode sid-qualified edge");
821 assert_eq!(e, Edge::new(7, "A", "B"));
822
823 let no_sid = json!({
824 "sender": "A",
825 "receiver": "B"
826 });
827 let e2 = decode_edge_json(&no_sid, Some(11)).expect("decode edge with sid hint");
828 assert_eq!(e2, Edge::new(11, "A", "B"));
829
830 let legacy = json!({
831 "from": "A",
832 "to": "B",
833 "sid": 11
834 });
835 let err = decode_edge_json(&legacy, None).expect_err("legacy edge shape must be rejected");
836 assert!(err.contains("invalid edge json"), "unexpected error: {err}");
837 }
838}