diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index e6cd5b0..0eb8b21 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -114,7 +114,7 @@ impl Pcap { &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 Pcap { 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, diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs index 1951e7e..4d63655 100644 --- a/crates/lib/src/capture/writer.rs +++ b/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 Writer for Pcap { @@ -36,4 +36,14 @@ impl Writer for Pcap { Ok(()) } + + fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()> { + match packet.protocol { + Protocol::TCP => { + self.send_tcp_fin(packet); + } + Protocol::UDP => {} + } + Ok(()) + } } diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 0ee26f6..0ba35ee 100644 --- a/crates/lib/src/socket.rs +++ b/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 { + pub(crate) struct WrappedCaptureSocket { inner: I, remote_address: SocketAddr, _protocol: PhantomData

, @@ -369,6 +369,26 @@ pub mod capture { } } + // this seems a bad way to do this, but its safe + impl Drop for WrappedCaptureSocket { + 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. ///