Implement TSO on linux tun device

gotatun/src/packet/ipv4/mod.rs

@@ -18,7 +18,10 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, TryFromBytes, Unali
 mod protocol;
 pub use protocol::*;
 
-use crate::packet::{DecodeError, Decoder, PseudoHeaderV4, Tcp, TcpDecoder, Udp, UdpDecoder};
+use crate::packet::{
+    DecodeError, Decoder, PseudoHeaderV4, Tcp, TcpDecoder, Udp, UdpDecoder,
+    util::checksum_ipv4_with_skip,
+};
 
 use super::util::size_must_be;
 
@@ -434,9 +437,9 @@ impl Ipv4Header {
         self.flags_and_fragment_offset.fragment_offset()
     }
 
-    /// Compute expected header checksum.
+    /// Compute expected header checksum. Only valid if IP header contains no optinal fields.
     pub fn compute_checksum(&self) -> u16 {
-        crate::packet::util::checksum_ipv4_with_skip(self.as_bytes())
+        checksum_ipv4_with_skip(self.as_bytes())
     }
 }
 
@@ -461,12 +464,35 @@ where
             .map_err(|_| eyre!("IPv4 packet was larger than {}", u16::MAX))?;
         Ok(())
     }
+
+    // /// Compute expected IP header checksum.
+    // pub fn calculate_ip_checksum(&self) -> u16 {
+    //     self.header.compute_checksum()
+    // }
+
+    // /// Update IP header checksum field.
+    // pub fn update_ip_checksum(&mut self) {
+    //     let checksum = self.header.compute_checksum();
+    //     self.header.header_checksum.set(checksum);
+    // }
 }
 
 impl<P> Ipv4<Ipv4Options<P>>
 where
     P: TryFromBytes + Immutable + KnownLayout,
 {
+    fn header_bytes(&self) -> eyre::Result<&[u8]> {
+        let header_len = usize::from(self.header.ihl()) * size_of::<u32>();
+
+        if header_len < Ipv4Header::LEN {
+            bail!("Invalid IHL");
+        };
+
+        self.as_bytes()
+            .get(..header_len)
+            .ok_or(eyre!("IHL larger than header"))
+    }
+
     fn options_and_payload_bytes(&self) -> eyre::Result<(&[u8], &[u8])> {
         let header_len = usize::from(self.header.ihl()) * size_of::<u32>();
 
@@ -495,6 +521,18 @@ where
         let payload = P::try_ref_from_bytes(bytes).map_err(|e| eyre!("{e}"))?;
         Ok(payload)
     }
+
+    /// Compute expected IP header checksum.
+    pub fn calculate_ip_checksum(&self) -> eyre::Result<u16> {
+        Ok(checksum_ipv4_with_skip(self.header_bytes()?))
+    }
+
+    /// Update IP header checksum field.
+    pub fn update_ip_checksum(&mut self) -> eyre::Result<()> {
+        let checksum = self.calculate_ip_checksum()?;
+        self.header.header_checksum.set(checksum);
+        Ok(())
+    }
 }
 
 impl Debug for Ipv4Header {

gotatun/src/tun/tun_async_device.rs

@@ -11,16 +11,23 @@
 
 //! Implementations of [`IpSend`] and [`IpRecv`] for the [`tun`] crate.
 
+mod linux;
+mod tso;
+mod virtio;
+
+use bytes::BytesMut;
 use tokio::{sync::watch, time::sleep};
+use tso::try_enable_tso;
 use tun::AbstractDevice;
+use zerocopy::IntoBytes;
 
 use crate::{
     packet::{Ip, Packet, PacketBufPool},
     task::Task,
     tun::{IpRecv, IpSend, MtuWatcher},
 };
 
-use std::{convert::Infallible, io, iter, sync::Arc, time::Duration};
+use std::{convert::Infallible, io, ops::Deref, sync::Arc, time::Duration};
 
 /// Error from [`TunDevice`].
 #[derive(Debug, thiserror::Error)]
@@ -57,6 +64,8 @@
 
     /// Task which monitors TUN device MTU. Aborted when dropped.
     _mtu_monitor: Task,
+
+    vnet_header: bool,
 }
 
 impl TunDevice {
@@ -79,10 +88,21 @@
         tun_config.platform_config(|p| {
             p.enable_routing(false);
         });
+
+        // HACK: always enable tso
+        #[cfg(target_os = "linux")]
+        tun_config.platform_config(|p| {
+            p.vnet_hdr(true);
+        });
+        let vnet_hdr = cfg!(target_os = "linux");
+
         // TODO: for wintun, must set path or enable signature check
         // we should upstream to `tun`
         let tun = tun::create_as_async(&tun_config).map_err(Error::OpenTun)?;
-        let tun = TunDevice::from_tun_device(tun)?;
+        // HACK: always enable tso
+        try_enable_tso(tun.deref()).unwrap();
+        let tun = TunDevice::from_tun_device(tun, vnet_hdr)?;
+
         Ok(tun)
     }
 
