diff --git a/pyproject.toml b/pyproject.toml index 88a51c8..c54b6e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,6 +4,7 @@ build-backend = "maturin" [project] name = "ipv8-rust-tunnels" +description = "IPv8 tunnel performance enhancements" requires-python = ">=3.8" classifiers = [ "Programming Language :: Rust", diff --git a/src/socks5.rs b/src/socks5.rs index bb15277..fe73f3e 100644 --- a/src/socks5.rs +++ b/src/socks5.rs @@ -78,17 +78,9 @@ async fn process_socks5_packet( let pkt = &buf[header.serialized_len()..].to_vec(); let address = &header.address; - let circuit_id = match addr_to_cid.get(address) { - Some(cid) => *cid, - None => { - let options = get_options(&associated_socket, &circuits); - let cid = match options.choose(&mut rand::thread_rng()) { - Some(cid) => *cid, - None => return None, - }; - addr_to_cid.insert(address.clone(), cid); - cid - } + let Some(circuit_id) = select_circuit(address, associated_socket, circuits, addr_to_cid) else { + warn!("No circuits available, dropping packet"); + return None; }; match circuits.lock().unwrap().get_mut(&circuit_id) { @@ -115,6 +107,33 @@ async fn process_socks5_packet( } } +fn select_circuit( + address: &Address, + socket: &Arc, + circuits: &Arc>>, + addr_to_cid: &mut HashMap, +) -> Option { + // Remove dead circuits from addr_to_cid + if let Some(cid) = addr_to_cid.get(address) { + if !circuits.lock().unwrap().contains_key(cid) { + debug!("Not sending packet for {} over dead circuit {}, removing link", address, cid); + addr_to_cid.remove(address); + } + } + + match addr_to_cid.get(address) { + Some(cid) => Some(*cid), + None => { + let options = get_options(&socket, &circuits); + let Some(&cid) = options.choose(&mut rand::thread_rng()) else { + return None; + }; + addr_to_cid.insert(address.clone(), cid); + Some(cid) + } + } +} + fn get_options(socket: &Arc, circuits: &Arc>>) -> Vec { let guard = circuits.lock().unwrap(); guard