arampoire/onetun
1
0
Fork 0
mirror of https://github.com/arampoire/onetun.git synced 2025-12-01 02:40:23 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use bytes

Browse source
This commit is contained in:
Aram 🍐 2023-01-12 01:40:04 -05:00
parent e62b7d30fe
commit 76b6a6e346
9 changed files with 61 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/tunnel/tcp.rs
View file

@ -1,17 +1,18 @@
use crate::config::{PortForwardConfig, PortProtocol};
use crate::virtual_iface::VirtualPort;
use anyhow::Context;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use std::ops::Range;
use std::sync::Arc;
use std::time::Duration;
use crate::events::{Bus, Event};
use anyhow::Context;
use bytes::BytesMut;
use rand::seq::SliceRandom;
use rand::thread_rng;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use crate::config::{PortForwardConfig, PortProtocol};
use crate::events::{Bus, Event};
use crate::virtual_iface::VirtualPort;
const MAX_PACKET: usize = 65536;
const MIN_PORT: u16 = 1000;
@ -81,7 +82,7 @@ async fn handle_tcp_proxy_connection(
let mut endpoint = bus.new_endpoint();
endpoint.send(Event::ClientConnectionInitiated(port_forward, virtual_port));
let mut buffer = Vec::with_capacity(MAX_PACKET);
let mut buffer = BytesMut::with_capacity(MAX_PACKET);
loop {
tokio::select! {
readable_result = socket.readable() => {
@ -90,7 +91,7 @@ async fn handle_tcp_proxy_connection(
match socket.try_read_buf(&mut buffer) {
Ok(size) if size > 0 => {
let data = Vec::from(&buffer[..size]);
endpoint.send(Event::LocalData(port_forward, virtual_port, data));
endpoint.send(Event::LocalData(port_forward, virtual_port, data.into()));
// Reset buffer
buffer.clear();
}

7
src/tunnel/udp.rs
View file

@ -4,14 +4,15 @@ use std::ops::Range;
use std::sync::Arc;
use std::time::Instant;
use crate::events::{Bus, Event};
use anyhow::Context;
use bytes::Bytes;
use priority_queue::double_priority_queue::DoublePriorityQueue;
use rand::seq::SliceRandom;
use rand::thread_rng;
use tokio::net::UdpSocket;
use crate::config::{PortForwardConfig, PortProtocol};
use crate::events::{Bus, Event};
use crate::virtual_iface::VirtualPort;
const MAX_PACKET: usize = 65536;
@ -98,7 +99,7 @@ async fn next_udp_datagram(
socket: &UdpSocket,
buffer: &mut [u8],
port_pool: UdpPortPool,
) -> anyhow::Result<Option<(VirtualPort, Vec<u8>)>> {
) -> anyhow::Result<Option<(VirtualPort, Bytes)>> {
let (size, peer_addr) = socket
.recv_from(buffer)
.await
@ -126,7 +127,7 @@ async fn next_udp_datagram(
port_pool.update_last_transmit(port).await;
let data = buffer[..size].to_vec();
Ok(Some((port, data)))
Ok(Some((port, data.into())))
}
/// A pool of virtual ports available for TCP connections.