mercuryworkshop-mirror/epoxy-tls
1
0
Fork 0
mirror of https://github.com/MercuryWorkshop/epoxy-tls.git synced 2025-05-12 14:00:01 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

add protocol header v2 to server and client

Browse source
This commit is contained in:
Toshit Chawda 2024-10-24 00:16:11 -07:00
parent cda7ed2190
commit 1ae3986a82
No known key found for this signature in database
GPG key ID: 91480ED99E2B3D9D
8 changed files with 84 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

17
server/src/handle/wisp/mod.rs
View file

@ -211,7 +211,7 @@ async fn handle_stream(
}
}
pub async fn handle_wisp(stream: WispResult, id: String) -> anyhow::Result<()> {
pub async fn handle_wisp(stream: WispResult, is_v2: bool, id: String) -> anyhow::Result<()> {
let (read, write) = stream;
cfg_if! {
if #[cfg(feature = "twisp")] {
@ -232,11 +232,16 @@ pub async fn handle_wisp(stream: WispResult, id: String) -> anyhow::Result<()> {
}
}
let (mux, fut) = ServerMux::create(read, write, buffer_size, extensions)
.await
.context("failed to create server multiplexor")?
.with_required_extensions(&required_extensions)
.await?;
let (mux, fut) = ServerMux::create(
read,
write,
buffer_size,
if is_v2 { extensions } else { None },
)
.await
.context("failed to create server multiplexor")?
.with_required_extensions(&required_extensions)
.await?;
let mux = Arc::new(mux);
debug!(

2
server/src/main.rs
View file

@ -208,7 +208,7 @@ fn handle_stream(stream: ServerRouteResult, id: String) {
tokio::spawn(async move {
CLIENTS.insert(id.clone(), (DashMap::new(), false));
let res = match stream {
ServerRouteResult::Wisp(stream) => handle_wisp(stream, id.clone()).await,
ServerRouteResult::Wisp(stream, is_v2) => handle_wisp(stream, is_v2, id.clone()).await,
ServerRouteResult::WsProxy(ws, path, udp) => {
handle_wsproxy(ws, id.clone(), path, udp).await
}

84
server/src/route.rs
View file

@ -5,8 +5,8 @@ use bytes::Bytes;
use fastwebsockets::{upgrade::UpgradeFut, FragmentCollector};
use http_body_util::Full;
use hyper::{
body::Incoming, server::conn::http1::Builder, service::service_fn, HeaderMap, Request,
Response, StatusCode,
body::Incoming, header::SEC_WEBSOCKET_PROTOCOL, server::conn::http1::Builder,
service::service_fn, HeaderMap, Request, Response, StatusCode,
};
use hyper_util::rt::TokioIo;
use log::{debug, error, trace};
@ -25,6 +25,25 @@ use crate::{
CONFIG,
};
pub type WispResult = (
Box<dyn WebSocketRead + Send>,
Box<dyn WebSocketWrite + Send>,
);
pub enum ServerRouteResult {
Wisp(WispResult, bool),
WsProxy(WebSocketStreamWrapper, String, bool),
}
impl Display for ServerRouteResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Wisp(..) => write!(f, "Wisp"),
Self::WsProxy(_, path, udp) => write!(f, "WsProxy path {:?} udp {:?}", path, udp),
}
}
}
type Body = Full<Bytes>;
fn non_ws_resp() -> anyhow::Result<Response<Body>> {
Ok(Response::builder()
@ -57,7 +76,7 @@ fn get_header(headers: &HeaderMap, header: &str) -> Option<String> {
}
enum HttpUpgradeResult {
Wisp,
Wisp(bool),
WsProxy(String, bool),
}
@ -90,7 +109,7 @@ where
let (resp, fut) = fastwebsockets::upgrade::upgrade(&mut req)?;
// replace body of Empty<Bytes> with Full<Bytes>
let resp = Response::from_parts(resp.into_parts().0, Body::new(Bytes::new()));
let mut resp = Response::from_parts(resp.into_parts().0, Body::new(Bytes::new()));
let headers = req.headers();
let ip_header = if CONFIG.server.use_real_ip_headers {
@ -99,25 +118,27 @@ where
None
};
if req
.uri()
.path()
.starts_with(&(CONFIG.wisp.prefix.clone() + "/"))
{
let ws_protocol = headers.get(SEC_WEBSOCKET_PROTOCOL);
let req_path = req.uri().path().to_string();
if req_path.starts_with(&(CONFIG.wisp.prefix.clone() + "/")) {
let has_ws_protocol = ws_protocol.is_some();
tokio::spawn(async move {
if let Err(err) = (callback)(fut, HttpUpgradeResult::Wisp, ip_header).await {
if let Err(err) =
(callback)(fut, HttpUpgradeResult::Wisp(has_ws_protocol), ip_header).await
{
error!("error while serving client: {:?}", err);
}
});
if let Some(protocol) = ws_protocol {
resp.headers_mut()
.append(SEC_WEBSOCKET_PROTOCOL, protocol.clone());
}
} else if CONFIG.wisp.allow_wsproxy {
let udp = req.uri().query().unwrap_or_default() == "?udp";
tokio::spawn(async move {
if let Err(err) = (callback)(
fut,
HttpUpgradeResult::WsProxy(req.uri().path().to_string(), udp),
ip_header,
)
.await
if let Err(err) =
(callback)(fut, HttpUpgradeResult::WsProxy(req_path, udp), ip_header).await
{
error!("error while serving client: {:?}", err);
}
@ -162,7 +183,7 @@ pub async fn route(
ws.set_auto_pong(false);
match res {
HttpUpgradeResult::Wisp => {
HttpUpgradeResult::Wisp(is_v2) => {
let (read, write) = ws.split(|x| {
let parts = x
.into_inner()
@ -173,10 +194,10 @@ pub async fn route(
});
(callback)(
ServerRouteResult::Wisp((
Box::new(read),
Box::new(write),
)),
ServerRouteResult::Wisp(
(Box::new(read), Box::new(write)),
is_v2,
),
maybe_ip,
)
}
@ -208,29 +229,10 @@ pub async fn route(
let write = GenericWebSocketWrite::new(FramedWrite::new(write, codec));
(callback)(
ServerRouteResult::Wisp((Box::new(read), Box::new(write))),
ServerRouteResult::Wisp((Box::new(read), Box::new(write)), true),
None,
);
}
}
Ok(())
}
pub type WispResult = (
Box<dyn WebSocketRead + Send>,
Box<dyn WebSocketWrite + Send>,
);
pub enum ServerRouteResult {
Wisp(WispResult),
WsProxy(WebSocketStreamWrapper, String, bool),
}
impl Display for ServerRouteResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Wisp(_) => write!(f, "Wisp"),
Self::WsProxy(_, path, udp) => write!(f, "WsProxy path {:?} udp {:?}", path, udp),
}
}
}