From 18cdefaefbe64ee683c5a43897f06f0b95cc4d72 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Fri, 24 Apr 2026 16:46:56 -0600 Subject: [PATCH] add realistic root knowledge mode to treetest --- treetest/src/app/mod.rs | 51 ++++- treetest/src/app/ui.rs | 232 ++++++++++++++++----- treetest/src/model.rs | 10 + treetest/src/sim.rs | 398 ++++++++++++++++++++++++++++++++++-- treetest/tests/knowledge.rs | 41 ++++ 5 files changed, 656 insertions(+), 76 deletions(-) create mode 100644 treetest/tests/knowledge.rs diff --git a/treetest/src/app/mod.rs b/treetest/src/app/mod.rs index cdba8d6..ada7367 100644 --- a/treetest/src/app/mod.rs +++ b/treetest/src/app/mod.rs @@ -11,7 +11,11 @@ use crossterm::{ }; use ratatui::DefaultTerminal; -use crate::{model::Selection, scenarios::built_in_scenarios, sim::Simulation}; +use crate::{ + model::{NodeId, Selection}, + scenarios::built_in_scenarios, + sim::Simulation, +}; /// Errors returned by the TUI application. #[derive(Debug, thiserror::Error)] @@ -113,6 +117,21 @@ impl App { KeyCode::Char('d') => self.perform_chat_data()?, KeyCode::Char('b') => self.perform_chat_bye()?, KeyCode::Char('f') => self.perform_invalid_fault_demo()?, + KeyCode::Char('g') => { + self.simulation.toggle_inspector_mode(); + self.refresh_selections(Some(self.selected().node_id())); + self.status = if self.simulation.is_realistic_mode() { + "Inspector switched to realistic mode.".to_owned() + } else { + "Inspector switched to ground truth mode.".to_owned() + }; + } + KeyCode::Char('m') => { + self.simulation.enable_realistic_mode_with_memory_reset(); + self.refresh_selections(Some(NodeId(0))); + self.status = + "Cleared root memory for deeper nodes and enabled realistic mode.".to_owned(); + } KeyCode::Char('s') => { let processed = self.simulation.step()?; self.status = if processed { @@ -133,12 +152,7 @@ impl App { fn load_scenario(&mut self, index: usize) -> Result<(), AppError> { self.scenario_index = index; self.simulation = Simulation::new(self.scenarios[index].clone())?; - self.selections = ui::build_selections(&self.simulation); - self.selection_index = self - .selections - .iter() - .position(|selection| *selection == self.simulation.initial_selection()) - .unwrap_or(0); + self.refresh_selections(Some(self.simulation.initial_selection().node_id())); self.status = format!("Loaded scenario: {}", self.scenarios[index].name); Ok(()) } @@ -152,6 +166,7 @@ impl App { Selection::Node(node_id) => { let result = self.simulation.call_endpoint_introspection(node_id)?; let steps = self.simulation.drain()?; + self.refresh_selections(Some(node_id)); self.status = format!("{} ({steps} steps)", result.label); } Selection::Leaf { node_id, leaf_name } => { @@ -159,6 +174,7 @@ impl App { .simulation .call_leaf_introspection(node_id, &leaf_name)?; let steps = self.simulation.drain()?; + self.refresh_selections(Some(node_id)); self.status = format!("{} ({steps} steps)", result.label); } } @@ -171,6 +187,7 @@ impl App { self.simulation .call_echo_leaf(node_id, &leaf_name, "demo echo from root")?; let steps = self.simulation.drain()?; + self.refresh_selections(Some(node_id)); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = "Select a leaf first, then press e.".to_owned(); @@ -193,6 +210,7 @@ impl App { b"ping".to_vec(), )?; let steps = self.simulation.drain()?; + self.refresh_selections(Some(node_id)); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = "Selected node has no endpoint procedures.".to_owned(); @@ -222,6 +240,7 @@ impl App { b"chunk please".to_vec(), )?; let steps = self.simulation.drain()?; + self.refresh_selections(Some(node_id)); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = "Selected node has no chunked procedure.".to_owned(); @@ -248,6 +267,7 @@ impl App { b"open chat".to_vec(), )?; let steps = self.simulation.drain()?; + self.refresh_selections(Some(node_id)); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = "Selected node has no chat procedure.".to_owned(); @@ -264,6 +284,7 @@ impl App { self.simulation .send_root_hook_data(hook_id, "hello from the root", false)?; let steps = self.simulation.drain()?; + self.refresh_selections(None); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = "No known hook yet. Press h to open chat first.".to_owned(); @@ -275,6 +296,7 @@ impl App { if let Some(hook_id) = self.simulation.hook_ids().last().copied() { let result = self.simulation.send_root_hook_data(hook_id, "bye", true)?; let steps = self.simulation.drain()?; + self.refresh_selections(None); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = "No known hook yet. Press h to open chat first.".to_owned(); @@ -284,9 +306,9 @@ impl App { fn perform_invalid_fault_demo(&mut self) -> Result<(), AppError> { if let Some(hook_id) = self.simulation.hook_ids().last().copied() { - let root_id = crate::model::NodeId(0); + let root_id = NodeId(0); if self.simulation.tree.nodes.len() > 1 { - let attacker = crate::model::NodeId(1); + let attacker = NodeId(1); let result = self.simulation.inject_invalid_peer_data( attacker, root_id, @@ -295,6 +317,7 @@ impl App { "spoofed data", )?; let steps = self.simulation.drain()?; + self.refresh_selections(None); self.status = format!("{} ({steps} steps)", result.label); } else { self.status = @@ -305,4 +328,14 @@ impl App { } Ok(()) } + + fn refresh_selections(&mut self, preferred_node: Option) { + let current = preferred_node.unwrap_or_else(|| self.selected().node_id()); + self.selections = ui::build_selections(&self.simulation); + self.selection_index = self + .selections + .iter() + .position(|selection| selection.node_id() == current) + .unwrap_or(0); + } } diff --git a/treetest/src/app/ui.rs b/treetest/src/app/ui.rs index 75010db..dc500e1 100644 --- a/treetest/src/app/ui.rs +++ b/treetest/src/app/ui.rs @@ -9,8 +9,8 @@ use ratatui::{ }; use crate::{ - model::{Selection, format_path}, - sim::{RecordedEvent, Simulation}, + model::{Selection, format_hook_ref, format_leaf_ref, format_path}, + sim::{InspectorMode, RecordedEvent, Simulation}, }; use super::App; @@ -32,11 +32,16 @@ impl App { } fn render_header(&self, frame: &mut Frame<'_>, area: Rect) { + let mode = match self.simulation.inspector_mode { + InspectorMode::GroundTruth => "ground truth", + InspectorMode::Realistic => "realistic", + }; let title = format!( - "treetest | scenario {} / {}: {}", + "treetest | scenario {} / {}: {} | {}", self.scenario_index + 1, self.scenarios.len(), - self.scenarios[self.scenario_index].name + self.scenarios[self.scenario_index].name, + mode ); frame.render_widget( Paragraph::new(title).block(Block::default().borders(Borders::ALL).title("Scenario")), @@ -107,18 +112,15 @@ impl App { node.display_path() ) } - Selection::Leaf { node_id, leaf_name } => { - format!( - "{} {} :: {}", - if index == self.selection_index { - ">" - } else { - " " - }, - self.simulation.node(*node_id).display_path(), - leaf_name - ) - } + Selection::Leaf { node_id, leaf_name } => format!( + "{} {}", + if index == self.selection_index { + ">" + } else { + " " + }, + format_leaf_ref(&self.simulation.node(*node_id).path, leaf_name) + ), }; ListItem::new(label) }) @@ -131,7 +133,21 @@ impl App { fn render_inspector(&self, frame: &mut Frame<'_>, area: Rect) { let selection = self.selected(); - let body = match selection { + let body = match self.simulation.inspector_mode { + InspectorMode::GroundTruth => self.render_ground_truth_inspector(selection), + InspectorMode::Realistic => self.render_realistic_inspector(selection), + }; + + frame.render_widget( + Paragraph::new(body) + .block(Block::default().borders(Borders::ALL).title("Inspector")) + .wrap(Wrap { trim: true }), + area, + ); + } + + fn render_ground_truth_inspector(&self, selection: &Selection) -> Text<'static> { + match selection { Selection::Node(node_id) => { let node = self.simulation.node(*node_id); let mut lines = vec![ @@ -147,16 +163,15 @@ impl App { Line::default(), Line::from("Endpoint procedures:"), ]; - lines.extend( - node.endpoint_procedures - .iter() - .map(|procedure| Line::from(format!("- {}", procedure.procedure_id))), - ); - lines.extend( - node.leaves - .iter() - .map(|leaf| Line::from(format!("- leaf {}", leaf.name))), - ); + lines.extend(node.endpoint_procedures.iter().map(|procedure| { + Line::from(format!( + "- {}: {}", + procedure.procedure_id, procedure.description + )) + })); + lines.extend(node.leaves.iter().map(|leaf| { + Line::from(format!("- {}", format_leaf_ref(&node.path, &leaf.name))) + })); Text::from(lines) } Selection::Leaf { node_id, leaf_name } => { @@ -166,21 +181,118 @@ impl App { .iter() .find(|leaf| &leaf.name == leaf_name) .expect("selection should stay valid"); - Text::from(vec![ - Line::from(format!("Leaf {}", leaf.name)).bold(), + let mut lines = vec![ + Line::from(format_leaf_ref(&node.path, &leaf.name)).bold(), Line::from(leaf.description.clone()), Line::from(format!("Node: {}", node.display_path())), - Line::from(format!("Procedures: {}", leaf.procedures.join(", "))), - ]) + Line::from("Procedures:"), + ]; + lines.extend( + leaf.procedures + .iter() + .map(|procedure| Line::from(format!("- {}", procedure))), + ); + Text::from(lines) } - }; + } + } - frame.render_widget( - Paragraph::new(body) - .block(Block::default().borders(Borders::ALL).title("Inspector")) - .wrap(Wrap { trim: true }), - area, - ); + fn render_realistic_inspector(&self, selection: &Selection) -> Text<'static> { + match selection { + Selection::Node(node_id) => { + let node = self.simulation.node(*node_id); + if let Some(learned) = self.simulation.root_knowledge.node(&node.path) { + let mut lines = vec![ + Line::from(learned.title.clone().unwrap_or_else(|| node.display_path())) + .bold(), + Line::from( + learned + .description + .clone() + .unwrap_or_else(|| "No learned description yet.".to_owned()), + ), + Line::from(format!("Path: {}", format_path(&learned.path))), + Line::from(format!("Known direct child: {}", learned.direct_child)), + Line::from(format!( + "Endpoint introspected: {}", + learned.endpoint_introspected + )), + Line::default(), + Line::from("Known endpoint procedures:"), + ]; + if learned.endpoint_procedures.is_empty() { + lines.push(Line::from("- none learned")); + } else { + lines.extend(learned.endpoint_procedures.iter().map(|procedure| { + Line::from(match &procedure.description { + Some(description) => { + format!("- {}: {}", procedure.procedure_id, description) + } + None => format!("- {}", procedure.procedure_id), + }) + })); + } + lines.push(Line::default()); + lines.push(Line::from("Known leaves:")); + if learned.leaves.is_empty() { + lines.push(Line::from("- none learned")); + } else { + lines.extend(learned.leaves.iter().map(|leaf| { + Line::from(format!( + "- {}", + format_leaf_ref(&learned.path, &leaf.leaf_name) + )) + })); + } + Text::from(lines) + } else { + Text::from(vec![ + Line::from(node.display_path()).bold(), + Line::from( + "The root host has not learned anything about this endpoint yet.", + ), + ]) + } + } + Selection::Leaf { node_id, leaf_name } => { + let node = self.simulation.node(*node_id); + if let Some(learned) = self.simulation.root_knowledge.node(&node.path) + && let Some(leaf) = learned + .leaves + .iter() + .find(|leaf| &leaf.leaf_name == leaf_name) + { + let mut lines = vec![ + Line::from(format_leaf_ref(&node.path, &leaf.leaf_name)).bold(), + Line::from( + leaf.description + .clone() + .unwrap_or_else(|| "No learned description yet.".to_owned()), + ), + Line::from(format!("Node: {}", node.display_path())), + Line::from("Known procedures:"), + ]; + if leaf.procedures.is_empty() { + lines.push(Line::from("- none learned")); + } else { + lines.extend(leaf.procedures.iter().map(|procedure| { + Line::from(match &procedure.description { + Some(description) => { + format!("- {}: {}", procedure.procedure_id, description) + } + None => format!("- {}", procedure.procedure_id), + }) + })); + } + Text::from(lines) + } else { + Text::from(vec![ + Line::from(format_leaf_ref(&node.path, leaf_name)).bold(), + Line::from("The root host has not learned this leaf yet."), + ]) + } + } + } } fn render_trace(&self, frame: &mut Frame<'_>, area: Rect) { @@ -211,9 +323,8 @@ impl App { .map(|hook| { let status = if hook.closed { "closed" } else { "open" }; ListItem::new(format!( - "#{} {} -> {} [{}] {}", - hook.hook_id, - format_path(&hook.host_path), + "{} -> {} [{}] {}", + format_hook_ref(&hook.host_path, hook.hook_id), format_path(&hook.peer_path), status, hook.last_message, @@ -230,7 +341,7 @@ impl App { let help = vec![ Line::from(self.status.clone()).style(Style::default().add_modifier(Modifier::BOLD)), Line::from( - "Keys: arrows move selection/scenario | i introspect | e echo leaf | p ping | c chunked | h open chat | d chat data | b chat bye | f invalid peer | s step | a autoplay | q quit", + "Keys: arrows move selection/scenario | i introspect | e echo leaf | p ping | c chunked | h open chat | d chat data | b chat bye | f invalid peer | g toggle ground-truth/realistic | m clear deeper memory + realistic | s step | a autoplay | q quit", ), Line::from(format!( "Current selection: {}", @@ -276,13 +387,38 @@ impl App { pub(super) fn build_selections(simulation: &Simulation) -> Vec { let mut selections = Vec::new(); - for node in &simulation.tree.nodes { + let node_ids: Vec<_> = match simulation.inspector_mode { + InspectorMode::GroundTruth => simulation.tree.nodes.iter().map(|node| node.id).collect(), + InspectorMode::Realistic => simulation + .root_knowledge + .known_paths() + .into_iter() + .filter_map(|path| simulation.tree.find_by_path(&path)) + .collect(), + }; + + for node_id in node_ids { + let node = simulation.node(node_id); selections.push(Selection::Node(node.id)); - for leaf in &node.leaves { - selections.push(Selection::Leaf { - node_id: node.id, - leaf_name: leaf.name.clone(), - }); + match simulation.inspector_mode { + InspectorMode::GroundTruth => { + for leaf in &node.leaves { + selections.push(Selection::Leaf { + node_id: node.id, + leaf_name: leaf.name.clone(), + }); + } + } + InspectorMode::Realistic => { + if let Some(learned) = simulation.root_knowledge.node(&node.path) { + for leaf in &learned.leaves { + selections.push(Selection::Leaf { + node_id: node.id, + leaf_name: leaf.leaf_name.clone(), + }); + } + } + } } } selections diff --git a/treetest/src/model.rs b/treetest/src/model.rs index c53948c..efc7734 100644 --- a/treetest/src/model.rs +++ b/treetest/src/model.rs @@ -182,3 +182,13 @@ pub fn format_path(path: &[String]) -> String { format!("/{}", path.join("/")) } } + +/// Formats a leaf reference using the protocol document's descriptive syntax. +pub fn format_leaf_ref(path: &[String], leaf_name: &str) -> String { + format!("{} {{ leaf: {} }}", format_path(path), leaf_name) +} + +/// Formats a hook reference using the protocol document's descriptive syntax. +pub fn format_hook_ref(path: &[String], hook_id: u64) -> String { + format!("{} {{ hook: {} }}", format_path(path), hook_id) +} diff --git a/treetest/src/sim.rs b/treetest/src/sim.rs index 27bed7a..805d132 100644 --- a/treetest/src/sim.rs +++ b/treetest/src/sim.rs @@ -12,14 +12,55 @@ use unshell::protocol::tree::{ ChildRoute, ConnectionState, Endpoint, Ingress, LeafBehavior, LocalEvent, ProtocolEndpoint, }; use unshell::protocol::{ - CallMessage, DataMessage, FaultMessage, FrameBytes, PacketHeader, PacketType, decode_frame, + CallMessage, DataMessage, EndpointIntrospection, FaultMessage, FrameBytes, LeafIntrospection, + PacketHeader, PacketType, decode_frame, deserialize_archived_bytes, }; use crate::model::{ DemoTree, EndpointProcedureKind, EndpointProcedureSpec, LeafKind, NodeId, ScenarioDefinition, - Selection, format_path, + Selection, format_hook_ref, format_leaf_ref, format_path, }; +/// Root inspector mode. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InspectorMode { + GroundTruth, + Realistic, +} + +/// Learned procedure metadata stored by the root host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LearnedProcedure { + pub procedure_id: String, + pub description: Option, +} + +/// Learned leaf metadata stored by the root host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LearnedLeaf { + pub leaf_name: String, + pub description: Option, + pub procedures: Vec, +} + +/// Learned endpoint metadata stored by the root host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LearnedNode { + pub path: Vec, + pub title: Option, + pub description: Option, + pub direct_child: bool, + pub endpoint_procedures: Vec, + pub leaves: Vec, + pub endpoint_introspected: bool, +} + +/// Root-host knowledge accumulated from local configuration and observed traffic. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RootKnowledge { + pub nodes: BTreeMap, LearnedNode>, +} + /// User-facing outcome of a root-originated action. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ActionResult { @@ -34,6 +75,7 @@ pub struct HookSnapshot { pub host_path: Vec, pub peer_path: Vec, pub procedure_id: String, + pub target_leaf: Option, pub closed: bool, pub last_message: String, } @@ -121,9 +163,208 @@ pub struct Simulation { pub trace: VecDeque, pub recorded_events: Vec, pub hooks: BTreeMap, + pub inspector_mode: InspectorMode, + pub root_knowledge: RootKnowledge, chat_sessions: BTreeMap, } +impl RootKnowledge { + fn new(tree: &DemoTree) -> Self { + let mut knowledge = Self { + nodes: BTreeMap::new(), + }; + for node in &tree.nodes { + if node.path.is_empty() || node.path.len() == 1 { + let direct_child = node.path.len() == 1; + let mut learned = LearnedNode { + path: node.path.clone(), + title: Some(node.title.clone()), + description: Some(node.description.clone()), + direct_child, + endpoint_procedures: Vec::new(), + leaves: Vec::new(), + endpoint_introspected: node.path.is_empty(), + }; + + if node.path.is_empty() { + learned.endpoint_procedures = node + .endpoint_procedures + .iter() + .map(|procedure| LearnedProcedure { + procedure_id: procedure.procedure_id.clone(), + description: Some(procedure.description.clone()), + }) + .collect(); + learned.leaves = node + .leaves + .iter() + .map(|leaf| LearnedLeaf { + leaf_name: leaf.name.clone(), + description: Some(leaf.description.clone()), + procedures: leaf + .procedures + .iter() + .map(|procedure_id| LearnedProcedure { + procedure_id: procedure_id.clone(), + description: Some(leaf.description.clone()), + }) + .collect(), + }) + .collect(); + } + + knowledge.nodes.insert(node.path.clone(), learned); + } + } + knowledge + } + + fn ensure_node(&mut self, demo_node: &crate::model::DemoNode) -> &mut LearnedNode { + let direct_child = demo_node.path.len() == 1; + self.nodes + .entry(demo_node.path.clone()) + .or_insert_with(|| LearnedNode { + path: demo_node.path.clone(), + title: Some(demo_node.title.clone()), + description: Some(demo_node.description.clone()), + direct_child, + endpoint_procedures: Vec::new(), + leaves: Vec::new(), + endpoint_introspected: false, + }) + } + + fn remember_endpoint_procedure( + &mut self, + demo_node: &crate::model::DemoNode, + procedure: &EndpointProcedureSpec, + ) { + let learned_node = self.ensure_node(demo_node); + push_procedure( + &mut learned_node.endpoint_procedures, + procedure.procedure_id.clone(), + Some(procedure.description.clone()), + ); + } + + fn remember_leaf_from_spec( + &mut self, + demo_node: &crate::model::DemoNode, + leaf_spec: &crate::model::LeafSpec, + ) { + let learned_node = self.ensure_node(demo_node); + let leaf = ensure_leaf( + &mut learned_node.leaves, + leaf_spec.name.clone(), + Some(leaf_spec.description.clone()), + ); + for procedure_id in &leaf_spec.procedures { + push_procedure( + &mut leaf.procedures, + procedure_id.clone(), + Some(leaf_spec.description.clone()), + ); + } + } + + fn remember_endpoint_introspection( + &mut self, + demo_node: &crate::model::DemoNode, + introspection: &EndpointIntrospection, + ) { + let learned_node = self.ensure_node(demo_node); + learned_node.endpoint_introspected = true; + for summary in &introspection.leaves { + let description = demo_node + .leaves + .iter() + .find(|leaf| leaf.name == summary.leaf_name) + .map(|leaf| leaf.description.clone()); + let leaf = ensure_leaf( + &mut learned_node.leaves, + summary.leaf_name.clone(), + description, + ); + for procedure_id in &summary.procedures { + push_procedure(&mut leaf.procedures, procedure_id.clone(), None); + } + } + } + + fn remember_leaf_introspection( + &mut self, + demo_node: &crate::model::DemoNode, + introspection: &LeafIntrospection, + ) { + let learned_node = self.ensure_node(demo_node); + let description = demo_node + .leaves + .iter() + .find(|leaf| leaf.name == introspection.leaf_name) + .map(|leaf| leaf.description.clone()); + let leaf = ensure_leaf( + &mut learned_node.leaves, + introspection.leaf_name.clone(), + description, + ); + for procedure_id in &introspection.procedures { + push_procedure(&mut leaf.procedures, procedure_id.clone(), None); + } + } + + fn clear_deeper_than_one_hop(&mut self) { + self.nodes.retain(|path, _| path.len() <= 1); + } + + pub fn node(&self, path: &[String]) -> Option<&LearnedNode> { + self.nodes.get(path) + } + + pub fn known_paths(&self) -> Vec> { + self.nodes.keys().cloned().collect() + } +} + +fn ensure_leaf<'a>( + leaves: &'a mut Vec, + leaf_name: String, + description: Option, +) -> &'a mut LearnedLeaf { + if let Some(index) = leaves.iter().position(|leaf| leaf.leaf_name == leaf_name) { + if leaves[index].description.is_none() { + leaves[index].description = description; + } + return &mut leaves[index]; + } + + leaves.push(LearnedLeaf { + leaf_name, + description, + procedures: Vec::new(), + }); + leaves.last_mut().expect("just pushed") +} + +fn push_procedure( + procedures: &mut Vec, + procedure_id: String, + description: Option, +) { + if let Some(existing) = procedures + .iter_mut() + .find(|procedure| procedure.procedure_id == procedure_id) + { + if existing.description.is_none() { + existing.description = description; + } + return; + } + procedures.push(LearnedProcedure { + procedure_id, + description, + }); +} + impl Simulation { /// Creates a fresh simulation from a scenario definition. pub fn new(scenario: ScenarioDefinition) -> Result { @@ -172,6 +413,8 @@ impl Simulation { }); } + let root_knowledge = RootKnowledge::new(&tree); + Ok(Self { scenario, tree, @@ -181,6 +424,8 @@ impl Simulation { trace: VecDeque::new(), recorded_events: Vec::new(), hooks: BTreeMap::new(), + inspector_mode: InspectorMode::GroundTruth, + root_knowledge, chat_sessions: BTreeMap::new(), }) } @@ -195,6 +440,25 @@ impl Simulation { self.tree.node(id) } + /// Clears deeper root memory and switches the inspector into realistic mode. + pub fn enable_realistic_mode_with_memory_reset(&mut self) { + self.root_knowledge.clear_deeper_than_one_hop(); + self.inspector_mode = InspectorMode::Realistic; + } + + /// Toggles the inspector between learned state and ground truth. + pub fn toggle_inspector_mode(&mut self) { + self.inspector_mode = match self.inspector_mode { + InspectorMode::GroundTruth => InspectorMode::Realistic, + InspectorMode::Realistic => InspectorMode::GroundTruth, + }; + } + + /// Returns whether the inspector is using learned state. + pub fn is_realistic_mode(&self) -> bool { + self.inspector_mode == InspectorMode::Realistic + } + /// Builds and routes an endpoint introspection call from the root. pub fn call_endpoint_introspection( &mut self, @@ -215,11 +479,18 @@ impl Simulation { leaf_name: &str, ) -> Result { let node_path = self.tree.node(node_id).path.clone(); - let node_display = self.tree.node(node_id).display_path(); self.require_leaf(node_id, leaf_name)?; + let node = self.tree.node(node_id).clone(); + if let Some(leaf_spec) = node.leaves.iter().find(|leaf| leaf.name == leaf_name) { + self.root_knowledge + .remember_leaf_from_spec(&node, leaf_spec); + } self.dispatch_root_call(node_path, Some(leaf_name.to_owned()), "", Vec::new())?; Ok(ActionResult { - label: format!("Inspect leaf {} on {}", leaf_name, node_display), + label: format!( + "Inspect {}", + format_leaf_ref(&self.node(node_id).path, leaf_name) + ), hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id), }) } @@ -233,9 +504,18 @@ impl Simulation { ) -> Result { let node_path = self.tree.node(node_id).path.clone(); let node_display = self.tree.node(node_id).display_path(); - let leaf = self.require_leaf(node_id, leaf_name)?; + let node = self.tree.node(node_id).clone(); + let procedures = self.require_leaf(node_id, leaf_name)?.procedures.clone(); + if let Some(leaf_spec) = node + .leaves + .iter() + .find(|known_leaf| known_leaf.name == leaf_name) + { + self.root_knowledge + .remember_leaf_from_spec(&node, leaf_spec); + } let procedure_id = - leaf.procedures + procedures .first() .cloned() .ok_or_else(|| SimError::UnknownProcedure { @@ -249,7 +529,10 @@ impl Simulation { text.as_bytes().to_vec(), )?; Ok(ActionResult { - label: format!("Echo via {leaf_name} on {}", node_display), + label: format!( + "Echo via {}", + format_leaf_ref(&self.node(node_id).path, leaf_name) + ), hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id), }) } @@ -264,6 +547,15 @@ impl Simulation { let node_path = self.tree.node(node_id).path.clone(); let node_display = self.tree.node(node_id).display_path(); self.require_endpoint_procedure(node_id, procedure_id)?; + let node = self.tree.node(node_id).clone(); + if let Some(procedure) = node + .endpoint_procedures + .iter() + .find(|known_procedure| known_procedure.procedure_id == procedure_id) + { + self.root_knowledge + .remember_endpoint_procedure(&node, procedure); + } self.dispatch_root_call(node_path, None, procedure_id, data)?; Ok(ActionResult { label: format!("Call {procedure_id} on {}", node_display), @@ -293,7 +585,10 @@ impl Simulation { }, node_display, dst_leaf - .map(|leaf_name| format!(" leaf {leaf_name}")) + .map(|leaf_name| format!( + " {}", + format_leaf_ref(&self.node(node_id).path, leaf_name) + )) .unwrap_or_default() ), hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id), @@ -324,7 +619,10 @@ impl Simulation { .map_err(|error| SimError::Protocol(error.to_string()))?; self.record_trace( self.root_id, - format!("root queued hook data for hook #{hook_id}: {text}"), + format!( + "root queued hook data for {}: {text}", + format_hook_ref(self.node(self.root_id).path.as_slice(), hook_id) + ), ); self.process_local_frame(self.root_id, frame)?; Ok(ActionResult { @@ -362,13 +660,17 @@ impl Simulation { self.record_trace( from_node_id, format!( - "injected invalid peer data toward {} for hook #{hook_id}", - format_path(&to_path) + "injected invalid peer data toward {} for {}", + format_path(&to_path), + format_hook_ref(self.node(to_node_id).path.as_slice(), hook_id) ), ); self.process_local_frame(from_node_id, frame)?; Ok(ActionResult { - label: format!("Inject invalid peer data for hook #{hook_id}"), + label: format!( + "Inject invalid peer data for {}", + format_hook_ref(self.node(to_node_id).path.as_slice(), hook_id) + ), hook_id: Some(hook_id), }) } @@ -432,6 +734,7 @@ impl Simulation { host_path: Vec::new(), peer_path: dst_path.clone(), procedure_id: procedure_id.to_owned(), + target_leaf: dst_leaf.clone(), closed: false, last_message: format!("created for {}", format_path(&dst_path)), }, @@ -448,7 +751,7 @@ impl Simulation { format_path(&dst_path), dst_leaf .as_ref() - .map(|leaf| format!(" leaf {leaf}")) + .map(|leaf| format!(" {}", format_leaf_ref(&dst_path, leaf))) .unwrap_or_default() ), ); @@ -543,8 +846,11 @@ impl Simulation { self.record_trace( node_id, format!( - "local Data on hook #{}: {text}", - header.hook_id.unwrap_or(0) + "local Data on {}: {text}", + format_hook_ref( + self.node(node_id).path.as_slice(), + header.hook_id.unwrap_or(0) + ) ), ); if let Some(hook_id) = header.hook_id { @@ -558,6 +864,10 @@ impl Simulation { snapshot.closed = true; } } + + if node_id == self.root_id { + self.learn_from_root_data(hook_id, &message); + } } if let Some(session) = self @@ -606,8 +916,11 @@ impl Simulation { self.record_trace( node_id, format!( - "local Fault on hook #{}: 0x{:02X}", - header.hook_id.unwrap_or(0), + "local Fault on {}: 0x{:02X}", + format_hook_ref( + self.node(node_id).path.as_slice(), + header.hook_id.unwrap_or(0) + ), message.fault.0 ), ); @@ -633,7 +946,7 @@ impl Simulation { header .dst_leaf .as_ref() - .map(|leaf| format!("leaf {leaf}")) + .map(|leaf| format_leaf_ref(&header.dst_path, leaf)) .unwrap_or_else(|| "endpoint".to_owned()) ), ); @@ -835,8 +1148,55 @@ impl Simulation { format!("{}: {}", node.display_path(), node.title) } Selection::Leaf { node_id, leaf_name } => { - format!("{} leaf {}", self.node(*node_id).display_path(), leaf_name) + format_leaf_ref(&self.node(*node_id).path, leaf_name) } } } + + fn learn_from_root_data(&mut self, hook_id: u64, message: &DataMessage) { + let Some(snapshot) = self.hooks.get(&hook_id).cloned() else { + return; + }; + let Some(node_id) = self.tree.find_by_path(&snapshot.peer_path) else { + return; + }; + let demo_node = self.node(node_id).clone(); + + if snapshot.procedure_id.is_empty() { + if snapshot.target_leaf.is_some() { + if let Ok(introspection) = deserialize_archived_bytes::< + unshell::protocol::introspection::ArchivedLeafIntrospection, + LeafIntrospection, + >(&message.data) + { + self.root_knowledge + .remember_leaf_introspection(&demo_node, &introspection); + } + } else if let Ok(introspection) = deserialize_archived_bytes::< + unshell::protocol::introspection::ArchivedEndpointIntrospection, + EndpointIntrospection, + >(&message.data) + { + self.root_knowledge + .remember_endpoint_introspection(&demo_node, &introspection); + } + return; + } + + if let Some(procedure) = demo_node + .endpoint_procedures + .iter() + .find(|procedure| procedure.procedure_id == snapshot.procedure_id) + { + self.root_knowledge + .remember_endpoint_procedure(&demo_node, procedure); + } + + if let Some(leaf_name) = &snapshot.target_leaf + && let Some(leaf_spec) = demo_node.leaves.iter().find(|leaf| &leaf.name == leaf_name) + { + self.root_knowledge + .remember_leaf_from_spec(&demo_node, leaf_spec); + } + } } diff --git a/treetest/tests/knowledge.rs b/treetest/tests/knowledge.rs new file mode 100644 index 0000000..741f75b --- /dev/null +++ b/treetest/tests/knowledge.rs @@ -0,0 +1,41 @@ +use treetest::{model::NodeId, scenarios::built_in_scenarios, sim::Simulation}; + +#[test] +fn realistic_mode_memory_reset_forgets_deeper_nodes() { + let scenarios = built_in_scenarios(); + let mut simulation = Simulation::new(scenarios[2].clone()).expect("scenario should build"); + + simulation + .call_echo_leaf(NodeId(2), "echo", "learn this nested leaf") + .expect("echo call should start"); + simulation.drain().expect("network should drain"); + + assert!( + simulation + .root_knowledge + .node(&simulation.node(NodeId(2)).path) + .is_some() + ); + + simulation.enable_realistic_mode_with_memory_reset(); + + assert!(simulation.is_realistic_mode()); + assert!( + simulation + .root_knowledge + .node(&simulation.node(NodeId(2)).path) + .is_none() + ); + assert!( + simulation + .root_knowledge + .node(&simulation.node(NodeId(1)).path) + .is_some() + ); + assert!( + simulation + .root_knowledge + .node(&simulation.node(NodeId(0)).path) + .is_some() + ); +}