Skip to content

Commit d02961b

Browse files
whitequarkwhitequark
committedDec 18, 2016
Implement TCP packet support.
·
v0.12.0v0.1.0
1 parent c7492b3 commit d02961b

File tree

3 files changed

+444
-1
lines changed

3 files changed

+444
-1
lines changed
 

‎src/wire/mod.rs

+3
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ mod ip;
8686
mod ipv4;
8787
mod icmpv4;
8888
mod udp;
89+
mod tcp;
8990

9091
pub use self::pretty_print::PrettyPrinter;
9192

@@ -116,3 +117,5 @@ pub use self::icmpv4::Repr as Icmpv4Repr;
116117

117118
pub use self::udp::Packet as UdpPacket;
118119
pub use self::udp::Repr as UdpRepr;
120+
121+
pub use self::tcp::Packet as TcpPacket;

‎src/wire/tcp.rs

+440
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,440 @@
1+
use byteorder::{ByteOrder, NetworkEndian};
2+
3+
use Error;
4+
use super::{InternetProtocolType, InternetAddress};
5+
use super::ip::checksum;
6+
7+
/// A read/write wrapper around an Transmission Control Protocol packet buffer.
8+
#[derive(Debug)]
9+
pub struct Packet<T: AsRef<[u8]>> {
10+
buffer: T
11+
}
12+
13+
mod field {
14+
#![allow(non_snake_case)]
15+
16+
use wire::field::*;
17+
18+
pub const SRC_PORT: Field = 0..2;
19+
pub const DST_PORT: Field = 2..4;
20+
pub const SEQ_NUM: Field = 4..8;
21+
pub const ACK_NUM: Field = 8..12;
22+
pub const FLAGS: Field = 12..14;
23+
pub const WIN_SIZE: Field = 14..16;
24+
pub const CHECKSUM: Field = 16..18;
25+
pub const URGENT: Field = 18..20;
26+
27+
pub const FLG_FIN: u16 = 0x001;
28+
pub const FLG_SYN: u16 = 0x002;
29+
pub const FLG_RST: u16 = 0x004;
30+
pub const FLG_PSH: u16 = 0x008;
31+
pub const FLG_ACK: u16 = 0x010;
32+
pub const FLG_URG: u16 = 0x020;
33+
pub const FLG_ECE: u16 = 0x040;
34+
pub const FLG_CWR: u16 = 0x080;
35+
pub const FLG_NS: u16 = 0x100;
36+
}
37+
38+
impl<T: AsRef<[u8]>> Packet<T> {
39+
/// Wrap a buffer with a TCP packet. Returns an error if the buffer
40+
/// is too small to contain one.
41+
pub fn new(buffer: T) -> Result<Packet<T>, Error> {
42+
let len = buffer.as_ref().len();
43+
if len < field::URGENT.end {
44+
Err(Error::Truncated)
45+
} else {
46+
Ok(Packet { buffer: buffer })
47+
}
48+
}
49+
50+
/// Consumes the packet, returning the underlying buffer.
51+
pub fn into_inner(self) -> T {
52+
self.buffer
53+
}
54+
55+
/// Return the source port field.
56+
#[inline(always)]
57+
pub fn src_port(&self) -> u16 {
58+
let data = self.buffer.as_ref();
59+
NetworkEndian::read_u16(&data[field::SRC_PORT])
60+
}
61+
62+
/// Return the destination port field.
63+
#[inline(always)]
64+
pub fn dst_port(&self) -> u16 {
65+
let data = self.buffer.as_ref();
66+
NetworkEndian::read_u16(&data[field::DST_PORT])
67+
}
68+
69+
/// Return the sequence number field.
70+
#[inline(always)]
71+
pub fn seq_number(&self) -> u32 {
72+
let data = self.buffer.as_ref();
73+
NetworkEndian::read_u32(&data[field::SEQ_NUM])
74+
}
75+
76+
/// Return the acknowledgement number field.
77+
#[inline(always)]
78+
pub fn ack_number(&self) -> u32 {
79+
let data = self.buffer.as_ref();
80+
NetworkEndian::read_u32(&data[field::ACK_NUM])
81+
}
82+
83+
/// Return the FIN flag.
84+
#[inline(always)]
85+
pub fn fin(&self) -> bool {
86+
let data = self.buffer.as_ref();
87+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
88+
raw & field::FLG_FIN != 0
89+
}
90+
91+
/// Return the SYN flag.
92+
#[inline(always)]
93+
pub fn syn(&self) -> bool {
94+
let data = self.buffer.as_ref();
95+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
96+
raw & field::FLG_SYN != 0
97+
}
98+
99+
/// Return the RST flag.
100+
#[inline(always)]
101+
pub fn rst(&self) -> bool {
102+
let data = self.buffer.as_ref();
103+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
104+
raw & field::FLG_RST != 0
105+
}
106+
107+
/// Return the PSH flag.
108+
#[inline(always)]
109+
pub fn psh(&self) -> bool {
110+
let data = self.buffer.as_ref();
111+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
112+
raw & field::FLG_PSH != 0
113+
}
114+
115+
/// Return the ACK flag.
116+
#[inline(always)]
117+
pub fn ack(&self) -> bool {
118+
let data = self.buffer.as_ref();
119+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
120+
raw & field::FLG_ACK != 0
121+
}
122+
123+
/// Return the URG flag.
124+
#[inline(always)]
125+
pub fn urg(&self) -> bool {
126+
let data = self.buffer.as_ref();
127+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
128+
raw & field::FLG_URG != 0
129+
}
130+
131+
/// Return the ECE flag.
132+
#[inline(always)]
133+
pub fn ece(&self) -> bool {
134+
let data = self.buffer.as_ref();
135+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
136+
raw & field::FLG_ECE != 0
137+
}
138+
139+
/// Return the CWR flag.
140+
#[inline(always)]
141+
pub fn cwr(&self) -> bool {
142+
let data = self.buffer.as_ref();
143+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
144+
raw & field::FLG_CWR != 0
145+
}
146+
147+
/// Return the NS flag.
148+
#[inline(always)]
149+
pub fn ns(&self) -> bool {
150+
let data = self.buffer.as_ref();
151+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
152+
raw & field::FLG_NS != 0
153+
}
154+
155+
/// Return the header length, in octets.
156+
#[inline(always)]
157+
pub fn header_len(&self) -> u8 {
158+
let data = self.buffer.as_ref();
159+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
160+
((raw >> 12) * 4) as u8
161+
}
162+
163+
/// Return the window size field.
164+
#[inline(always)]
165+
pub fn window_len(&self) -> u16 {
166+
let data = self.buffer.as_ref();
167+
NetworkEndian::read_u16(&data[field::WIN_SIZE])
168+
}
169+
170+
/// Return the checksum field.
171+
#[inline(always)]
172+
pub fn checksum(&self) -> u16 {
173+
let data = self.buffer.as_ref();
174+
NetworkEndian::read_u16(&data[field::CHECKSUM])
175+
}
176+
177+
/// Return the urgent pointer field.
178+
#[inline(always)]
179+
pub fn urgent_at(&self) -> u16 {
180+
let data = self.buffer.as_ref();
181+
NetworkEndian::read_u16(&data[field::URGENT])
182+
}
183+
184+
/// Validate the packet checksum.
185+
///
186+
/// # Panics
187+
/// This function panics unless `src_addr` and `dst_addr` belong to the same family,
188+
/// and that family is IPv4 or IPv6.
189+
pub fn verify_checksum(&self, src_addr: &InternetAddress, dst_addr: &InternetAddress) -> bool {
190+
let data = self.buffer.as_ref();
191+
checksum::combine(&[
192+
checksum::pseudo_header(src_addr, dst_addr, InternetProtocolType::Tcp,
193+
data.len() as u32),
194+
checksum::data(data)
195+
]) == !0
196+
}
197+
}
198+
199+
impl<'a, T: AsRef<[u8]> + ?Sized> Packet<&'a T> {
200+
/// Return a pointer to the payload.
201+
#[inline(always)]
202+
pub fn payload(&self) -> &'a [u8] {
203+
let header_len = self.header_len() as usize;
204+
let data = self.buffer.as_ref();
205+
&data[header_len..]
206+
}
207+
}
208+
209+
impl<T: AsRef<[u8]> + AsMut<[u8]>> Packet<T> {
210+
/// Set the source port field.
211+
#[inline(always)]
212+
pub fn set_src_port(&mut self, value: u16) {
213+
let mut data = self.buffer.as_mut();
214+
NetworkEndian::write_u16(&mut data[field::SRC_PORT], value)
215+
}
216+
217+
/// Set the destination port field.
218+
#[inline(always)]
219+
pub fn set_dst_port(&mut self, value: u16) {
220+
let mut data = self.buffer.as_mut();
221+
NetworkEndian::write_u16(&mut data[field::DST_PORT], value)
222+
}
223+
224+
/// Set the sequence number field.
225+
#[inline(always)]
226+
pub fn set_seq_number(&mut self, value: u32) {
227+
let mut data = self.buffer.as_mut();
228+
NetworkEndian::write_u32(&mut data[field::SEQ_NUM], value)
229+
}
230+
231+
/// Set the acknowledgement number field.
232+
#[inline(always)]
233+
pub fn set_ack_number(&mut self, value: u32) {
234+
let mut data = self.buffer.as_mut();
235+
NetworkEndian::write_u32(&mut data[field::ACK_NUM], value)
236+
}
237+
238+
/// Set the FIN flag.
239+
#[inline(always)]
240+
pub fn set_fin(&mut self, value: bool) {
241+
let mut data = self.buffer.as_mut();
242+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
243+
let raw = if value { raw | field::FLG_FIN } else { raw & !field::FLG_FIN };
244+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
245+
}
246+
247+
/// Set the SYN flag.
248+
#[inline(always)]
249+
pub fn set_syn(&mut self, value: bool) {
250+
let mut data = self.buffer.as_mut();
251+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
252+
let raw = if value { raw | field::FLG_SYN } else { raw & !field::FLG_SYN };
253+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
254+
}
255+
256+
/// Set the RST flag.
257+
#[inline(always)]
258+
pub fn set_rst(&mut self, value: bool) {
259+
let mut data = self.buffer.as_mut();
260+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
261+
let raw = if value { raw | field::FLG_RST } else { raw & !field::FLG_RST };
262+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
263+
}
264+
265+
/// Set the PSH flag.
266+
#[inline(always)]
267+
pub fn set_psh(&mut self, value: bool) {
268+
let mut data = self.buffer.as_mut();
269+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
270+
let raw = if value { raw | field::FLG_PSH } else { raw & !field::FLG_PSH };
271+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
272+
}
273+
274+
/// Set the ACK flag.
275+
#[inline(always)]
276+
pub fn set_ack(&mut self, value: bool) {
277+
let mut data = self.buffer.as_mut();
278+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
279+
let raw = if value { raw | field::FLG_ACK } else { raw & !field::FLG_ACK };
280+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
281+
}
282+
283+
/// Set the URG flag.
284+
#[inline(always)]
285+
pub fn set_urg(&mut self, value: bool) {
286+
let mut data = self.buffer.as_mut();
287+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
288+
let raw = if value { raw | field::FLG_URG } else { raw & !field::FLG_URG };
289+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
290+
}
291+
292+
/// Set the ECE flag.
293+
#[inline(always)]
294+
pub fn set_ece(&mut self, value: bool) {
295+
let mut data = self.buffer.as_mut();
296+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
297+
let raw = if value { raw | field::FLG_ECE } else { raw & !field::FLG_ECE };
298+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
299+
}
300+
301+
/// Set the CWR flag.
302+
#[inline(always)]
303+
pub fn set_cwr(&mut self, value: bool) {
304+
let mut data = self.buffer.as_mut();
305+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
306+
let raw = if value { raw | field::FLG_CWR } else { raw & !field::FLG_CWR };
307+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
308+
}
309+
310+
/// Set the NS flag.
311+
#[inline(always)]
312+
pub fn set_ns(&mut self, value: bool) {
313+
let mut data = self.buffer.as_mut();
314+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
315+
let raw = if value { raw | field::FLG_NS } else { raw & !field::FLG_NS };
316+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
317+
}
318+
319+
/// Set the header length, in octets.
320+
#[inline(always)]
321+
pub fn set_header_len(&mut self, value: u8) {
322+
let mut data = self.buffer.as_mut();
323+
let raw = NetworkEndian::read_u16(&data[field::FLAGS]);
324+
let raw = (raw & !0xf000) | ((value as u16) / 4) << 12;
325+
NetworkEndian::write_u16(&mut data[field::FLAGS], raw)
326+
}
327+
328+
/// Return the window size field.
329+
#[inline(always)]
330+
pub fn set_window_len(&mut self, value: u16) {
331+
let mut data = self.buffer.as_mut();
332+
NetworkEndian::write_u16(&mut data[field::WIN_SIZE], value)
333+
}
334+
335+
/// Set the checksum field.
336+
#[inline(always)]
337+
pub fn set_checksum(&mut self, value: u16) {
338+
let mut data = self.buffer.as_mut();
339+
NetworkEndian::write_u16(&mut data[field::CHECKSUM], value)
340+
}
341+
342+
/// Set the urgent pointer field.
343+
#[inline(always)]
344+
pub fn set_urgent_at(&mut self, value: u16) {
345+
let mut data = self.buffer.as_mut();
346+
NetworkEndian::write_u16(&mut data[field::URGENT], value)
347+
}
348+
349+
/// Compute and fill in the header checksum.
350+
///
351+
/// # Panics
352+
/// This function panics unless `src_addr` and `dst_addr` belong to the same family,
353+
/// and that family is IPv4 or IPv6.
354+
pub fn fill_checksum(&mut self, src_addr: &InternetAddress, dst_addr: &InternetAddress) {
355+
self.set_checksum(0);
356+
let checksum = {
357+
let data = self.buffer.as_ref();
358+
!checksum::combine(&[
359+
checksum::pseudo_header(src_addr, dst_addr, InternetProtocolType::Tcp,
360+
data.len() as u32),
361+
checksum::data(data)
362+
])
363+
};
364+
self.set_checksum(checksum)
365+
}
366+
}
367+
368+
impl<'a, T: AsRef<[u8]> + AsMut<[u8]> + ?Sized> Packet<&'a mut T> {
369+
/// Return a mutable pointer to the payload data.
370+
#[inline(always)]
371+
pub fn payload_mut(&mut self) -> &mut [u8] {
372+
let header_len = self.header_len() as usize;
373+
let mut data = self.buffer.as_mut();
374+
&mut data[header_len..]
375+
}
376+
}
377+
378+
#[cfg(test)]
379+
mod test {
380+
use wire::Ipv4Address;
381+
use super::*;
382+
383+
const SRC_ADDR: Ipv4Address = Ipv4Address([192, 168, 1, 1]);
384+
const DST_ADDR: Ipv4Address = Ipv4Address([192, 168, 1, 2]);
385+
386+
static PACKET_BYTES: [u8; 24] =
387+
[0xbf, 0x00, 0x00, 0x50,
388+
0x01, 0x23, 0x45, 0x67,
389+
0x89, 0xab, 0xcd, 0xef,
390+
0x50, 0x35, 0x01, 0x23,
391+
0x20, 0xbe, 0x02, 0x01,
392+
0xaa, 0x00, 0x00, 0xff];
393+
394+
static PAYLOAD_BYTES: [u8; 4] =
395+
[0xaa, 0x00, 0x00, 0xff];
396+
397+
#[test]
398+
fn test_deconstruct() {
399+
let packet = Packet::new(&PACKET_BYTES[..]).unwrap();
400+
assert_eq!(packet.src_port(), 48896);
401+
assert_eq!(packet.dst_port(), 80);
402+
assert_eq!(packet.seq_number(), 0x01234567);
403+
assert_eq!(packet.ack_number(), 0x89abcdef);
404+
assert_eq!(packet.header_len(), 20);
405+
assert_eq!(packet.fin(), true);
406+
assert_eq!(packet.syn(), false);
407+
assert_eq!(packet.rst(), true);
408+
assert_eq!(packet.psh(), false);
409+
assert_eq!(packet.ack(), true);
410+
assert_eq!(packet.urg(), true);
411+
assert_eq!(packet.window_len(), 0x0123);
412+
assert_eq!(packet.urgent_at(), 0x0201);
413+
assert_eq!(packet.checksum(), 0x20be);
414+
assert_eq!(packet.payload(), &PAYLOAD_BYTES[..]);
415+
assert_eq!(packet.verify_checksum(&SRC_ADDR.into(), &DST_ADDR.into()), true);
416+
}
417+
418+
#[test]
419+
fn test_construct() {
420+
let mut bytes = vec![0; 24];
421+
let mut packet = Packet::new(&mut bytes).unwrap();
422+
packet.set_src_port(48896);
423+
packet.set_dst_port(80);
424+
packet.set_seq_number(0x01234567);
425+
packet.set_ack_number(0x89abcdef);
426+
packet.set_header_len(20);
427+
packet.set_fin(true);
428+
packet.set_syn(false);
429+
packet.set_rst(true);
430+
packet.set_psh(false);
431+
packet.set_ack(true);
432+
packet.set_urg(true);
433+
packet.set_window_len(0x0123);
434+
packet.set_urgent_at(0x0201);
435+
packet.set_checksum(0xEEEE);
436+
packet.payload_mut().copy_from_slice(&PAYLOAD_BYTES[..]);
437+
packet.fill_checksum(&SRC_ADDR.into(), &DST_ADDR.into());
438+
assert_eq!(&packet.into_inner()[..], &PACKET_BYTES[..]);
439+
}
440+
}

‎src/wire/udp.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ impl<T: AsRef<[u8]> + AsMut<[u8]>> Packet<T> {
150150
}
151151

152152
impl<'a, T: AsRef<[u8]> + AsMut<[u8]> + ?Sized> Packet<&'a mut T> {
153-
/// Return a mutable pointer to the type-specific data.
153+
/// Return a mutable pointer to the payload.
154154
#[inline(always)]
155155
pub fn payload_mut(&mut self) -> &mut [u8] {
156156
let length = self.len();

0 commit comments

Comments
 (0)
Please sign in to comment.