fix(lib): route on the full four-byte message code

Dispatch read one byte of the code, so a code with a non-zero high byte
was handled as its low byte. Message::get_message_code now reads the
whole little-endian uint32 and handlers register by u32. PeerInit and
PierceFirewall, the only frames with a one-byte code, get their own
accessor for the listener's first frame.
This commit is contained in:
Michel de Graaf
2026-09-07 15:03:49 +02:00
parent 40a8eb1dd5
commit f1aa94fecc
40 changed files with 99 additions and 55 deletions
+2 -2
View File
@@ -481,7 +481,7 @@ impl SeederState {
let Ok(init) = read_framed(&mut p) else { let Ok(init) = read_framed(&mut p) else {
return; return;
}; };
if init.get_message_code() != 1 { if init.get_init_code() != 1 {
return; return;
} }
} }
@@ -735,7 +735,7 @@ fn accept_file_connection(
let Ok(mut init) = read_framed(&mut stream) else { let Ok(mut init) = read_framed(&mut stream) else {
continue; continue;
}; };
if init.get_message_code() != 1 { if init.get_init_code() != 1 {
continue; continue;
} }
init.set_pointer(5); init.set_pointer(5);
+1 -1
View File
@@ -535,7 +535,7 @@ impl PeerActor {
message message
.get_message_name( .get_message_name(
MessageType::Peer, MessageType::Peer,
u32::from(message.get_message_code()) message.get_message_code()
) )
.map_err(|e| e.to_string()) .map_err(|e| e.to_string())
); );
+1 -1
View File
@@ -803,7 +803,7 @@ impl ServerActor {
message message
.get_message_name( .get_message_name(
MessageType::Server, MessageType::Server,
u32::from(message.get_message_code()) message.get_message_code()
) )
.map_err(|e| e.to_string()) .map_err(|e| e.to_string())
); );
+37
View File
@@ -37,3 +37,40 @@ impl<Op> MessageDispatcher<Op> {
} }
} }
} }
#[cfg(test)]
mod tests {
use super::*;
use crate::message::MessageHandler;
use std::sync::mpsc::channel;
struct Echo(u32);
impl MessageHandler<u32> for Echo {
fn get_code(&self) -> u32 {
self.0
}
fn handle(&self, _: &mut Message, sender: Sender<u32>) {
let _ = sender.send(self.0);
}
}
fn framed(code: u32) -> Message {
Message::new_with_data(Message::new().write_int32(code).get_buffer())
}
#[test]
fn a_code_above_255_is_not_confused_with_its_low_byte() {
let (sender, received) = channel();
let mut handlers = Handlers::new();
handlers.register_handler(Echo(1));
handlers.register_handler(Echo(1001));
let dispatcher =
MessageDispatcher::new("test".to_string(), sender, handlers);
dispatcher.dispatch(&mut framed(257));
dispatcher.dispatch(&mut framed(1001));
assert_eq!(received.try_iter().collect::<Vec<_>>(), [1001]);
}
}
+3 -3
View File
@@ -4,11 +4,11 @@ use crate::message::Message;
use std::sync::mpsc::Sender; use std::sync::mpsc::Sender;
pub trait MessageHandler<Op>: Send { pub trait MessageHandler<Op>: Send {
fn get_code(&self) -> u8; fn get_code(&self) -> u32;
fn handle(&self, message: &mut Message, sender: Sender<Op>); fn handle(&self, message: &mut Message, sender: Sender<Op>);
} }
pub struct Handlers<Op> { pub struct Handlers<Op> {
handlers: HashMap<u8, Box<dyn MessageHandler<Op> + Send>>, handlers: HashMap<u32, Box<dyn MessageHandler<Op> + Send>>,
} }
impl<Op> Default for Handlers<Op> { impl<Op> Default for Handlers<Op> {
@@ -35,7 +35,7 @@ impl<Op> Handlers<Op> {
#[must_use] #[must_use]
pub fn get_handler( pub fn get_handler(
&self, &self,
code: u8, code: u32,
) -> Option<&(dyn MessageHandler<Op> + Send)> { ) -> Option<&(dyn MessageHandler<Op> + Send)> {
self.handlers.get(&code).map(|v| &**v) self.handlers.get(&code).map(|v| &**v)
} }
+16 -9
View File
@@ -51,21 +51,19 @@ impl Message {
} }
} }
#[allow(dead_code)]
#[must_use] #[must_use]
pub fn get_message_code_u32(&self) -> u32 { pub fn get_message_code(&self) -> u32 {
if self.data.len() < 8 { if self.data.len() < 8 {
return 0; return 0;
} }
self.data[4..8].try_into().map_or(0, u32::from_le_bytes) self.data[4..8].try_into().map_or(0, u32::from_le_bytes)
} }
/// The one-byte code of a peer init frame (PierceFirewall 0, PeerInit 1),
/// the only messages whose code is not four bytes wide.
#[must_use] #[must_use]
pub fn get_message_code(&self) -> u8 { pub fn get_init_code(&self) -> u8 {
if self.data.len() <= 4 { self.data.get(4).copied().unwrap_or(0)
return 0;
}
self.data[4]
} }
#[must_use] #[must_use]
@@ -388,10 +386,19 @@ fn read_string_with_missing_length_prefix_does_not_panic() {
#[test] #[test]
fn get_message_code_on_short_message_does_not_panic() { fn get_message_code_on_short_message_does_not_panic() {
// A zero-length frame drains to exactly the 4-byte length prefix, leaving // A zero-length frame drains to exactly the 4-byte length prefix, leaving
// no room for a message code at data[4]. // no room for a code after it.
let msg = Message::new_with_data(vec![0, 0, 0, 0]); let msg = Message::new_with_data(vec![0, 0, 0, 0]);
assert_eq!(msg.get_message_code(), 0); assert_eq!(msg.get_message_code(), 0);
assert_eq!(msg.get_message_code_u32(), 0); assert_eq!(msg.get_init_code(), 0);
}
#[test]
fn an_init_frame_reads_its_one_byte_code() {
let init = Message::new_with_data(
Message::new().write_int8(1).write_string("ab").get_buffer(),
);
assert_eq!(init.get_init_code(), 1);
assert_ne!(init.get_message_code(), 1);
} }
#[test] #[test]
@@ -50,7 +50,7 @@ pub fn build_file_search_response(
pub struct FileSearchResponse; pub struct FileSearchResponse;
impl MessageHandler<PeerMessage> for FileSearchResponse { impl MessageHandler<PeerMessage> for FileSearchResponse {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
9 9
} }
fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) {
@@ -8,7 +8,7 @@ use std::sync::mpsc::Sender;
/// owns the shares) builds the real SharedFileListResponse in reply. /// owns the shares) builds the real SharedFileListResponse in reply.
pub struct GetShareFileList; pub struct GetShareFileList;
impl MessageHandler<PeerMessage> for GetShareFileList { impl MessageHandler<PeerMessage> for GetShareFileList {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
4 4
} }
fn handle(&self, _message: &mut Message, sender: Sender<PeerMessage>) { fn handle(&self, _message: &mut Message, sender: Sender<PeerMessage>) {
@@ -8,7 +8,7 @@ use crate::{
pub struct PeerInit; pub struct PeerInit;
impl MessageHandler<PeerMessage> for PeerInit { impl MessageHandler<PeerMessage> for PeerInit {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
1 1
} }
@@ -12,7 +12,7 @@ use std::sync::mpsc::Sender;
pub struct PlaceInQueueRequest; pub struct PlaceInQueueRequest;
impl MessageHandler<PeerMessage> for PlaceInQueueRequest { impl MessageHandler<PeerMessage> for PlaceInQueueRequest {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
51 51
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct PlaceInQueueResponse; pub struct PlaceInQueueResponse;
impl MessageHandler<PeerMessage> for PlaceInQueueResponse { impl MessageHandler<PeerMessage> for PlaceInQueueResponse {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
44 44
} }
@@ -8,7 +8,7 @@ use std::sync::mpsc::Sender;
pub struct QueueUploadHandler; pub struct QueueUploadHandler;
impl MessageHandler<PeerMessage> for QueueUploadHandler { impl MessageHandler<PeerMessage> for QueueUploadHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
43 43
} }
@@ -17,7 +17,7 @@ pub struct SharedDirectory {
/// Receives a peer's `SharedFileListResponse` (peer code 5) when browsing them. /// Receives a peer's `SharedFileListResponse` (peer code 5) when browsing them.
pub struct SharedFileListResponseHandler; pub struct SharedFileListResponseHandler;
impl MessageHandler<PeerMessage> for SharedFileListResponseHandler { impl MessageHandler<PeerMessage> for SharedFileListResponseHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
5 5
} }
fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) {
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct TransferRequest; pub struct TransferRequest;
impl MessageHandler<PeerMessage> for TransferRequest { impl MessageHandler<PeerMessage> for TransferRequest {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
40 40
} }
fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) {
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct TransferResponse; pub struct TransferResponse;
impl MessageHandler<PeerMessage> for TransferResponse { impl MessageHandler<PeerMessage> for TransferResponse {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
41 41
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct UploadDeniedHandler; pub struct UploadDeniedHandler;
impl MessageHandler<PeerMessage> for UploadDeniedHandler { impl MessageHandler<PeerMessage> for UploadDeniedHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
50 50
} }
fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) {
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct UploadFailedHandler; pub struct UploadFailedHandler;
impl MessageHandler<PeerMessage> for UploadFailedHandler { impl MessageHandler<PeerMessage> for UploadFailedHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
46 46
} }
fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<PeerMessage>) {
@@ -5,7 +5,7 @@ use std::sync::mpsc::Sender;
pub struct ConnectToPeerHandler; pub struct ConnectToPeerHandler;
impl MessageHandler<ServerMessage> for ConnectToPeerHandler { impl MessageHandler<ServerMessage> for ConnectToPeerHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
18 18
} }
fn handle(&self, message: &mut Message, sender: Sender<ServerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<ServerMessage>) {
@@ -8,7 +8,7 @@ use std::sync::mpsc::Sender;
pub struct ExcludedSearchPhrasesHandler; pub struct ExcludedSearchPhrasesHandler;
impl MessageHandler<ServerMessage> for ExcludedSearchPhrasesHandler { impl MessageHandler<ServerMessage> for ExcludedSearchPhrasesHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
160 160
} }
@@ -9,7 +9,7 @@ use crate::{
pub struct FileSearchHandler; pub struct FileSearchHandler;
impl MessageHandler<ServerMessage> for FileSearchHandler { impl MessageHandler<ServerMessage> for FileSearchHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
26 26
} }
fn handle(&self, message: &mut Message, sender: Sender<ServerMessage>) { fn handle(&self, message: &mut Message, sender: Sender<ServerMessage>) {
@@ -5,7 +5,7 @@ use std::sync::mpsc::Sender;
pub struct GetPeerAddressHandler; pub struct GetPeerAddressHandler;
impl MessageHandler<ServerMessage> for GetPeerAddressHandler { impl MessageHandler<ServerMessage> for GetPeerAddressHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
3 3
} }
@@ -8,7 +8,7 @@ use std::sync::mpsc::Sender;
pub struct JoinRoomHandler; pub struct JoinRoomHandler;
impl MessageHandler<ServerMessage> for JoinRoomHandler { impl MessageHandler<ServerMessage> for JoinRoomHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
14 14
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct LeaveRoomHandler; pub struct LeaveRoomHandler;
impl MessageHandler<ServerMessage> for LeaveRoomHandler { impl MessageHandler<ServerMessage> for LeaveRoomHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
15 15
} }
+1 -1
View File
@@ -8,7 +8,7 @@ use crate::message::MessageHandler;
pub struct LoginHandler; pub struct LoginHandler;
impl MessageHandler<ServerMessage> for LoginHandler { impl MessageHandler<ServerMessage> for LoginHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
1 1
} }
@@ -336,7 +336,7 @@ fn a_transfer_denial_carries_the_token_and_the_reason() {
let message = let message =
MessageFactory::build_transfer_denial_message(555, "Cancelled"); MessageFactory::build_transfer_denial_message(555, "Cancelled");
let mut decoded = Message::new_with_data(message.get_buffer()); let mut decoded = Message::new_with_data(message.get_buffer());
assert_eq!(decoded.get_message_code_u32(), 41); assert_eq!(decoded.get_message_code(), 41);
decoded.set_pointer(8); decoded.set_pointer(8);
assert_eq!(decoded.read_int32(), 555); assert_eq!(decoded.read_int32(), 555);
assert!(!decoded.read_bool()); assert!(!decoded.read_bool());
@@ -8,7 +8,7 @@ use std::sync::mpsc::Sender;
pub struct MessageUser; pub struct MessageUser;
impl MessageHandler<ServerMessage> for MessageUser { impl MessageHandler<ServerMessage> for MessageUser {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
22 22
} }
@@ -9,7 +9,7 @@ use crate::{
pub struct ParentMinSpeedHandler; pub struct ParentMinSpeedHandler;
impl MessageHandler<ServerMessage> for ParentMinSpeedHandler { impl MessageHandler<ServerMessage> for ParentMinSpeedHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
83 83
} }
@@ -10,7 +10,7 @@ pub struct ParentSpeedRatioHandler;
// The server sends us a speed ratio determining the number of children we can have in the distributed network. The maximum number of children is our upload speed divided by the speed ratio. // The server sends us a speed ratio determining the number of children we can have in the distributed network. The maximum number of children is our upload speed divided by the speed ratio.
impl MessageHandler<ServerMessage> for ParentSpeedRatioHandler { impl MessageHandler<ServerMessage> for ParentSpeedRatioHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
84 84
} }
@@ -14,7 +14,7 @@ use crate::{
pub struct PrivilegedUsersHandler; pub struct PrivilegedUsersHandler;
impl MessageHandler<ServerMessage> for PrivilegedUsersHandler { impl MessageHandler<ServerMessage> for PrivilegedUsersHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
69 69
} }
@@ -37,7 +37,7 @@ impl MessageHandler<ServerMessage> for PrivilegedUsersHandler {
pub struct CheckPrivilegesHandler; pub struct CheckPrivilegesHandler;
impl MessageHandler<ServerMessage> for CheckPrivilegesHandler { impl MessageHandler<ServerMessage> for CheckPrivilegesHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
92 92
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct ReloggedHandler; pub struct ReloggedHandler;
impl MessageHandler<ServerMessage> for ReloggedHandler { impl MessageHandler<ServerMessage> for ReloggedHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
41 41
} }
@@ -8,7 +8,7 @@ use std::sync::mpsc::Sender;
pub struct RoomListHandler; pub struct RoomListHandler;
impl MessageHandler<ServerMessage> for RoomListHandler { impl MessageHandler<ServerMessage> for RoomListHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
64 64
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct SayChatroomHandler; pub struct SayChatroomHandler;
impl MessageHandler<ServerMessage> for SayChatroomHandler { impl MessageHandler<ServerMessage> for SayChatroomHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
13 13
} }
@@ -12,7 +12,7 @@ use std::sync::mpsc::Sender;
pub struct GetUserStatusHandler; pub struct GetUserStatusHandler;
impl MessageHandler<ServerMessage> for GetUserStatusHandler { impl MessageHandler<ServerMessage> for GetUserStatusHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
7 7
} }
@@ -33,7 +33,7 @@ impl MessageHandler<ServerMessage> for GetUserStatusHandler {
pub struct GetUserStatsHandler; pub struct GetUserStatsHandler;
impl MessageHandler<ServerMessage> for GetUserStatsHandler { impl MessageHandler<ServerMessage> for GetUserStatsHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
36 36
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct UserJoinedRoomHandler; pub struct UserJoinedRoomHandler;
impl MessageHandler<ServerMessage> for UserJoinedRoomHandler { impl MessageHandler<ServerMessage> for UserJoinedRoomHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
16 16
} }
@@ -7,7 +7,7 @@ use std::sync::mpsc::Sender;
pub struct UserLeftRoomHandler; pub struct UserLeftRoomHandler;
impl MessageHandler<ServerMessage> for UserLeftRoomHandler { impl MessageHandler<ServerMessage> for UserLeftRoomHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
17 17
} }
@@ -21,7 +21,7 @@ use std::sync::mpsc::Sender;
pub struct WatchUserHandler; pub struct WatchUserHandler;
impl MessageHandler<ServerMessage> for WatchUserHandler { impl MessageHandler<ServerMessage> for WatchUserHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
5 5
} }
@@ -13,7 +13,7 @@ use crate::{
pub struct WishListIntervalHandler; pub struct WishListIntervalHandler;
impl MessageHandler<ServerMessage> for WishListIntervalHandler { impl MessageHandler<ServerMessage> for WishListIntervalHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u32 {
104 104
} }
+1 -1
View File
@@ -360,7 +360,7 @@ fn handle_incoming_connection(
// A firewalled peer brokered through the server connects back with a // A firewalled peer brokered through the server connects back with a
// PierceFirewall (code 0) instead of a PeerInit (code 1). // PierceFirewall (code 0) instead of a PeerInit (code 1).
if message.get_message_code() == 0 { if message.get_init_code() == 0 {
handle_pierce_firewall( handle_pierce_firewall(
message, stream, reader, &context, &peer_ip, peer_port, message, stream, reader, &context, &peer_ip, peer_port,
); );
+1 -1
View File
@@ -191,7 +191,7 @@ fn test_new_from_message2() {
let mut message = Message::new_with_data(data); let mut message = Message::new_with_data(data);
message.set_pointer(8); message.set_pointer(8);
println!("code: {}", message.get_message_code_u32()); println!("code: {}", message.get_message_code());
let peer = Peer::new_from_message(&mut message).unwrap(); let peer = Peer::new_from_message(&mut message).unwrap();
+4 -4
View File
@@ -750,7 +750,7 @@ fn login_raw(
/// passes. Returns the matching message. /// passes. Returns the matching message.
fn read_until_code( fn read_until_code(
stream: &mut TcpStream, stream: &mut TcpStream,
code: u8, code: u32,
timeout: Duration, timeout: Duration,
) -> Option<Message> { ) -> Option<Message> {
let deadline = Instant::now() + timeout; let deadline = Instant::now() + timeout;
@@ -767,7 +767,7 @@ fn read_until_code(
/// Like [`read_until_code`] but turns a miss into an `io::Error` for `?`. /// Like [`read_until_code`] but turns a miss into an `io::Error` for `?`.
fn expect_code( fn expect_code(
stream: &mut TcpStream, stream: &mut TcpStream,
code: u8, code: u32,
timeout: Duration, timeout: Duration,
) -> std::io::Result<Message> { ) -> std::io::Result<Message> {
read_until_code(stream, code, timeout).ok_or_else(|| { read_until_code(stream, code, timeout).ok_or_else(|| {
@@ -1379,7 +1379,7 @@ fn run_mock_direct_peer(cfg: &MockDirectUpload) -> std::io::Result<()> {
p.set_nonblocking(false)?; p.set_nonblocking(false)?;
p.set_read_timeout(Some(Duration::from_secs(10)))?; p.set_read_timeout(Some(Duration::from_secs(10)))?;
let mut init = read_framed(&mut p)?; let mut init = read_framed(&mut p)?;
assert_eq!(init.get_message_code(), 1, "expected inbound PeerInit"); assert_eq!(init.get_init_code(), 1, "expected inbound PeerInit");
init.set_pointer(5); init.set_pointer(5);
assert_eq!(init.read_string(), cfg.downloader_username, "PeerInit user"); assert_eq!(init.read_string(), cfg.downloader_username, "PeerInit user");
assert_eq!(init.read_string(), "P", "PeerInit connection type"); assert_eq!(init.read_string(), "P", "PeerInit connection type");
@@ -2079,7 +2079,7 @@ fn a_third_party_client_browses_our_shares_via_the_server_broker() {
let mut pierce = read_framed(&mut p).expect("a PierceFirewall frame"); let mut pierce = read_framed(&mut p).expect("a PierceFirewall frame");
assert_eq!( assert_eq!(
pierce.get_message_code(), pierce.get_init_code(),
0, 0,
"the brokered connect-back must start with a PierceFirewall" "the brokered connect-back must start with a PierceFirewall"
); );