RustDesk/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2026-02-17 22:10:17 +08:00
Code Issues Packages Projects Releases Wiki Activity

v1.1.12

Browse Source
This commit is contained in:
rustdesk
2024-10-07 16:21:36 +08:00
parent 3b386b6b54
commit 6f18a97644
6 changed files with 44 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
src/relay_server.rs
View File

@@ -392,19 +392,30 @@ async fn handle_connection(
async fn make_pair(
stream: TcpStream,
addr: SocketAddr,
mut addr: SocketAddr,
key: &str,
limiter: Limiter,
ws: bool,
) -> ResultType<()> {
if ws {
make_pair_(
tokio_tungstenite::accept_async(stream).await?,
addr,
key,
limiter,
)
.await;
use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
let callback = |req: &Request, response: Response| {
let headers = req.headers();
let real_ip = headers
.get("X-Real-IP")
.or_else(|| headers.get("X-Forwarded-For"))
.and_then(|header_value| header_value.to_str().ok());
if let Some(ip) = real_ip {
if ip.contains('.') {
addr = format!("{ip}:0").parse().unwrap_or(addr);
} else {
addr = format!("[{ip}]:0").parse().unwrap_or(addr);
}
}
Ok(response)
};
let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
make_pair_(ws_stream, addr, key, limiter).await;
} else {
make_pair_(FramedStream::from(stream, addr), addr, key, limiter).await;
}

20
src/rendezvous_server.rs
View File

@@ -1106,13 +1106,29 @@ impl RendezvousServer {
async fn handle_listener_inner(
&mut self,
stream: TcpStream,
addr: SocketAddr,
mut addr: SocketAddr,
key: &str,
ws: bool,
) -> ResultType<()> {
let mut sink;
if ws {
let ws_stream = tokio_tungstenite::accept_async(stream).await?;
use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
let callback = |req: &Request, response: Response| {
let headers = req.headers();
let real_ip = headers
.get("X-Real-IP")
.or_else(|| headers.get("X-Forwarded-For"))
.and_then(|header_value| header_value.to_str().ok());
if let Some(ip) = real_ip {
if ip.contains('.') {
addr = format!("{ip}:0").parse().unwrap_or(addr);
} else {
addr = format!("[{ip}]:0").parse().unwrap_or(addr);
}
}
Ok(response)
};
let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
let (a, mut b) = ws_stream.split();
sink = Some(Sink::Ws(a));
while let Ok(Some(Ok(msg))) = timeout(30_000, b.next()).await {