xref: /webrtc/util/src/vnet/chunk.rs (revision 7ceeeeb0) 
1 #[cfg(test)]
2 mod chunk_test;
3 
4 use super::net::*;
5 use crate::error::Result;
6 
7 use std::fmt;
8 use std::net::{IpAddr, SocketAddr};
9 use std::ops::{BitAnd, BitOr};
10 use std::str::FromStr;
11 use std::sync::atomic::{AtomicU64, Ordering};
12 use std::time::SystemTime;
13 
14 lazy_static! {
15     static ref TAG_CTR: AtomicU64 = AtomicU64::new(0);
16 }
17 
18 /// Encodes a u64 value to a lowercase base 36 string.
19 pub fn base36(value: impl Into<u64>) -> String {
20     let mut digits: Vec<u8> = vec![];
21 
22     let mut value = value.into();
23     while value > 0 {
24         let digit = (value % 36) as usize;
25         value /= 36;
26 
27         digits.push(b"0123456789abcdefghijklmnopqrstuvwxyz"[digit]);
28     }
29 
30     digits.reverse();
31     format!("{:0>8}", String::from_utf8(digits).unwrap())
32 }
33 
34 // Generate a base36-encoded unique tag
35 // See: https://play.golang.org/p/0ZaAID1q-HN
36 fn assign_chunk_tag() -> String {
37     let n = TAG_CTR.fetch_add(1, Ordering::SeqCst);
38     base36(n)
39 }
40 
41 #[derive(Copy, Clone, PartialEq, Debug)]
42 pub(crate) struct TcpFlag(pub(crate) u8);
43 
44 pub(crate) const TCP_FLAG_ZERO: TcpFlag = TcpFlag(0x00);
45 pub(crate) const TCP_FLAG_FIN: TcpFlag = TcpFlag(0x01);
46 pub(crate) const TCP_FLAG_SYN: TcpFlag = TcpFlag(0x02);
47 pub(crate) const TCP_FLAG_RST: TcpFlag = TcpFlag(0x04);
48 pub(crate) const TCP_FLAG_PSH: TcpFlag = TcpFlag(0x08);
49 pub(crate) const TCP_FLAG_ACK: TcpFlag = TcpFlag(0x10);
50 
51 impl BitOr for TcpFlag {
52     type Output = Self;
53 
54     // rhs is the "right-hand side" of the expression `a | b`
55     fn bitor(self, rhs: Self) -> Self::Output {
56         Self(self.0 | rhs.0)
57     }
58 }
59 
60 impl BitAnd for TcpFlag {
61     type Output = Self;
62 
63     // rhs is the "right-hand side" of the expression `a & b`
64     fn bitand(self, rhs: Self) -> Self::Output {
65         Self(self.0 & rhs.0)
66     }
67 }
68 
69 impl fmt::Display for TcpFlag {
70     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71         let mut sa = vec![];
72         if *self & TCP_FLAG_FIN != TCP_FLAG_ZERO {
73             sa.push("FIN");
74         }
75         if *self & TCP_FLAG_SYN != TCP_FLAG_ZERO {
76             sa.push("SYN");
77         }
78         if *self & TCP_FLAG_RST != TCP_FLAG_ZERO {
79             sa.push("RST");
80         }
81         if *self & TCP_FLAG_PSH != TCP_FLAG_ZERO {
82             sa.push("PSH");
83         }
84         if *self & TCP_FLAG_ACK != TCP_FLAG_ZERO {
85             sa.push("ACK");
86         }
87 
88         write!(f, "{}", sa.join("-"))
89     }
90 }
91 
92 // Chunk represents a packet passed around in the vnet
93 pub trait Chunk: fmt::Display + fmt::Debug {
94     fn set_timestamp(&mut self) -> SystemTime; // used by router
95     fn get_timestamp(&self) -> SystemTime; // used by router
96     fn get_source_ip(&self) -> IpAddr; // used by routee
97     fn get_destination_ip(&self) -> IpAddr; // used by router
98     fn set_source_addr(&mut self, address: &str) -> Result<()>; // used by nat
99     fn set_destination_addr(&mut self, address: &str) -> Result<()>; // used by nat
100 
101     fn source_addr(&self) -> SocketAddr;
102     fn destination_addr(&self) -> SocketAddr;
103     fn user_data(&self) -> Vec<u8>;
104     fn tag(&self) -> String;
105     fn network(&self) -> String; // returns "udp" or "tcp"
106     fn clone_to(&self) -> Box<dyn Chunk + Send + Sync>;
107 }
108 
109 #[derive(PartialEq, Debug)]
110 pub(crate) struct ChunkIp {
111     pub(crate) timestamp: SystemTime,
112     pub(crate) source_ip: IpAddr,
113     pub(crate) destination_ip: IpAddr,
114     pub(crate) tag: String,
115 }
116 
117 impl ChunkIp {
118     fn set_timestamp(&mut self) -> SystemTime {
119         self.timestamp = SystemTime::now();
120         self.timestamp
121     }
122 
123     fn get_timestamp(&self) -> SystemTime {
124         self.timestamp
125     }
126 
127     fn get_destination_ip(&self) -> IpAddr {
128         self.destination_ip
129     }
130 
131     fn get_source_ip(&self) -> IpAddr {
132         self.source_ip
133     }
134 
135     fn tag(&self) -> String {
136         self.tag.clone()
137     }
138 }
139 
140 #[derive(PartialEq, Debug)]
141 pub(crate) struct ChunkUdp {
142     pub(crate) chunk_ip: ChunkIp,
143     pub(crate) source_port: u16,
144     pub(crate) destination_port: u16,
145     pub(crate) user_data: Vec<u8>,
146 }
147 
148 impl fmt::Display for ChunkUdp {
149     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150         write!(
151             f,
152             "{} chunk {} {} => {}",
153             self.network(),
154             self.tag(),
155             self.source_addr(),
156             self.destination_addr(),
157         )
158     }
159 }
160 
161 impl Chunk for ChunkUdp {
162     fn set_timestamp(&mut self) -> SystemTime {
163         self.chunk_ip.set_timestamp()
164     }
165 
166     fn get_timestamp(&self) -> SystemTime {
167         self.chunk_ip.get_timestamp()
168     }
169 
170     fn get_destination_ip(&self) -> IpAddr {
171         self.chunk_ip.get_destination_ip()
172     }
173 
174     fn get_source_ip(&self) -> IpAddr {
175         self.chunk_ip.get_source_ip()
176     }
177 
178     fn tag(&self) -> String {
179         self.chunk_ip.tag()
180     }
181 
182     fn source_addr(&self) -> SocketAddr {
183         SocketAddr::new(self.chunk_ip.source_ip, self.source_port)
184     }
185 
186     fn destination_addr(&self) -> SocketAddr {
187         SocketAddr::new(self.chunk_ip.destination_ip, self.destination_port)
188     }
189 
190     fn user_data(&self) -> Vec<u8> {
191         self.user_data.clone()
192     }
193 
194     fn clone_to(&self) -> Box<dyn Chunk + Send + Sync> {
195         Box::new(ChunkUdp {
196             chunk_ip: ChunkIp {
197                 timestamp: self.chunk_ip.timestamp,
198                 source_ip: self.chunk_ip.source_ip,
199                 destination_ip: self.chunk_ip.destination_ip,
200                 tag: self.chunk_ip.tag.clone(),
201             },
202             source_port: self.source_port,
203             destination_port: self.destination_port,
204             user_data: self.user_data.clone(),
205         })
206     }
207 
208     fn network(&self) -> String {
209         UDP_STR.to_owned()
210     }
211 
212     fn set_source_addr(&mut self, address: &str) -> Result<()> {
213         let addr = SocketAddr::from_str(address)?;
214         self.chunk_ip.source_ip = addr.ip();
215         self.source_port = addr.port();
216         Ok(())
217     }
218 
219     fn set_destination_addr(&mut self, address: &str) -> Result<()> {
220         let addr = SocketAddr::from_str(address)?;
221         self.chunk_ip.destination_ip = addr.ip();
222         self.destination_port = addr.port();
223         Ok(())
224     }
225 }
226 
227 impl ChunkUdp {
228     pub(crate) fn new(src_addr: SocketAddr, dst_addr: SocketAddr) -> Self {
229         ChunkUdp {
230             chunk_ip: ChunkIp {
231                 timestamp: SystemTime::now(),
232                 source_ip: src_addr.ip(),
233                 destination_ip: dst_addr.ip(),
234                 tag: assign_chunk_tag(),
235             },
236             source_port: src_addr.port(),
237             destination_port: dst_addr.port(),
238             user_data: vec![],
239         }
240     }
241 }
242 
243 #[derive(PartialEq, Debug)]
244 pub(crate) struct ChunkTcp {
245     chunk_ip: ChunkIp,
246     source_port: u16,
247     destination_port: u16,
248     flags: TcpFlag, // control bits
249     user_data: Vec<u8>, // only with PSH flag
250                     // seq             :u32,  // always starts with 0
251                     // ack             :u32,  // always starts with 0
252 }
253 
254 impl fmt::Display for ChunkTcp {
255     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
256         write!(
257             f,
258             "{} {} chunk {} {} => {}",
259             self.network(),
260             self.flags,
261             self.chunk_ip.tag,
262             self.source_addr(),
263             self.destination_addr(),
264         )
265     }
266 }
267 
268 impl Chunk for ChunkTcp {
269     fn set_timestamp(&mut self) -> SystemTime {
270         self.chunk_ip.set_timestamp()
271     }
272 
273     fn get_timestamp(&self) -> SystemTime {
274         self.chunk_ip.get_timestamp()
275     }
276 
277     fn get_destination_ip(&self) -> IpAddr {
278         self.chunk_ip.get_destination_ip()
279     }
280 
281     fn get_source_ip(&self) -> IpAddr {
282         self.chunk_ip.get_source_ip()
283     }
284 
285     fn tag(&self) -> String {
286         self.chunk_ip.tag()
287     }
288 
289     fn source_addr(&self) -> SocketAddr {
290         SocketAddr::new(self.chunk_ip.source_ip, self.source_port)
291     }
292 
293     fn destination_addr(&self) -> SocketAddr {
294         SocketAddr::new(self.chunk_ip.destination_ip, self.destination_port)
295     }
296 
297     fn user_data(&self) -> Vec<u8> {
298         self.user_data.clone()
299     }
300 
301     fn clone_to(&self) -> Box<dyn Chunk + Send + Sync> {
302         Box::new(ChunkTcp {
303             chunk_ip: ChunkIp {
304                 timestamp: self.chunk_ip.timestamp,
305                 source_ip: self.chunk_ip.source_ip,
306                 destination_ip: self.chunk_ip.destination_ip,
307                 tag: self.chunk_ip.tag.clone(),
308             },
309             source_port: self.source_port,
310             destination_port: self.destination_port,
311             flags: self.flags,
312             user_data: self.user_data.clone(),
313         })
314     }
315 
316     fn network(&self) -> String {
317         "tcp".to_owned()
318     }
319 
320     fn set_source_addr(&mut self, address: &str) -> Result<()> {
321         let addr = SocketAddr::from_str(address)?;
322         self.chunk_ip.source_ip = addr.ip();
323         self.source_port = addr.port();
324         Ok(())
325     }
326 
327     fn set_destination_addr(&mut self, address: &str) -> Result<()> {
328         let addr = SocketAddr::from_str(address)?;
329         self.chunk_ip.destination_ip = addr.ip();
330         self.destination_port = addr.port();
331         Ok(())
332     }
333 }
334 
335 impl ChunkTcp {
336     pub(crate) fn new(src_addr: SocketAddr, dst_addr: SocketAddr, flags: TcpFlag) -> Self {
337         ChunkTcp {
338             chunk_ip: ChunkIp {
339                 timestamp: SystemTime::now(),
340                 source_ip: src_addr.ip(),
341                 destination_ip: dst_addr.ip(),
342                 tag: assign_chunk_tag(),
343             },
344             source_port: src_addr.port(),
345             destination_port: dst_addr.port(),
346             flags,
347             user_data: vec![],
348         }
349     }
350 }
351