GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

socks5: add type for returned data and index

Browse Source
This commit is contained in:
Jon Häggblad
2022-10-24 00:47:28 +02:00
parent 81f6c1b52b
commit 7f1e19ebe0
3 changed files with 47 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/socks5/ordered-buffer/src/buffer.rs
+38 -25
View File
@@ -13,6 +13,13 @@ pub struct OrderedMessageBuffer {
messages: HashMap<u64, OrderedMessage>,
}
/// Data returned from `OrderedMessageBuffer` on a successful read of gapless ordered data.
#[derive(Debug, PartialEq, Eq)]
pub struct ReadContiguousData {
pub data: Vec<u8>,
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<Vec<u8>>, u64) {
pub fn read(&mut self) -> Option<ReadContiguousData> {
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)
}
}
common/socks5/ordered-buffer/src/lib.rs
+1 -1
View File
@@ -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;
common/socks5/proxy-helpers/src/connection_controller.rs
+8 -10
View File
@@ -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<ConnectionSender>,
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<Vec<u8>>, u64) {
fn read_from_buf(&mut self) -> Option<ReadContiguousData> {
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,
})
{