From 82c4d218d7b505a97cf0e8357c0f208021fabc41 Mon Sep 17 00:00:00 2001
From: JettChenT <jettchen12345@gmail.com>
Date: Mon, 29 May 2023 10:36:10 +0800
Subject: [PATCH] Add read and write functions for TunInterface

This adds read and write functionality for TunInterface.
---
 Cargo.lock                |  1 +
 tun/Cargo.toml            |  1 +
 tun/src/unix/apple/mod.rs | 39 +++++++++++++++++++++++++--
 tun/src/unix/linux/mod.rs | 13 ++++++++-
 tun/src/unix/mod.rs       | 55 ++++++++++++++++++++++++++++++++++++++-
 5 files changed, 105 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c779b81..4edb2ec 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1298,6 +1298,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "bindgen",
+ "byteorder",
  "fehler",
  "lazy_static",
  "libc",
diff --git a/tun/Cargo.toml b/tun/Cargo.toml
index e0a5e45..33446d1 100644
--- a/tun/Cargo.toml
+++ b/tun/Cargo.toml
@@ -9,6 +9,7 @@ fehler = "1.0"
 nix = { version = "0.26", features = ["ioctl"] }
 socket2 = "0.4"
 tokio = { version = "1.28", features = [] }
+byteorder = "1.4"
 
 [target.'cfg(windows)'.dependencies]
 lazy_static = "1.4"
diff --git a/tun/src/unix/apple/mod.rs b/tun/src/unix/apple/mod.rs
index 4230c41..a2a5674 100644
--- a/tun/src/unix/apple/mod.rs
+++ b/tun/src/unix/apple/mod.rs
@@ -1,9 +1,14 @@
+use byteorder::{ByteOrder, NetworkEndian};
 use fehler::throws;
-use libc::c_char;
+use libc::{c_char, iovec, writev, AF_INET, AF_INET6};
 use socket2::{Domain, SockAddr, Socket, Type};
+use std::io::{IoSlice, Write};
 use std::net::{Ipv4Addr, SocketAddrV4};
 use std::os::fd::{AsRawFd, RawFd};
-use std::{io::Error, mem};
+use std::{
+    io::{Error, Read},
+    mem,
+};
 
 mod kern_control;
 mod sys;
@@ -122,6 +127,36 @@ impl TunInterface {
     }
 }
 
+impl Write for TunInterface {
+    #[throws]
+    fn write(&mut self, buf: &[u8]) -> usize {
+        use std::io::ErrorKind;
+        let proto = match buf[0] >> 4 {
+            6 => Ok(AF_INET6),
+            4 => Ok(AF_INET),
+            _ => Err(Error::new(ErrorKind::InvalidInput, "Invalid IP version")),
+        }?;
+        let mut pbuf = [0; 4];
+        NetworkEndian::write_i32(&mut pbuf, proto);
+
+        let bufs = [IoSlice::new(&pbuf), IoSlice::new(buf)];
+        let bytes_written: isize = unsafe {
+            writev(
+                self.as_raw_fd(),
+                bufs.as_ptr() as *const iovec,
+                bufs.len() as i32,
+            )
+        };
+        bytes_written
+            .try_into()
+            .map_err(|_| Error::new(ErrorKind::Other, "Conversion error"))?
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/tun/src/unix/linux/mod.rs b/tun/src/unix/linux/mod.rs
index 20eabaa..8467731 100644
--- a/tun/src/unix/linux/mod.rs
+++ b/tun/src/unix/linux/mod.rs
@@ -2,7 +2,7 @@ use fehler::throws;
 
 use socket2::{Domain, SockAddr, Socket, Type};
 use std::fs::OpenOptions;
-use std::io::Error;
+use std::io::{Error, Write};
 use std::mem;
 use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
 use std::os::fd::RawFd;
@@ -140,6 +140,17 @@ impl TunInterface {
     }
 }
 
+impl Write for TunInterface {
+    #[throws]
+    fn write(&mut self, buf: &[u8]) -> usize {
+        self.socket.send(buf)?
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::TunInterface;
diff --git a/tun/src/unix/mod.rs b/tun/src/unix/mod.rs
index 6ee3bf4..28b4ec8 100644
--- a/tun/src/unix/mod.rs
+++ b/tun/src/unix/mod.rs
@@ -1,4 +1,7 @@
-use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
+use std::{
+    io::{Error, Read},
+    os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd},
+};
 
 mod queue;
 
@@ -10,6 +13,7 @@ mod imp;
 #[path = "linux/mod.rs"]
 mod imp;
 
+use fehler::throws;
 pub use imp::TunInterface;
 pub use queue::TunQueue;
 
@@ -32,6 +36,19 @@ impl IntoRawFd for TunInterface {
         self.socket.into_raw_fd()
     }
 }
+
+impl TunInterface {
+    // #[throws]
+    // pub fn write(&self, buf: &[u8]) -> usize {
+    //     self.socket.send(buf)?
+    // }
+
+    #[throws]
+    pub fn read(&mut self, buf: &mut [u8]) -> usize {
+        self.socket.read(buf)?
+    }
+}
+
 pub fn ifname_to_string(buf: [libc::c_char; libc::IFNAMSIZ]) -> String {
     // TODO: Switch to `CStr::from_bytes_until_nul` when stabilized
     unsafe {
@@ -48,3 +65,39 @@ pub fn string_to_ifname(name: &str) -> [libc::c_char; libc::IFNAMSIZ] {
     buf[..len].copy_from_slice(unsafe { &*(name.as_bytes() as *const _ as *const [libc::c_char]) });
     buf
 }
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+    use std::io::Write;
+    use std::io::{self, BufRead};
+    use std::net::Ipv4Addr;
+
+    #[throws]
+    #[test]
+    fn tst_read() {
+        // This test is interactive, you need to send a packet to any server through 192.168.1.10
+        // EG. `sudo route add 8.8.8.8 192.168.1.10`,
+        //`dig @8.8.8.8 hackclub.com`
+        let mut tun = TunInterface::new()?;
+        println!("tun name: {:?}", tun.name()?);
+        tun.set_ipv4_addr(Ipv4Addr::from([192, 168, 1, 10]))?;
+        println!("tun ip: {:?}", tun.ipv4_addr()?);
+        println!("Waiting for a packet...");
+        let buf = &mut [0u8; 1500];
+        let res = tun.read(buf);
+        println!("Received!");
+        assert!(res.is_ok());
+    }
+
+    #[test]
+    #[throws]
+    fn write_packets() {
+        let mut tun = TunInterface::new()?;
+        let mut buf = [0u8; 1500];
+        buf[0] = 6 << 4;
+        let bytes_written = tun.write(&buf)?;
+        assert_eq!(bytes_written, 1504);
+    }
+}