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

cleanup

Browse source
This commit is contained in:
Σrebe - Romain GERARD 2023-12-04 08:50:47 +01:00
parent 32e7cdcb8e
commit 8495870bff
No known key found for this signature in database
GPG key ID: 7A42B4B97E0332F4
3 changed files with 17 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/tunnel/client.rs
View file

@ -2,6 +2,7 @@ use super::{to_host_port, JwtTunnelConfig, JWT_KEY};
use crate::{LocalToRemote, WsClientConfig}; use crate::{LocalToRemote, WsClientConfig};
use anyhow::{anyhow, Context}; use anyhow::{anyhow, Context};
use base64::Engine;
use fastwebsockets::WebSocket; use fastwebsockets::WebSocket;
use futures_util::pin_mut; use futures_util::pin_mut;
use hyper::header::{AUTHORIZATION, COOKIE, SEC_WEBSOCKET_VERSION, UPGRADE}; use hyper::header::{AUTHORIZATION, COOKIE, SEC_WEBSOCKET_VERSION, UPGRADE};
@ -186,7 +187,8 @@ where
.and_then(|h| { .and_then(|h| {
h.to_str() h.to_str()
.ok() .ok()
.and_then(|s| Url::parse(s).ok()) .and_then(|s| base64::engine::general_purpose::STANDARD.decode(s).ok())
.and_then(|s| Url::parse(&String::from_utf8_lossy(&s)).ok())
.and_then(|url| match (url.host(), url.port()) { .and_then(|url| match (url.host(), url.port()) {
(Some(h), Some(p)) => Some((h.to_owned(), p)), (Some(h), Some(p)) => Some((h.to_owned(), p)),
_ => None, _ => None,

8
src/tunnel/io.rs
View file

@ -1,5 +1,5 @@
use fastwebsockets::{Frame, OpCode, Payload, WebSocketError, WebSocketRead, WebSocketWrite}; use fastwebsockets::{Frame, OpCode, Payload, WebSocketError, WebSocketRead, WebSocketWrite};
use futures_util::pin_mut; use futures_util::{pin_mut, FutureExt};
use hyper::upgrade::Upgraded; use hyper::upgrade::Upgraded;
use std::time::Duration; use std::time::Duration;
@ -28,8 +28,10 @@ pub(super) async fn propagate_read(
let frequency = ping_frequency.unwrap_or(Duration::from_secs(3600 * 24)); let frequency = ping_frequency.unwrap_or(Duration::from_secs(3600 * 24));
let start_at = Instant::now().checked_add(frequency).unwrap_or(Instant::now()); let start_at = Instant::now().checked_add(frequency).unwrap_or(Instant::now());
let timeout = tokio::time::interval_at(start_at, frequency); let timeout = tokio::time::interval_at(start_at, frequency);
pin_mut!(timeout); let should_close = close_tx.closed().fuse();
pin_mut!(timeout);
pin_mut!(should_close);
pin_mut!(local_rx); pin_mut!(local_rx);
loop { loop {
let read_len = select! { let read_len = select! {
@ -37,7 +39,7 @@ pub(super) async fn propagate_read(
read_len = local_rx.read(&mut buffer) => read_len, read_len = local_rx.read(&mut buffer) => read_len,
_ = close_tx.closed() => break, _ = &mut should_close => break,
_ = timeout.tick(), if ping_frequency.is_some() => { _ = timeout.tick(), if ping_frequency.is_some() => {
debug!("sending ping to keep websocket connection alive"); debug!("sending ping to keep websocket connection alive");

12
src/tunnel/server.rs
View file

@ -1,5 +1,6 @@
use ahash::{HashMap, HashMapExt}; use ahash::{HashMap, HashMapExt};
use anyhow::anyhow; use anyhow::anyhow;
use base64::Engine;
use futures_util::{pin_mut, Stream, StreamExt}; use futures_util::{pin_mut, Stream, StreamExt};
use std::cmp::min; use std::cmp::min;
use std::fmt::Debug; use std::fmt::Debug;
@ -264,9 +265,14 @@ async fn server_upgrade(
.instrument(Span::current()), .instrument(Span::current()),
); );
response if protocol == LocalProtocol::ReverseSocks5 {
.headers_mut() response.headers_mut().insert(
.insert(COOKIE, HeaderValue::from_str(&format!("fake://{}:{}", dest, port)).unwrap()); COOKIE,
HeaderValue::from_str(
&base64::engine::general_purpose::STANDARD.encode(format!("fake://{}:{}", dest, port)),
)?,
);
}
Ok(response) Ok(response)
} }