impl: tcp fin

crates/lib/src/capture/pcap.rs

@@ -114,7 +114,7 @@ impl<W: Write> Pcap<W> {
                     &info,
                     IpNextHeaderProtocols::Tcp,
                     &buf[..buf_size],
-                    vec![EnhancedPacketOption::Comment("Generated TCP ack".into())],
+                    vec![EnhancedPacketOption::Comment("Generated TCP ACK".into())],
                 );
             }
             Protocol::UDP => {
@@ -285,8 +285,50 @@ impl<W: Write> Pcap<W> {
         self.state.has_sent_handshake = true;
     }
 
-    /// Take a transport layer packet as a buffer and write it after encoding
-    /// all the layers under it.
+    pub(crate) fn send_tcp_fin(&mut self, info: &CapturePacket) {
+        let mut buf = vec![0; BUFFER_SIZE];
+        let (source_port, dest_port) = info.ports_by_direction();
+
+        let buf_size = {
+            let mut tcp = MutableTcpPacket::new(&mut buf).unwrap();
+            tcp.set_source(source_port);
+            tcp.set_destination(dest_port);
+            tcp.set_data_offset(5);
+            tcp.set_window(43440);
+
+            match info.direction {
+                Direction::Send => {
+                    tcp.set_sequence(self.state.send_seq);
+                    tcp.set_acknowledgement(self.state.rec_seq);
+                }
+                Direction::Receive => {
+                    tcp.set_sequence(self.state.rec_seq);
+                    tcp.set_acknowledgement(self.state.send_seq);
+                }
+            }
+
+            tcp.set_flags(TcpFlags::FIN | TcpFlags::ACK);
+            tcp.packet_size()
+        };
+
+        self.write_transport_payload(
+            info,
+            IpNextHeaderProtocols::Tcp,
+            &buf[..buf_size],
+            vec![EnhancedPacketOption::Comment("Generated TCP FIN".into())],
+        );
+
+        // Update sequence number
+        match info.direction {
+            Direction::Send => {
+                self.state.send_seq = self.state.send_seq.wrapping_add(1);
+            }
+            Direction::Receive => {
+                self.state.rec_seq = self.state.rec_seq.wrapping_add(1);
+            }
+        }
+    }
+
     fn write_transport_payload(
         &mut self,
         info: &CapturePacket,

crates/lib/src/capture/writer.rs

@@ -14,7 +14,7 @@ lazy_static! {
 pub(crate) trait Writer {
     fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> GDResult<()>;
     fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()>;
-    //TODO: Close connection
+    fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()>;
 }
 
 impl<W: Write> Writer for Pcap<W> {
@@ -36,4 +36,14 @@ impl<W: Write> Writer for Pcap<W> {
 
         Ok(())
     }
+
+    fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()> {
+        match packet.protocol {
+            Protocol::TCP => {
+                self.send_tcp_fin(packet);
+            }
+            Protocol::UDP => {}
+        }
+        Ok(())
+    }
 }

crates/lib/src/socket.rs

@@ -239,7 +239,7 @@ pub mod capture {
     /// * `I` - The inner socket type.
     /// * `P` - The protocol provider.
     #[derive(Clone, Debug)]
-    pub(crate) struct WrappedCaptureSocket<I, P> {
+    pub(crate) struct WrappedCaptureSocket<I: Socket, P: ProtocolProvider> {
         inner: I,
         remote_address: SocketAddr,
         _protocol: PhantomData<P>,
@@ -369,6 +369,26 @@ pub mod capture {
         }
     }
 
+    // this seems a bad way to do this, but its safe
+    impl<I: Socket, P: ProtocolProvider> Drop for WrappedCaptureSocket<I, P> {
+        fn drop(&mut self) {
+            // Construct the CapturePacket info
+            let info = CapturePacket {
+                direction: Direction::Send,
+                protocol: P::protocol(),
+                remote_address: &self.remote_address,
+                local_address: &self
+                    .local_addr()
+                    .unwrap_or_else(|_| SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), 0)),
+            };
+
+            // If a capture writer is set, close the connection and capture the packet.
+            if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() {
+                let _ = writer.close_connection(&info);
+            }
+        }
+    }
+
     /// A specialized `WrappedCaptureSocket` for UDP, using `UdpSocketImpl` as the inner socket
     /// and `ProtocolUDP` as the protocol provider.
     ///