Remove probability of collision in connection IDs by switching to u128 in Proxy

src/qos_core/src/io/stream.rs

@@ -382,9 +382,7 @@ mod test {
 		os::{fd::AsRawFd, unix::net::UnixListener},
 		path::Path,
 		str::from_utf8,
-		sync::mpsc,
 		thread,
-		time::Duration,
 	};
 
 	use super::*;

src/qos_net/src/error.rs

@@ -23,10 +23,14 @@ pub enum QosNetError {
 	ParseError(String),
 	/// DNS Resolution error
 	DNSResolutionError(String),
-	/// Attempt to save a connection with a duplicate ID
-	DuplicateConnectionId(u32),
+	/// Attempt to save a connection with a duplicate connection ID
+	DuplicateConnectionId(u128),
+	/// Error that should not happen: saving a connection overrode another
+	/// We take great care to ensure code prior to saving the connection checks
+	/// that it's not present in the connection map (and if it is, we return `DuplicateConnectionId`)
+	ConnectionOverridden(u128),
 	/// Attempt to send a message to a remote connection, but ID isn't found
-	ConnectionIdNotFound(u32),
+	ConnectionIdNotFound(u128),
 	/// Happens when a socket `read` returns too much data for the provided
 	/// buffer and the data doesn't fit. The first `usize` is the size of the
 	/// received data, the second `usize` is the size of the buffer.

src/qos_net/src/proxy.rs

@@ -1,5 +1,8 @@
 //! Protocol proxy for our remote QOS net proxy
-use std::io::{Read, Write};
+use std::{
+	collections::HashMap,
+	io::{Read, Write},
+};
 
 use borsh::BorshDeserialize;
 use qos_core::server;
@@ -9,6 +12,7 @@ use crate::{
 	proxy_connection::{self, ProxyConnection},
 	proxy_msg::ProxyMsg,
 };
+use rand::Rng;
 
 const MEGABYTE: usize = 1024 * 1024;
 const MAX_ENCODED_MSG_LEN: usize = 128 * MEGABYTE;
@@ -17,7 +21,7 @@ pub const DEFAULT_MAX_CONNECTION_SIZE: usize = 512;
 
 /// Socket<>TCP proxy to enable remote connections
 pub struct Proxy {
-	connections: Vec<ProxyConnection>,
+	connections: HashMap<u128, ProxyConnection>,
 	max_connections: usize,
 }
 
@@ -32,49 +36,63 @@ impl Proxy {
 	#[must_use]
 	pub fn new() -> Self {
 		Self {
-			connections: vec![],
+			connections: HashMap::new(),
 			max_connections: DEFAULT_MAX_CONNECTION_SIZE,
 		}
 	}
 
 	#[must_use]
 	pub fn new_with_max_connections(max_connections: usize) -> Self {
-		Self { connections: vec![], max_connections }
+		Self { connections: HashMap::new(), max_connections }
 	}
 
+	/// Save the connection in the proxy and assigns a connection ID
 	fn save_connection(
 		&mut self,
 		connection: ProxyConnection,
-	) -> Result<(), QosNetError> {
-		if self.connections.iter().any(|c| c.id == connection.id) {
-			Err(QosNetError::DuplicateConnectionId(connection.id))
-		} else {
-			if self.connections.len() >= self.max_connections {
-				return Err(QosNetError::TooManyConnections(
-					self.max_connections,
-				));
-			}
-			self.connections.push(connection);
-			Ok(())
+	) -> Result<u128, QosNetError> {
+		if self.connections.len() >= self.max_connections {
+			return Err(QosNetError::TooManyConnections(self.max_connections));
+		}
+		let connection_id = self.next_id();
+		if self.connections.contains_key(&connection_id) {
+			// This should never happen because "next_id" auto-increments
+			// Still, out of an abundance of caution, we error out here.
+			return Err(QosNetError::DuplicateConnectionId(connection_id));
 		}
+
+		match self.connections.insert(connection_id, connection) {
+			// Should never, ever happen because we checked above that the connection id was not present before proceeding.
+			// We explicitly handle this case here out of paranoia. If this happens, it means saving this connection
+			// overrode another. This is _very_ concerning.
+			Some(_) => Err(QosNetError::ConnectionOverridden(connection_id)),
+			// Normal case: no value was present before
+			None => Ok(connection_id),
+		}
+	}
+
+	// Simple convenience method to get the next connection ID
+	// We use a simple strategy here: pick a random u128.
+	fn next_id(&mut self) -> u128 {
+		rand::thread_rng().gen::<u128>()
 	}
 
-	fn remove_connection(&mut self, id: u32) -> Result<(), QosNetError> {
-		match self.connections.iter().position(|c| c.id == id) {
-			Some(i) => {
-				self.connections.remove(i);
+	fn remove_connection(&mut self, id: u128) -> Result<(), QosNetError> {
+		match self.get_connection(id) {
+			Some(_) => {
+				self.connections.remove(&id);
 				Ok(())
 			}
 			None => Err(QosNetError::ConnectionIdNotFound(id)),
 		}
 	}
 
-	fn get_connection(&mut self, id: u32) -> Option<&mut ProxyConnection> {
-		self.connections.iter_mut().find(|c| c.id == id)
+	fn get_connection(&mut self, id: u128) -> Option<&mut ProxyConnection> {
+		self.connections.get_mut(&id)
 	}
 
 	/// Close a connection by its ID
-	pub fn close(&mut self, connection_id: u32) -> ProxyMsg {
+	pub fn close(&mut self, connection_id: u128) -> ProxyMsg {
 		match self.shutdown_and_remove_connection(connection_id) {
 			Ok(_) => ProxyMsg::CloseResponse { connection_id },
 			Err(e) => ProxyMsg::ProxyError(e),
@@ -83,7 +101,7 @@ impl Proxy {
 
 	fn shutdown_and_remove_connection(
 		&mut self,
-		id: u32,
+		id: u128,
 	) -> Result<(), QosNetError> {
 		let conn = self
 			.get_connection(id)
@@ -113,10 +131,9 @@ impl Proxy {
 			dns_port,
 		) {
 			Ok(conn) => {
-				let connection_id = conn.id;
 				let remote_ip = conn.ip.clone();
 				match self.save_connection(conn) {
-					Ok(()) => {
+					Ok(connection_id) => {
 						println!(
 							"Connection to {hostname} established and saved as ID {connection_id}"
 						);
@@ -140,10 +157,9 @@ impl Proxy {
 	pub fn connect_by_ip(&mut self, ip: String, port: u16) -> ProxyMsg {
 		match proxy_connection::ProxyConnection::new_from_ip(ip.clone(), port) {
 			Ok(conn) => {
-				let connection_id = conn.id;
 				let remote_ip = conn.ip.clone();
 				match self.save_connection(conn) {
-					Ok(()) => {
+					Ok(connection_id) => {
 						println!("Connection to {ip} established and saved as ID {connection_id}");
 						ProxyMsg::ConnectResponse { connection_id, remote_ip }
 					}
@@ -161,7 +177,7 @@ impl Proxy {
 	}
 
 	/// Performs a Read on a connection
-	pub fn read(&mut self, connection_id: u32, size: usize) -> ProxyMsg {
+	pub fn read(&mut self, connection_id: u128, size: usize) -> ProxyMsg {
 		if let Some(conn) = self.get_connection(connection_id) {
 			let mut buf: Vec<u8> = vec![0; size];
 			match conn.read(&mut buf) {
@@ -203,7 +219,7 @@ impl Proxy {
 	}
 
 	/// Performs a Write on an existing connection
-	pub fn write(&mut self, connection_id: u32, data: Vec<u8>) -> ProxyMsg {
+	pub fn write(&mut self, connection_id: u128, data: Vec<u8>) -> ProxyMsg {
 		if let Some(conn) = self.get_connection(connection_id) {
 			match conn.write(&data) {
 				Ok(size) => ProxyMsg::WriteResponse { connection_id, size },
@@ -217,7 +233,7 @@ impl Proxy {
 	}
 
 	/// Performs a Flush on an existing TCP connection
-	pub fn flush(&mut self, connection_id: u32) -> ProxyMsg {
+	pub fn flush(&mut self, connection_id: u128) -> ProxyMsg {
 		if let Some(conn) = self.get_connection(connection_id) {
 			match conn.flush() {
 				Ok(_) => ProxyMsg::FlushResponse { connection_id },

src/qos_net/src/proxy_connection.rs

@@ -9,14 +9,11 @@ use hickory_resolver::{
 	config::{NameServerConfigGroup, ResolverConfig, ResolverOpts},
 	Resolver,
 };
-use rand::Rng;
 
 use crate::error::QosNetError;
 
 /// Struct representing a TCP connection held on our proxy
 pub struct ProxyConnection {
-	/// Unsigned integer with the connection ID (random positive int)
-	pub id: u32,
 	/// IP address of the remote host
 	pub ip: String,
 	/// TCP stream object
@@ -33,19 +30,10 @@ impl ProxyConnection {
 		dns_port: u16,
 	) -> Result<ProxyConnection, QosNetError> {
 		let ip = resolve_hostname(hostname, dns_resolvers, dns_port)?;
-
-		// Generate a new random u32 to get an ID. We'll use it to name our
-		// socket. This will be our connection ID.
-		let mut rng = rand::thread_rng();
-		let connection_id: u32 = rng.gen::<u32>();
-
 		let tcp_addr = SocketAddr::new(ip, port);
 		let tcp_stream = TcpStream::connect(tcp_addr)?;
-		Ok(ProxyConnection {
-			id: connection_id,
-			ip: ip.to_string(),
-			tcp_stream,
-		})
+
+		Ok(ProxyConnection { ip: ip.to_string(), tcp_stream })
 	}
 
 	/// Create a new `ProxyConnection` from an IP address. This results in a
@@ -54,16 +42,11 @@ impl ProxyConnection {
 		ip: String,
 		port: u16,
 	) -> Result<ProxyConnection, QosNetError> {
-		// Generate a new random u32 to get an ID. We'll use it to name our
-		// socket. This will be our connection ID.
-		let mut rng = rand::thread_rng();
-		let connection_id: u32 = rng.gen::<u32>();
-
 		let ip_addr = ip.parse()?;
 		let tcp_addr = SocketAddr::new(ip_addr, port);
 		let tcp_stream = TcpStream::connect(tcp_addr)?;
 
-		Ok(ProxyConnection { id: connection_id, ip, tcp_stream })
+		Ok(ProxyConnection { ip, tcp_stream })
 	}
 
 	/// Closes the underlying TCP connection (`Shutdown::Both`)

src/qos_net/src/proxy_msg.rs

@@ -37,31 +37,31 @@ pub enum ProxyMsg {
 	ConnectResponse {
 		/// Connection ID to reference the opened connection in later messages
 		/// (`Read`, `Write`, `Flush`)
-		connection_id: u32,
+		connection_id: u128,
 		/// The remote host IP, e.g. "1.2.3.4"
 		remote_ip: String,
 	},
 	/// Request from the enclave app to close the connection
 	CloseRequest {
 		/// Connection ID
-		connection_id: u32,
+		connection_id: u128,
 	},
 	/// Response for `CloseRequest`
 	CloseResponse {
 		/// Connection ID
-		connection_id: u32,
+		connection_id: u128,
 	},
 	/// Read from a remote connection
 	ReadRequest {
 		/// A connection ID from `ConnectResponse`
-		connection_id: u32,
+		connection_id: u128,
 		/// number of bytes to read
 		size: usize,
 	},
 	/// Response to `ReadRequest` containing read data
 	ReadResponse {
 		/// A connection ID from `ConnectResponse`
-		connection_id: u32,
+		connection_id: u128,
 		/// number of bytes read
 		data: Vec<u8>,
 		/// buffer after mutation from `read`. The first `size` bytes contain
@@ -71,28 +71,28 @@ pub enum ProxyMsg {
 	/// Write to a remote connection
 	WriteRequest {
 		/// A connection ID from `ConnectResponse`
-		connection_id: u32,
+		connection_id: u128,
 		/// Data to be sent
 		data: Vec<u8>,
 	},
 	/// Response to `WriteRequest` containing the number of successfully
 	/// written bytes.
 	WriteResponse {
 		/// Connection ID from `ConnectResponse`
-		connection_id: u32,
+		connection_id: u128,
 		/// Number of bytes written successfully
 		size: usize,
 	},
 	/// Write to a remote connection
 	FlushRequest {
 		/// A connection ID from `ConnectResponse`
-		connection_id: u32,
+		connection_id: u128,
 	},
 	/// Response to `FlushRequest`
 	/// The response only contains the connection ID. Success is implicit: if
 	/// the flush response fails, a `ProxyError` will be sent instead.
 	FlushResponse {
 		/// Connection ID from `ConnectResponse`
-		connection_id: u32,
+		connection_id: u128,
 	},
 }

src/qos_net/src/proxy_stream.rs

@@ -20,7 +20,7 @@ pub struct ProxyStream {
 	/// Once a connection is established (successful `ConnectByName` or
 	/// `ConnectByIp` request), this connection ID is set to the u32 in
 	/// `ConnectResponse`.
-	pub connection_id: u32,
+	pub connection_id: u128,
 	/// The remote host this connection points to
 	pub remote_hostname: Option<String>,
 	/// The remote IP this connection points to
@@ -437,7 +437,7 @@ mod test {
 	/// Useful in tests! :)
 	struct LocalStream {
 		proxy: Box<Proxy>,
-		pub connection_id: u32,
+		pub connection_id: u128,
 		pub remote_hostname: Option<String>,
 	}