connection changes and buffers changes
This commit is contained in:
@@ -2,3 +2,5 @@ pub mod dns;
|
|||||||
pub mod ip_store;
|
pub mod ip_store;
|
||||||
pub mod tcp_connection;
|
pub mod tcp_connection;
|
||||||
pub mod udp_connection;
|
pub mod udp_connection;
|
||||||
|
|
||||||
|
pub const CHANNEL_CAPACITY: usize = 2048;
|
||||||
|
|||||||
@@ -18,24 +18,33 @@ pub enum ConnectionState {
|
|||||||
pub struct TcpConnection {
|
pub struct TcpConnection {
|
||||||
handle: SocketHandle,
|
handle: SocketHandle,
|
||||||
state: ConnectionState,
|
state: ConnectionState,
|
||||||
tx: mpsc::UnboundedSender<Vec<u8>>,
|
// UPLOAD: Ограниченный канал (Bounded) для Backpressure (128)
|
||||||
|
tx: mpsc::Sender<Vec<u8>>,
|
||||||
|
// DOWNLOAD: Безлимитный канал (Unbounded), чтобы Муксер НИКОГДА не зависал
|
||||||
rx: mpsc::UnboundedReceiver<Vec<u8>>,
|
rx: mpsc::UnboundedReceiver<Vec<u8>>,
|
||||||
pending_data: Vec<u8>,
|
pending_data: Vec<u8>,
|
||||||
handshake_rx: Option<oneshot::Receiver<()>>,
|
handshake_rx: Option<oneshot::Receiver<()>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
const MAX_PENDING: usize = 2 * 1024 * 1024;
|
// Подняли лимит буфера для загрузок до 32 Мегабайт.
|
||||||
|
// Этого хватит для Спидтеста на скорости ~200-300 Мбит/с без разрыва соединения.
|
||||||
|
const MAX_PENDING: usize = 32 * 1024 * 1024;
|
||||||
|
|
||||||
impl TcpConnection {
|
impl TcpConnection {
|
||||||
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
|
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
|
||||||
let (tx_to_mux, mut rx_from_smol) = mpsc::unbounded_channel::<Vec<u8>>();
|
// UPLOAD: Ограничиваем очередь
|
||||||
|
let (tx_to_mux, mut rx_from_smol) = mpsc::channel::<Vec<u8>>(128);
|
||||||
|
|
||||||
|
// DOWNLOAD: Безлимитный канал до синхронного tick()
|
||||||
let (tx_to_smol, rx_from_proxy) = mpsc::unbounded_channel::<Vec<u8>>();
|
let (tx_to_smol, rx_from_proxy) = mpsc::unbounded_channel::<Vec<u8>>();
|
||||||
let (handshake_tx, handshake_rx) = oneshot::channel();
|
let (handshake_tx, handshake_rx) = oneshot::channel();
|
||||||
|
|
||||||
let stream_id = muxer.next_id();
|
let stream_id = muxer.next_id();
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let (v_tx, mut v_rx) = mpsc::channel(64);
|
// ИСПРАВЛЕНИЕ: Даем Муксеру ограниченный канал, как он и просит (тип mpsc::Sender)
|
||||||
|
// Делаем его достаточно вместительным (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).await;
|
||||||
|
|
||||||
let connect_payload = target_addr.to_string();
|
let connect_payload = target_addr.to_string();
|
||||||
@@ -83,6 +92,9 @@ impl TcpConnection {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// МАГИЯ ЗДЕСЬ: Мы мгновенно читаем из ограниченного v_rx и переливаем
|
||||||
|
// в безлимитный tx_to_smol. send() у безлимитного канала никогда не блокируется.
|
||||||
|
// Поэтому v_rx всегда пустой, и Муксер никогда не зависнет!
|
||||||
let from_proxy = async {
|
let from_proxy = async {
|
||||||
while let Some(data) = v_rx.recv().await {
|
while let Some(data) = v_rx.recv().await {
|
||||||
if data.is_empty() {
|
if data.is_empty() {
|
||||||
@@ -113,7 +125,7 @@ impl TcpConnection {
|
|||||||
handle,
|
handle,
|
||||||
state: ConnectionState::Handshaking,
|
state: ConnectionState::Handshaking,
|
||||||
tx: tx_to_mux,
|
tx: tx_to_mux,
|
||||||
rx: rx_from_proxy,
|
rx: rx_from_proxy, // Это UnboundedReceiver, здесь ничего менять не надо
|
||||||
pending_data: vec![],
|
pending_data: vec![],
|
||||||
handshake_rx: Some(handshake_rx),
|
handshake_rx: Some(handshake_rx),
|
||||||
}
|
}
|
||||||
@@ -156,19 +168,17 @@ impl TcpConnection {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
ConnectionState::Closed => {
|
ConnectionState::Closed => {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
|
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn is_finished(&self, socket: &tcp::Socket) -> bool {
|
pub fn is_finished(&self, socket: &tcp::Socket) -> bool {
|
||||||
use tcp::State;
|
use tcp::State;
|
||||||
|
|
||||||
matches!(socket.state(), State::Closed | State::TimeWait)
|
matches!(socket.state(), State::Closed | State::TimeWait)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -177,25 +187,69 @@ impl TcpConnection {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn poll_and_process(&mut self, socket: &mut tcp::Socket) {
|
fn poll_and_process(&mut self, socket: &mut tcp::Socket) {
|
||||||
|
// 1. UPLOAD: Чанкинг и Backpressure (Защита от краша сервера)
|
||||||
if socket.can_recv() {
|
if socket.can_recv() {
|
||||||
let mut total_read = 0;
|
while socket.can_recv() {
|
||||||
|
let mut channel_full = false;
|
||||||
|
let mut channel_closed = false;
|
||||||
|
|
||||||
let _ = socket.recv(|data| {
|
let _ = socket.recv(|data| {
|
||||||
if !data.is_empty() {
|
if data.is_empty() {
|
||||||
if self.tx.send(data.to_vec()).is_ok() {
|
return (0, ());
|
||||||
total_read = data.len();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
(total_read, ())
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let chunk_size = std::cmp::min(data.len(), 16000);
|
||||||
|
let chunk = &data[..chunk_size];
|
||||||
|
|
||||||
|
match self.tx.try_send(chunk.to_vec()) {
|
||||||
|
Ok(_) => (chunk_size, ()),
|
||||||
|
Err(mpsc::error::TrySendError::Full(_)) => {
|
||||||
|
channel_full = true;
|
||||||
|
(0, ()) // Оставляем данные в smoltcp
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
channel_closed = true;
|
||||||
|
(0, ())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
if channel_full || channel_closed {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. DOWNLOAD: Сброс безлимитного канала в буфер
|
||||||
|
while let Ok(data) = self.rx.try_recv() {
|
||||||
|
if self.pending_data.is_empty() && socket.can_send() {
|
||||||
|
match socket.send_slice(&data) {
|
||||||
|
Ok(n) if n < data.len() => {
|
||||||
|
self.pending_data.extend_from_slice(&data[n..]);
|
||||||
|
}
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(_) => {
|
||||||
|
self.pending_data.extend_from_slice(&data);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
self.pending_data.extend_from_slice(&data);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Защита от OOM. 32MB более чем достаточно.
|
||||||
if self.pending_data.len() > MAX_PENDING {
|
if self.pending_data.len() > MAX_PENDING {
|
||||||
netrunner_logger::error!(%self.handle, "TCP Buffer overflow ({} bytes). Dropping connection.", self.pending_data.len());
|
netrunner_logger::error!(
|
||||||
|
%self.handle,
|
||||||
|
"TCP Buffer overflow ({} bytes). Dropping connection.",
|
||||||
|
self.pending_data.len()
|
||||||
|
);
|
||||||
socket.abort();
|
socket.abort();
|
||||||
self.state = ConnectionState::Closed;
|
self.state = ConnectionState::Closed;
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. DOWNLOAD: Отправка буфера в smoltcp
|
||||||
if !self.pending_data.is_empty() && socket.can_send() {
|
if !self.pending_data.is_empty() && socket.can_send() {
|
||||||
match socket.send_slice(&self.pending_data) {
|
match socket.send_slice(&self.pending_data) {
|
||||||
Ok(n) => {
|
Ok(n) => {
|
||||||
@@ -204,19 +258,5 @@ impl TcpConnection {
|
|||||||
Err(_) => {}
|
Err(_) => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if self.pending_data.is_empty() && socket.can_send() {
|
|
||||||
if let Ok(data) = self.rx.try_recv() {
|
|
||||||
match socket.send_slice(&data) {
|
|
||||||
Ok(n) if n < data.len() => {
|
|
||||||
self.pending_data.extend_from_slice(&data[n..]);
|
|
||||||
}
|
|
||||||
Ok(_) => {}
|
|
||||||
Err(_) => {
|
|
||||||
self.pending_data = data;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,11 +11,13 @@ use netrunner_core::{
|
|||||||
proxy::connection::muxer::{MuxMessage, Muxer},
|
proxy::connection::muxer::{MuxMessage, Muxer},
|
||||||
};
|
};
|
||||||
|
|
||||||
|
use crate::connections::CHANNEL_CAPACITY;
|
||||||
|
|
||||||
pub struct UdpConnection {
|
pub struct UdpConnection {
|
||||||
pub handle: SocketHandle,
|
pub handle: SocketHandle,
|
||||||
stream_id: u32,
|
stream_id: u32,
|
||||||
tx_to_net: mpsc::Sender<MuxMessage>,
|
tx_to_net: mpsc::Sender<MuxMessage>,
|
||||||
rx_from_net: mpsc::Receiver<Bytes>,
|
rx_from_net: mpsc::UnboundedReceiver<Bytes>, // <-- ТУТ
|
||||||
client_endpoint: Option<IpEndpoint>,
|
client_endpoint: Option<IpEndpoint>,
|
||||||
last_activity: Instant,
|
last_activity: Instant,
|
||||||
token: CancellationToken,
|
token: CancellationToken,
|
||||||
@@ -23,8 +25,6 @@ pub struct UdpConnection {
|
|||||||
|
|
||||||
const UDP_TIMEOUT: Duration = Duration::from_secs(60);
|
const UDP_TIMEOUT: Duration = Duration::from_secs(60);
|
||||||
|
|
||||||
const CHANNEL_CAPACITY: usize = 128;
|
|
||||||
|
|
||||||
impl UdpConnection {
|
impl UdpConnection {
|
||||||
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
|
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
|
||||||
let stream_id = muxer.next_id();
|
let stream_id = muxer.next_id();
|
||||||
@@ -33,11 +33,14 @@ impl UdpConnection {
|
|||||||
|
|
||||||
let (tx_to_net, mut rx_from_smol) = mpsc::channel::<MuxMessage>(CHANNEL_CAPACITY);
|
let (tx_to_net, mut rx_from_smol) = mpsc::channel::<MuxMessage>(CHANNEL_CAPACITY);
|
||||||
|
|
||||||
let (v_tx, v_rx) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY);
|
// ИСПРАВЛЕНИЕ: Внутренний безлимитный канал
|
||||||
|
let (internal_tx, internal_rx) = mpsc::unbounded_channel::<Bytes>();
|
||||||
|
|
||||||
let m_clone = muxer.clone();
|
let m_clone = muxer.clone();
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
|
// Даем Муксеру ограниченный канал (тип Sender), чтобы компилятор пропустил
|
||||||
|
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).await;
|
||||||
|
|
||||||
let _ = m_clone
|
let _ = m_clone
|
||||||
@@ -56,8 +59,18 @@ impl UdpConnection {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Мгновенно выкачиваем данные из Муксера в безлимитный внутренний канал
|
||||||
|
let from_proxy = async {
|
||||||
|
while let Some(data) = v_rx.recv().await {
|
||||||
|
if internal_tx.send(data).is_err() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = to_proxy => {}
|
_ = to_proxy => {}
|
||||||
|
_ = from_proxy => {} // Не забудь добавить этот блок в select!
|
||||||
_ = task_token.cancelled() => {}
|
_ = task_token.cancelled() => {}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -75,7 +88,7 @@ impl UdpConnection {
|
|||||||
handle,
|
handle,
|
||||||
stream_id,
|
stream_id,
|
||||||
tx_to_net,
|
tx_to_net,
|
||||||
rx_from_net: v_rx,
|
rx_from_net: internal_rx, // Сохраняем безлимитный ресивер
|
||||||
client_endpoint: None,
|
client_endpoint: None,
|
||||||
last_activity: Instant::now(),
|
last_activity: Instant::now(),
|
||||||
token,
|
token,
|
||||||
|
|||||||
@@ -358,7 +358,7 @@ impl ConnectionManager {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
|
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
|
||||||
let mut sockets = SocketSet::new(Vec::with_capacity(256));
|
let mut sockets = SocketSet::new(Vec::with_capacity(1024));
|
||||||
|
|
||||||
for _ in 0..n_icmp {
|
for _ in 0..n_icmp {
|
||||||
sockets.add(Self::create_icmp_socket());
|
sockets.add(Self::create_icmp_socket());
|
||||||
|
|||||||
Reference in New Issue
Block a user