From 7f1e19ebe039c8cd34ad8cdc6d8c7a123486f7fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Mon, 24 Oct 2022 00:47:28 +0200 Subject: [PATCH] socks5: add type for returned data and index --- common/socks5/ordered-buffer/src/buffer.rs | 63 +++++++++++-------- common/socks5/ordered-buffer/src/lib.rs | 2 +- .../src/connection_controller.rs | 18 +++--- 3 files changed, 47 insertions(+), 36 deletions(-) diff --git a/common/socks5/ordered-buffer/src/buffer.rs b/common/socks5/ordered-buffer/src/buffer.rs index bc6d8c8619..33c4d69142 100644 --- a/common/socks5/ordered-buffer/src/buffer.rs +++ b/common/socks5/ordered-buffer/src/buffer.rs @@ -13,6 +13,13 @@ pub struct OrderedMessageBuffer { messages: HashMap, } +/// Data returned from `OrderedMessageBuffer` on a successful read of gapless ordered data. +#[derive(Debug, PartialEq, Eq)] +pub struct ReadContiguousData { + pub data: Vec, + pub last_index: u64, +} + impl OrderedMessageBuffer { pub fn new() -> OrderedMessageBuffer { OrderedMessageBuffer { @@ -42,9 +49,9 @@ impl OrderedMessageBuffer { /// a read will return the bytes of messages 0, 1, 2. Subsequent reads will /// return `None` until message 3 comes in, at which point 3, 4, and any /// further contiguous messages which have arrived will be returned. - pub fn read(&mut self) -> (Option>, u64) { + pub fn read(&mut self) -> Option { if !self.messages.contains_key(&self.next_index) { - return (None, self.next_index); + return None; } let mut contiguous_messages = Vec::new(); @@ -66,7 +73,10 @@ impl OrderedMessageBuffer { .collect(); trace!("Returning {} bytes from ordered message buffer", data.len()); - (Some(data), index) + Some(ReadContiguousData { + data, + last_index: index, + }) } } @@ -102,14 +112,14 @@ mod test_chunking_and_reassembling { }; buffer.write(first_message); - let first_read = buffer.read().0.unwrap(); + let first_read = buffer.read().unwrap().data; assert_eq!(vec![1, 2, 3, 4], first_read); buffer.write(second_message); - let second_read = buffer.read().0.unwrap(); + let second_read = buffer.read().unwrap().data; assert_eq!(vec![5, 6, 7, 8], second_read); - assert_eq!(None, buffer.read().0); // second read on fully ordered result set is empty + assert_eq!(None, buffer.read()); // second read on fully ordered result set is empty } #[test] @@ -128,8 +138,8 @@ mod test_chunking_and_reassembling { buffer.write(first_message); buffer.write(second_message); let second_read = buffer.read(); - assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], second_read.0.unwrap()); - assert_eq!(None, buffer.read().0); // second read on fully ordered result set is empty + assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], second_read.unwrap().data); + assert_eq!(None, buffer.read()); // second read on fully ordered result set is empty } #[test] @@ -147,9 +157,9 @@ mod test_chunking_and_reassembling { buffer.write(second_message); buffer.write(first_message); - let read = buffer.read().0; - assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], read.unwrap()); - assert_eq!(None, buffer.read().0); // second read on fully ordered result set is empty + let read = buffer.read().unwrap().data; + assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], read); + assert_eq!(None, buffer.read()); // second read on fully ordered result set is empty } } @@ -182,11 +192,11 @@ mod test_chunking_and_reassembling { #[test] fn everything_up_to_the_indexing_gap_is_returned() { let mut buffer = setup(); - let ordered_bytes = buffer.read().0.unwrap(); + let ordered_bytes = buffer.read().unwrap().data; assert_eq!([0, 0, 0, 0, 1, 1, 1, 1].to_vec(), ordered_bytes); // we shouldn't get any more from a second attempt if nothing is added - assert_eq!(None, buffer.read().0); + assert_eq!(None, buffer.read()); // let's add another message, leaving a gap in place at index 2 let five_message = OrderedMessage { @@ -194,7 +204,7 @@ mod test_chunking_and_reassembling { index: 5, }; buffer.write(five_message); - assert_eq!(None, buffer.read().0); + assert_eq!(None, buffer.read()); } #[test] @@ -208,7 +218,7 @@ mod test_chunking_and_reassembling { }; buffer.write(two_message); - let more_ordered_bytes = buffer.read().0.unwrap(); + let more_ordered_bytes = buffer.read().unwrap().data; assert_eq!([2, 2, 2, 2, 3, 3, 3, 3].to_vec(), more_ordered_bytes); // let's add another message @@ -218,7 +228,7 @@ mod test_chunking_and_reassembling { }; buffer.write(five_message); - assert_eq!(None, buffer.read().0); + assert_eq!(None, buffer.read()); // let's fill in the gap of 4s now and read again let four_message = OrderedMessage { @@ -227,10 +237,13 @@ mod test_chunking_and_reassembling { }; buffer.write(four_message); - assert_eq!([4, 4, 4, 4, 5, 5, 5, 5].to_vec(), buffer.read().0.unwrap()); + assert_eq!( + [4, 4, 4, 4, 5, 5, 5, 5].to_vec(), + buffer.read().unwrap().data + ); // at this point we should again get back nothing if we try a read - assert_eq!(None, buffer.read().0); + assert_eq!(None, buffer.read()); } #[test] @@ -250,11 +263,11 @@ mod test_chunking_and_reassembling { }; buffer.write(zero_message); - assert!(buffer.read().0.is_some()); // burn the buffer + assert!(buffer.read().is_some()); // burn the buffer buffer.write(two_message); buffer.write(one_message); - assert!(buffer.read().0.is_some()); + assert!(buffer.read().is_some()); assert_eq!(buffer.next_index, 3); } @@ -282,23 +295,23 @@ mod test_chunking_and_reassembling { index: 4, }; buffer.write(zero_message); - assert!(buffer.read().0.is_some()); + assert!(buffer.read().is_some()); assert_eq!(buffer.next_index, 1); buffer.write(four_message); - assert!(buffer.read().0.is_none()); + assert!(buffer.read().is_none()); assert_eq!(buffer.next_index, 1); buffer.write(three_message); - assert!(buffer.read().0.is_none()); + assert!(buffer.read().is_none()); assert_eq!(buffer.next_index, 1); buffer.write(two_message); - assert!(buffer.read().0.is_none()); + assert!(buffer.read().is_none()); assert_eq!(buffer.next_index, 1); buffer.write(one_message); - assert!(buffer.read().0.is_some()); + assert!(buffer.read().is_some()); assert_eq!(buffer.next_index, 5) } } diff --git a/common/socks5/ordered-buffer/src/lib.rs b/common/socks5/ordered-buffer/src/lib.rs index 7a3c9e6501..9ed9d9c867 100644 --- a/common/socks5/ordered-buffer/src/lib.rs +++ b/common/socks5/ordered-buffer/src/lib.rs @@ -2,7 +2,7 @@ mod buffer; mod message; mod sender; -pub use buffer::OrderedMessageBuffer; +pub use buffer::{OrderedMessageBuffer, ReadContiguousData}; pub use message::MessageError; pub use message::OrderedMessage; pub use sender::OrderedMessageSender; diff --git a/common/socks5/proxy-helpers/src/connection_controller.rs b/common/socks5/proxy-helpers/src/connection_controller.rs index 139c124be3..a21131e82e 100644 --- a/common/socks5/proxy-helpers/src/connection_controller.rs +++ b/common/socks5/proxy-helpers/src/connection_controller.rs @@ -4,7 +4,7 @@ use futures::channel::mpsc; use futures::StreamExt; use log::*; -use ordered_buffer::{OrderedMessage, OrderedMessageBuffer}; +use ordered_buffer::{OrderedMessage, OrderedMessageBuffer, ReadContiguousData}; use socks5_requests::ConnectionId; use std::collections::{HashMap, HashSet}; use task::ShutdownListener; @@ -38,7 +38,7 @@ pub enum ControllerCommand { struct ActiveConnection { is_closed: bool, - closed_at: u64, + closed_at_index: u64, connection_sender: Option, ordered_buffer: OrderedMessageBuffer, } @@ -53,12 +53,12 @@ impl ActiveConnection { } }; if is_closed { - self.closed_at = ordered_message.index; + self.closed_at_index = ordered_message.index; } self.ordered_buffer.write(ordered_message); } - fn read_from_buf(&mut self) -> (Option>, u64) { + fn read_from_buf(&mut self) -> Option { self.ordered_buffer.read() } } @@ -103,7 +103,7 @@ impl Controller { is_closed: false, connection_sender: Some(connection_sender), ordered_buffer: OrderedMessageBuffer::new(), - closed_at: u64::MAX, + closed_at_index: u64::MAX, }; if let Some(_active_conn) = self.active_connections.insert(conn_id, active_connection) { error!("Received a duplicate 'Connect'!") @@ -136,11 +136,9 @@ impl Controller { } else if !is_closed { error!("Tried to write an empty message to a not-closing connection. Please let us know if you see this message"); } - // if messages get unordered, make sure we don't lose information about - // remote socket getting closed! - if let (Some(payload), last_index) = active_connection.read_from_buf() { - if last_index > active_connection.closed_at { + if let Some(payload) = active_connection.read_from_buf() { + if payload.last_index > active_connection.closed_at_index { active_connection.is_closed = true; } if let Err(err) = active_connection @@ -148,7 +146,7 @@ impl Controller { .as_mut() .unwrap() .unbounded_send(ConnectionMessage { - payload, + payload: payload.data, socket_closed: active_connection.is_closed, }) {