|
|
@@ -392,19 +392,30 @@ async fn handle_connection(
|
|
392
|
392
|
|
|
393
|
393
|
async fn make_pair(
|
|
394
|
394
|
stream: TcpStream,
|
|
395
|
|
- addr: SocketAddr,
|
|
|
395
|
+ mut addr: SocketAddr,
|
|
396
|
396
|
key: &str,
|
|
397
|
397
|
limiter: Limiter,
|
|
398
|
398
|
ws: bool,
|
|
399
|
399
|
) -> ResultType<()> {
|
|
400
|
400
|
if ws {
|
|
401
|
|
- make_pair_(
|
|
402
|
|
- tokio_tungstenite::accept_async(stream).await?,
|
|
403
|
|
- addr,
|
|
404
|
|
- key,
|
|
405
|
|
- limiter,
|
|
406
|
|
- )
|
|
407
|
|
- .await;
|
|
|
401
|
+ use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
|
|
|
402
|
+ let callback = |req: &Request, response: Response| {
|
|
|
403
|
+ let headers = req.headers();
|
|
|
404
|
+ let real_ip = headers
|
|
|
405
|
+ .get("X-Real-IP")
|
|
|
406
|
+ .or_else(|| headers.get("X-Forwarded-For"))
|
|
|
407
|
+ .and_then(|header_value| header_value.to_str().ok());
|
|
|
408
|
+ if let Some(ip) = real_ip {
|
|
|
409
|
+ if ip.contains('.') {
|
|
|
410
|
+ addr = format!("{ip}:0").parse().unwrap_or(addr);
|
|
|
411
|
+ } else {
|
|
|
412
|
+ addr = format!("[{ip}]:0").parse().unwrap_or(addr);
|
|
|
413
|
+ }
|
|
|
414
|
+ }
|
|
|
415
|
+ Ok(response)
|
|
|
416
|
+ };
|
|
|
417
|
+ let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
|
|
|
418
|
+ make_pair_(ws_stream, addr, key, limiter).await;
|
|
408
|
419
|
} else {
|
|
409
|
420
|
make_pair_(FramedStream::from(stream, addr), addr, key, limiter).await;
|
|
410
|
421
|
}
|