From 633253647fd9e021e28609703d52b0711fa1a09f Mon Sep 17 00:00:00 2001
From: rustdesk <info@rustdesk.com>
Date: Mon, 26 Dec 2022 16:41:33 +0800
Subject: [PATCH] ipv6 mangle

---
 libs/hbb_common/src/config.rs |  9 +++++++--
 libs/hbb_common/src/lib.rs    | 38 ++++++++++++++++++++++++++++++-----
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/libs/hbb_common/src/config.rs b/libs/hbb_common/src/config.rs
index d4701aada..e2592bbea 100644
--- a/libs/hbb_common/src/config.rs
+++ b/libs/hbb_common/src/config.rs
@@ -1,7 +1,7 @@
 use std::{
     collections::HashMap,
     fs,
-    net::{IpAddr, Ipv4Addr, SocketAddr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
     path::{Path, PathBuf},
     sync::{Arc, Mutex, RwLock},
     time::SystemTime,
@@ -512,7 +512,12 @@ impl Config {
 
     #[inline]
     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 {
diff --git a/libs/hbb_common/src/lib.rs b/libs/hbb_common/src/lib.rs
index 0f9f7824c..5a6d24b91 100644
--- a/libs/hbb_common/src/lib.rs
+++ b/libs/hbb_common/src/lib.rs
@@ -10,7 +10,7 @@ pub use protos::rendezvous as rendezvous_proto;
 use std::{
     fs::File,
     io::{self, BufRead},
-    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
     path::Path,
     time::{self, SystemTime, UNIX_EPOCH},
 };
@@ -66,7 +66,7 @@ macro_rules! allow_err {
         } else {
         }
     };
-    
+
     ($e:expr, $($arg:tt)*) => {
         if let Err(err) = $e {
             log::debug!(
@@ -117,13 +117,31 @@ impl AddrMangle {
                 }
                 bytes[..(16 - n_padding)].to_vec()
             }
-            _ => {
-                panic!("Only support ipv4");
+            SocketAddr::V6(addr_v6) => {
+                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 {
+        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];
         padded[..bytes.len()].copy_from_slice(&bytes);
         let number = u128::from_le_bytes(padded);
@@ -264,11 +282,21 @@ mod tests {
     fn test_mangle() {
         let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(192, 168, 16, 32), 21116));
         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]
     fn test_allow_err() {
         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"
+        );
     }
 }