hex/wstunnel
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

use BytesMut instead of vec

Browse source
This commit is contained in:
Σrebe - Romain GERARD 2023-10-27 09:15:15 +02:00
parent f813d925d6
commit b70d547370
No known key found for this signature in database
GPG key ID: 7A42B4B97E0332F4
5 changed files with 43 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/tunnel/client.rs
View file

@ -58,7 +58,7 @@ pub async fn connect(
&client_cfg.http_upgrade_path_prefix,
jsonwebtoken::encode(alg, &data, secret).unwrap_or_default(),
))
.header(HOST, &client_cfg.host_http_header)
.header(HOST, &client_cfg.http_header_host)
.header(UPGRADE, "websocket")
.header(CONNECTION, "upgrade")
.header(SEC_WEBSOCKET_KEY, fastwebsockets::handshake::generate_key())

25
src/tunnel/io.rs
View file

@ -1,6 +1,9 @@
use fastwebsockets::{Frame, OpCode, Payload, WebSocketError, WebSocketRead, WebSocketWrite};
use futures_util::pin_mut;
use hyper::upgrade::Upgraded;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio::select;
@ -22,21 +25,21 @@ pub(super) async fn propagate_read(
let mut buffer = vec![0u8; 8 * 1024];
pin_mut!(local_rx);
loop {
let read = select! {
let read_len = select! {
biased;
read_len = local_rx.read(buffer.as_mut_slice()) => read_len,
read_len = local_rx.read(&mut buffer) => read_len,
_ = close_tx.closed() => break,
_ = timeout(ping_frequency, futures_util::future::pending::<()>()) => {
debug!("sending ping to keep websocket connection alive");
ws_tx.write_frame(Frame::new(true, OpCode::Ping, None, Payload::Borrowed(&[]))).await?;
ws_tx.write_frame(Frame::new(true, OpCode::Ping, None, Payload::BorrowedMut(&mut []))).await?;
continue;
}
};
let read_len = match read {
let read_len = match read_len {
Ok(read_len) if read_len > 0 => read_len,
Ok(_) => break,
Err(err) => {
@ -50,7 +53,7 @@ pub(super) async fn propagate_read(
trace!("read {} bytes", read_len);
match ws_tx
.write_frame(Frame::binary(Payload::Borrowed(&buffer[..read_len])))
.write_frame(Frame::binary(Payload::BorrowedMut(&mut buffer[..read_len])))
.await
{
Ok(_) => {}
@ -60,8 +63,10 @@ pub(super) async fn propagate_read(
}
}
if read_len == buffer.len() {
buffer.resize(read_len * 2, 0);
if buffer.capacity() == read_len {
buffer.clear();
info!("capa: {} read:{}", buffer.capacity(), read_len);
buffer.resize(buffer.capacity() * 2, 0);
}
}
@ -85,14 +90,14 @@ pub(super) async fn propagate_write(
pin_mut!(local_tx);
loop {
let ret = select! {
let msg = select! {
biased;
ret = ws_rx.read_frame(&mut x) => ret,
msg = ws_rx.read_frame(&mut x) => msg,
_ = &mut close_rx => break,
};
let msg = match ret {
let msg = match msg {
Ok(msg) => msg,
Err(err) => {
error!("error while reading from websocket rx {}", err);