@@ -104,7 +124,7 @@
     }
 
     /// Construct from a [`tun::AsyncDevice`].
-    pub fn from_tun_device(tun: tun::AsyncDevice) -> Result<Self, Error> {
+    pub fn from_tun_device(tun: tun::AsyncDevice, vnet_header: bool) -> Result<Self, Error> {
         #[cfg(target_os = "linux")]
         if tun.packet_information() {
             return Err(Error::UnsupportedFeature("packet_information".to_string()));
@@ -140,6 +160,7 @@
             state: Arc::new(TunDeviceState {
                 mtu: rx.into(),
                 _mtu_monitor: mtu_monitor,
+                vnet_header,
             }),
         })
     }
@@ -152,11 +173,28 @@
 
 impl IpSend for TunDevice {
     async fn send(&mut self, packet: Packet<Ip>) -> io::Result<()> {
+        let mut packet = packet.into_bytes();
+        if self.state.vnet_header {
+            let header = virtio::VirtioNetHeader {
+                flags: virtio::Flags::new(),
+                gso_type: virtio::GsoType::VIRTIO_NET_HDR_GSO_NONE,
+                hdr_len: 0,
+                gso_size: 0,
+                csum_start: 0,
+                csum_offset: 0,
+            };
+            let mut buf =
+                BytesMut::with_capacity(header.as_bytes().len() + packet.as_bytes().len());
+            buf.extend_from_slice(header.as_bytes());
+            buf.extend_from_slice(packet.as_bytes());
+            *packet.buf_mut() = buf;
+        }
         self.tun.send(&packet.into_bytes()).await?;
         Ok(())
     }
 }
 
+#[cfg(not(any(target_os = "linux", target_os = "android")))]
 impl IpRecv for TunDevice {
     async fn recv<'a>(
         &'a mut self,
@@ -166,7 +204,7 @@
         let n = self.tun.recv(&mut packet).await?;
         packet.truncate(n);
         match packet.try_into_ip() {
             Ok(packet) => Ok(iter::once(packet)),
             Err(e) => Err(io::Error::other(e.to_string())),
         }
     }
@@ -175,3 +213,55 @@
         self.state.mtu.clone()
     }
 }
+
+#[cfg(any(target_os = "linux", target_os = "android"))]
+impl IpRecv for TunDevice {
+    async fn recv<'a>(
+        &'a mut self,
+        pool: &mut PacketBufPool,
+    ) -> io::Result<impl Iterator<Item = Packet<Ip>> + 'a> {
+        use bytes::BytesMut;
+        use either::Either;
+        use zerocopy::FromBytes;
+
+        use crate::tun::tun_async_device::virtio::VirtioNetHeader;
+
+        // FIXME: pool buffers have a cap of 4096, but we need more
+        //let mut packet = pool.get();
+        let _ = pool;
+
+        let mut buf = BytesMut::zeroed(usize::from(u16::MAX));
+        let n = self.tun.recv(&mut buf).await?;
+        buf.truncate(n);
+
+        if self.state.vnet_header {
+            let vnet_hdr = buf.split_to(size_of::<VirtioNetHeader>());
+            let vnet_hdr = *VirtioNetHeader::ref_from_bytes(&vnet_hdr).unwrap();
+
+            let packet = Packet::from_bytes(buf)
+                .try_into_ipvx()
+                .map_err(|e| io::Error::other(e.to_string()))?;
+
+            // TODO
+            let mtu = 1200;
+
+            // TODO: if segmentation and checksum offload is disabled,
+            // we could take a more efficient branch where we do not need to check
+            // packet length, and whether it's an IP/TCP packet.
+            match packet {
+                Either::Left(ipv4_packet) => {
+                    tso::new_tso_iter_ipv4(ipv4_packet, usize::from(vnet_hdr.gso_size))
+                }
+                Either::Right(ipv6_packet) => {
+                    tso::new_tso_iter_ipv6(ipv6_packet, usize::from(vnet_hdr.gso_size))
+                }
+            }
+        } else {
+            todo!("vnet_hdr is disabled")
+        }
+    }
+
+    fn mtu(&self) -> MtuWatcher {
+        self.state.mtu.clone()
+    }
+}