nineap/netrunner-proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dash map in muxer

Browse Source
This commit is contained in:
Kirill
2026-03-21 20:27:02 +07:00
parent 2777e7fd24
commit e39fa755db
8 changed files with 56 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+21
View File
@@ -597,6 +597,20 @@ dependencies = [
"syn 2.0.116", "syn 2.0.116",
] ]
[[package]]
name = "dashmap"
version = "6.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf"
dependencies = [
"cfg-if",
"crossbeam-utils",
"hashbrown 0.14.5",
"lock_api",
"once_cell",
"parking_lot_core",
]
[[package]] [[package]]
name = "data-encoding" name = "data-encoding"
version = "2.10.0" version = "2.10.0"
@@ -984,6 +998,12 @@ dependencies = [
"byteorder", "byteorder",
] ]
[[package]]
name = "hashbrown"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
[[package]] [[package]]
name = "hashbrown" name = "hashbrown"
version = "0.15.5" version = "0.15.5"
@@ -1525,6 +1545,7 @@ dependencies = [
"async-trait", "async-trait",
"bytes", "bytes",
"chacha20poly1305", "chacha20poly1305",
"dashmap",
"hex", "hex",
"hkdf", "hkdf",
"hmac", "hmac",
client/src/connections/tcp_connection.rs
+5 -5
View File
@@ -45,7 +45,7 @@ impl TcpConnection {
// ИСПРАВЛЕНИЕ: Даем Муксеру ограниченный канал, как он и просит (тип mpsc::Sender) // ИСПРАВЛЕНИЕ: Даем Муксеру ограниченный канал, как он и просит (тип mpsc::Sender)
// Делаем его достаточно вместительным (1024) // Делаем его достаточно вместительным (1024)
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024); let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024);
muxer.register_stream(stream_id, v_tx).await; muxer.register_stream(stream_id, v_tx);
let connect_payload = target_addr.to_string(); let connect_payload = target_addr.to_string();
if muxer if muxer
@@ -57,7 +57,7 @@ impl TcpConnection {
.await .await
.is_err() .is_err()
{ {
muxer.remove_stream(stream_id).await; muxer.remove_stream(stream_id);
return; return;
} }
@@ -68,13 +68,13 @@ impl TcpConnection {
let _ = handshake_tx.send(()); let _ = handshake_tx.send(());
} else { } else {
netrunner_logger::warn!(stream_id, "Server rejected TCP connection"); netrunner_logger::warn!(stream_id, "Server rejected TCP connection");
muxer.remove_stream(stream_id).await; muxer.remove_stream(stream_id);
return; return;
} }
} }
_ => { _ => {
netrunner_logger::error!(stream_id, "Timeout waiting for proxy response"); netrunner_logger::error!(stream_id, "Timeout waiting for proxy response");
muxer.remove_stream(stream_id).await; muxer.remove_stream(stream_id);
return; return;
} }
} }
@@ -118,7 +118,7 @@ impl TcpConnection {
data: Bytes::new(), data: Bytes::new(),
}) })
.await; .await;
muxer.remove_stream(stream_id).await; muxer.remove_stream(stream_id);
}); });
Self { Self {
client/src/connections/udp_connection.rs
+2 -2
View File
@@ -41,7 +41,7 @@ impl UdpConnection {
tokio::spawn(async move { tokio::spawn(async move {
// Даем Муксеру ограниченный канал (тип Sender), чтобы компилятор пропустил // Даем Муксеру ограниченный канал (тип Sender), чтобы компилятор пропустил
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024); let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024);
m_clone.register_stream(stream_id, v_tx).await; m_clone.register_stream(stream_id, v_tx);
let _ = m_clone let _ = m_clone
.send_to_netwrok(MuxMessage { .send_to_netwrok(MuxMessage {
@@ -81,7 +81,7 @@ impl UdpConnection {
data: Bytes::new(), data: Bytes::new(),
}) })
.await; .await;
m_clone.remove_stream(stream_id).await; m_clone.remove_stream(stream_id);
}); });
Self { Self {
core/Cargo.toml
+1
View File
@@ -17,4 +17,5 @@ rand = "0.10.0"
hex = "0.4.3" hex = "0.4.3"
tokio-util = { version = "0.7", features = ["rt"] } tokio-util = { version = "0.7", features = ["rt"] }
async-trait = "0.1.89" async-trait = "0.1.89"
dashmap = "6.1.0"
netrunner-logger = { path = "../tools/log" } netrunner-logger = { path = "../tools/log" }
core/src/proxy/connection/bridge.rs
+11 -4
View File
@@ -5,6 +5,7 @@ use bytes::{Bytes, BytesMut};
use netrunner_logger::{debug, error}; use netrunner_logger::{debug, error};
use tokio::net::UdpSocket; use tokio::net::UdpSocket;
use tokio::sync::mpsc; use tokio::sync::mpsc;
pub async fn run_proxy_bridge<R, W>( pub async fn run_proxy_bridge<R, W>(
stream_id: u32, stream_id: u32,
mut reader: R, mut reader: R,
@@ -58,6 +59,7 @@ pub async fn run_proxy_bridge<R, W>(
} }
} }
// Отправляем Close фрейм перед выходом
let _ = muxer let _ = muxer
.send_to_netwrok(MuxMessage { .send_to_netwrok(MuxMessage {
stream_id, stream_id,
@@ -65,8 +67,12 @@ pub async fn run_proxy_bridge<R, W>(
data: Bytes::new(), data: Bytes::new(),
}) })
.await; .await;
// Небольшая пауза, чтобы Close фрейм успел уйти в сеть
tokio::time::sleep(std::time::Duration::from_millis(500)).await; tokio::time::sleep(std::time::Duration::from_millis(500)).await;
muxer.remove_stream(stream_id).await;
// ИЗМЕНЕНИЕ: remove_stream теперь синхронный
muxer.remove_stream(stream_id);
} }
pub async fn run_udp_bridge( pub async fn run_udp_bridge(
@@ -79,7 +85,6 @@ pub async fn run_udp_bridge(
loop { loop {
tokio::select! { tokio::select! {
res = socket.recv(&mut buf) => { res = socket.recv(&mut buf) => {
match res { match res {
Ok(0) => break, Ok(0) => break,
@@ -98,7 +103,6 @@ pub async fn run_udp_bridge(
} }
} }
maybe_data = v_rx.recv() => { maybe_data = v_rx.recv() => {
match maybe_data { match maybe_data {
Some(data) => { Some(data) => {
@@ -124,6 +128,9 @@ pub async fn run_udp_bridge(
data: Bytes::new(), data: Bytes::new(),
}) })
.await; .await;
tokio::time::sleep(std::time::Duration::from_millis(500)).await; tokio::time::sleep(std::time::Duration::from_millis(500)).await;
muxer.remove_stream(stream_id).await;
// ИЗМЕНЕНИЕ: remove_stream теперь синхронный
muxer.remove_stream(stream_id);
} }
core/src/proxy/connection/connection.rs
+1 -1
View File
@@ -206,7 +206,7 @@ impl TunnelHandler for ClientHandler {
} => { } => {
let stream_id = self.muxer.next_id(); let stream_id = self.muxer.next_id();
let (v_tx, mut v_rx) = mpsc::channel::<bytes::Bytes>(1024); let (v_tx, mut v_rx) = mpsc::channel::<bytes::Bytes>(1024);
self.muxer.register_stream(stream_id, v_tx).await; self.muxer.register_stream(stream_id, v_tx);
self.muxer self.muxer
.send_to_netwrok(MuxMessage { .send_to_netwrok(MuxMessage {
core/src/proxy/connection/handler.rs
+4 -4
View File
@@ -42,7 +42,7 @@ impl StreamHandler {
let muxer = self.muxer.clone(); let muxer = self.muxer.clone();
let (v_tx, v_rx) = tokio::sync::mpsc::channel(512); let (v_tx, v_rx) = tokio::sync::mpsc::channel(512);
muxer.register_stream(stream_id, v_tx).await; muxer.register_stream(stream_id, v_tx);
tokio::spawn(async move { tokio::spawn(async move {
let start = std::time::Instant::now(); let start = std::time::Instant::now();
@@ -96,7 +96,7 @@ impl StreamHandler {
let muxer = self.muxer.clone(); let muxer = self.muxer.clone();
let (v_tx, v_rx) = tokio::sync::mpsc::channel(512); let (v_tx, v_rx) = tokio::sync::mpsc::channel(512);
muxer.register_stream(stream_id, v_tx).await; muxer.register_stream(stream_id, v_tx);
tokio::spawn(async move { tokio::spawn(async move {
info!(stream_id, target = %target_str, "Attempting remote UDP connection"); info!(stream_id, target = %target_str, "Attempting remote UDP connection");
@@ -142,7 +142,7 @@ impl StreamHandler {
} }
async fn send_error_reply(muxer: &Muxer, stream_id: u32, code: u8, frame_type: FrameType) { async fn send_error_reply(muxer: &Muxer, stream_id: u32, code: u8, frame_type: FrameType) {
muxer.remove_stream(stream_id).await; muxer.remove_stream(stream_id);
let mut reply_buf = BytesMut::with_capacity(10); let mut reply_buf = BytesMut::with_capacity(10);
let reply = SocksReply::ConnectResult { let reply = SocksReply::ConnectResult {
reply_code: code, reply_code: code,
@@ -162,6 +162,6 @@ impl StreamHandler {
async fn on_close(&self, stream_id: u32) { async fn on_close(&self, stream_id: u32) {
self.muxer.dispatch_to_local(stream_id, Bytes::new()).await; self.muxer.dispatch_to_local(stream_id, Bytes::new()).await;
self.muxer.remove_stream(stream_id).await; self.muxer.remove_stream(stream_id);
} }
} }
core/src/proxy/connection/muxer.rs
+11 -16
View File
@@ -1,11 +1,10 @@
use crate::protocol::codec::frame::FrameType; use crate::protocol::codec::frame::FrameType;
use bytes::Bytes; use bytes::Bytes;
use std::collections::HashMap; use dashmap::DashMap; // Добавь в Cargo.toml: dashmap = "6.0"
use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::mpsc::error::SendError; use tokio::sync::mpsc::error::SendError;
use tokio::sync::mpsc::Sender; use tokio::sync::mpsc::Sender;
use tokio::sync::RwLock;
pub struct IdGenerator { pub struct IdGenerator {
counter: AtomicU32, counter: AtomicU32,
@@ -33,7 +32,7 @@ pub struct MuxMessage {
#[derive(Clone)] #[derive(Clone)]
pub struct Muxer { pub struct Muxer {
to_network: Sender<MuxMessage>, to_network: Sender<MuxMessage>,
streams: Arc<RwLock<HashMap<u32, Sender<Bytes>>>>, streams: Arc<DashMap<u32, Sender<Bytes>>>,
id_gen: Arc<IdGenerator>, id_gen: Arc<IdGenerator>,
} }
@@ -41,7 +40,7 @@ impl Muxer {
pub fn new(to_network: Sender<MuxMessage>, is_client: bool) -> Self { pub fn new(to_network: Sender<MuxMessage>, is_client: bool) -> Self {
Self { Self {
to_network, to_network,
streams: Arc::new(RwLock::new(HashMap::new())), streams: Arc::new(DashMap::new()),
id_gen: Arc::new(IdGenerator::new(is_client)), id_gen: Arc::new(IdGenerator::new(is_client)),
} }
} }
@@ -54,19 +53,17 @@ impl Muxer {
self.to_network.send(message).await self.to_network.send(message).await
} }
pub async fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) { pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
let mut lock = self.streams.write().await; self.streams.insert(stream_id, tx);
lock.insert(stream_id, tx);
netrunner_logger::debug!( netrunner_logger::debug!(
stream_id, stream_id,
total_active = lock.len(), total_active = self.streams.len(),
"MUXER: [REGISTER] Stream added" "MUXER: [REGISTER] Stream added"
); );
} }
pub async fn remove_stream(&self, stream_id: u32) { pub fn remove_stream(&self, stream_id: u32) {
let mut lock = self.streams.write().await; self.streams.remove(&stream_id);
lock.remove(&stream_id);
} }
pub async fn send_control( pub async fn send_control(
@@ -86,10 +83,8 @@ impl Muxer {
} }
pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) { pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
let tx = { // DashMap позволяет получить доступ к элементу без явного RwLock
let lock = self.streams.read().await; let tx = self.streams.get(&stream_id).map(|r| r.value().clone());
lock.get(&stream_id).cloned()
};
if let Some(tx) = tx { if let Some(tx) = tx {
if data.is_empty() { if data.is_empty() {
@@ -107,7 +102,7 @@ impl Muxer {
stream_id, stream_id,
"MUXER: [WARN] Local channel closed, dropping packet" "MUXER: [WARN] Local channel closed, dropping packet"
); );
self.remove_stream(stream_id).await; self.remove_stream(stream_id);
} }
} else { } else {
netrunner_logger::trace!( netrunner_logger::trace!(