From f496a7919bb1647bb8143c398fa59afe48f00f7d Mon Sep 17 00:00:00 2001
From: Datong Sun <dndx@idndx.com>
Date: Fri, 13 May 2022 08:02:47 -0700
Subject: [PATCH] feat(phantun) new option `--handshake-packet` that allows
 additional packet to be sent to the other end after TCP connection
 establishment

---
 phantun/src/bin/client.rs | 28 ++++++++++++++++++++++++++--
 phantun/src/bin/server.rs | 24 ++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/phantun/src/bin/client.rs b/phantun/src/bin/client.rs
index 9832b83..6bb1f14 100644
--- a/phantun/src/bin/client.rs
+++ b/phantun/src/bin/client.rs
@@ -4,6 +4,7 @@ use fake_tcp::{Socket, Stack};
 use log::{debug, error, info};
 use phantun::utils::{assign_ipv6_address, new_udp_reuseport};
 use std::collections::HashMap;
+use std::fs;
 use std::io;
 use std::net::{Ipv4Addr, SocketAddr};
 use std::sync::Arc;
@@ -97,6 +98,17 @@ async fn main() -> io::Result<()> {
                 .default_value("fcc8::2")
                 .takes_value(true),
         )
+        .arg(
+            Arg::new("handshake_packet")
+                .long("handshake-packet")
+                .required(false)
+                .value_name("PATH")
+                .help("Specify a file, which, after TCP handshake, its content will be sent as the \
+                      first data packet to the server.\n\
+                      Note: ensure this file's size does not exceed the MTU of the outgoing interface. \
+                      The content is always sent out in a single packet and will not be further segmented")
+                .takes_value(true),
+        )
         .get_matches();
 
     let local_addr: SocketAddr = matches
@@ -139,6 +151,10 @@ async fn main() -> io::Result<()> {
     };
 
     let tun_name = matches.value_of("tun").unwrap();
+    let handshake_packet: Option<Vec<u8>> = matches
+        .value_of("handshake_packet")
+        .map(fs::read)
+        .transpose()?;
 
     let num_cpus = num_cpus::get();
     info!("{} cores available", num_cpus);
@@ -186,9 +202,17 @@ async fn main() -> io::Result<()> {
             }
 
             let sock = Arc::new(sock.unwrap());
+            if let Some(ref p) = handshake_packet {
+                if sock.send(p).await.is_none() {
+                    error!("Failed to send handshake packet to remote, closing connection.");
+                    continue;
+                }
+
+                debug!("Sent handshake packet to: {}", sock);
+            }
+
             // send first packet
-            let res = sock.send(&buf_r[..size]).await;
-            if res.is_none() {
+            if sock.send(&buf_r[..size]).await.is_none() {
                 continue;
             }
 
diff --git a/phantun/src/bin/server.rs b/phantun/src/bin/server.rs
index 59bd6b3..38c96ae 100644
--- a/phantun/src/bin/server.rs
+++ b/phantun/src/bin/server.rs
@@ -3,6 +3,7 @@ use fake_tcp::packet::MAX_PACKET_LEN;
 use fake_tcp::Stack;
 use log::{debug, error, info};
 use phantun::utils::{assign_ipv6_address, new_udp_reuseport};
+use std::fs;
 use std::io;
 use std::net::Ipv4Addr;
 use std::sync::Arc;
@@ -97,6 +98,17 @@ async fn main() -> io::Result<()> {
                 .default_value("fcc9::2")
                 .takes_value(true),
         )
+        .arg(
+            Arg::new("handshake_packet")
+                .long("handshake-packet")
+                .required(false)
+                .value_name("PATH")
+                .help("Specify a file, which, after TCP handshake, its content will be sent as the \
+                      first data packet to the client.\n\
+                      Note: ensure this file's size does not exceed the MTU of the outgoing interface. \
+                      The content is always sent out in a single packet and will not be further segmented")
+                .takes_value(true),
+        )
         .get_matches();
 
     let local_port: u16 = matches
@@ -137,6 +149,10 @@ async fn main() -> io::Result<()> {
     };
 
     let tun_name = matches.value_of("tun").unwrap();
+    let handshake_packet: Option<Vec<u8>> = matches
+        .value_of("handshake_packet")
+        .map(fs::read)
+        .transpose()?;
 
     let num_cpus = num_cpus::get();
     info!("{} cores available", num_cpus);
@@ -169,6 +185,14 @@ async fn main() -> io::Result<()> {
         loop {
             let sock = Arc::new(stack.accept().await);
             info!("New connection: {}", sock);
+            if let Some(ref p) = handshake_packet {
+                if sock.send(p).await.is_none() {
+                    error!("Failed to send handshake packet to remote, closing connection.");
+                    continue;
+                }
+
+                debug!("Sent handshake packet to: {}", sock);
+            }
 
             let packet_received = Arc::new(Notify::new());
             let quit = CancellationToken::new();