todo rewrite client
This commit is contained in:
@@ -0,0 +1,130 @@
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use netrunner_common::proxy::connection::muxer::{MuxMessage, Muxer};
|
||||
use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet, SocketStorage};
|
||||
use smoltcp::phy::{DeviceCapabilities, Medium};
|
||||
use smoltcp::socket::tcp::{Socket as SmolTcpSocket, SocketBuffer};
|
||||
use smoltcp::time::Instant;
|
||||
use smoltcp::wire::IpListenEndpoint;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
|
||||
use tracing::{debug, info, trace};
|
||||
|
||||
use crate::tun::virt_device::{TokenBuffer, VirtTunDevice};
|
||||
|
||||
pub enum BridgeState {
|
||||
WaitingHandshake,
|
||||
WaitingConnect,
|
||||
DataTransferring {
|
||||
tx_to_muxer: tokio::sync::mpsc::Sender<MuxMessage>,
|
||||
stream_id: u32,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct NetStack {
|
||||
interface: Interface,
|
||||
sockets: SocketSet<'static>,
|
||||
socket_buffers: HashMap<SocketHandle, BytesMut>,
|
||||
bridges: HashMap<SocketHandle, BridgeState>,
|
||||
muxer: Muxer,
|
||||
input: UnboundedSender<TokenBuffer>,
|
||||
output: UnboundedReceiver<TokenBuffer>,
|
||||
tun_available: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref START_TIME: std::time::Instant = std::time::Instant::now();
|
||||
}
|
||||
|
||||
fn current_smoltcp_time() -> Instant {
|
||||
let nanos = START_TIME.elapsed().as_micros() as i64;
|
||||
smoltcp::time::Instant::from_micros(nanos)
|
||||
}
|
||||
|
||||
impl NetStack {
|
||||
pub fn new(muxer: Muxer) -> Self {
|
||||
let now = current_smoltcp_time();
|
||||
|
||||
let mut caps = DeviceCapabilities::default();
|
||||
caps.medium = Medium::Ip;
|
||||
caps.max_transmission_unit = 1500;
|
||||
|
||||
let (virt_device, iface_output, iface_input, in_buf_avail) = VirtTunDevice::new(caps);
|
||||
let device = Box::leak(Box::new(virt_device));
|
||||
|
||||
let config = Config::new(smoltcp::wire::HardwareAddress::Ip);
|
||||
let mut interface = Interface::new(config, device, now);
|
||||
interface.set_any_ip(true);
|
||||
|
||||
interface.update_ip_addrs(|addrs| {
|
||||
addrs
|
||||
.push(smoltcp::wire::IpCidr::new(
|
||||
smoltcp::wire::IpAddress::v4(10, 0, 0, 2),
|
||||
24,
|
||||
))
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
interface
|
||||
.routes_mut()
|
||||
.add_default_ipv4_route(smoltcp::wire::Ipv4Address::new(10, 0, 0, 2))
|
||||
.unwrap();
|
||||
|
||||
let rx_data = Box::leak(vec![0u8; 4096].into_boxed_slice());
|
||||
let tx_data = Box::leak(vec![0u8; 4096].into_boxed_slice());
|
||||
|
||||
let mut socket = SmolTcpSocket::new(SocketBuffer::new(rx_data), SocketBuffer::new(tx_data));
|
||||
|
||||
// 4. Теперь listen() сработает, так как у интерфейса есть адрес!
|
||||
let endpoint = IpListenEndpoint {
|
||||
addr: None, // ВАЖНО: Принимаем пакеты для ЛЮБОГО IP назначения
|
||||
port: 443, // Для HTTPS
|
||||
};
|
||||
socket.listen(endpoint).unwrap();
|
||||
|
||||
let storage: Vec<SocketStorage> = (0..16).map(|_| SocketStorage::EMPTY).collect();
|
||||
let mut sockets = SocketSet::new(Box::leak(storage.into_boxed_slice()));
|
||||
sockets.add(socket);
|
||||
|
||||
Self {
|
||||
interface,
|
||||
sockets,
|
||||
socket_buffers: HashMap::new(),
|
||||
bridges: HashMap::new(),
|
||||
muxer,
|
||||
input: iface_input,
|
||||
output: iface_output,
|
||||
tun_available: in_buf_avail,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn process_tun_input(&mut self, data: &[u8]) {
|
||||
let mut token = TokenBuffer::with_capacity(data.len());
|
||||
token.extend_from_slice(data);
|
||||
|
||||
if self.input.send(token).is_ok() {
|
||||
self.tun_available
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn next_outbound_packet(&mut self) -> Option<TokenBuffer> {
|
||||
self.output.recv().await
|
||||
}
|
||||
|
||||
pub fn poll_delay(&mut self) -> tokio::time::Sleep {
|
||||
let timestamp = current_smoltcp_time();
|
||||
let ms = self
|
||||
.interface
|
||||
.poll_delay(timestamp, &self.sockets)
|
||||
.map(|d| d.total_millis())
|
||||
.unwrap_or(10); // Порог отзывчивости стека
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(ms))
|
||||
}
|
||||
|
||||
pub fn poll(&mut self) {
|
||||
let timestamp = current_smoltcp_time();
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -1,2 +1,2 @@
|
||||
pub mod stack;
|
||||
pub mod interface;
|
||||
pub mod tun;
|
||||
|
||||
+45
-36
@@ -1,29 +1,27 @@
|
||||
mod stack;
|
||||
mod tun;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use netrunner_common::{
|
||||
logger_init, proxy::connection::connection::ConnectionRole, proxy::network::Network,
|
||||
use netrunner_client::{
|
||||
interface::NetStack,
|
||||
tun::{
|
||||
linux_tun_device::create_linux_tun,
|
||||
virt_device::{TokenBuffer, VirtTunDevice},
|
||||
},
|
||||
};
|
||||
use stack::interface::NetStack;
|
||||
use netrunner_common::{
|
||||
logger_init,
|
||||
proxy::{connection::connection::ConnectionRole, network::Network},
|
||||
};
|
||||
use smoltcp::phy::{DeviceCapabilities, Medium};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tracing::{error, info};
|
||||
use tun::{desktop::create_linux_tun, device::TunDevice};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
// 1. Инициализация логов
|
||||
let mut tun = create_linux_tun();
|
||||
logger_init();
|
||||
info!("Starting NetRunner VPN Bridge...");
|
||||
|
||||
// 2. Настройка физического уровня (TUN)
|
||||
let tun_dev = create_linux_tun();
|
||||
let my_phy = TunDevice::new(tun_dev, 1500);
|
||||
|
||||
// 3. Создаем объект Network
|
||||
let net = Network::new(8080, ConnectionRole::Client, Some("0.0.0.0:4443".into()));
|
||||
|
||||
// 4. Инициализируем ЕДИНЫЙ туннель для всего приложения.
|
||||
// Этот метод внутри создает TLS-подключение, Muxer и запускает TunnelEngine.
|
||||
info!("Initializing global TLS tunnel to proxy...");
|
||||
let muxer = match net.initialize_client_tunnel().await {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
@@ -32,13 +30,7 @@ async fn main() {
|
||||
}
|
||||
};
|
||||
|
||||
// 5. Создаем стек, передавая ему РАБОЧИЙ муксер.
|
||||
// Теперь данные из TUN будут уходить в реальный TLS-туннель.
|
||||
let mut stack = NetStack::new(my_phy, muxer.clone());
|
||||
|
||||
// 6. Запускаем SOCKS-сервер (Network), чтобы он слушал порт 8080
|
||||
// и использовал тот же самый муксер для обычных прокси-запросов.
|
||||
let net_handle = {
|
||||
let mut net_handle = {
|
||||
let muxer_for_net = muxer.clone();
|
||||
tokio::spawn(async move {
|
||||
info!("SOCKS5 server starting on 127.0.0.1:8080");
|
||||
@@ -46,22 +38,39 @@ async fn main() {
|
||||
})
|
||||
};
|
||||
|
||||
info!("VPN BRIDGE IS RUNNING");
|
||||
let mut net_stack = NetStack::new(muxer);
|
||||
|
||||
// 7. Запускаем цикл обработки стека (блокирующий поток)
|
||||
let stack_loop = tokio::task::spawn_blocking(move || loop {
|
||||
stack.poll();
|
||||
});
|
||||
loop {
|
||||
// 1. Сначала ВСЕГДА даем стеку поработать (обработать то, что уже пришло)
|
||||
net_stack.poll();
|
||||
|
||||
// Ждем завершения (по сути бесконечно)
|
||||
tokio::select! {
|
||||
res = stack_loop => {
|
||||
if let Err(e) = res {
|
||||
error!("Stack loop panicked: {:?}", e);
|
||||
let mut buf = [0u8; 1600];
|
||||
|
||||
tokio::select! {
|
||||
// Читаем из реального мира (TUN) и закидываем в очередь стека
|
||||
tun_res = tun.read(&mut buf) => {
|
||||
if let Ok(n) = tun_res {
|
||||
// Прямой вызов обработки (без лишних каналов, если это один поток)
|
||||
// Или через твой метод, если логика разделена:
|
||||
net_stack.process_tun_input(&buf[..n]);
|
||||
}
|
||||
}
|
||||
|
||||
// Читаем из очереди стека и отдаем в реальный мир (TUN)
|
||||
Some(packet_to_tun) = net_stack.next_outbound_packet() => {
|
||||
let _ = tun.write_all(&packet_to_tun).await;
|
||||
}
|
||||
|
||||
// Ждем, пока стек сам попросит проснуться (таймеры TCP)
|
||||
_ = net_stack.poll_delay() => {
|
||||
// Просто просыпаемся. На следующей итерации вызовется poll()
|
||||
}
|
||||
|
||||
// Ошибка прокси
|
||||
net_res = &mut net_handle => {
|
||||
error!("SOCKS5 server stopped: {:?}", net_res);
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ = net_handle => {
|
||||
error!("Network server stopped unexpectedly");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,353 +0,0 @@
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use netrunner_common::protocol::codec::frame::FrameType;
|
||||
use netrunner_common::protocol::codec::socks::{SocksReply, SocksRequest};
|
||||
use netrunner_common::protocol::parser::parser::Parser;
|
||||
use netrunner_common::proxy::connection::muxer::{MuxMessage, Muxer};
|
||||
use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet, SocketStorage};
|
||||
use smoltcp::socket::tcp::{Socket as SmolTcpSocket, SocketBuffer};
|
||||
use smoltcp::socket::AnySocket;
|
||||
use smoltcp::time::Instant;
|
||||
use smoltcp::wire::IpListenEndpoint;
|
||||
use std::collections::HashMap;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use tracing::{debug, info, trace};
|
||||
|
||||
pub enum BridgeState {
|
||||
WaitingHandshake,
|
||||
WaitingConnect,
|
||||
DataTransferring {
|
||||
tx_to_muxer: tokio::sync::mpsc::Sender<MuxMessage>,
|
||||
stream_id: u32,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct NetStack<D: smoltcp::phy::Device + AsRawFd + 'static> {
|
||||
interface: Interface,
|
||||
device: &'static mut D,
|
||||
sockets: SocketSet<'static>,
|
||||
socket_buffers: HashMap<SocketHandle, BytesMut>,
|
||||
bridges: HashMap<SocketHandle, BridgeState>,
|
||||
muxer: Muxer,
|
||||
outbound_rx: tokio::sync::mpsc::UnboundedReceiver<(SocketHandle, Bytes)>,
|
||||
outbound_tx: tokio::sync::mpsc::UnboundedSender<(SocketHandle, Bytes)>,
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref START_TIME: std::time::Instant = std::time::Instant::now();
|
||||
}
|
||||
|
||||
fn _current_smoltcp_time() -> Instant {
|
||||
let nanos = START_TIME.elapsed().as_micros() as i64;
|
||||
smoltcp::time::Instant::from_micros(nanos)
|
||||
}
|
||||
|
||||
impl<D: smoltcp::phy::Device + AsRawFd> NetStack<D> {
|
||||
pub fn new(device_obj: D, muxer: Muxer) -> Self {
|
||||
let now = _current_smoltcp_time();
|
||||
let device: &'static mut D = Box::leak(Box::new(device_obj));
|
||||
|
||||
let config = Config::new(smoltcp::wire::HardwareAddress::Ip);
|
||||
let mut interface = Interface::new(config, device, now);
|
||||
interface.set_any_ip(true);
|
||||
|
||||
interface.update_ip_addrs(|addrs| {
|
||||
addrs
|
||||
.push(smoltcp::wire::IpCidr::new(
|
||||
smoltcp::wire::IpAddress::v4(10, 0, 0, 2),
|
||||
24,
|
||||
))
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
interface
|
||||
.routes_mut()
|
||||
.add_default_ipv4_route(smoltcp::wire::Ipv4Address::new(10, 0, 0, 1))
|
||||
.unwrap();
|
||||
|
||||
let rx_data = Box::leak(vec![0u8; 4096].into_boxed_slice());
|
||||
let tx_data = Box::leak(vec![0u8; 4096].into_boxed_slice());
|
||||
|
||||
let mut socket = SmolTcpSocket::new(SocketBuffer::new(rx_data), SocketBuffer::new(tx_data));
|
||||
|
||||
// 4. Теперь listen() сработает, так как у интерфейса есть адрес!
|
||||
let endpoint = IpListenEndpoint {
|
||||
addr: None, // ВАЖНО: Принимаем пакеты для ЛЮБОГО IP назначения
|
||||
port: 443, // Для HTTPS
|
||||
};
|
||||
socket.listen(endpoint).unwrap();
|
||||
|
||||
let storage: Vec<SocketStorage> = (0..16).map(|_| SocketStorage::EMPTY).collect();
|
||||
let mut sockets = SocketSet::new(Box::leak(storage.into_boxed_slice()));
|
||||
sockets.add(socket);
|
||||
|
||||
let (outbound_tx, outbound_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
|
||||
Self {
|
||||
interface,
|
||||
device,
|
||||
sockets,
|
||||
socket_buffers: HashMap::new(),
|
||||
bridges: HashMap::new(),
|
||||
muxer,
|
||||
outbound_rx,
|
||||
outbound_tx,
|
||||
}
|
||||
}
|
||||
|
||||
fn drive_interface(&mut self, timestamp: Instant) {
|
||||
match self
|
||||
.interface
|
||||
.poll(timestamp, self.device, &mut self.sockets)
|
||||
{
|
||||
smoltcp::iface::PollResult::None => {}
|
||||
res => debug!("Interface activity: {:?}", res),
|
||||
}
|
||||
let delay = self
|
||||
.interface
|
||||
.poll_delay(timestamp, &self.sockets)
|
||||
.map(|d| d.total_millis() as i32)
|
||||
.unwrap_or(10);
|
||||
let raw_fd = self.device.as_raw_fd();
|
||||
let mut fds = [libc::pollfd {
|
||||
fd: raw_fd,
|
||||
events: libc::POLLIN, // Ждем данные на вход
|
||||
revents: 0,
|
||||
}];
|
||||
|
||||
unsafe {
|
||||
libc::poll(fds.as_mut_ptr(), 1, delay);
|
||||
}
|
||||
}
|
||||
|
||||
fn manage_socket_lifecycle(&mut self) {
|
||||
for (_, any_socket) in self.sockets.iter_mut() {
|
||||
let socket = SmolTcpSocket::downcast_mut(any_socket).unwrap();
|
||||
|
||||
// Если сокет закрылся или "отвисел" в TimeWait, возвращаем его в Listen
|
||||
if !socket.is_active() && !socket.is_listening() {
|
||||
let endpoint = IpListenEndpoint {
|
||||
addr: None,
|
||||
port: 443,
|
||||
};
|
||||
socket.listen(endpoint).unwrap();
|
||||
println!("DEBUG: Сокет готов к новому соединению");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn process_socks_logic(&mut self, handle: SocketHandle) {
|
||||
let buf = self.socket_buffers.get_mut(&handle).unwrap();
|
||||
let state = self
|
||||
.bridges
|
||||
.entry(handle)
|
||||
.or_insert(BridgeState::WaitingHandshake);
|
||||
|
||||
trace!(handle = ?handle, buf_len = buf.len(), "Processing SOCKS logic");
|
||||
match state {
|
||||
BridgeState::WaitingHandshake => {
|
||||
if let Ok(Some(SocksRequest::Handshake { .. })) = SocksRequest::parse(buf) {
|
||||
let mut reply = BytesMut::with_capacity(2);
|
||||
SocksReply::HandshakeSelect { method: 0x00 }.write_to(&mut reply);
|
||||
|
||||
let socket = self.sockets.get_mut::<SmolTcpSocket>(handle);
|
||||
socket.send_slice(&reply).unwrap();
|
||||
|
||||
*state = BridgeState::WaitingConnect;
|
||||
return; // <--- ВАЖНО: Выходим, чтобы smoltcp отправил этот пакет отдельно
|
||||
}
|
||||
}
|
||||
BridgeState::WaitingConnect => {
|
||||
// 1. Парсим запрос (парсер должен сам отрезать байты через advance/split_to)
|
||||
if let Ok(Some(SocksRequest::Connect { target, .. })) = SocksRequest::parse(buf) {
|
||||
let stream_id = self.muxer.next_id();
|
||||
let target_str = target.to_string();
|
||||
info!(handle = ?handle, target = %target_str, stream_id, "SOCKS5 Connect request");
|
||||
|
||||
let (v_tx, mut v_rx) = tokio::sync::mpsc::channel::<Bytes>(1024);
|
||||
|
||||
// 2. Пытаемся зарегистрировать стрим СИНХРОННО
|
||||
// Это гарантирует, что Muxer узнает об ID до того, как придет ответ из сети
|
||||
if self.muxer.try_register_stream(stream_id, v_tx.clone()) {
|
||||
debug!(stream_id, "Stream registered synchronously via try_write");
|
||||
|
||||
// Сразу шлем Connect в сеть
|
||||
let _ = self.muxer.to_network.try_send(MuxMessage {
|
||||
stream_id,
|
||||
frame_type: FrameType::Connect,
|
||||
data: Bytes::from(target_str),
|
||||
});
|
||||
} else {
|
||||
// ПЛАН Б: Если лок занят, спавним асинхронную задачу.
|
||||
// ВАЖНО: Мы НЕ шлем Connect здесь, его пришлет сама задача ПОСЛЕ регистрации.
|
||||
let muxer_clone = self.muxer.clone();
|
||||
let target_bytes = Bytes::from(target_str);
|
||||
let v_tx_clone = v_tx.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
muxer_clone.register_stream(stream_id, v_tx_clone).await;
|
||||
let _ = muxer_clone
|
||||
.to_network
|
||||
.send(MuxMessage {
|
||||
stream_id,
|
||||
frame_type: FrameType::Connect,
|
||||
data: target_bytes,
|
||||
})
|
||||
.await;
|
||||
debug!(
|
||||
stream_id,
|
||||
"Stream registered asynchronously (lock was busy)"
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
// 3. Запускаем задачу проброса данных из сети в smoltcp
|
||||
let outbound_tx = self.outbound_tx.clone();
|
||||
let handle_clone = handle;
|
||||
tokio::spawn(async move {
|
||||
while let Some(data) = v_rx.recv().await {
|
||||
if outbound_tx.send((handle_clone, data)).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
buf.clear();
|
||||
// 4. Отвечаем клиенту (curl)
|
||||
let mut reply = BytesMut::new();
|
||||
SocksReply::ConnectResult {
|
||||
reply_code: 0x00,
|
||||
atyp: 0x01,
|
||||
addr: [0, 0, 0, 0],
|
||||
port: 0,
|
||||
}
|
||||
.write_to(&mut reply);
|
||||
|
||||
let socket = self.sockets.get_mut::<SmolTcpSocket>(handle);
|
||||
socket
|
||||
.send_slice(&reply)
|
||||
.expect("Failed to send SOCKS reply");
|
||||
|
||||
// ПЕРЕХОД
|
||||
*state = BridgeState::DataTransferring {
|
||||
tx_to_muxer: self.muxer.to_network.clone(),
|
||||
stream_id,
|
||||
};
|
||||
return;
|
||||
}
|
||||
}
|
||||
BridgeState::DataTransferring {
|
||||
tx_to_muxer,
|
||||
stream_id,
|
||||
} => {
|
||||
if !buf.is_empty() {
|
||||
// Забираем все накопленные данные из буфера
|
||||
let data = buf.split().freeze();
|
||||
trace!(handle = ?handle, stream_id = *stream_id, bytes = data.len(), "Forwarding data smoltcp -> muxer");
|
||||
// Отправляем в Muxer
|
||||
let _ = tx_to_muxer.try_send(MuxMessage {
|
||||
stream_id: *stream_id,
|
||||
frame_type: FrameType::Data,
|
||||
data,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn process_payloads(&mut self) {
|
||||
let handles: Vec<_> = self.sockets.iter().map(|(h, _)| h).collect();
|
||||
|
||||
for handle in handles {
|
||||
let socket = self.sockets.get_mut::<SmolTcpSocket>(handle);
|
||||
|
||||
// 1. СОСТОЯНИЕ: Соединение только что установилось
|
||||
if socket.is_active() && !self.bridges.contains_key(&handle) {
|
||||
if let Some(target) = socket.local_endpoint() {
|
||||
let stream_id = self.muxer.next_id();
|
||||
info!("TUN Intercept: Auto-initiating SOCKS for target {}", target);
|
||||
|
||||
// ВАЖНО: Мы САМИ инициируем процесс для локального прокси
|
||||
// Регистрируем стрим в муксере
|
||||
let (v_tx, mut v_rx) = tokio::sync::mpsc::channel::<Bytes>(1024);
|
||||
self.muxer.try_register_stream(stream_id, v_tx);
|
||||
|
||||
// Шлем CONNECT (твой муксер/прокси поймет это как SOCKS-запрос)
|
||||
let _ = self.muxer.to_network.try_send(MuxMessage {
|
||||
stream_id,
|
||||
frame_type: FrameType::Connect,
|
||||
data: Bytes::from(target.to_string()),
|
||||
});
|
||||
|
||||
// Запускаем стандартную задачу проброса данных ИЗ сети в smoltcp
|
||||
let outbound_tx = self.outbound_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(data) = v_rx.recv().await {
|
||||
let _ = outbound_tx.send((handle, data));
|
||||
}
|
||||
});
|
||||
|
||||
// Переходим сразу в режим передачи данных
|
||||
// Теперь всё, что пришлет Firefox, мы будем просто гнать в Muxer как FrameType::Data
|
||||
self.bridges.insert(
|
||||
handle,
|
||||
BridgeState::DataTransferring {
|
||||
stream_id,
|
||||
tx_to_muxer: self.muxer.to_network.clone(),
|
||||
},
|
||||
);
|
||||
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// 2. СОСТОЯНИЕ: Передача данных (DataTransferring)
|
||||
if let Some(BridgeState::DataTransferring {
|
||||
stream_id,
|
||||
tx_to_muxer,
|
||||
}) = self.bridges.get(&handle)
|
||||
{
|
||||
if socket.can_recv() {
|
||||
let mut temp_buf = vec![0u8; 2048];
|
||||
if let Ok(size) = socket.recv_slice(&mut temp_buf) {
|
||||
// Это РЕАЛЬНЫЕ данные от Firefox (например, TLS Client Hello)
|
||||
// Мы их просто оборачиваем в твой протокол и шлем в муксер
|
||||
let _ = tx_to_muxer.try_send(MuxMessage {
|
||||
stream_id: *stream_id,
|
||||
frame_type: FrameType::Data,
|
||||
data: Bytes::copy_from_slice(&temp_buf[..size]),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn poll(&mut self) {
|
||||
let timestamp = _current_smoltcp_time();
|
||||
|
||||
self.drive_interface(timestamp);
|
||||
|
||||
self.process_payloads();
|
||||
|
||||
while let Ok((handle, data)) = self.outbound_rx.try_recv() {
|
||||
if let Some(socket_state) = self.bridges.get(&handle) {
|
||||
// ПИШЕМ ТОЛЬКО ЕСЛИ МЫ УЖЕ В DATA TRANSFER
|
||||
if let BridgeState::DataTransferring { .. } = socket_state {
|
||||
if self.sockets.iter().any(|(h, _)| h == handle) {
|
||||
let socket = self.sockets.get_mut::<SmolTcpSocket>(handle);
|
||||
if socket.can_send() {
|
||||
debug!(handle = ?handle, len = data.len(), head = ?&data[..std::cmp::min(data.len(), 10)], "Writing to curl");
|
||||
let _ = socket.send_slice(&data);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Если мы еще в WaitingConnect, возвращаем данные в очередь
|
||||
// или просто игнорируем (они придут следующим тиком)
|
||||
debug!(handle = ?handle, "Data arrived before SOCKS handshake finished, delaying...");
|
||||
// Можно использовать self.outbound_tx.send(...) чтобы вернуть в хвост
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.manage_socket_lifecycle();
|
||||
}
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
pub mod interface;
|
||||
@@ -1,14 +0,0 @@
|
||||
use tun::Device;
|
||||
|
||||
pub fn create_linux_tun() -> Device {
|
||||
let mut config = tun::Configuration::default();
|
||||
config
|
||||
.tun_name("netr0")
|
||||
.address((10, 0, 0, 1)) // IP нашего "моста"
|
||||
.netmask((255, 255, 255, 0)) // Маска подсети
|
||||
.up(); // Сразу включаем интерфейс
|
||||
|
||||
// Это создаст в системе интерфейс tun0
|
||||
let dev = tun::create(&config).expect("Нужны права root или CAP_NET_ADMIN!");
|
||||
dev
|
||||
}
|
||||
@@ -1,105 +0,0 @@
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use smoltcp::phy::{ChecksumCapabilities, Device, DeviceCapabilities, Medium, RxToken, TxToken};
|
||||
use std::io::{Read, Write};
|
||||
use std::os::fd::RawFd;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
pub struct TunDevice<T: Read + Write + AsRawFd> {
|
||||
pub io: T, // Сюда мы запихнем или File (Android) или Tun-либу
|
||||
mtu: usize,
|
||||
read_buffer: BytesMut,
|
||||
}
|
||||
|
||||
impl<T: Read + Write + AsRawFd> TunDevice<T> {
|
||||
pub fn new(io: T, mtu: usize) -> Self {
|
||||
let fd = io.as_raw_fd();
|
||||
unsafe {
|
||||
let flags = libc::fcntl(fd, libc::F_GETFL);
|
||||
if flags == -1 {
|
||||
panic!("Не удалось получить флаги файла");
|
||||
}
|
||||
libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK);
|
||||
}
|
||||
Self {
|
||||
io,
|
||||
mtu,
|
||||
read_buffer: BytesMut::with_capacity(mtu),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Read + Write + AsRawFd> Device for TunDevice<T> {
|
||||
type RxToken<'a>
|
||||
= TunRxToken
|
||||
where
|
||||
Self: 'a;
|
||||
type TxToken<'a>
|
||||
= TunTxToken<'a, T>
|
||||
where
|
||||
Self: 'a;
|
||||
|
||||
fn capabilities(&self) -> DeviceCapabilities {
|
||||
let mut caps = DeviceCapabilities::default();
|
||||
caps.medium = Medium::Ip;
|
||||
caps.max_transmission_unit = self.mtu;
|
||||
caps.checksum = ChecksumCapabilities::ignored();
|
||||
caps
|
||||
}
|
||||
|
||||
fn receive(
|
||||
&mut self,
|
||||
_timestamp: smoltcp::time::Instant,
|
||||
) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
|
||||
let mut tmp_buf = [0u8; 2048];
|
||||
match self.io.read(&mut tmp_buf) {
|
||||
Ok(n) if n > 0 => {
|
||||
let rx_data = Bytes::copy_from_slice(&tmp_buf[..n]);
|
||||
|
||||
let rx = TunRxToken { buffer: rx_data };
|
||||
let tx = TunTxToken { io: &mut self.io };
|
||||
Some((rx, tx))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn transmit(&mut self, _timestamp: smoltcp::time::Instant) -> Option<Self::TxToken<'_>> {
|
||||
Some(TunTxToken { io: &mut self.io })
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TunRxToken {
|
||||
buffer: Bytes,
|
||||
}
|
||||
|
||||
impl RxToken for TunRxToken {
|
||||
fn consume<R, F>(mut self, f: F) -> R
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R,
|
||||
{
|
||||
f(&mut self.buffer)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TunTxToken<'a, T: Write> {
|
||||
io: &'a mut T,
|
||||
}
|
||||
|
||||
impl<'a, T: Write> TxToken for TunTxToken<'a, T> {
|
||||
fn consume<R, F>(self, len: usize, f: F) -> R
|
||||
where
|
||||
F: FnOnce(&mut [u8]) -> R,
|
||||
{
|
||||
// ПРЯМАЯ ЗАПИСЬ БЕЗ ДОБАВЛЕНИЯ PI-ЗАГОЛОВКА
|
||||
let mut buffer = vec![0u8; len];
|
||||
let result = f(&mut buffer);
|
||||
let _ = self.io.write_all(&buffer);
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Read + Write + AsRawFd> AsRawFd for TunDevice<T> {
|
||||
fn as_raw_fd(&self) -> RawFd {
|
||||
self.io.as_raw_fd()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
//! IP packet encapsulation
|
||||
|
||||
use std::net::IpAddr;
|
||||
|
||||
use smoltcp::wire::{IpProtocol, IpVersion, Ipv4Packet, Ipv6Packet};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum IpPacket<T: AsRef<[u8]>> {
|
||||
Ipv4(Ipv4Packet<T>),
|
||||
Ipv6(Ipv6Packet<T>),
|
||||
}
|
||||
|
||||
impl<T: AsRef<[u8]> + Copy> IpPacket<T> {
|
||||
pub fn new_checked(packet: T) -> smoltcp::wire::Result<Option<Self>> {
|
||||
let buffer = packet.as_ref();
|
||||
match IpVersion::of_packet(buffer)? {
|
||||
IpVersion::Ipv4 => Ok(Some(Self::Ipv4(Ipv4Packet::new_checked(packet)?))),
|
||||
IpVersion::Ipv6 => Ok(Some(Self::Ipv6(Ipv6Packet::new_checked(packet)?))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn src_addr(&self) -> IpAddr {
|
||||
match *self {
|
||||
Self::Ipv4(ref packet) => IpAddr::from(packet.src_addr()),
|
||||
Self::Ipv6(ref packet) => IpAddr::from(packet.src_addr()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dst_addr(&self) -> IpAddr {
|
||||
match *self {
|
||||
Self::Ipv4(ref packet) => IpAddr::from(packet.dst_addr()),
|
||||
Self::Ipv6(ref packet) => IpAddr::from(packet.dst_addr()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn protocol(&self) -> IpProtocol {
|
||||
match *self {
|
||||
Self::Ipv4(ref packet) => packet.next_header(),
|
||||
Self::Ipv6(ref packet) => packet.next_header(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T: AsRef<[u8]> + ?Sized> IpPacket<&'a T> {
|
||||
/// Return a pointer to the payload.
|
||||
#[inline]
|
||||
pub fn payload(&self) -> &'a [u8] {
|
||||
match *self {
|
||||
IpPacket::Ipv4(ref packet) => packet.payload(),
|
||||
IpPacket::Ipv6(ref packet) => packet.payload(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
use tun::AsyncDevice;
|
||||
|
||||
//tun device
|
||||
pub fn create_linux_tun() -> AsyncDevice {
|
||||
let mut config = tun::Configuration::default();
|
||||
config
|
||||
.tun_name("netr0")
|
||||
.address((10, 0, 0, 1))
|
||||
.netmask((255, 255, 255, 0))
|
||||
.up();
|
||||
|
||||
let dev = tun::create_as_async(&config).expect("Нужны права root или CAP_NET_ADMIN!");
|
||||
dev
|
||||
}
|
||||
+204
-2
@@ -1,2 +1,204 @@
|
||||
pub mod desktop;
|
||||
pub mod device;
|
||||
use std::os::unix::io::RawFd;
|
||||
use std::{io, mem, net::IpAddr, time::Duration};
|
||||
|
||||
use byte_string::ByteStr;
|
||||
use ipnet::IpNet;
|
||||
use log::{error, info, trace, warn};
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
use tun::{
|
||||
create_as_async, AbstractDevice, AsyncDevice, Configuration as TunConfiguration,
|
||||
Error as TunError, Layer,
|
||||
};
|
||||
|
||||
use crate::tun::ip_packet::IpPacket;
|
||||
use crate::tun::virt_device::TokenBuffer;
|
||||
|
||||
mod ip_packet;
|
||||
pub mod linux_tun_device;
|
||||
pub mod virt_device;
|
||||
|
||||
/// Tun service builder
|
||||
pub struct TunBuilder {
|
||||
tun_config: TunConfiguration,
|
||||
udp_expiry_duration: Option<Duration>,
|
||||
udp_capacity: Option<usize>,
|
||||
}
|
||||
|
||||
/// TunConfiguration contains a HANDLE, which is a *mut c_void on Windows.
|
||||
unsafe impl Send for TunBuilder {}
|
||||
|
||||
impl TunBuilder {
|
||||
/// Create a Tun service builder
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tun_config: TunConfiguration::default(),
|
||||
udp_expiry_duration: None,
|
||||
udp_capacity: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn address(&mut self, addr: IpNet) {
|
||||
self.tun_config.address(addr.addr()).netmask(addr.netmask());
|
||||
}
|
||||
|
||||
pub fn destination(&mut self, addr: IpNet) {
|
||||
self.tun_config.destination(addr.addr());
|
||||
}
|
||||
|
||||
pub fn name(&mut self, name: &str) {
|
||||
self.tun_config.tun_name(name);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
pub fn file_descriptor(&mut self, fd: RawFd) {
|
||||
self.tun_config.raw_fd(fd);
|
||||
}
|
||||
|
||||
pub fn udp_expiry_duration(&mut self, udp_expiry_duration: Duration) {
|
||||
self.udp_expiry_duration = Some(udp_expiry_duration);
|
||||
}
|
||||
|
||||
pub fn udp_capacity(&mut self, udp_capacity: usize) {
|
||||
self.udp_capacity = Some(udp_capacity);
|
||||
}
|
||||
|
||||
/// Build Tun server
|
||||
pub async fn build(mut self) -> io::Result<Tun> {
|
||||
self.tun_config.layer(Layer::L3).up();
|
||||
|
||||
let device = match create_as_async(&self.tun_config) {
|
||||
Ok(d) => d,
|
||||
Err(TunError::Io(err)) => return Err(err),
|
||||
Err(err) => return Err(io::Error::other(err)),
|
||||
};
|
||||
|
||||
Ok(Tun { device })
|
||||
}
|
||||
}
|
||||
|
||||
/// Tun service
|
||||
pub struct Tun {
|
||||
device: AsyncDevice,
|
||||
}
|
||||
|
||||
impl Tun {
|
||||
/// Start serving
|
||||
pub async fn run(mut self) -> io::Result<()> {
|
||||
info!(
|
||||
"tun device {}",
|
||||
self.device
|
||||
.tun_name()
|
||||
.or_else(|r| Ok::<_, ()>(r.to_string()))
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
let address = match self.device.address() {
|
||||
Ok(a) => a,
|
||||
Err(err) => {
|
||||
error!("[TUN] failed to get device address, error: {}", err);
|
||||
return Err(io::Error::other(err));
|
||||
}
|
||||
};
|
||||
|
||||
let netmask = match self.device.netmask() {
|
||||
Ok(n) => n,
|
||||
Err(err) => {
|
||||
error!("[TUN] failed to get device netmask, error: {}", err);
|
||||
return Err(io::Error::other(err));
|
||||
}
|
||||
};
|
||||
|
||||
let address_net = match IpNet::with_netmask(address, netmask) {
|
||||
Ok(n) => n,
|
||||
Err(err) => {
|
||||
error!(
|
||||
"[TUN] invalid address {}, netmask {}, error: {}",
|
||||
address, netmask, err
|
||||
);
|
||||
return Err(io::Error::other(err));
|
||||
}
|
||||
};
|
||||
|
||||
trace!(
|
||||
"[TUN] tun device network: {} (address: {}, netmask: {})",
|
||||
address_net,
|
||||
address,
|
||||
netmask
|
||||
);
|
||||
|
||||
let address_broadcast = address_net.broadcast();
|
||||
|
||||
let create_packet_buffer = || {
|
||||
const PACKET_BUFFER_SIZE: usize = 65536;
|
||||
let mut packet_buffer = TokenBuffer::with_capacity(PACKET_BUFFER_SIZE);
|
||||
unsafe {
|
||||
packet_buffer.set_len(PACKET_BUFFER_SIZE);
|
||||
}
|
||||
packet_buffer
|
||||
};
|
||||
|
||||
let mut packet_buffer = create_packet_buffer();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// tun device
|
||||
n = self.device.read(&mut packet_buffer) => {
|
||||
let n = n?;
|
||||
|
||||
let mut packet_buffer = mem::replace(&mut packet_buffer, create_packet_buffer());
|
||||
unsafe {
|
||||
packet_buffer.set_len(n);
|
||||
}
|
||||
|
||||
trace!("[TUN] received IP packet {:?}", ByteStr::new(&packet_buffer));
|
||||
|
||||
if let Err(err) = self.handle_tun_frame(&address_broadcast, packet_buffer).await {
|
||||
error!("[TUN] handle IP frame failed, error: {}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_tun_frame(
|
||||
&mut self,
|
||||
device_broadcast_addr: &IpAddr,
|
||||
frame: TokenBuffer,
|
||||
) -> smoltcp::wire::Result<()> {
|
||||
let packet = match IpPacket::new_checked(frame.as_ref())? {
|
||||
Some(packet) => packet,
|
||||
None => {
|
||||
warn!("unrecognized IP packet {:?}", ByteStr::new(&frame));
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
trace!("[TUN] {:?}", packet);
|
||||
|
||||
let src_ip_addr = packet.src_addr();
|
||||
let dst_ip_addr = packet.dst_addr();
|
||||
let src_non_unicast = src_ip_addr == *device_broadcast_addr
|
||||
|| match src_ip_addr {
|
||||
IpAddr::V4(v4) => v4.is_broadcast() || v4.is_multicast() || v4.is_unspecified(),
|
||||
IpAddr::V6(v6) => v6.is_multicast() || v6.is_unspecified(),
|
||||
};
|
||||
let dst_non_unicast = dst_ip_addr == *device_broadcast_addr
|
||||
|| match dst_ip_addr {
|
||||
IpAddr::V4(v4) => v4.is_broadcast() || v4.is_multicast() || v4.is_unspecified(),
|
||||
IpAddr::V6(v6) => v6.is_multicast() || v6.is_unspecified(),
|
||||
};
|
||||
|
||||
if src_non_unicast || dst_non_unicast {
|
||||
trace!(
|
||||
"[TUN] IP packet {} (unicast? {}) -> {} (unicast? {}) throwing away",
|
||||
src_ip_addr,
|
||||
!src_non_unicast,
|
||||
dst_ip_addr,
|
||||
!dst_non_unicast
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
//! Virtual Device for receiving packets from tun
|
||||
|
||||
use std::{
|
||||
marker::PhantomData,
|
||||
mem,
|
||||
ops::{Deref, DerefMut},
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc, LazyLock, Mutex,
|
||||
},
|
||||
};
|
||||
|
||||
use bytes::BytesMut;
|
||||
use smoltcp::{
|
||||
phy::{self, Device, DeviceCapabilities},
|
||||
time::Instant,
|
||||
};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
pub struct VirtTunDevice {
|
||||
capabilities: DeviceCapabilities,
|
||||
in_buf: mpsc::UnboundedReceiver<TokenBuffer>,
|
||||
out_buf: mpsc::UnboundedSender<TokenBuffer>,
|
||||
in_buf_avail: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl VirtTunDevice {
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub fn new(
|
||||
capabilities: DeviceCapabilities,
|
||||
) -> (
|
||||
Self,
|
||||
mpsc::UnboundedReceiver<TokenBuffer>,
|
||||
mpsc::UnboundedSender<TokenBuffer>,
|
||||
Arc<AtomicBool>,
|
||||
) {
|
||||
let (iface_tx, iface_output) = mpsc::unbounded_channel();
|
||||
let (iface_input, iface_rx) = mpsc::unbounded_channel();
|
||||
let in_buf_avail = Arc::new(AtomicBool::new(false));
|
||||
|
||||
(
|
||||
Self {
|
||||
capabilities,
|
||||
in_buf: iface_rx,
|
||||
out_buf: iface_tx,
|
||||
in_buf_avail: in_buf_avail.clone(),
|
||||
},
|
||||
iface_output,
|
||||
iface_input,
|
||||
in_buf_avail,
|
||||
)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn recv_available(&self) -> bool {
|
||||
self.in_buf_avail.load(Ordering::Acquire)
|
||||
}
|
||||
}
|
||||
|
||||
impl Device for VirtTunDevice {
|
||||
type RxToken<'a> = VirtRxToken<'a>;
|
||||
type TxToken<'a> = VirtTxToken<'a>;
|
||||
|
||||
fn receive(&mut self, _timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
|
||||
if let Ok(buffer) = self.in_buf.try_recv() {
|
||||
let rx = Self::RxToken {
|
||||
buffer,
|
||||
phantom_device: PhantomData,
|
||||
};
|
||||
let tx = VirtTxToken(self);
|
||||
return Some((rx, tx));
|
||||
}
|
||||
self.in_buf_avail.store(false, Ordering::Release);
|
||||
None
|
||||
}
|
||||
|
||||
fn transmit(&mut self, _timestamp: Instant) -> Option<Self::TxToken<'_>> {
|
||||
Some(VirtTxToken(self))
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> DeviceCapabilities {
|
||||
self.capabilities.clone()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct VirtRxToken<'a> {
|
||||
buffer: TokenBuffer,
|
||||
phantom_device: PhantomData<&'a VirtTunDevice>,
|
||||
}
|
||||
|
||||
impl phy::RxToken for VirtRxToken<'_> {
|
||||
fn consume<R, F>(self, f: F) -> R
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R,
|
||||
{
|
||||
f(&self.buffer)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct VirtTxToken<'a>(&'a mut VirtTunDevice);
|
||||
|
||||
impl phy::TxToken for VirtTxToken<'_> {
|
||||
fn consume<R, F>(self, len: usize, f: F) -> R
|
||||
where
|
||||
F: FnOnce(&mut [u8]) -> R,
|
||||
{
|
||||
let mut buffer = TokenBuffer::with_capacity(len);
|
||||
unsafe {
|
||||
buffer.set_len(len);
|
||||
}
|
||||
|
||||
let result = f(&mut buffer);
|
||||
self.0
|
||||
.out_buf
|
||||
.send(buffer)
|
||||
.expect("channel closed unexpectedly");
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
// Maximum number of TokenBuffer cached globally.
|
||||
//
|
||||
// Each of them has capacity 65536 (defined in tun/mod.rs), so 64 * 65536 = 4MB.
|
||||
const TOKEN_BUFFER_LIST_MAX_SIZE: usize = 64;
|
||||
static TOKEN_BUFFER_LIST: LazyLock<Mutex<Vec<BytesMut>>> = LazyLock::new(|| Mutex::new(Vec::new()));
|
||||
|
||||
pub struct TokenBuffer {
|
||||
buffer: BytesMut,
|
||||
}
|
||||
|
||||
impl Drop for TokenBuffer {
|
||||
fn drop(&mut self) {
|
||||
let mut list = TOKEN_BUFFER_LIST.lock().unwrap();
|
||||
if list.len() >= TOKEN_BUFFER_LIST_MAX_SIZE {
|
||||
return;
|
||||
}
|
||||
|
||||
let empty_buffer = BytesMut::new();
|
||||
let mut buffer = mem::replace(&mut self.buffer, empty_buffer);
|
||||
buffer.clear();
|
||||
|
||||
list.push(buffer);
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenBuffer {
|
||||
pub fn with_capacity(cap: usize) -> Self {
|
||||
let mut list = TOKEN_BUFFER_LIST.lock().unwrap();
|
||||
if let Some(mut buffer) = list.pop() {
|
||||
buffer.reserve(cap);
|
||||
return Self { buffer };
|
||||
}
|
||||
Self {
|
||||
buffer: BytesMut::with_capacity(cap),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for TokenBuffer {
|
||||
type Target = BytesMut;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.buffer
|
||||
}
|
||||
}
|
||||
|
||||
impl DerefMut for TokenBuffer {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
&mut self.buffer
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user