hongshaorou/rustdesk
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ipv6 mangle

Browse Source
This commit is contained in:
rustdesk 2022-12-26 16:41:33 +08:00
parent a848fa0f9b
commit 633253647f
2 changed files with 40 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
libs/hbb_common/src/config.rs
View File

@ -1,7 +1,7 @@
use std::{ use std::{
collections::HashMap, collections::HashMap,
fs, fs,
net::{IpAddr, Ipv4Addr, SocketAddr}, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
path::{Path, PathBuf}, path::{Path, PathBuf},
sync::{Arc, Mutex, RwLock}, sync::{Arc, Mutex, RwLock},
time::SystemTime, time::SystemTime,
@ -512,7 +512,12 @@ impl Config {
#[inline] #[inline]
pub fn get_any_listen_addr() -> SocketAddr { pub fn get_any_listen_addr() -> SocketAddr {
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0) SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
}
#[inline]
pub fn get_any_listen_addr_v6() -> SocketAddr {
SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)
} }
pub fn get_rendezvous_server() -> String { pub fn get_rendezvous_server() -> String {

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

@ -10,7 +10,7 @@ pub use protos::rendezvous as rendezvous_proto;
use std::{ use std::{
fs::File, fs::File,
io::{self, BufRead}, io::{self, BufRead},
net::{Ipv4Addr, SocketAddr, SocketAddrV4}, net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
path::Path, path::Path,
time::{self, SystemTime, UNIX_EPOCH}, time::{self, SystemTime, UNIX_EPOCH},
}; };
@ -117,13 +117,31 @@ impl AddrMangle {
} }
bytes[..(16 - n_padding)].to_vec() bytes[..(16 - n_padding)].to_vec()
} }
_ => { SocketAddr::V6(addr_v6) => {
panic!("Only support ipv4"); let mut x = addr_v6.ip().octets().to_vec();
let port: [u8; 2] = addr_v6.port().to_le_bytes();
x.push(port[0]);
x.push(port[1]);
x
} }
} }
} }
pub fn decode(bytes: &[u8]) -> SocketAddr { pub fn decode(bytes: &[u8]) -> SocketAddr {
if bytes.len() > 16 {
if bytes.len() != 18 {
return Config::get_any_listen_addr_v6();
}
#[allow(invalid_value)]
let mut tmp: [u8; 2] = unsafe { std::mem::MaybeUninit::uninit().assume_init() };
tmp.copy_from_slice(&bytes[16..]);
let port = u16::from_le_bytes(tmp);
#[allow(invalid_value)]
let mut tmp: [u8; 16] = unsafe { std::mem::MaybeUninit::uninit().assume_init() };
tmp.copy_from_slice(&bytes[..16]);
let ip = std::net::Ipv6Addr::from(tmp);
return SocketAddr::new(IpAddr::V6(ip), port);
}
let mut padded = [0u8; 16]; let mut padded = [0u8; 16];
padded[..bytes.len()].copy_from_slice(&bytes); padded[..bytes.len()].copy_from_slice(&bytes);
let number = u128::from_le_bytes(padded); let number = u128::from_le_bytes(padded);
@ -264,11 +282,21 @@ mod tests {
fn test_mangle() { fn test_mangle() {
let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(192, 168, 16, 32), 21116)); let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(192, 168, 16, 32), 21116));
assert_eq!(addr, AddrMangle::decode(&AddrMangle::encode(addr))); assert_eq!(addr, AddrMangle::decode(&AddrMangle::encode(addr)));
let addr = "[2001:db8::1]:8080".parse::<SocketAddr>().unwrap();
assert_eq!(addr, AddrMangle::decode(&AddrMangle::encode(addr)));
let addr = "[2001:db8:ff::1111]:80".parse::<SocketAddr>().unwrap();
assert_eq!(addr, AddrMangle::decode(&AddrMangle::encode(addr)));
} }
#[test] #[test]
fn test_allow_err() { fn test_allow_err() {
allow_err!(Err("test err") as Result<(), &str>); allow_err!(Err("test err") as Result<(), &str>);
allow_err!(Err("test err with msg") as Result<(), &str>, "prompt {}", "failed"); allow_err!(
Err("test err with msg") as Result<(), &str>,
"prompt {}",
"failed"
);
} }
} }