kratanet/nat/
processor.rs

1use crate::pkt::RecvPacket;
2use crate::pkt::RecvPacketIp;
3use anyhow::Result;
4use bytes::BytesMut;
5use etherparse::Icmpv4Header;
6use etherparse::Icmpv4Type;
7use etherparse::Icmpv6Header;
8use etherparse::Icmpv6Type;
9use etherparse::IpNumber;
10use etherparse::IpPayloadSlice;
11use etherparse::Ipv4Slice;
12use etherparse::Ipv6Slice;
13use etherparse::SlicedPacket;
14use etherparse::TcpHeaderSlice;
15use etherparse::UdpHeaderSlice;
16use log::warn;
17use log::{debug, trace};
18use smoltcp::wire::EthernetAddress;
19use smoltcp::wire::IpAddress;
20use smoltcp::wire::IpCidr;
21use smoltcp::wire::IpEndpoint;
22use std::collections::hash_map::Entry;
23use tokio::select;
24use tokio::sync::mpsc::channel;
25use tokio::sync::mpsc::Receiver;
26use tokio::sync::mpsc::Sender;
27use tokio::task::JoinHandle;
28
29use super::handler::NatHandler;
30use super::handler::NatHandlerContext;
31use super::handler::NatHandlerFactory;
32use super::key::NatKey;
33use super::key::NatKeyProtocol;
34use super::table::NatTable;
35
36const RECEIVE_CHANNEL_QUEUE_LEN: usize = 3000;
37const RECLAIM_CHANNEL_QUEUE_LEN: usize = 30;
38
39pub struct NatProcessor {
40    mtu: usize,
41    local_mac: EthernetAddress,
42    local_cidrs: Vec<IpCidr>,
43    table: NatTable,
44    factory: Box<dyn NatHandlerFactory>,
45    transmit_sender: Sender<BytesMut>,
46    reclaim_sender: Sender<NatKey>,
47    reclaim_receiver: Receiver<NatKey>,
48    receive_receiver: Receiver<BytesMut>,
49}
50
51enum NatProcessorSelect {
52    Reclaim(Option<NatKey>),
53    ReceivedPacket(Option<BytesMut>),
54}
55
56impl NatProcessor {
57    pub fn launch(
58        mtu: usize,
59        factory: Box<dyn NatHandlerFactory>,
60        local_mac: EthernetAddress,
61        local_cidrs: Vec<IpCidr>,
62        transmit_sender: Sender<BytesMut>,
63    ) -> Result<(Sender<BytesMut>, JoinHandle<()>)> {
64        let (reclaim_sender, reclaim_receiver) = channel(RECLAIM_CHANNEL_QUEUE_LEN);
65        let (receive_sender, receive_receiver) = channel(RECEIVE_CHANNEL_QUEUE_LEN);
66        let mut processor = Self {
67            mtu,
68            local_mac,
69            local_cidrs,
70            factory,
71            table: NatTable::new(),
72            transmit_sender,
73            reclaim_sender,
74            receive_receiver,
75            reclaim_receiver,
76        };
77
78        let handle = tokio::task::spawn(async move {
79            if let Err(error) = processor.process().await {
80                warn!("nat processing failed: {}", error);
81            }
82        });
83
84        Ok((receive_sender, handle))
85    }
86
87    pub async fn process(&mut self) -> Result<()> {
88        loop {
89            let selection = select! {
90                x = self.reclaim_receiver.recv() => NatProcessorSelect::Reclaim(x),
91                x = self.receive_receiver.recv() => NatProcessorSelect::ReceivedPacket(x),
92            };
93
94            match selection {
95                NatProcessorSelect::Reclaim(Some(key)) => {
96                    if self.table.inner.remove(&key).is_some() {
97                        debug!("reclaimed nat key: {}", key);
98                    }
99                }
100
101                NatProcessorSelect::ReceivedPacket(Some(packet)) => {
102                    if let Ok(slice) = SlicedPacket::from_ethernet(&packet) {
103                        let Ok(packet) = RecvPacket::new(&packet, &slice) else {
104                            continue;
105                        };
106
107                        self.process_packet(&packet).await?;
108                    }
109                }
110
111                NatProcessorSelect::ReceivedPacket(None) | NatProcessorSelect::Reclaim(None) => {
112                    break
113                }
114            }
115        }
116        Ok(())
117    }
118
119    pub async fn process_reclaim(&mut self) -> Result<Option<NatKey>> {
120        Ok(if let Some(key) = self.reclaim_receiver.recv().await {
121            if self.table.inner.remove(&key).is_some() {
122                debug!("reclaimed nat key: {}", key);
123                Some(key)
124            } else {
125                None
126            }
127        } else {
128            None
129        })
130    }
131
132    pub async fn process_packet<'a>(&mut self, packet: &RecvPacket<'a>) -> Result<()> {
133        let Some(ether) = packet.ether else {
134            return Ok(());
135        };
136
137        let mac = EthernetAddress(ether.destination());
138        if mac != self.local_mac {
139            trace!(
140                "received packet with destination {} which is not the local mac {}",
141                mac,
142                self.local_mac
143            );
144            return Ok(());
145        }
146
147        let key = match packet.ip {
148            Some(RecvPacketIp::Ipv4(ipv4)) => self.extract_key_ipv4(packet, ipv4)?,
149            Some(RecvPacketIp::Ipv6(ipv6)) => self.extract_key_ipv6(packet, ipv6)?,
150            _ => None,
151        };
152
153        let Some(key) = key else {
154            return Ok(());
155        };
156
157        for cidr in &self.local_cidrs {
158            if cidr.contains_addr(&key.external_ip.addr) {
159                return Ok(());
160            }
161        }
162
163        let context = NatHandlerContext {
164            mtu: self.mtu,
165            key,
166            transmit_sender: self.transmit_sender.clone(),
167            reclaim_sender: self.reclaim_sender.clone(),
168        };
169        let handler: Option<&mut Box<dyn NatHandler>> = match self.table.inner.entry(key) {
170            Entry::Occupied(entry) => Some(entry.into_mut()),
171            Entry::Vacant(entry) => {
172                if let Some(handler) = self.factory.nat(context).await {
173                    debug!("creating nat entry for key: {}", key);
174                    Some(entry.insert(handler))
175                } else {
176                    None
177                }
178            }
179        };
180
181        if let Some(handler) = handler {
182            if !handler.receive(packet.raw).await? {
183                self.reclaim_sender.try_send(key)?;
184            }
185        }
186        Ok(())
187    }
188
189    pub fn extract_key_ipv4<'a>(
190        &mut self,
191        packet: &RecvPacket<'a>,
192        ipv4: &Ipv4Slice<'a>,
193    ) -> Result<Option<NatKey>> {
194        let source_addr = IpAddress::Ipv4(ipv4.header().source_addr().into());
195        let dest_addr = IpAddress::Ipv4(ipv4.header().destination_addr().into());
196        Ok(match ipv4.header().protocol() {
197            IpNumber::TCP => {
198                self.extract_key_tcp(packet, source_addr, dest_addr, ipv4.payload())?
199            }
200
201            IpNumber::UDP => {
202                self.extract_key_udp(packet, source_addr, dest_addr, ipv4.payload())?
203            }
204
205            IpNumber::ICMP => {
206                self.extract_key_icmpv4(packet, source_addr, dest_addr, ipv4.payload())?
207            }
208
209            _ => None,
210        })
211    }
212
213    pub fn extract_key_ipv6<'a>(
214        &mut self,
215        packet: &RecvPacket<'a>,
216        ipv6: &Ipv6Slice<'a>,
217    ) -> Result<Option<NatKey>> {
218        let source_addr = IpAddress::Ipv6(ipv6.header().source_addr().into());
219        let dest_addr = IpAddress::Ipv6(ipv6.header().destination_addr().into());
220        Ok(match ipv6.header().next_header() {
221            IpNumber::TCP => {
222                self.extract_key_tcp(packet, source_addr, dest_addr, ipv6.payload())?
223            }
224
225            IpNumber::UDP => {
226                self.extract_key_udp(packet, source_addr, dest_addr, ipv6.payload())?
227            }
228
229            IpNumber::IPV6_ICMP => {
230                self.extract_key_icmpv6(packet, source_addr, dest_addr, ipv6.payload())?
231            }
232
233            _ => None,
234        })
235    }
236
237    pub fn extract_key_udp<'a>(
238        &mut self,
239        packet: &RecvPacket<'a>,
240        source_addr: IpAddress,
241        dest_addr: IpAddress,
242        payload: &IpPayloadSlice<'a>,
243    ) -> Result<Option<NatKey>> {
244        let Some(ether) = packet.ether else {
245            return Ok(None);
246        };
247        let header = UdpHeaderSlice::from_slice(payload.payload)?;
248        let source = IpEndpoint::new(source_addr, header.source_port());
249        let dest = IpEndpoint::new(dest_addr, header.destination_port());
250        Ok(Some(NatKey {
251            protocol: NatKeyProtocol::Udp,
252            client_mac: EthernetAddress(ether.source()),
253            local_mac: EthernetAddress(ether.destination()),
254            client_ip: source,
255            external_ip: dest,
256        }))
257    }
258
259    pub fn extract_key_icmpv4<'a>(
260        &mut self,
261        packet: &RecvPacket<'a>,
262        source_addr: IpAddress,
263        dest_addr: IpAddress,
264        payload: &IpPayloadSlice<'a>,
265    ) -> Result<Option<NatKey>> {
266        let Some(ether) = packet.ether else {
267            return Ok(None);
268        };
269        let (header, _) = Icmpv4Header::from_slice(payload.payload)?;
270        let Icmpv4Type::EchoRequest(_) = header.icmp_type else {
271            return Ok(None);
272        };
273        let source = IpEndpoint::new(source_addr, 0);
274        let dest = IpEndpoint::new(dest_addr, 0);
275        Ok(Some(NatKey {
276            protocol: NatKeyProtocol::Icmp,
277            client_mac: EthernetAddress(ether.source()),
278            local_mac: EthernetAddress(ether.destination()),
279            client_ip: source,
280            external_ip: dest,
281        }))
282    }
283
284    pub fn extract_key_icmpv6<'a>(
285        &mut self,
286        packet: &RecvPacket<'a>,
287        source_addr: IpAddress,
288        dest_addr: IpAddress,
289        payload: &IpPayloadSlice<'a>,
290    ) -> Result<Option<NatKey>> {
291        let Some(ether) = packet.ether else {
292            return Ok(None);
293        };
294        let (header, _) = Icmpv6Header::from_slice(payload.payload)?;
295        let Icmpv6Type::EchoRequest(_) = header.icmp_type else {
296            return Ok(None);
297        };
298        let source = IpEndpoint::new(source_addr, 0);
299        let dest = IpEndpoint::new(dest_addr, 0);
300        Ok(Some(NatKey {
301            protocol: NatKeyProtocol::Icmp,
302            client_mac: EthernetAddress(ether.source()),
303            local_mac: EthernetAddress(ether.destination()),
304            client_ip: source,
305            external_ip: dest,
306        }))
307    }
308
309    pub fn extract_key_tcp<'a>(
310        &mut self,
311        packet: &RecvPacket<'a>,
312        source_addr: IpAddress,
313        dest_addr: IpAddress,
314        payload: &IpPayloadSlice<'a>,
315    ) -> Result<Option<NatKey>> {
316        let Some(ether) = packet.ether else {
317            return Ok(None);
318        };
319        let header = TcpHeaderSlice::from_slice(payload.payload)?;
320        let source = IpEndpoint::new(source_addr, header.source_port());
321        let dest = IpEndpoint::new(dest_addr, header.destination_port());
322        Ok(Some(NatKey {
323            protocol: NatKeyProtocol::Tcp,
324            client_mac: EthernetAddress(ether.source()),
325            local_mac: EthernetAddress(ether.destination()),
326            client_ip: source,
327            external_ip: dest,
328        }))
329    }
330}
kratanet/nat/processor.rs

kratanet/nat/
processor.rs
