From 042f5af49f5dbdc488d98c543a9d0708d13d23dc Mon Sep 17 00:00:00 2001
From: Datong Sun <dndx@idndx.com>
Date: Sat, 30 Oct 2021 09:02:28 -0700
Subject: [PATCH] feat(phantun) add support for connecting to/from IPv6 based
 UDP endpoint

---
 phantun/src/bin/client.rs | 21 +++++++++++++++------
 phantun/src/bin/server.rs | 14 ++++++++++----
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/phantun/src/bin/client.rs b/phantun/src/bin/client.rs
index 59deebf..6be1645 100644
--- a/phantun/src/bin/client.rs
+++ b/phantun/src/bin/client.rs
@@ -16,8 +16,17 @@ use tokio_tun::TunBuilder;
 
 const UDP_TTL: Duration = Duration::from_secs(180);
 
-fn new_udp_reuseport(addr: SocketAddrV4) -> UdpSocket {
-    let udp_sock = socket2::Socket::new(socket2::Domain::IPV4, socket2::Type::DGRAM, None).unwrap();
+fn new_udp_reuseport(addr: SocketAddr) -> UdpSocket {
+    let udp_sock = socket2::Socket::new(
+        if addr.is_ipv4() {
+            socket2::Domain::IPV4
+        } else {
+            socket2::Domain::IPV6
+        },
+        socket2::Type::DGRAM,
+        None,
+    )
+    .unwrap();
     udp_sock.set_reuse_port(true).unwrap();
     // from tokio-rs/mio/blob/master/src/sys/unix/net.rs
     udp_sock.set_cloexec(true).unwrap();
@@ -40,7 +49,7 @@ async fn main() {
                 .long("local")
                 .required(true)
                 .value_name("IP:PORT")
-                .help("Sets the IP and port where Phantun Client listens for incoming UDP datagrams")
+                .help("Sets the IP and port where Phantun Client listens for incoming UDP datagrams, IPv6 address need to be specified as: \"[IPv6]:PORT\"")
                 .takes_value(true),
         )
         .arg(
@@ -83,7 +92,7 @@ async fn main() {
         )
         .get_matches();
 
-    let local_addr: SocketAddrV4 = matches
+    let local_addr: SocketAddr = matches
         .value_of("local")
         .unwrap()
         .parse()
@@ -117,7 +126,7 @@ async fn main() {
     info!("Created TUN device {}", tun[0].name());
 
     let udp_sock = Arc::new(new_udp_reuseport(local_addr));
-    let connections = Arc::new(RwLock::new(HashMap::<SocketAddrV4, Arc<Socket>>::new()));
+    let connections = Arc::new(RwLock::new(HashMap::<SocketAddr, Arc<Socket>>::new()));
 
     let mut stack = Stack::new(tun);
 
@@ -126,7 +135,7 @@ async fn main() {
 
         loop {
             tokio::select! {
-                Ok((size, SocketAddr::V4(addr))) = udp_sock.recv_from(&mut buf_r) => {
+                Ok((size, addr)) = udp_sock.recv_from(&mut buf_r) => {
                     // seen UDP packet to listening socket, this means:
                     // 1. It is a new UDP connection, or
                     // 2. It is some extra packets not filtered by more specific
diff --git a/phantun/src/bin/server.rs b/phantun/src/bin/server.rs
index 1e54dae..0484ba9 100644
--- a/phantun/src/bin/server.rs
+++ b/phantun/src/bin/server.rs
@@ -4,7 +4,7 @@ use clap::{crate_version, App, Arg};
 use fake_tcp::packet::MAX_PACKET_LEN;
 use fake_tcp::Stack;
 use log::{error, info};
-use std::net::{Ipv4Addr, SocketAddrV4};
+use std::net::{Ipv4Addr, SocketAddr};
 use tokio::net::UdpSocket;
 use tokio::time::{self, Duration};
 use tokio_tun::TunBuilder;
@@ -32,7 +32,7 @@ async fn main() {
                 .long("remote")
                 .required(true)
                 .value_name("IP:PORT")
-                .help("Sets the address and port where Phantun Server forwards UDP packets to")
+                .help("Sets the address and port where Phantun Server forwards UDP packets to, IPv6 address need to be specified as: \"[IPv6]:PORT\"")
                 .takes_value(true),
         )
         .arg(
@@ -71,7 +71,7 @@ async fn main() {
         .unwrap()
         .parse()
         .expect("bad local port");
-    let remote_addr: SocketAddrV4 = matches
+    let remote_addr: SocketAddr = matches
         .value_of("remote")
         .unwrap()
         .parse()
@@ -113,7 +113,13 @@ async fn main() {
             info!("New connection: {}", sock);
 
             tokio::spawn(async move {
-                let udp_sock = UdpSocket::bind("0.0.0.0:0").await.unwrap();
+                let udp_sock = UdpSocket::bind(if remote_addr.is_ipv4() {
+                    "0.0.0.0:0"
+                } else {
+                    "[::]:0"
+                })
+                .await
+                .unwrap();
                 udp_sock.connect(remote_addr).await.unwrap();
 
                 loop {