rustdesk/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2025-08-17 00:45:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

try_into_v4

Browse Source
This commit is contained in:
rustdesk 2023-01-10 16:09:25 +08:00
parent 338af1af9d
commit fc83fa0a04
2 changed files with 24 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
libs/hbb_common/src/lib.rs
View File

@ -95,11 +95,10 @@ pub type ResultType<F, E = anyhow::Error> = anyhow::Result<F, E>;
pub struct AddrMangle(); pub struct AddrMangle();
impl AddrMangle { #[inline]
pub fn encode(addr: SocketAddr) -> Vec<u8> { pub fn try_into_v4(addr: SocketAddr) -> SocketAddr {
// not work with [:1]:<port> match addr {
let addr = match addr { SocketAddr::V6(v6) if !addr.ip().is_loopback() => {
SocketAddr::V6(v6) => {
if let Some(v4) = v6.ip().to_ipv4() { if let Some(v4) = v6.ip().to_ipv4() {
SocketAddr::new(IpAddr::V4(v4), addr.port()) SocketAddr::new(IpAddr::V4(v4), addr.port())
} else { } else {
@ -107,7 +106,13 @@ impl AddrMangle {
} }
} }
_ => addr, _ => addr,
}; }
}
impl AddrMangle {
pub fn encode(addr: SocketAddr) -> Vec<u8> {
// not work with [:1]:<port>
let addr = try_into_v4(addr);
match addr { match addr {
SocketAddr::V4(addr_v4) => { SocketAddr::V4(addr_v4) => {
let tm = (SystemTime::now() let tm = (SystemTime::now()
@ -353,5 +358,7 @@ mod test {
assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v4)), addr_v4); assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v4)), addr_v4);
let addr_v6 = "[ef::fe]:8080".parse().unwrap(); let addr_v6 = "[ef::fe]:8080".parse().unwrap();
assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v6)), addr_v6); assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v6)), addr_v6);
let addr_v6 = "[::1]:8080".parse().unwrap();
assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v6)), addr_v6);
} }
} }

12
src/rendezvous_server.rs
View File

@ -4,6 +4,7 @@ use hbb_common::{
allow_err, allow_err,
bytes::{Bytes, BytesMut}, bytes::{Bytes, BytesMut},
bytes_codec::BytesCodec, bytes_codec::BytesCodec,
config,
futures::future::join_all, futures::future::join_all,
futures_util::{ futures_util::{
sink::SinkExt, sink::SinkExt,
@ -25,6 +26,7 @@ use hbb_common::{
time::{interval, Duration}, time::{interval, Duration},
}, },
tokio_util::codec::Framed, tokio_util::codec::Framed,
try_into_v4,
udp::FramedSocket, udp::FramedSocket,
AddrMangle, ResultType, AddrMangle, ResultType,
}; };
@ -475,7 +477,7 @@ impl RendezvousServer {
Some(rendezvous_message::Union::PunchHoleRequest(ph)) => { Some(rendezvous_message::Union::PunchHoleRequest(ph)) => {
// there maybe several attempt, so sink can be none // there maybe several attempt, so sink can be none
if let Some(sink) = sink.take() { if let Some(sink) = sink.take() {
self.tcp_punch.lock().await.insert(addr, sink); self.tcp_punch.lock().await.insert(try_into_v4(addr), sink);
} }
allow_err!(self.handle_tcp_punch_hole_request(addr, ph, key, ws).await); allow_err!(self.handle_tcp_punch_hole_request(addr, ph, key, ws).await);
return true; return true;
@ -483,7 +485,7 @@ impl RendezvousServer {
Some(rendezvous_message::Union::RequestRelay(mut rf)) => { Some(rendezvous_message::Union::RequestRelay(mut rf)) => {
// there maybe several attempt, so sink can be none // there maybe several attempt, so sink can be none
if let Some(sink) = sink.take() { if let Some(sink) = sink.take() {
self.tcp_punch.lock().await.insert(addr, sink); self.tcp_punch.lock().await.insert(try_into_v4(addr), sink);
} }
if let Some(peer) = self.pm.get_in_memory(&rf.id).await { if let Some(peer) = self.pm.get_in_memory(&rf.id).await {
let mut msg_out = RendezvousMessage::new(); let mut msg_out = RendezvousMessage::new();
@ -1048,7 +1050,7 @@ impl RendezvousServer {
async fn handle_listener2(&self, stream: TcpStream, addr: SocketAddr) { async fn handle_listener2(&self, stream: TcpStream, addr: SocketAddr) {
let mut rs = self.clone(); let mut rs = self.clone();
if addr.ip().to_string() == "127.0.0.1" { if addr.ip().is_loopback() {
tokio::spawn(async move { tokio::spawn(async move {
let mut stream = stream; let mut stream = stream;
let mut buffer = [0; 64]; let mut buffer = [0; 64];
@ -1203,7 +1205,7 @@ async fn check_relay_servers(rs0: Arc<RelayServers>, tx: Sender) {
for x in rs0.iter() { for x in rs0.iter() {
let mut host = x.to_owned(); let mut host = x.to_owned();
if !host.contains(':') { if !host.contains(':') {
host = format!("{}:{}", host, hbb_common::config::RELAY_PORT); host = format!("{}:{}", host, config::RELAY_PORT);
} }
let rs = rs.clone(); let rs = rs.clone();
let x = x.clone(); let x = x.clone();
@ -1226,7 +1228,7 @@ async fn check_relay_servers(rs0: Arc<RelayServers>, tx: Sender) {
// temp solution to solve udp socket failure // temp solution to solve udp socket failure
async fn test_hbbs(addr: SocketAddr) -> ResultType<()> { async fn test_hbbs(addr: SocketAddr) -> ResultType<()> {
let mut socket = FramedSocket::new("0.0.0.0:0").await?; let mut socket = FramedSocket::new(config::Config::get_any_listen_addr(addr.is_ipv4())).await?;
let mut msg_out = RendezvousMessage::new(); let mut msg_out = RendezvousMessage::new();
msg_out.set_register_peer(RegisterPeer { msg_out.set_register_peer(RegisterPeer {
id: "(:test_hbbs:)".to_owned(), id: "(:test_hbbs:)".to_owned(),