From 864e8090ba1b473479966a1419630109ffd64210 Mon Sep 17 00:00:00 2001
From: Qiming HOU <hqm03ster@gmail.com>
Date: Tue, 14 Feb 2023 16:36:42 +0800
Subject: [PATCH] add --wireguard mode

---
 common.cpp  | 17 +++++++++++++++++
 common.h    |  2 ++
 misc.cpp    | 13 +++++++++++++
 network.cpp |  8 ++++++--
 network.h   |  2 ++
 5 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/common.cpp b/common.cpp
index ffcddde..27362f7 100644
--- a/common.cpp
+++ b/common.cpp
@@ -11,6 +11,7 @@
 
 #include <random>
 #include <cmath>
+#include <stdint.h>
 
 // static int random_number_fd=-1;
 int force_socket_buf = 0;
@@ -190,6 +191,8 @@ int address_t::from_sockaddr(sockaddr *addr, socklen_t slen) {
     return 0;
 }
 
+int g_randomize_local_addr = 0;
+static uint32_t g_lo_ip = 0x7f010001u;
 int address_t::new_connected_udp_fd() {
     int new_udp_fd;
     new_udp_fd = socket(get_type(), SOCK_DGRAM, IPPROTO_UDP);
@@ -200,6 +203,20 @@ int address_t::new_connected_udp_fd() {
     setnonblocking(new_udp_fd);
     set_buf_size(new_udp_fd, socket_buf_size);
 
+    struct sockaddr_in *paddr_inet = (struct sockaddr_in *)&inner;
+    if (paddr_inet->sin_family == AF_INET && g_randomize_local_addr && 
+    (ntohl(paddr_inet->sin_addr.s_addr) & 0xff000000u) == 0x7f000000u) {
+        // wireguard allows only one port number per address, so change source address on reconnection
+        struct sockaddr_in addr_bound;
+        memset(&addr_bound, 0, sizeof(addr_bound));
+        addr_bound.sin_family = AF_INET;
+        addr_bound.sin_addr.s_addr = htonl(g_lo_ip);
+        g_lo_ip += 0x2u;
+        mylog(log_debug, "randomizing local address when connecting to localhost, binding local ip %s\n", my_ntoa(g_lo_ip));
+        if (bind(new_udp_fd, (struct sockaddr *)&addr_bound, sizeof(addr_bound)) != 0) {
+            mylog(log_warn, "lo addr: bind failed\n");
+        }
+    }
     mylog(log_debug, "created new udp_fd %d\n", new_udp_fd);
     int ret = connect(new_udp_fd, (struct sockaddr *)&inner, get_len());
     if (ret != 0) {
diff --git a/common.h b/common.h
index b6c559f..1df436e 100644
--- a/common.h
+++ b/common.h
@@ -163,6 +163,8 @@ extern int force_socket_buf;
 
 extern int g_fix_gro;
 
+extern int g_randomize_local_addr;
+
 /*
 struct ip_port_t
 {
diff --git a/misc.cpp b/misc.cpp
index 737c9be..66ab8e6 100644
--- a/misc.cpp
+++ b/misc.cpp
@@ -296,6 +296,9 @@ void process_arg(int argc, char *argv[])  // process all options
             {"no-pcap-mutex", no_argument, 0, 1},
 #endif
             {"fix-gro", no_argument, 0, 1},
+            {"do-fragment", no_argument, 0, 1},
+            {"rand-addr", no_argument, 0, 1},
+            {"wireguard", no_argument, 0, 1},
             {NULL, 0, 0, 0}};
 
     process_log_level(argc, argv);
@@ -677,6 +680,16 @@ void process_arg(int argc, char *argv[])  // process all options
                 } else if (strcmp(long_options[option_index].name, "fix-gro") == 0) {
                     mylog(log_info, "--fix-gro enabled\n");
                     g_fix_gro = 1;
+                } else if (strcmp(long_options[option_index].name, "do-fragment") == 0) {
+                    mylog(log_info, "--do-fragment enabled\n");
+                    g_should_fragment = 1;
+                } else if (strcmp(long_options[option_index].name, "rand-addr") == 0) {
+                    mylog(log_info, "--rand-addr enabled\n");
+                    g_randomize_local_addr = 1;
+                } else if (strcmp(long_options[option_index].name, "wireguard") == 0) {
+                    mylog(log_info, "--wireguard mode enabled, turning on --do-fragment and --rand-addr\n");
+                    g_should_fragment = 1;
+                    g_randomize_local_addr = 1;
                 } else {
                     mylog(log_warn, "ignored unknown long option ,option_index:%d code:<%x>\n", option_index, optopt);
                 }
diff --git a/network.cpp b/network.cpp
index adb8d9b..f1e527c 100644
--- a/network.cpp
+++ b/network.cpp
@@ -1159,6 +1159,7 @@ printf("pcap send!\n");*/
 }
 #endif
 
+int g_should_fragment = 0;
 int send_raw_ip(raw_info_t &raw_info, const char *payload, int payloadlen) {
     const packet_info_t &send_info = raw_info.send_info;
     const packet_info_t &recv_info = raw_info.recv_info;
@@ -1188,8 +1189,11 @@ int send_raw_ip(raw_info_t &raw_info, const char *payload, int payloadlen) {
             // iph->id = 0; //Id of this packet  ,kernel will auto fill this if id is zero  ,or really?????// todo //seems like there is a problem
         }
 
-        iph->frag_off = htons(0x4000);  // DF set,others are zero
-        // iph->frag_off = htons(0x0000); //DF set,others are zero
+        if (g_should_fragment) {
+            iph->frag_off = htons(0x0000);  //DF cleared,others are zero
+        } else {
+            iph->frag_off = htons(0x4000);  // DF set,others are zero
+        }
         iph->ttl = (unsigned char)ttl_value;
         iph->protocol = send_info.protocol;
         iph->check = 0;                        // Set to 0 before calculating checksum
diff --git a/network.h b/network.h
index dbeb0bf..6ccc215 100644
--- a/network.h
+++ b/network.h
@@ -56,6 +56,8 @@ struct icmphdr {
 };
 #endif
 
+extern int g_should_fragment;
+
 struct my_iphdr {
 #ifdef UDP2RAW_LITTLE_ENDIAN
     unsigned char ihl : 4;