diff --git a/base62/src/aes.rs b/base62/src/aes.rs index 83507a1..4a8cab5 100644 --- a/base62/src/aes.rs +++ b/base62/src/aes.rs @@ -3,71 +3,72 @@ use aes::cipher::{BlockModeDecrypt, BlockModeEncrypt, KeyIvInit}; use cbc::cipher::block_padding::Pkcs7; use regex::Regex; +/// Returns the next AES block-sized buffer length for PKCS#7 padding. fn pkcs7_padded_length(input_len: usize) -> usize { let block_size = 16; ((input_len / block_size) + 1) * block_size } +fn xor_fold(bytes: &[u8]) -> u8 { + bytes.iter().fold(0, |salt, &byte| salt ^ byte) +} + +fn xor_key_with_salt(key: &mut [u8; 32], salt: u8) { + for byte in key.iter_mut() { + *byte ^= salt; + } +} + +/// Encrypts `plaintext` with AES-256-CBC and base62-encodes the result. pub fn encrypt_aes(plaintext: &str, key_str: &str, iv: [u8; 16]) -> String { let plaintext = plaintext.as_bytes(); - // Hash the env key to get a 32-byte (256-bit) AES key + // Hash the environment key into the fixed-width AES-256 key material expected below. let key = hash(key_str.as_bytes()); - // Generate a pseudo-random salt byte based on the plaintext - // I hope this does not break the encryption. - let mut salt = 0; + // This crate intentionally perturbs the key with a single plaintext-derived byte so + // similar inputs diverge quickly after encryption while remaining reversible. + let salt = xor_fold(plaintext); - for byte in plaintext { - salt ^= byte; - } - - let mut key_salted = key.clone(); - - // Salt the key by XORing the salt byte with all the key bytes. - // This ensures that the "hash" generated from the plaintext will - // make the encrypted result extremely different. - for i in 0..32 { - key_salted[i] ^= salt; - } + let mut salted_key = key; + xor_key_with_salt(&mut salted_key, salt); let buf_len = pkcs7_padded_length(plaintext.len()); let mut buf = vec![0u8; buf_len]; let pt_len = plaintext.len(); - buf[..pt_len].copy_from_slice(&plaintext); + buf[..pt_len].copy_from_slice(plaintext); - let mut ct = cbc::Encryptor::::new(&key_salted.into(), &iv.into()) + let mut ciphertext = cbc::Encryptor::::new(&salted_key.into(), &iv.into()) .encrypt_padded::(&mut buf, pt_len) .unwrap() .to_vec(); - // Add the salt byte to the key byte, - ct.insert(0, salt); + // Prefix the salt so decryption can rebuild the same salted key. + ciphertext.insert(0, salt); - // Encode result in base62 - Base62::encode_full(&ct, &key) + Base62::encode_full(&ciphertext, &key) } +/// Wraps [`encrypt_aes`] output in `_..._` markers for line-wise replacement. pub fn encrypt_aes_lines(plaintext: &str, key_str: &str, iv: [u8; 16]) -> String { format!("_{}_", encrypt_aes(plaintext, key_str, iv)) } +/// Decodes base62 input and decrypts the payload with AES-256-CBC. pub fn decrypt_aes(input: &str, key_str: &str, iv: [u8; 16]) -> Result { - // Hash the env key to get a 32-byte (256-bit) AES key let mut key = hash(key_str.as_bytes()); - let mut cipher_bytes = Base62::decode_full(input, &key).unwrap(); + let mut cipher_bytes = Base62::decode_full(input, &key)?; + + if cipher_bytes.is_empty() { + return Err("decryption failed".to_string()); + } let salt = cipher_bytes.remove(0); - // XOR the salt bytes with the key bytes - // This replicates - for i in 0..32 { - key[i] ^= salt; - } + xor_key_with_salt(&mut key, salt); - // Create buffer for result let buf_len = cipher_bytes.len(); let mut buf: Vec = vec![0; buf_len]; buf[..cipher_bytes.len()].copy_from_slice(&cipher_bytes); @@ -79,26 +80,25 @@ pub fn decrypt_aes(input: &str, key_str: &str, iv: [u8; 16]) -> Result String { let mut decrypted_result = input.to_string(); - let mut total_offset = 0; + let mut replacement_offset: isize = 0; - // Split input by segments of base62 chars, denoted by two _'s, and attempt to decode - for aes_block in Regex::new(r"_([0-9a-zA-Z]*?)_").unwrap().find_iter(&input) { + // Walk the original input and patch into a separate mutable buffer. The offset keeps the + // current match aligned after prior replacements change the string length. + for aes_block in Regex::new(r"_([0-9a-zA-Z]*?)_").unwrap().find_iter(input) { let range = aes_block.range(); let aes_block = aes_block.as_str()[1..(aes_block.len() - 1)].to_string(); - // If the decryption is successful, offset the current offset position if let Ok(decrypted_block) = decrypt_aes(&aes_block, key_str, iv) { - let range = (range.start + total_offset as usize)..(range.end + total_offset as usize); + let adjusted_start = (range.start as isize + replacement_offset) as usize; + let adjusted_end = (range.end as isize + replacement_offset) as usize; + let adjusted_range = adjusted_start..adjusted_end; - // Offset range by the difference between the decrypted block length and the original range length - total_offset += decrypted_block.len().clone() - (range.end - range.start); + replacement_offset += decrypted_block.len() as isize - adjusted_range.len() as isize; - decrypted_result.replace_range(range, &decrypted_block); - } else { - // If the decode is unsuccessful, leave the underscore-denoted region as is - continue; + decrypted_result.replace_range(adjusted_range, &decrypted_block); } } diff --git a/base62/src/base62.rs b/base62/src/base62.rs index 902b96f..a270c12 100644 --- a/base62/src/base62.rs +++ b/base62/src/base62.rs @@ -1,6 +1,6 @@ use crate::{STATIC_BYTE_MAP, hash}; -// Randomly mapped Base62 characters +/// Base-62 encoder/decoder with a deterministic per-key character permutation. pub struct Base62 { charset: [char; 62], } @@ -12,34 +12,32 @@ pub const BASE62_CHARS: [char; 62] = [ 'v', 'w', 'x', 'y', 'z', ]; -// Const for ratio +/// `8.0 / log2(62.0)`, used to estimate encoded length from a byte length. const ENCODING_RATIO: f64 = 8.0 / 5.954196310386875; // 8.0 / log2(62.0) impl Base62 { + /// Builds the charset permutation for `key` and `nonce`. pub fn new(key: &[u8], nonce: usize) -> Self { - // Hash key again, for the chance that this random function can be used to derive the key - // My solution to not being good at cryptography lol + // Re-hash the caller-provided key so charset generation always runs on a fixed-width input. let key = hash(key); let mut charset: [char; 62] = [0 as char; 62]; - // Create a vector of indices from 0 to 61 - let mut current_indices = (0..62).map(|i| i as usize).collect::>(); + let mut available_positions = (0..62).collect::>(); - // Loop through each byte in the key until all chars are filled - for i in 0..62 as usize { - let rand = STATIC_BYTE_MAP[(key[i as usize % key.len()] as usize + nonce) % 255]; + for (char_index, ch) in BASE62_CHARS.iter().copied().enumerate() { + let random_byte = STATIC_BYTE_MAP[(key[char_index % key.len()] as usize + nonce) % 255]; - let index_index = rand as usize % current_indices.len(); - let put_index = current_indices.remove(index_index); + let choice_index = random_byte as usize % available_positions.len(); + let charset_index = available_positions.remove(choice_index); - charset[put_index] = BASE62_CHARS[i]; + charset[charset_index] = ch; } - return Self { charset }; + Self { charset } } - // Convert character to base-62 value using custom charset + /// Converts a character to its base-62 value in this instance's charset. fn char_to_value(&self, ch: char) -> Result { self.charset .iter() @@ -49,7 +47,7 @@ impl Base62 { } /// Encodes a byte slice into a base-62 string using a custom character set - /// Supports arbitrary length input by using big integer arithmetic + /// while preserving leading zero bytes. pub fn encode(&self, data: &[u8]) -> String { if data.is_empty() { return String::new(); @@ -68,7 +66,7 @@ impl Base62 { let mut result = Vec::new(); let mut num = data.to_vec(); - // Convert to base-62 using division + // Repeated division keeps the implementation independent from bigint crates. while !is_zero(&num) { let remainder = div_mod_62(&mut num); result.push(self.charset[remainder]); @@ -85,7 +83,7 @@ impl Base62 { } /// Decodes a base-62 string back into bytes using a custom character set - /// Supports arbitrary length output + /// while preserving leading zero bytes. pub fn decode(&self, encoded: &str) -> Result, String> { if encoded.is_empty() { return Ok(Vec::new()); @@ -102,7 +100,7 @@ impl Base62 { return Ok(vec![0; leading_zeros]); } - // Convert base-62 string to bytes using multiplication + // Rebuild the big-endian integer via repeated multiply-add. let mut num = vec![0u8]; for ch in encoded.chars() { @@ -117,42 +115,41 @@ impl Base62 { Ok(result) } + /// Encodes `data` using the nonce convention shared with [`decode_full`]. pub fn encode_full(data: &[u8], key: &[u8]) -> String { - // Predict the length of the encoded data - let length = predict_base62_len(data); + let predicted_len = predict_base62_len(data); - let base = Base62::new(&key, length % 255); + let base = Base62::new(key, predicted_len % 255); let encoded = base.encode(data); - // For the case that the encoded length is not equal to the predicted length - // The nonce must be derived from this length, so this needs to be ensured - // - // Re-encode with the correct length - if encoded.len() != length { - let len = encoded.len(); - let base = Base62::new(&key, len % 255); + // The charset nonce is derived from the final encoded length, so a misprediction must + // trigger one more pass with the actual length-derived nonce. + if encoded.len() != predicted_len { + let actual_len = encoded.len(); + let base = Base62::new(key, actual_len % 255); let encoded = base.encode(data); - assert_eq!(encoded.len(), len); + assert_eq!(encoded.len(), actual_len); encoded } else { encoded } } + + /// Decodes a string previously produced by [`encode_full`]. pub fn decode_full(data: &str, key: &[u8]) -> Result, String> { - let base = Base62::new(&key, data.len() % 255); + let base = Base62::new(key, data.len() % 255); base.decode(data) } } -// Helper: Check if big integer (as bytes) is zero +/// Returns whether the big-endian integer represented by `num` is zero. fn is_zero(num: &[u8]) -> bool { num.iter().all(|&b| b == 0) } -// Helper: Divide big integer by 62 and return remainder -// Modifies num in place to be the quotient +/// Divides an in-place big-endian integer by `62`, returning the remainder. fn div_mod_62(num: &mut Vec) -> usize { let mut remainder = 0u16; let mut all_zero = true; @@ -166,7 +163,7 @@ fn div_mod_62(num: &mut Vec) -> usize { } } - // Remove leading zeros from quotient + // Keep a canonical representation so the next loop iteration can stop at `[0]`. if all_zero { num.clear(); num.push(0); @@ -180,8 +177,7 @@ fn div_mod_62(num: &mut Vec) -> usize { remainder as usize } -// Helper: Multiply big integer by 62 and add a value -// Modifies num in place +/// Multiplies an in-place big-endian integer by `multiplier` and adds `add`. fn mul_add(num: &mut Vec, multiplier: u16, add: u8) { let mut carry = add as u16; @@ -191,7 +187,6 @@ fn mul_add(num: &mut Vec, multiplier: u16, add: u8) { carry = product >> 8; } - // Add remaining carry bytes while carry > 0 { num.insert(0, (carry & 0xFF) as u8); carry >>= 8; @@ -205,22 +200,13 @@ pub fn predict_base62_len(input_bytes: &[u8]) -> usize { return 0; } - // 1. Count leading zero bytes. let num_leading_zeros = input_bytes.iter().take_while(|&&b| b == 0).count(); - - // 2. Calculate length of the rest of the bytes. let num_rest_bytes = input_bytes.len() - num_leading_zeros; if num_rest_bytes == 0 { - // If all bytes were zeros, the length is just the number of zeros. num_leading_zeros } else { - // 3. Calculate the mathematical upper bound for the non-zero part. - // This is ceil(num_rest_bytes * 8_bits / log2(62)) - // which is ceil(num_rest_bytes * log_62(256)) let rest_len = (num_rest_bytes as f64 * ENCODING_RATIO).ceil(); - - // 4. Total length is zeros + rest_len num_leading_zeros + rest_len as usize } } diff --git a/base62/src/lib.rs b/base62/src/lib.rs index f1ed8b3..e7aac55 100644 --- a/base62/src/lib.rs +++ b/base62/src/lib.rs @@ -1,10 +1,12 @@ +//! Base62 encoding helpers plus the AES wrapper used by `ush-obfuscate`. + mod aes; mod base62; -// Exports pub use aes::{decrypt_aes, decrypt_aes_lines, encrypt_aes, encrypt_aes_lines}; pub use base62::Base62; +/// Static IV shared by the proc-macro crate and the runtime decoder. pub const STATIC_IV: [u8; 16] = [ 0x6d, 0x79, 0x5f, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x69, 0x76, 0x5f, 0x30, 0x31, 0x32, ]; @@ -27,12 +29,14 @@ pub const STATIC_BYTE_MAP: [u8; 256] = [ use sha2::{Digest, Sha256}; +/// Returns the SHA-256 digest of `input`. pub fn hash(input: &[u8]) -> [u8; 32] { let mut hasher = Sha256::new(); hasher.update(input); hasher.finalize().into() } +/// Encodes a `usize` as a big-endian byte slice without leading zero bytes. pub fn encode_usize(value: usize) -> Vec { if value == 0 { return vec![0]; @@ -42,6 +46,9 @@ pub fn encode_usize(value: usize) -> Vec { bytes[leading_zeros..].to_vec() } +/// Decodes a big-endian `usize` previously produced by [`encode_usize`]. +/// +/// The caller must ensure `bytes.len() <= size_of::()`. pub fn decode_usize(bytes: &[u8]) -> usize { let mut buf = [0u8; size_of::()]; let offset = buf.len() - bytes.len(); diff --git a/src/logger/global.rs b/src/logger/global.rs index 3324101..6cd2a0f 100644 --- a/src/logger/global.rs +++ b/src/logger/global.rs @@ -58,6 +58,8 @@ pub fn set_logger(logger: &'static dyn Logger) { } /// Returns the currently installed global logger. +/// +/// Until [`set_logger`] runs, this returns the internal null sink. #[must_use] pub fn global_logger() -> &'static dyn Logger { GLOBAL_LOGGER.get() @@ -66,6 +68,8 @@ pub fn global_logger() -> &'static dyn Logger { /// Sends a single record through the installed global logger. /// /// Most code should prefer the exported logging macros. +/// This helper exists for integrations that already have a preformatted message +/// and explicit source context. pub fn log(level: LogLevel, message: &str, file: Option<&'static str>, line: Option) { global_logger().log(&Record::new(level, message, file, line)); } diff --git a/src/logger/sink.rs b/src/logger/sink.rs index 7fb3614..5648680 100644 --- a/src/logger/sink.rs +++ b/src/logger/sink.rs @@ -25,6 +25,12 @@ impl CompatibilityLogger { Self { min_level } } + /// Returns the minimum severity that passes the filter. + #[must_use] + pub const fn min_level(&self) -> LogLevel { + self.min_level + } + /// Returns whether a record at `level` would be accepted. /// /// # Examples diff --git a/src/protocol/codec.rs b/src/protocol/codec.rs index 3f0f462..4968497 100644 --- a/src/protocol/codec.rs +++ b/src/protocol/codec.rs @@ -51,30 +51,44 @@ pub struct ParsedFrame<'a> { } impl<'a> ParsedFrame<'a> { + /// Returns the deserialized packet header. + /// + /// The header is owned by `ParsedFrame` because decoding must validate it + /// before any routing decision is made. + #[must_use] pub fn header(&self) -> &PacketHeader { &self.header } + /// Returns the header packet type for quick dispatch. + #[must_use] pub fn packet_type(&self) -> PacketType { self.header.packet_type } + /// Returns the raw archived payload section. + #[must_use] pub fn payload_bytes(&self) -> &'a [u8] { self.payload_bytes } + /// Clones the decoded header out of the parsed frame. + #[must_use] pub fn deserialize_header(&self) -> PacketHeader { self.header.clone() } + /// Deserializes the payload as a [`CallMessage`]. pub fn deserialize_call(&self) -> Result { deserialize_archived_bytes::(self.payload_bytes) } + /// Deserializes the payload as a [`DataMessage`]. pub fn deserialize_data(&self) -> Result { deserialize_archived_bytes::(self.payload_bytes) } + /// Deserializes the payload as a [`FaultMessage`]. pub fn deserialize_fault(&self) -> Result { deserialize_archived_bytes::(self.payload_bytes) } @@ -211,10 +225,9 @@ where } fn align_section(bytes: &[u8]) -> AlignedVec { - if bytes.as_ptr().align_offset(16) == 0 { - // Still need to return AlignedVec for the API, but maybe we can avoid - // some overhead. Actually, AlignedVec is just a wrapper around Vec. - } + // The framed wire format prefixes each archived section with a 4-byte length, + // so callers cannot rely on the borrowed slice meeting rkyv's alignment. + // Copying into `AlignedVec` keeps the alignment fix local and predictable. let mut aligned = AlignedVec::with_capacity(bytes.len()); aligned.extend_from_slice(bytes); aligned diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index f2d8229..a7ae061 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -22,6 +22,10 @@ pub use types::{ }; pub use validation::{ValidationError, validate_call, validate_header, validate_procedure_id}; +/// Encodes a header and payload with the crate's default frame codec. +/// +/// This is a convenience wrapper around [`RkyvCodec`] for callers that do not +/// need to choose a codec explicitly. pub fn encode_packet

