diff --git a/src/dispatch.rs b/src/dispatch.rs new file mode 100644 index 0000000..129ed22 --- /dev/null +++ b/src/dispatch.rs @@ -0,0 +1,39 @@ +//! Classify decrypted application plaintext before it hits the message log. + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Kind { + Chat, + Loc, + Profile, + Invoice, + Receipt, + Ping, + Drop, +} + +pub fn classify(pt: &[u8]) -> Kind { + if pt.starts_with(b"loc ") { + return Kind::Loc; + } + if pt.starts_with(b"prf ") { + return Kind::Profile; + } + if pt.starts_with(b"inv ") { + return Kind::Invoice; + } + if pt.starts_with(b"rcp ") { + return Kind::Receipt; + } + if pt.starts_with(b"png ") { + return Kind::Ping; + } + if pt.len() >= 4 + && pt[0].is_ascii_lowercase() + && pt[1].is_ascii_lowercase() + && pt[2].is_ascii_lowercase() + && pt[3] == b' ' + { + return Kind::Drop; + } + Kind::Chat +} diff --git a/src/lib.rs b/src/lib.rs index 8d22191..c560e94 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,4 @@ +pub mod dispatch; pub mod frame; pub mod hs; pub mod loc; diff --git a/src/node.rs b/src/node.rs index 5f2b29f..8a11dfe 100644 --- a/src/node.rs +++ b/src/node.rs @@ -9,6 +9,7 @@ use futures::io::{AsyncRead, AsyncWrite}; use tor_cell::relaycell::msg::Connected; use tor_hsservice::{RunningOnionService, handle_rend_requests}; +use crate::dispatch::{self, Kind}; use crate::frame; use crate::hs::{self, Client, HS_PORT}; use crate::loc; @@ -308,26 +309,33 @@ impl Node { .map_err(session_err)?; let ct = frame::read_frame(stream).await.map_err(|e| e.to_string())?; let pt = sess.decrypt(&ct).map_err(session_err)?; - if let Some(loc) = loc::decode(&pt) { - let applied = store.lock().map_err(|e| e.to_string())?.apply_loc( - &sess.peer_identity, - &loc.onion, - loc.ts, - &loc.sig, - ); - match applied { - Ok(true) => {} - Ok(false) => eprintln!("loc dropped (bad sig, stale ts, or unknown friend)"), - Err(e) => return Err(e.to_string()), + match dispatch::classify(&pt) { + Kind::Loc => { + if let Some(loc) = loc::decode(&pt) { + let applied = store.lock().map_err(|e| e.to_string())?.apply_loc( + &sess.peer_identity, + &loc.onion, + loc.ts, + &loc.sig, + ); + match applied { + Ok(true) => {} + Ok(false) => eprintln!("loc dropped (bad sig, stale ts, or unknown friend)"), + Err(e) => return Err(e.to_string()), + } + } + Ok(()) } - return Ok(()); + Kind::Chat => { + store + .lock() + .map_err(|e| e.to_string())? + .append_message(&sess.peer_identity, "in", &pt) + .map_err(|e| e.to_string())?; + Ok(()) + } + Kind::Profile | Kind::Invoice | Kind::Receipt | Kind::Ping | Kind::Drop => Ok(()), } - store - .lock() - .map_err(|e| e.to_string())? - .append_message(&sess.peer_identity, "in", &pt) - .map_err(|e| e.to_string())?; - Ok(()) } } diff --git a/tests/dispatch.rs b/tests/dispatch.rs new file mode 100644 index 0000000..e66bb78 --- /dev/null +++ b/tests/dispatch.rs @@ -0,0 +1,39 @@ +//! M6: typed-frame dispatcher — classify before append_message. + +use onionwire::dispatch::{classify, Kind}; + +#[test] +fn chat_is_chat() { + assert_eq!(classify(b"hello wire"), Kind::Chat); +} + +#[test] +fn loc_still_loc() { + assert_eq!(classify(b"loc abc.onion\n1\n00"), Kind::Loc); +} + +#[test] +fn unknown_typed_prefix_is_drop() { + assert_eq!(classify(b"zzz not a real type"), Kind::Drop); +} + +#[test] +fn short_or_binary_without_prefix_is_chat() { + assert_eq!(classify(b"hi"), Kind::Chat); + assert_eq!(classify(&[0xff, 0xfe]), Kind::Chat); +} + +#[test] +fn known_prefixes_are_typed() { + assert_eq!(classify(b"prf name"), Kind::Profile); + assert_eq!(classify(b"inv 1"), Kind::Invoice); + assert_eq!(classify(b"rcp tx"), Kind::Receipt); + assert_eq!(classify(b"png 1"), Kind::Ping); +} + +#[test] +fn uppercase_or_digit_prefix_is_chat() { + assert_eq!(classify(b"ZZZ not typed"), Kind::Chat); + assert_eq!(classify(b"ab1 leftover"), Kind::Chat); + assert_eq!(classify(b"abcd"), Kind::Chat); +}