(header: &PacketHeader, payload: &P) -> Result where P: for<'a> rkyv::Serialize< @@ -35,6 +39,7 @@ where codec::encode_packet(header, payload) } +/// Decodes a framed packet with the crate's default frame codec. pub fn decode_frame(bytes: &[u8]) -> Result, FrameError> { codec::decode_frame(bytes) } diff --git a/src/protocol/traits.rs b/src/protocol/traits.rs index 2ad31d5..6f66bb1 100644 --- a/src/protocol/traits.rs +++ b/src/protocol/traits.rs @@ -25,14 +25,31 @@ impl RouteResolution for T where T: RouteProvider + ?Sized {} /// Hook storage contract for pending and active protocol flows. pub trait HookStore { + /// Allocates a hook identifier scoped to `return_path`. fn allocate_hook_id(&mut self, return_path: &[String]) -> u64; + + /// Inserts a hook created by an incoming call before the peer is confirmed. fn insert_pending(&mut self, pending: PendingHook) -> Result<(), HookConflict>; + + /// Inserts an already-established hook flow. fn insert_active(&mut self, active: ActiveHook) -> Result<(), HookConflict>; + + /// Promotes a pending hook once the responding peer is known. fn activate_pending(&mut self, key: &HookKey, peer_path: Vec) -> Option<()>; + + /// Removes a pending hook. fn remove_pending(&mut self, key: &HookKey) -> Option; + + /// Removes an active hook. fn remove_active(&mut self, key: &HookKey) -> Option; + + /// Returns immutable access to a pending hook. fn pending(&self, key: &HookKey) -> Option<&PendingHook>; + + /// Returns immutable access to an active hook. fn active(&self, key: &HookKey) -> Option<&ActiveHook>; + + /// Returns mutable access to an active hook. fn active_mut(&mut self, key: &HookKey) -> Option<&mut ActiveHook>; } @@ -76,9 +93,13 @@ impl HookStore for HookTable { /// Leaf metadata contract used for protocol discovery payloads. pub trait LeafMetadata { + /// Returns the leaf name exposed in routing and introspection. fn leaf_name(&self) -> &str; + + /// Returns the supported canonical procedure identifiers. fn procedures(&self) -> &[String]; + /// Builds the compact endpoint-wide discovery record for this leaf. fn summary(&self) -> LeafIntrospectionSummary { LeafIntrospectionSummary { leaf_name: self.leaf_name().into(), @@ -86,6 +107,7 @@ pub trait LeafMetadata { } } + /// Builds the full leaf-specific discovery payload. fn introspection(&self) -> LeafIntrospection { LeafIntrospection { leaf_name: self.leaf_name().into(), @@ -116,7 +138,10 @@ impl LeafMetadata for LeafNode { /// Packet processor and local runtime contract for framed protocol traffic. pub trait PacketProcessor { + /// Returns the endpoint path that owns this processor. fn path(&self) -> &[String]; + + /// Receives one framed packet from the given ingress side. fn receive( &mut self, ingress: &Ingress, diff --git a/src/protocol/tree/endpoint/builders.rs b/src/protocol/tree/endpoint/builders.rs index 57f89b7..e265f05 100644 --- a/src/protocol/tree/endpoint/builders.rs +++ b/src/protocol/tree/endpoint/builders.rs @@ -23,6 +23,7 @@ impl ProtocolEndpoint { /// let endpoint = ProtocolEndpoint::new(Vec::new(), None, Vec::new(), Vec::new()); /// assert!(endpoint.path().is_empty()); /// ``` + #[must_use] pub fn new( path: Vec, parent_path: Option>, @@ -54,6 +55,7 @@ impl ProtocolEndpoint { } /// Allocates a locally unique hook id. + #[must_use] pub fn allocate_hook_id(&mut self) -> u64 { self.hooks.allocate_hook_id(&self.path) } diff --git a/src/protocol/tree/endpoint/core.rs b/src/protocol/tree/endpoint/core.rs index 1831477..7bdc086 100644 --- a/src/protocol/tree/endpoint/core.rs +++ b/src/protocol/tree/endpoint/core.rs @@ -21,19 +21,24 @@ use super::super::{HookTable, RouteDecision}; /// Local connection state used for child route eligibility. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ConnectionState { + /// The child exists in the static topology but is not currently routable. Unregistered, + /// The child may receive routed traffic. Registered, } /// Child path plus current registration state. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ChildRoute { + /// Absolute child endpoint path. pub path: Vec, + /// Whether the child currently participates in routing. pub state: ConnectionState, } impl ChildRoute { /// Convenience constructor for the common registered-child case. + #[must_use] pub fn registered(path: Vec) -> Self { Self { path, @@ -45,14 +50,18 @@ impl ChildRoute { /// Test leaf behavior implemented by the endpoint runtime. #[derive(Debug, Clone, PartialEq, Eq)] pub enum LeafBehavior { + /// Mirrors the incoming payload back over the declared response hook. Echo, } /// Static leaf metadata used for procedure dispatch and introspection. #[derive(Debug, Clone, PartialEq, Eq)] pub struct LeafSpec { + /// Stable local leaf name. pub name: String, + /// Procedures supported by the leaf. pub procedures: Vec, + /// Built-in behavior used by the lightweight test runtime. pub behavior: LeafBehavior, } @@ -62,22 +71,28 @@ pub struct LeafSpec { /// `PROTOCOL.md` routing and call sections. #[derive(Debug, Clone, PartialEq, Eq)] pub enum Ingress { + /// Received from the parent link. Parent, + /// Received from the child at the given absolute path. Child(Vec), + /// Injected locally by code running on this endpoint. Local, } /// Locally delivered protocol events. #[derive(Debug, Clone, PartialEq, Eq)] pub enum LocalEvent { + /// A call reached this endpoint runtime. Call { header: PacketHeader, message: CallMessage, }, + /// Hook data reached this endpoint runtime. Data { header: PacketHeader, message: DataMessage, }, + /// A protocol fault reached this endpoint runtime. Fault { header: PacketHeader, message: FaultMessage, @@ -87,15 +102,20 @@ pub enum LocalEvent { /// Result of processing one framed packet. #[derive(Debug, Default)] pub struct EndpointOutcome { + /// Forwarding actions to perform after local processing. pub forwards: Vec<(RouteDecision, FrameBytes)>, + /// Events delivered to local runtime consumers. pub events: Vec, + /// Whether the packet was intentionally dropped with no other side effects. pub dropped: bool, } /// Errors returned while decoding or validating a packet. #[derive(Debug)] pub enum EndpointError { + /// The frame could not be decoded. Frame(FrameError), + /// The decoded packet violated protocol invariants. Validation(ValidationError), } @@ -124,7 +144,10 @@ impl From for EndpointError { /// Public packet-processing trait exposed by the tree runtime. pub trait Endpoint { + /// Returns the absolute endpoint path. fn path(&self) -> &[String]; + + /// Processes one incoming frame from the given ingress side. fn receive( &mut self, ingress: &Ingress, diff --git a/src/protocol/tree/endpoint/hooks.rs b/src/protocol/tree/endpoint/hooks.rs index 23e3adc..8f14b95 100644 --- a/src/protocol/tree/endpoint/hooks.rs +++ b/src/protocol/tree/endpoint/hooks.rs @@ -61,6 +61,9 @@ impl ProtocolEndpoint { message: DataMessage, ) -> Result { let hook_id = header.hook_id.expect("validated"); + // The hook host can address its hook directly with `self.path + hook_id`. + // A non-host peer only knows the hook id it was given earlier, so it must + // recover the host-scoped key from active state using its validated path. let key = self .hooks .active(&HookKey::new(self.path.clone(), hook_id)) diff --git a/src/protocol/tree/hook.rs b/src/protocol/tree/hook.rs index aa469fc..a272fea 100644 --- a/src/protocol/tree/hook.rs +++ b/src/protocol/tree/hook.rs @@ -12,6 +12,8 @@ pub struct HookKey { } impl HookKey { + /// Creates a host-scoped key from the return path and hook identifier. + #[must_use] pub fn new(return_path: Vec, hook_id: u64) -> Self { Self { return_path, @@ -64,12 +66,18 @@ impl Default for HookTable { } impl HookTable { + /// Allocates the next locally unique hook identifier. + /// + /// Hook IDs are scoped by return path, so this counter only needs to be + /// unique within one endpoint runtime. + #[must_use] pub fn allocate_hook_id(&mut self, _return_path: &[String]) -> u64 { let id = self.next_id; self.next_id = self.next_id.wrapping_add(1); id } + /// Inserts a pending hook created by an inbound call. pub fn insert_pending(&mut self, pending: PendingHook) -> Result<(), HookConflict> { let key = HookKey::new(pending.return_path.clone(), pending.hook_id); if self.pending.contains_key(&key) || self.active.contains_key(&key) { @@ -79,6 +87,7 @@ impl HookTable { Ok(()) } + /// Inserts an already-active hook flow. pub fn insert_active(&mut self, active: ActiveHook) -> Result<(), HookConflict> { let key = HookKey::new(active.return_path.clone(), active.hook_id); if self.pending.contains_key(&key) || self.active.contains_key(&key) { @@ -88,6 +97,7 @@ impl HookTable { Ok(()) } + /// Promotes a pending hook into the active table once its peer is known. pub fn activate_pending(&mut self, key: &HookKey, peer_path: Vec) -> Option<()> { let pending = self.pending.remove(key)?; self.active.insert( @@ -104,22 +114,29 @@ impl HookTable { Some(()) } + /// Removes a pending hook entry. pub fn remove_pending(&mut self, key: &HookKey) -> Option { self.pending.remove(key) } + /// Removes an active hook entry. pub fn remove_active(&mut self, key: &HookKey) -> Option { self.active.remove(key) } + /// Returns a pending hook by its host-scoped key. + #[must_use] pub fn pending(&self, key: &HookKey) -> Option<&PendingHook> { self.pending.get(key) } + /// Returns an active hook by its host-scoped key. + #[must_use] pub fn active(&self, key: &HookKey) -> Option<&ActiveHook> { self.active.get(key) } + /// Returns mutable access to an active hook by its host-scoped key. pub fn active_mut(&mut self, key: &HookKey) -> Option<&mut ActiveHook> { self.active.get_mut(key) } @@ -130,23 +147,27 @@ impl HookTable { /// cannot derive the full key from the packet header alone. The peer uses /// its already-validated active state to recover the host-scoped key. pub fn find_active_key_by_peer(&self, hook_id: u64, peer_path: &[String]) -> Option { - let mut matches = self + let mut matching_keys = self .active .iter() .filter(|(_key, active)| active.hook_id == hook_id && active.peer_path == peer_path) .map(|(key, _)| key.clone()); - let first = matches.next()?; - if matches.next().is_some() { + let key = matching_keys.next()?; + if matching_keys.next().is_some() { return None; } - Some(first) + Some(key) } + /// Returns the number of pending hooks. + #[must_use] pub fn pending_len(&self) -> usize { self.pending.len() } + /// Returns the number of active hooks. + #[must_use] pub fn active_len(&self) -> usize { self.active.len() } diff --git a/src/protocol/validation.rs b/src/protocol/validation.rs index 29a983b..08a52c4 100644 --- a/src/protocol/validation.rs +++ b/src/protocol/validation.rs @@ -32,6 +32,9 @@ impl fmt::Display for ValidationError { impl core::error::Error for ValidationError {} /// Validates packet header invariants from the protocol. +/// +/// This checks only the header fields themselves. Payload-dependent rules belong +/// in helpers such as [`validate_call`]. pub fn validate_header(header: &PacketHeader) -> Result<(), ValidationError> { match header.packet_type { PacketType::Call => { diff --git a/treetest/src/app/actions.rs b/treetest/src/app/actions.rs index 09ea098..f289dc3 100644 --- a/treetest/src/app/actions.rs +++ b/treetest/src/app/actions.rs @@ -8,6 +8,19 @@ use super::{App, AppError, NodeId, Selection}; impl App { + /// Drains queued simulator work, refreshes the visible selection list, and + /// reports the action result in the footer. + fn finish_action( + &mut self, + preferred_node: Option, + label: impl Into, + ) -> Result<(), AppError> { + let steps = self.simulation.drain()?; + self.refresh_selections(preferred_node); + self.status = format!("{} ({steps} steps)", label.into()); + Ok(()) + } + /// Performs protocol introspection for the current selection. /// /// Rationale: node and leaf introspection share one key because the protocol @@ -18,18 +31,14 @@ impl App { // Route the blank procedure to endpoint-wide introspection. let result = self.simulation.call_endpoint_introspection(node_id)?; // Drain immediately so the inspector reflects the learned state. - let steps = self.simulation.drain()?; - self.refresh_selections(Some(node_id)); - self.status = format!("{} ({steps} steps)", result.label); + self.finish_action(Some(node_id), result.label)?; } Selection::Leaf { node_id, leaf_name } => { // Route the blank procedure to one specific leaf. let result = self .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); + self.finish_action(Some(node_id), result.label)?; } } Ok(()) @@ -40,102 +49,99 @@ impl App { /// Rationale: the payload is fixed so the demo highlights packet flow rather /// than turning the TUI into a line editor. pub(super) fn perform_echo(&mut self) -> Result<(), AppError> { - if let Selection::Leaf { node_id, leaf_name } = self.selected().clone() { - let result = - 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 { + let Selection::Leaf { node_id, leaf_name } = self.selected().clone() else { self.status = "Select a leaf first, then press e.".to_owned(); - } + return Ok(()); + }; + + let result = self + .simulation + .call_echo_leaf(node_id, &leaf_name, "demo echo from root")?; + self.finish_action(Some(node_id), result.label)?; Ok(()) } /// Calls the first endpoint-level procedure on the selected node. pub(super) fn perform_ping(&mut self) -> Result<(), AppError> { - if let Selection::Node(node_id) = self.selected().clone() { - if let Some(procedure_id) = self - .simulation - .node(node_id) - .endpoint_procedures - .first() - .map(|procedure| procedure.procedure_id.clone()) - { - let result = self.simulation.call_endpoint_procedure( - node_id, - &procedure_id, - 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(); - } - } else { + let Selection::Node(node_id) = self.selected().clone() else { self.status = "Select a node first, then press p.".to_owned(); - } + return Ok(()); + }; + + let Some(procedure_id) = self + .simulation + .node(node_id) + .endpoint_procedures + .first() + .map(|procedure| procedure.procedure_id.clone()) + else { + self.status = "Selected node has no endpoint procedures.".to_owned(); + return Ok(()); + }; + + let result = + self.simulation + .call_endpoint_procedure(node_id, &procedure_id, b"ping".to_vec())?; + self.finish_action(Some(node_id), result.label)?; Ok(()) } /// Calls the chunked-response procedure on the selected node. pub(super) fn perform_chunked(&mut self) -> Result<(), AppError> { - if let Selection::Node(node_id) = self.selected().clone() { - if let Some(procedure_id) = self - .simulation - .node(node_id) - .endpoint_procedures - .iter() - .find(|procedure| { - procedure.description.contains("chunk") - || procedure.procedure_id.contains("chunked") - }) - .map(|procedure| procedure.procedure_id.clone()) - { - let result = self.simulation.call_endpoint_procedure( - node_id, - &procedure_id, - 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(); - } - } else { + let Selection::Node(node_id) = self.selected().clone() else { self.status = "Select a node first, then press c.".to_owned(); - } + return Ok(()); + }; + + let Some(procedure_id) = self + .simulation + .node(node_id) + .endpoint_procedures + .iter() + .find(|procedure| { + procedure.description.contains("chunk") + || procedure.procedure_id.contains("chunked") + }) + .map(|procedure| procedure.procedure_id.clone()) + else { + self.status = "Selected node has no chunked procedure.".to_owned(); + return Ok(()); + }; + + let result = self.simulation.call_endpoint_procedure( + node_id, + &procedure_id, + b"chunk please".to_vec(), + )?; + self.finish_action(Some(node_id), result.label)?; Ok(()) } /// Opens a long-lived chat hook on the selected node. pub(super) fn perform_chat_call(&mut self) -> Result<(), AppError> { - if let Selection::Node(node_id) = self.selected().clone() { - if let Some(procedure_id) = self - .simulation - .node(node_id) - .endpoint_procedures - .iter() - .find(|procedure| procedure.procedure_id.contains("chat")) - .map(|procedure| procedure.procedure_id.clone()) - { - let result = self.simulation.call_endpoint_procedure( - node_id, - &procedure_id, - 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(); - } - } else { + let Selection::Node(node_id) = self.selected().clone() else { self.status = "Select a node first, then press h.".to_owned(); - } + return Ok(()); + }; + + let Some(procedure_id) = self + .simulation + .node(node_id) + .endpoint_procedures + .iter() + .find(|procedure| procedure.procedure_id.contains("chat")) + .map(|procedure| procedure.procedure_id.clone()) + else { + self.status = "Selected node has no chat procedure.".to_owned(); + return Ok(()); + }; + + let result = self.simulation.call_endpoint_procedure( + node_id, + &procedure_id, + b"open chat".to_vec(), + )?; + self.finish_action(Some(node_id), result.label)?; Ok(()) } @@ -144,57 +150,54 @@ impl App { /// Rationale: using the latest hook keeps the demo simple while still /// exposing bidirectional hook behavior. pub(super) fn perform_chat_data(&mut self) -> Result<(), AppError> { - if let Some(hook_id) = self.simulation.hook_ids().last().copied() { - let result = - 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 { + let Some(hook_id) = self.simulation.hook_ids().last().copied() else { self.status = "No known hook yet. Press h to open chat first.".to_owned(); - } + return Ok(()); + }; + + let result = self + .simulation + .send_root_hook_data(hook_id, "hello from the root", false)?; + self.finish_action(None, result.label)?; Ok(()) } /// Ends the newest known chat hook from the root side. pub(super) fn perform_chat_bye(&mut self) -> Result<(), AppError> { - 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 { + let Some(hook_id) = self.simulation.hook_ids().last().copied() else { self.status = "No known hook yet. Press h to open chat first.".to_owned(); - } + return Ok(()); + }; + + let result = self.simulation.send_root_hook_data(hook_id, "bye", true)?; + self.finish_action(None, result.label)?; Ok(()) } /// Injects intentionally invalid hook data to exercise fault handling. pub(super) fn perform_invalid_fault_demo(&mut self) -> Result<(), AppError> { - if let Some(hook_id) = self.simulation.hook_ids().last().copied() { - // The root is always node zero in every built-in scenario. - let root_id = NodeId(0); - if self.simulation.tree.nodes.len() > 1 { - // The first child is enough to spoof a wrong peer path. - let attacker = NodeId(1); - let result = self.simulation.inject_invalid_peer_data( - attacker, - root_id, - hook_id, - "demo.endpoint.v1.chat.session", - "spoofed data", - )?; - let steps = self.simulation.drain()?; - self.refresh_selections(None); - self.status = format!("{} ({steps} steps)", result.label); - } else { - self.status = - "This scenario has no second node for invalid-peer traffic.".to_owned(); - } - } else { + let Some(hook_id) = self.simulation.hook_ids().last().copied() else { self.status = "Open a hook first before injecting invalid traffic.".to_owned(); + return Ok(()); + }; + + if self.simulation.tree.nodes.len() <= 1 { + self.status = "This scenario has no second node for invalid-peer traffic.".to_owned(); + return Ok(()); } + + // The root is always node zero in every built-in scenario. + let root_id = NodeId(0); + // The first child is enough to spoof a wrong peer path. + let attacker = NodeId(1); + let result = self.simulation.inject_invalid_peer_data( + attacker, + root_id, + hook_id, + "demo.endpoint.v1.chat.session", + "spoofed data", + )?; + self.finish_action(None, result.label)?; Ok(()) } } diff --git a/treetest/src/app/mod.rs b/treetest/src/app/mod.rs index 8ed009d..174ed1c 100644 --- a/treetest/src/app/mod.rs +++ b/treetest/src/app/mod.rs @@ -11,7 +11,7 @@ mod ui; use ratatui::DefaultTerminal; use crate::{ - model::{NodeId, Selection}, + model::{NodeId, ScenarioDefinition, Selection}, scenarios::built_in_scenarios, sim::Simulation, }; @@ -19,8 +19,10 @@ use crate::{ /// Errors returned by the TUI application. #[derive(Debug, thiserror::Error)] pub enum AppError { + /// Terminal setup, teardown, or input/output failed. #[error(transparent)] Io(#[from] std::io::Error), + /// The simulator rejected an operation or could not advance. #[error(transparent)] Sim(#[from] crate::sim::SimError), } @@ -32,7 +34,7 @@ pub fn run() -> Result<(), AppError> { #[derive(Debug)] struct App { - scenarios: Vec, + scenarios: Vec, scenario_index: usize, simulation: Simulation, selection_index: usize, diff --git a/treetest/src/app/shell.rs b/treetest/src/app/shell.rs index 433268a..c6d6419 100644 --- a/treetest/src/app/shell.rs +++ b/treetest/src/app/shell.rs @@ -68,11 +68,19 @@ impl App { terminal.draw(|frame| self.render(frame))?; // Poll with a timeout so redraws stay responsive without busy-spinning. - if event::poll(Duration::from_millis(100))? - && let Event::Key(key) = event::read()? - && key.kind == KeyEventKind::Press - && !self.handle_key(key.code)? - { + if !event::poll(Duration::from_millis(100))? { + continue; + } + + let Event::Key(key) = event::read()? else { + continue; + }; + + if key.kind != KeyEventKind::Press { + continue; + } + + if !self.handle_key(key.code)? { break; } } @@ -176,14 +184,14 @@ impl App { /// so selection repair needs to happen in one dedicated place. pub(super) fn refresh_selections(&mut self, preferred_node: Option) { // Prefer an explicit node if the caller knows what should stay selected. - let current = preferred_node.unwrap_or_else(|| self.selected().node_id()); + let selected_node_id = preferred_node.unwrap_or_else(|| self.selected().node_id()); self.selections = ui::build_selections(&self.simulation); // Fall back to the first row when the previous node disappeared. self.selection_index = self .selections .iter() - .position(|selection| selection.node_id() == current) + .position(|selection| selection.node_id() == selected_node_id) .unwrap_or(0); } } diff --git a/treetest/src/app/ui/panels/chrome.rs b/treetest/src/app/ui/panels/chrome.rs index d583870..5a55790 100644 --- a/treetest/src/app/ui/panels/chrome.rs +++ b/treetest/src/app/ui/panels/chrome.rs @@ -20,7 +20,7 @@ impl App { pub(crate) fn render(&self, frame: &mut Frame<'_>) { // Split the screen into a small header, a large working area, and a // persistent status/footer region. - let chunks = Layout::default() + let rows = Layout::default() .direction(Direction::Vertical) .constraints([ Constraint::Length(3), @@ -29,9 +29,9 @@ impl App { ]) .split(frame.area()); - self.render_header(frame, chunks[0]); - self.render_body(frame, chunks[1]); - self.render_footer(frame, chunks[2]); + self.render_header(frame, rows[0]); + self.render_body(frame, rows[1]); + self.render_footer(frame, rows[2]); } /// Renders the scenario header bar. diff --git a/treetest/src/app/ui/panels/lists.rs b/treetest/src/app/ui/panels/lists.rs index 1cc6abe..43773c3 100644 --- a/treetest/src/app/ui/panels/lists.rs +++ b/treetest/src/app/ui/panels/lists.rs @@ -60,7 +60,7 @@ impl App { /// while ground-truth mode intentionally exposes the entire scenario tree. pub(crate) fn build_selections(simulation: &Simulation) -> Vec { let mut selections = Vec::new(); - let node_ids: Vec<_> = match simulation.inspector_mode { + let visible_node_ids: Vec<_> = match simulation.inspector_mode { InspectorMode::GroundTruth => simulation.tree.nodes.iter().map(|node| node.id).collect(), InspectorMode::Realistic => simulation .root_knowledge @@ -70,7 +70,7 @@ pub(crate) fn build_selections(simulation: &Simulation) -> Vec { .collect(), }; - for node_id in node_ids { + for node_id in visible_node_ids { let node = simulation.node(node_id); selections.push(Selection::Node(node.id)); match simulation.inspector_mode { diff --git a/treetest/src/sim/knowledge.rs b/treetest/src/sim/knowledge.rs index ee4c137..18f3620 100644 --- a/treetest/src/sim/knowledge.rs +++ b/treetest/src/sim/knowledge.rs @@ -13,40 +13,56 @@ use crate::model::EndpointProcedureSpec; /// Root inspector mode. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum InspectorMode { + /// Render the full scenario definition, including information the root has + /// not yet learned through traffic or introspection. GroundTruth, + /// Render only the subset of state the root host could plausibly know. Realistic, } /// Learned procedure metadata stored by the root host. #[derive(Debug, Clone, PartialEq, Eq)] pub struct LearnedProcedure { + /// Stable protocol identifier for the learned procedure. pub procedure_id: String, + /// Optional human-readable description learned from config or introspection. pub description: Option, } /// Learned leaf metadata stored by the root host. #[derive(Debug, Clone, PartialEq, Eq)] pub struct LearnedLeaf { + /// Leaf name relative to the endpoint path. pub leaf_name: String, + /// Optional human-readable description for the leaf. pub description: Option, + /// Procedures currently known on the leaf. pub procedures: Vec, } /// Learned endpoint metadata stored by the root host. #[derive(Debug, Clone, PartialEq, Eq)] pub struct LearnedNode { + /// Absolute node path from the root. pub path: Vec, + /// Optional display title shown in the inspector. pub title: Option, + /// Optional endpoint description shown in the inspector. pub description: Option, + /// Whether the node is a direct child of the root. pub direct_child: bool, + /// Endpoint-level procedures known on the node itself. pub endpoint_procedures: Vec, + /// Leaf metadata currently known for the node. pub leaves: Vec, + /// Whether endpoint introspection definitely ran against this node. pub endpoint_introspected: bool, } /// Root-host knowledge accumulated from local configuration and observed traffic. #[derive(Debug, Clone, PartialEq, Eq)] pub struct RootKnowledge { + /// Learned nodes keyed by their absolute path. pub nodes: BTreeMap, LearnedNode>, } @@ -58,49 +74,9 @@ impl RootKnowledge { }; for node in &tree.nodes { if node.path.is_empty() || node.path.len() == 1 { - // Realistic mode intentionally starts with root plus direct children, - // not the full transitive tree. - 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() { - // The root always knows its own procedures and leaves because - // those are locally configured, not discovered remotely. - 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 + .nodes + .insert(node.path.clone(), initial_learned_node(node)); } } knowledge @@ -223,12 +199,56 @@ impl RootKnowledge { } } +/// Builds the root's initial record for one statically known node. +fn initial_learned_node(node: &crate::model::DemoNode) -> LearnedNode { + let mut learned = LearnedNode { + path: node.path.clone(), + title: Some(node.title.clone()), + description: Some(node.description.clone()), + direct_child: node.path.len() == 1, + endpoint_procedures: Vec::new(), + leaves: Vec::new(), + endpoint_introspected: node.path.is_empty(), + }; + + if node.path.is_empty() { + // The root always knows its own procedures and leaves because those are + // locally configured, not discovered remotely. + 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(); + } + + learned +} + /// Returns one learned leaf entry, creating it if necessary. -fn ensure_leaf<'a>( - leaves: &'a mut Vec, +fn ensure_leaf( + leaves: &mut Vec, leaf_name: String, description: Option, -) -> &'a mut LearnedLeaf { +) -> &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; diff --git a/treetest/src/sim/runtime/events/application.rs b/treetest/src/sim/runtime/events/application.rs index b20373f..41377ef 100644 --- a/treetest/src/sim/runtime/events/application.rs +++ b/treetest/src/sim/runtime/events/application.rs @@ -32,16 +32,14 @@ impl Simulation { match procedure.kind { EndpointProcedureKind::Ping => { let reply = format!("pong from {}", self.node(node_id).display_path()); - let frame = self.nodes[node_id.0] - .endpoint - .make_data( - hook.return_path.clone(), - hook.hook_id, - procedure.procedure_id.clone(), - reply.clone().into_bytes(), - true, - ) - .map_err(|error| SimError::Protocol(error.to_string()))?; + let frame = self.make_endpoint_data_frame( + node_id, + hook.return_path.clone(), + hook.hook_id, + procedure.procedure_id.clone(), + reply.clone().into_bytes(), + true, + )?; self.record_trace(node_id, format!("endpoint sent ping reply: {reply}")); self.process_local_frame(node_id, frame)?; } @@ -54,16 +52,14 @@ impl Simulation { .iter() .enumerate() { - let frame = self.nodes[node_id.0] - .endpoint - .make_data( - hook.return_path.clone(), - hook.hook_id, - procedure.procedure_id.clone(), - text.as_bytes().to_vec(), - index == 2, - ) - .map_err(|error| SimError::Protocol(error.to_string()))?; + let frame = self.make_endpoint_data_frame( + node_id, + hook.return_path.clone(), + hook.hook_id, + procedure.procedure_id.clone(), + text.as_bytes().to_vec(), + index == 2, + )?; self.record_trace(node_id, format!("endpoint sent chunk {}", index + 1)); self.process_local_frame(node_id, frame)?; } @@ -80,16 +76,14 @@ impl Simulation { procedure_id: procedure.procedure_id.clone(), }, ); - let frame = self.nodes[node_id.0] - .endpoint - .make_data( - hook.return_path.clone(), - hook.hook_id, - procedure.procedure_id.clone(), - b"chat ready".to_vec(), - false, - ) - .map_err(|error| SimError::Protocol(error.to_string()))?; + let frame = self.make_endpoint_data_frame( + node_id, + hook.return_path.clone(), + hook.hook_id, + procedure.procedure_id.clone(), + b"chat ready".to_vec(), + false, + )?; self.record_trace(node_id, "chat handler opened session".to_owned()); self.process_local_frame(node_id, frame)?; } @@ -97,6 +91,23 @@ impl Simulation { Ok(()) } + /// Builds one endpoint-originated data frame after application logic decides + /// what to send back on an already-validated hook. + fn make_endpoint_data_frame( + &mut self, + node_id: NodeId, + return_path: Vec, + hook_id: u64, + procedure_id: String, + data: Vec, + end_hook: bool, + ) -> Result { + self.nodes[node_id.0] + .endpoint + .make_data(return_path, hook_id, procedure_id, data, end_hook) + .map_err(|error| SimError::Protocol(error.to_string())) + } + /// Resolves one endpoint procedure from the ground-truth node metadata. pub(super) fn lookup_endpoint_procedure( &self, diff --git a/treetest/src/sim/runtime/events/local.rs b/treetest/src/sim/runtime/events/local.rs index 804f3c6..5da32b7 100644 --- a/treetest/src/sim/runtime/events/local.rs +++ b/treetest/src/sim/runtime/events/local.rs @@ -20,51 +20,24 @@ impl Simulation { match event { LocalEvent::Data { header, message } => { let text = String::from_utf8_lossy(&message.data).to_string(); - self.record_trace( - node_id, - format!( - "local Data on {}: {text}", - format_hook_ref( - self.node(node_id).path.as_slice(), - header.hook_id.unwrap_or(0) - ) - ), + let hook_ref = format_hook_ref( + self.node(node_id).path.as_slice(), + header.hook_id.unwrap_or(0), ); + self.record_trace(node_id, format!("local Data on {hook_ref}: {text}")); if let Some(hook_id) = header.hook_id { - if let Some(snapshot) = self.hooks.get_mut(&hook_id) { - // Keep the most recent human-readable payload in the UI. - snapshot.last_message = if text.is_empty() { - format!("binary payload ({} bytes)", message.data.len()) - } else { - text.clone() - }; - if message.end_hook { - snapshot.closed = true; - } - } - + self.update_hook_snapshot(hook_id, &text, message.data.len(), message.end_hook); if node_id == self.root_id { self.learn_from_root_data(hook_id, &message); } } - if let Some(session) = self - .chat_sessions - .get(&header.hook_id.unwrap_or(0)) - .cloned() - .filter(|session| session.node_id == node_id) - { + if let Some(session) = self.chat_session_for_event(node_id, header.hook_id) { // Rationale: chat responses are implemented here instead of in the // core endpoint so the protocol crate stays generic. The simulator // acts as the application layer sitting above validated hook traffic. - let reply = if text.eq_ignore_ascii_case("bye") { - Some(("chat session closed".to_owned(), true)) - } else if !text.is_empty() { - Some((format!("chat ack: {}", text.to_uppercase()), false)) - } else { - None - }; + let reply = chat_reply_for_text(&text); if let Some((reply, end_hook)) = reply { let frame = self.nodes[session.node_id.0] @@ -92,16 +65,13 @@ impl Simulation { }); } LocalEvent::Fault { header, message } => { + let hook_ref = format_hook_ref( + self.node(node_id).path.as_slice(), + header.hook_id.unwrap_or(0), + ); self.record_trace( node_id, - format!( - "local Fault on {}: 0x{:02X}", - format_hook_ref( - self.node(node_id).path.as_slice(), - header.hook_id.unwrap_or(0) - ), - message.fault.0 - ), + format!("local Fault on {hook_ref}: 0x{:02X}", message.fault.0), ); if let Some(hook_id) = header.hook_id { if let Some(snapshot) = self.hooks.get_mut(&hook_id) { @@ -140,4 +110,45 @@ impl Simulation { } Ok(()) } + + fn update_hook_snapshot( + &mut self, + hook_id: u64, + text: &str, + payload_len: usize, + end_hook: bool, + ) { + if let Some(snapshot) = self.hooks.get_mut(&hook_id) { + // Keep the most recent human-readable payload in the UI. + snapshot.last_message = if text.is_empty() { + format!("binary payload ({payload_len} bytes)") + } else { + text.to_owned() + }; + if end_hook { + snapshot.closed = true; + } + } + } + + fn chat_session_for_event( + &self, + node_id: NodeId, + hook_id: Option, + ) -> Option { + self.chat_sessions + .get(&hook_id.unwrap_or(0)) + .cloned() + .filter(|session| session.node_id == node_id) + } +} + +fn chat_reply_for_text(text: &str) -> Option<(String, bool)> { + if text.eq_ignore_ascii_case("bye") { + return Some(("chat session closed".to_owned(), true)); + } + if text.is_empty() { + return None; + } + Some((format!("chat ack: {}", text.to_uppercase()), false)) } diff --git a/treetest/src/sim/runtime/learning.rs b/treetest/src/sim/runtime/learning.rs index bcf76a5..41f5a99 100644 --- a/treetest/src/sim/runtime/learning.rs +++ b/treetest/src/sim/runtime/learning.rs @@ -1,4 +1,8 @@ //! Root-side knowledge learning from returned data. +//! +//! The simulator learns only from data that arrives back at the root on a known +//! hook. This keeps the realistic inspector aligned with what the UI-triggered +//! action actually observed. use unshell::protocol::{ DataMessage, EndpointIntrospection, LeafIntrospection, deserialize_archived_bytes, @@ -17,23 +21,7 @@ impl Simulation { 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); - } + self.learn_from_root_introspection(&snapshot, &demo_node, message); return; } @@ -54,3 +42,33 @@ impl Simulation { } } } + +impl Simulation { + fn learn_from_root_introspection( + &mut self, + snapshot: &super::super::types::HookSnapshot, + demo_node: &crate::model::DemoNode, + message: &DataMessage, + ) { + 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); + } + return; + } + + if let Ok(introspection) = deserialize_archived_bytes::< + unshell::protocol::introspection::ArchivedEndpointIntrospection, + EndpointIntrospection, + >(&message.data) + { + self.root_knowledge + .remember_endpoint_introspection(demo_node, &introspection); + } + } +} diff --git a/ush-obfuscate/src/env.rs b/ush-obfuscate/src/env.rs index d79861b..0f86355 100644 --- a/ush-obfuscate/src/env.rs +++ b/ush-obfuscate/src/env.rs @@ -1,6 +1,9 @@ const ENV_KEY_NAME: &str = "OBFUSCATION_KEY"; const BACKUP_ENV_KEY: &str = "OBFUSCATION_KEY_DO_NOT_USE"; +/// Returns the obfuscation key used by the proc macros. +/// +/// The fallback keeps macro expansion deterministic when the environment variable is absent. pub fn get_encryption_key() -> String { if let Ok(key) = std::env::var(ENV_KEY_NAME) { key diff --git a/ush-obfuscate/src/format_helper.rs b/ush-obfuscate/src/format_helper.rs index 083e5fa..fc1ac67 100644 --- a/ush-obfuscate/src/format_helper.rs +++ b/ush-obfuscate/src/format_helper.rs @@ -3,6 +3,7 @@ use quote::quote; use syn::parse::{Parse, ParseStream}; use syn::{Expr, Lit, Token, parse_macro_input}; +/// Expands `sym_format!` into a string builder that obfuscates static segments only. pub fn sym_format(input: TokenStream) -> TokenStream { let PrintlnArgs { format_str, args } = parse_macro_input!(input as PrintlnArgs); @@ -42,9 +43,6 @@ pub fn sym_format(input: TokenStream) -> TokenStream { quote! { #full_spec } }; - // quote! { - // println!(#fmt_spec, #arg); - // } parts.push(quote! { format!(#fmt_spec, #arg) }); @@ -105,9 +103,13 @@ impl Parse for PrintlnArgs { #[derive(Debug)] enum FormatSegment { Static(String), - Dynamic(String, usize), // format spec, arg index + Dynamic(String, usize), } +/// Splits a Rust formatting string into literal and replacement segments. +/// +/// This only handles the subset needed by `sym_format!`: positional replacements in order, +/// plus escaped braces. fn parse_format_string(fmt: &str) -> Vec { let mut segments = Vec::new(); let mut current_static = String::new(); @@ -122,13 +124,11 @@ fn parse_format_string(fmt: &str) -> Vec { continue; } - // Save current static segment if !current_static.is_empty() { segments.push(FormatSegment::Static(current_static.clone())); current_static.clear(); } - // Parse format spec let mut spec = String::new(); while let Some(&next_ch) = chars.peek() { if next_ch == '}' { diff --git a/ush-obfuscate/src/proc_impl_switcher.rs b/ush-obfuscate/src/proc_impl_switcher.rs index c727686..76fce6f 100644 --- a/ush-obfuscate/src/proc_impl_switcher.rs +++ b/ush-obfuscate/src/proc_impl_switcher.rs @@ -1,5 +1,5 @@ /// Call some other function -macro_rules! passtrough { +macro_rules! passthrough { ($name:tt, $ref:expr) => { pub fn $name(input: TokenStream) -> TokenStream { $ref(input) @@ -42,25 +42,27 @@ pub mod proc_impl { unwrap_string!(sym_fn); } -#[cfg(feature = "obfuscate_aes")] +#[cfg(all(feature = "obfuscate_aes", not(feature = "obfuscate_ref")))] pub mod proc_impl { use proc_macro::TokenStream; - passtrough!(xor, crate::obfuscate::xor); - passtrough!(junk_asm, crate::obfuscate::junk_asm); + passthrough!(xor, crate::obfuscate::xor); + passthrough!(junk_asm, crate::obfuscate::junk_asm); - passtrough!(sym, crate::symbolic_aes::aes_str); - passtrough!(sym_fn, crate::symbolic_aes::aes_fn_name); + passthrough!(sym, crate::symbolic_aes::aes_str); + passthrough!(sym_fn, crate::symbolic_aes::aes_fn_name); } #[cfg(feature = "obfuscate_ref")] pub mod proc_impl { use proc_macro::TokenStream; - use syn::{LitStr, parse_macro_input}; - passtrough!(xor, crate::obfuscate::xor); - passtrough!(junk_asm, crate::obfuscate::junk_asm); + passthrough!(xor, crate::obfuscate::xor); + passthrough!(junk_asm, crate::obfuscate::junk_asm); - passtrough!(sym, crate::symbolic_ref::sym_ref); - passtrough!(sym_fn, crate::symbolic_ref::sym_ref_fn); + // `sym` and `sym_fn` need one concrete strategy. When both feature flags are enabled, + // prefer symbolic references so `cargo clippy --all-features` still selects a single, + // deterministic implementation. + passthrough!(sym, crate::symbolic_ref::sym_ref); + passthrough!(sym_fn, crate::symbolic_ref::sym_ref_fn); } diff --git a/ush-obfuscate/src/symbolic_ref/mod.rs b/ush-obfuscate/src/symbolic_ref/mod.rs index cf7eb42..f37c969 100644 --- a/ush-obfuscate/src/symbolic_ref/mod.rs +++ b/ush-obfuscate/src/symbolic_ref/mod.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; - use base62::{Base62, hash}; use proc_macro::TokenStream; use quote::quote; @@ -10,11 +8,13 @@ use crate::env::get_encryption_key; static mut SYM_COUNTER: Vec = Vec::new(); #[allow(static_mut_refs)] +/// Returns how many unique symbols have been registered in this macro process. pub fn get_symbol_number() -> usize { unsafe { SYM_COUNTER.len() } } #[allow(static_mut_refs)] +/// Returns the stable numeric ID for `text`, inserting it on first use. pub fn get_symbol(text: &str) -> usize { unsafe { if let Some(n) = SYM_COUNTER.iter().position(|r| r == text) { @@ -27,47 +27,40 @@ pub fn get_symbol(text: &str) -> usize { } } -fn ref_string(input: String) -> String { - let n = get_symbol(&input); +fn encode_symbol_reference(symbol: String) -> String { + let symbol_index = get_symbol(&symbol); - let data = base62::encode_usize(n); - let key = hash(&get_encryption_key().as_bytes()); + let data = base62::encode_usize(symbol_index); + let key = hash(get_encryption_key().as_bytes()); let encoded = format!("_{}_", Base62::encode_full(&data, &key)); - println!("Aliased '{}' as '{encoded}'", input); + // Macro expansion logs make it easier to correlate exported symbols with their aliases. + println!("Aliased '{}' as '{encoded}'", symbol); encoded } +/// Replaces a string literal with its symbolic reference alias. pub fn sym_ref(input: TokenStream) -> TokenStream { - // Parse the input as a string literal let lit_str = parse_macro_input!(input as LitStr); let original_name = lit_str.value(); - let encoded = ref_string(original_name); + let encoded = encode_symbol_reference(original_name); - // Expand to a static string literal TokenStream::from(quote! { #encoded }) } +/// Re-exports a function under a symbolic reference alias. pub fn sym_ref_fn(input: TokenStream) -> TokenStream { - // Parse the input function let func = parse_macro_input!(input as ItemFn); - - // Get the original function name let fn_name = func.sig.ident.to_string(); - // Generate the new, obfuscated name - let obfuscated_name = ref_string(fn_name); - - // Create a new string literal for the name + let obfuscated_name = encode_symbol_reference(fn_name); let new_name_lit = LitStr::new(&obfuscated_name, func.sig.ident.span()); - // Re-build the function, but add #[no_mangle] - // and rename the *exported* symbol via #[export_name] TokenStream::from(quote! { #[unsafe(export_name = #new_name_lit)] #func