stackforge_core/layer/tcp/
builder.rs1use std::net::{Ipv4Addr, Ipv6Addr};
26
27use super::checksum::{tcp_checksum_ipv4, tcp_checksum_ipv6};
28use super::flags::TcpFlags;
29use super::header::{TCP_MIN_HEADER_LEN, TcpLayer, offsets};
30use super::options::{TcpAoValue, TcpOption, TcpOptions, TcpOptionsBuilder, TcpSackBlock};
31use crate::layer::field::FieldError;
32
33#[derive(Debug, Clone)]
35pub struct TcpBuilder {
36 src_port: u16,
38 dst_port: u16,
39 seq: u32,
40 ack: u32,
41 data_offset: Option<u8>,
42 reserved: u8,
43 flags: TcpFlags,
44 window: u16,
45 checksum: Option<u16>,
46 urgent_ptr: u16,
47
48 options: TcpOptions,
50
51 payload: Vec<u8>,
53
54 auto_checksum: bool,
56 auto_data_offset: bool,
57
58 src_ip: Option<IpAddr>,
60 dst_ip: Option<IpAddr>,
61
62 flags_explicitly_set: bool,
65}
66
67#[derive(Debug, Clone, Copy)]
69pub enum IpAddr {
70 V4(Ipv4Addr),
71 V6(Ipv6Addr),
72}
73
74impl From<Ipv4Addr> for IpAddr {
75 fn from(addr: Ipv4Addr) -> Self {
76 IpAddr::V4(addr)
77 }
78}
79
80impl From<Ipv6Addr> for IpAddr {
81 fn from(addr: Ipv6Addr) -> Self {
82 IpAddr::V6(addr)
83 }
84}
85
86impl Default for TcpBuilder {
87 fn default() -> Self {
88 Self {
89 src_port: 20,
90 dst_port: 80,
91 seq: 0,
92 ack: 0,
93 data_offset: None,
94 reserved: 0,
95 flags: TcpFlags::from_u16(0x02), window: 8192,
97 checksum: None,
98 urgent_ptr: 0,
99 options: TcpOptions::new(),
100 payload: Vec::new(),
101 auto_checksum: true,
102 auto_data_offset: true,
103 src_ip: None,
104 dst_ip: None,
105 flags_explicitly_set: false,
106 }
107 }
108}
109
110impl TcpBuilder {
111 #[must_use]
113 pub fn new() -> Self {
114 Self::default()
115 }
116
117 pub fn from_bytes(data: &[u8]) -> Result<Self, FieldError> {
119 let layer = TcpLayer::at_offset_dynamic(data, 0)?;
120
121 let mut builder = Self::new();
122 builder.src_port = layer.src_port(data)?;
123 builder.dst_port = layer.dst_port(data)?;
124 builder.seq = layer.seq(data)?;
125 builder.ack = layer.ack(data)?;
126 builder.data_offset = Some(layer.data_offset(data)?);
127 builder.reserved = layer.reserved(data)?;
128 builder.flags = layer.flags(data)?;
129 builder.window = layer.window(data)?;
130 builder.checksum = Some(layer.checksum(data)?);
131 builder.urgent_ptr = layer.urgent_ptr(data)?;
132
133 if layer.options_len(data) > 0 {
135 builder.options = layer.options(data)?;
136 }
137
138 let header_len = layer.calculate_header_len(data);
140 if data.len() > header_len {
141 builder.payload = data[header_len..].to_vec();
142 }
143
144 builder.auto_checksum = false;
146 builder.auto_data_offset = false;
147
148 Ok(builder)
149 }
150
151 #[must_use]
155 pub fn src_port(mut self, port: u16) -> Self {
156 self.src_port = port;
157 self
158 }
159
160 #[must_use]
162 pub fn sport(self, port: u16) -> Self {
163 self.src_port(port)
164 }
165
166 #[must_use]
168 pub fn dst_port(mut self, port: u16) -> Self {
169 self.dst_port = port;
170 self
171 }
172
173 #[must_use]
175 pub fn dport(self, port: u16) -> Self {
176 self.dst_port(port)
177 }
178
179 #[must_use]
181 pub fn seq(mut self, seq: u32) -> Self {
182 self.seq = seq;
183 self
184 }
185
186 #[must_use]
188 pub fn ack_num(mut self, ack: u32) -> Self {
189 self.ack = ack;
190 self
191 }
192
193 #[must_use]
195 pub fn data_offset(mut self, offset: u8) -> Self {
196 self.data_offset = Some(offset);
197 self.auto_data_offset = false;
198 self
199 }
200
201 #[must_use]
203 pub fn dataofs(self, offset: u8) -> Self {
204 self.data_offset(offset)
205 }
206
207 #[must_use]
209 pub fn reserved(mut self, reserved: u8) -> Self {
210 self.reserved = reserved & 0x07;
211 self
212 }
213
214 #[must_use]
216 pub fn flags(mut self, flags: TcpFlags) -> Self {
217 self.flags = flags;
218 self.flags_explicitly_set = true;
219 self
220 }
221
222 #[must_use]
224 pub fn flags_str(mut self, s: &str) -> Self {
225 self.flags = TcpFlags::from_str(s);
226 self.flags_explicitly_set = true;
227 self
228 }
229
230 fn clear_defaults_if_needed(&mut self) {
235 if !self.flags_explicitly_set {
236 self.flags = TcpFlags::from_u16(0);
237 self.flags_explicitly_set = true;
238 }
239 }
240
241 #[must_use]
243 pub fn syn(mut self) -> Self {
244 self.clear_defaults_if_needed();
245 self.flags.syn = true;
246 self
247 }
248
249 #[must_use]
251 pub fn ack(mut self) -> Self {
252 self.clear_defaults_if_needed();
253 self.flags.ack = true;
254 self
255 }
256
257 #[must_use]
259 pub fn fin(mut self) -> Self {
260 self.clear_defaults_if_needed();
261 self.flags.fin = true;
262 self
263 }
264
265 #[must_use]
267 pub fn rst(mut self) -> Self {
268 self.clear_defaults_if_needed();
269 self.flags.rst = true;
270 self
271 }
272
273 #[must_use]
275 pub fn psh(mut self) -> Self {
276 self.clear_defaults_if_needed();
277 self.flags.psh = true;
278 self
279 }
280
281 #[must_use]
283 pub fn urg(mut self) -> Self {
284 self.clear_defaults_if_needed();
285 self.flags.urg = true;
286 self
287 }
288
289 #[must_use]
291 pub fn ece(mut self) -> Self {
292 self.clear_defaults_if_needed();
293 self.flags.ece = true;
294 self
295 }
296
297 #[must_use]
299 pub fn cwr(mut self) -> Self {
300 self.clear_defaults_if_needed();
301 self.flags.cwr = true;
302 self
303 }
304
305 #[must_use]
307 pub fn ns(mut self) -> Self {
308 self.clear_defaults_if_needed();
309 self.flags.ns = true;
310 self
311 }
312
313 #[must_use]
315 pub fn syn_ack(mut self) -> Self {
316 self.clear_defaults_if_needed();
317 self.flags.syn = true;
318 self.flags.ack = true;
319 self
320 }
321
322 #[must_use]
324 pub fn fin_ack(mut self) -> Self {
325 self.clear_defaults_if_needed();
326 self.flags.fin = true;
327 self.flags.ack = true;
328 self
329 }
330
331 #[must_use]
333 pub fn psh_ack(mut self) -> Self {
334 self.clear_defaults_if_needed();
335 self.flags.psh = true;
336 self.flags.ack = true;
337 self
338 }
339
340 #[must_use]
342 pub fn window(mut self, window: u16) -> Self {
343 self.window = window;
344 self
345 }
346
347 #[must_use]
349 pub fn checksum(mut self, checksum: u16) -> Self {
350 self.checksum = Some(checksum);
351 self.auto_checksum = false;
352 self
353 }
354
355 #[must_use]
357 pub fn chksum(self, checksum: u16) -> Self {
358 self.checksum(checksum)
359 }
360
361 #[must_use]
363 pub fn urgent_ptr(mut self, urgptr: u16) -> Self {
364 self.urgent_ptr = urgptr;
365 self
366 }
367
368 #[must_use]
370 pub fn urgptr(self, urgptr: u16) -> Self {
371 self.urgent_ptr(urgptr)
372 }
373
374 pub fn src_ip<T: Into<IpAddr>>(mut self, ip: T) -> Self {
378 self.src_ip = Some(ip.into());
379 self
380 }
381
382 pub fn dst_ip<T: Into<IpAddr>>(mut self, ip: T) -> Self {
384 self.dst_ip = Some(ip.into());
385 self
386 }
387
388 #[must_use]
390 pub fn ipv4_addrs(mut self, src: Ipv4Addr, dst: Ipv4Addr) -> Self {
391 self.src_ip = Some(IpAddr::V4(src));
392 self.dst_ip = Some(IpAddr::V4(dst));
393 self
394 }
395
396 #[must_use]
398 pub fn ipv6_addrs(mut self, src: Ipv6Addr, dst: Ipv6Addr) -> Self {
399 self.src_ip = Some(IpAddr::V6(src));
400 self.dst_ip = Some(IpAddr::V6(dst));
401 self
402 }
403
404 #[must_use]
408 pub fn options(mut self, options: TcpOptions) -> Self {
409 self.options = options;
410 self
411 }
412
413 #[must_use]
415 pub fn option(mut self, option: TcpOption) -> Self {
416 self.options.push(option);
417 self
418 }
419
420 pub fn with_options<F>(mut self, f: F) -> Self
422 where
423 F: FnOnce(TcpOptionsBuilder) -> TcpOptionsBuilder,
424 {
425 self.options = f(TcpOptionsBuilder::new()).build();
426 self
427 }
428
429 #[must_use]
431 pub fn mss(mut self, mss: u16) -> Self {
432 self.options.push(TcpOption::Mss(mss));
433 self
434 }
435
436 #[must_use]
438 pub fn wscale(mut self, scale: u8) -> Self {
439 self.options.push(TcpOption::WScale(scale));
440 self
441 }
442
443 #[must_use]
445 pub fn sack_ok(mut self) -> Self {
446 self.options.push(TcpOption::SackOk);
447 self
448 }
449
450 #[must_use]
452 pub fn sack(mut self, blocks: Vec<TcpSackBlock>) -> Self {
453 self.options.push(TcpOption::Sack(blocks));
454 self
455 }
456
457 #[must_use]
459 pub fn timestamp(mut self, ts_val: u32, ts_ecr: u32) -> Self {
460 self.options.push(TcpOption::timestamp(ts_val, ts_ecr));
461 self
462 }
463
464 #[must_use]
466 pub fn nop(mut self) -> Self {
467 self.options.push(TcpOption::Nop);
468 self
469 }
470
471 #[must_use]
473 pub fn eol(mut self) -> Self {
474 self.options.push(TcpOption::Eol);
475 self
476 }
477
478 #[must_use]
480 pub fn tfo(mut self, cookie: Option<Vec<u8>>) -> Self {
481 self.options.push(TcpOption::Tfo { cookie });
482 self
483 }
484
485 #[must_use]
487 pub fn md5(mut self, signature: [u8; 16]) -> Self {
488 self.options.push(TcpOption::Md5(signature));
489 self
490 }
491
492 #[must_use]
494 pub fn ao(mut self, key_id: u8, rnext_key_id: u8, mac: Vec<u8>) -> Self {
495 self.options
496 .push(TcpOption::Ao(TcpAoValue::new(key_id, rnext_key_id, mac)));
497 self
498 }
499
500 pub fn payload(mut self, payload: impl Into<Vec<u8>>) -> Self {
504 self.payload = payload.into();
505 self
506 }
507
508 #[must_use]
510 pub fn append_payload(mut self, data: &[u8]) -> Self {
511 self.payload.extend_from_slice(data);
512 self
513 }
514
515 #[must_use]
519 pub fn auto_checksum(mut self, enabled: bool) -> Self {
520 self.auto_checksum = enabled;
521 self
522 }
523
524 #[must_use]
526 pub fn auto_data_offset(mut self, enabled: bool) -> Self {
527 self.auto_data_offset = enabled;
528 self
529 }
530
531 #[must_use]
535 pub fn header_size(&self) -> usize {
536 if let Some(doff) = self.data_offset {
537 (doff as usize) * 4
538 } else {
539 let opts_len = self.options.padded_len();
540 TCP_MIN_HEADER_LEN + opts_len
541 }
542 }
543
544 #[must_use]
546 pub fn packet_size(&self) -> usize {
547 self.header_size() + self.payload.len()
548 }
549
550 #[must_use]
552 pub fn build(&self) -> Vec<u8> {
553 let total_size = self.packet_size();
554 let mut buf = vec![0u8; total_size];
555 self.build_into(&mut buf)
556 .expect("buffer is correctly sized");
557 buf
558 }
559
560 pub fn build_into(&self, buf: &mut [u8]) -> Result<usize, FieldError> {
562 let header_size = self.header_size();
563 let total_size = self.packet_size();
564
565 if buf.len() < total_size {
566 return Err(FieldError::BufferTooShort {
567 offset: 0,
568 need: total_size,
569 have: buf.len(),
570 });
571 }
572
573 let data_offset = if self.auto_data_offset {
575 (header_size / 4) as u8
576 } else {
577 self.data_offset.unwrap_or(5)
578 };
579
580 buf[offsets::SRC_PORT] = (self.src_port >> 8) as u8;
582 buf[offsets::SRC_PORT + 1] = (self.src_port & 0xFF) as u8;
583
584 buf[offsets::DST_PORT] = (self.dst_port >> 8) as u8;
586 buf[offsets::DST_PORT + 1] = (self.dst_port & 0xFF) as u8;
587
588 buf[offsets::SEQ..offsets::SEQ + 4].copy_from_slice(&self.seq.to_be_bytes());
590
591 buf[offsets::ACK..offsets::ACK + 4].copy_from_slice(&self.ack.to_be_bytes());
593
594 buf[offsets::DATA_OFFSET] =
596 ((data_offset & 0x0F) << 4) | ((self.reserved & 0x07) << 1) | self.flags.ns_bit();
597
598 buf[offsets::FLAGS] = self.flags.to_byte();
600
601 buf[offsets::WINDOW] = (self.window >> 8) as u8;
603 buf[offsets::WINDOW + 1] = (self.window & 0xFF) as u8;
604
605 buf[offsets::CHECKSUM] = 0;
607 buf[offsets::CHECKSUM + 1] = 0;
608
609 buf[offsets::URG_PTR] = (self.urgent_ptr >> 8) as u8;
611 buf[offsets::URG_PTR + 1] = (self.urgent_ptr & 0xFF) as u8;
612
613 if !self.options.is_empty() {
615 let opts_bytes = self.options.to_bytes();
616 let opts_end = offsets::OPTIONS + opts_bytes.len();
617 if opts_end <= header_size {
618 buf[offsets::OPTIONS..opts_end].copy_from_slice(&opts_bytes);
619 }
620 }
621
622 if !self.payload.is_empty() {
624 buf[header_size..header_size + self.payload.len()].copy_from_slice(&self.payload);
625 }
626
627 let checksum = if self.auto_checksum {
629 match (self.src_ip, self.dst_ip) {
630 (Some(IpAddr::V4(src)), Some(IpAddr::V4(dst))) => {
631 tcp_checksum_ipv4(src, dst, &buf[..total_size])
632 },
633 (Some(IpAddr::V6(src)), Some(IpAddr::V6(dst))) => {
634 tcp_checksum_ipv6(src, dst, &buf[..total_size])
635 },
636 _ => 0, }
638 } else {
639 self.checksum.unwrap_or(0)
640 };
641 buf[offsets::CHECKSUM] = (checksum >> 8) as u8;
642 buf[offsets::CHECKSUM + 1] = (checksum & 0xFF) as u8;
643
644 Ok(total_size)
645 }
646
647 #[must_use]
649 pub fn build_header(&self) -> Vec<u8> {
650 let header_size = self.header_size();
651 let mut buf = vec![0u8; header_size];
652
653 let builder = Self {
655 payload: Vec::new(),
656 ..self.clone()
657 };
658 builder
659 .build_into(&mut buf)
660 .expect("buffer is correctly sized");
661
662 buf
663 }
664}
665
666impl TcpBuilder {
669 #[must_use]
671 pub fn syn_packet() -> Self {
672 Self::new().syn().ack_num(0)
673 }
674
675 #[must_use]
677 pub fn syn_ack_packet() -> Self {
678 Self::new().syn_ack()
679 }
680
681 #[must_use]
683 pub fn ack_packet() -> Self {
684 Self::new().flags(TcpFlags::A)
685 }
686
687 #[must_use]
689 pub fn fin_ack_packet() -> Self {
690 Self::new().fin_ack()
691 }
692
693 #[must_use]
695 pub fn rst_packet() -> Self {
696 Self::new().flags(TcpFlags::R)
697 }
698
699 #[must_use]
701 pub fn rst_ack_packet() -> Self {
702 Self::new().flags(TcpFlags::RA)
703 }
704
705 #[must_use]
707 pub fn data_packet() -> Self {
708 Self::new().psh_ack()
709 }
710}
711
712#[cfg(feature = "rand")]
715impl TcpBuilder {
716 #[must_use]
718 pub fn random_seq(mut self) -> Self {
719 use rand::Rng;
720 self.seq = rand::rng().random();
721 self
722 }
723
724 #[must_use]
726 pub fn random_sport(mut self) -> Self {
727 use rand::Rng;
728 self.src_port = rand::rng().random_range(49152..=65535);
729 self
730 }
731}
732
733#[cfg(test)]
734mod tests {
735 use super::*;
736
737 #[test]
738 fn test_basic_build() {
739 let pkt = TcpBuilder::new()
740 .src_port(12345)
741 .dst_port(80)
742 .seq(1000)
743 .syn()
744 .window(65535)
745 .build();
746
747 assert_eq!(pkt.len(), 20); let layer = TcpLayer::at_offset(0);
750 assert_eq!(layer.src_port(&pkt).unwrap(), 12345);
751 assert_eq!(layer.dst_port(&pkt).unwrap(), 80);
752 assert_eq!(layer.seq(&pkt).unwrap(), 1000);
753 assert_eq!(layer.window(&pkt).unwrap(), 65535);
754
755 let flags = layer.flags(&pkt).unwrap();
756 assert!(flags.syn);
757 assert!(!flags.ack);
758 }
759
760 #[test]
761 fn test_with_options() {
762 let pkt = TcpBuilder::new()
763 .src_port(12345)
764 .dst_port(80)
765 .syn()
766 .mss(1460)
767 .wscale(7)
768 .sack_ok()
769 .nop()
770 .build();
771
772 let layer = TcpLayer::at_offset(0);
773 let opts = layer.options(&pkt).unwrap();
774
775 assert_eq!(opts.mss(), Some(1460));
776 assert_eq!(opts.wscale(), Some(7));
777 assert!(opts.sack_permitted());
778 }
779
780 #[test]
781 fn test_with_payload() {
782 let payload = b"Hello, TCP!";
783 let pkt = TcpBuilder::new()
784 .src_port(12345)
785 .dst_port(80)
786 .psh_ack()
787 .payload(payload.to_vec())
788 .build();
789
790 assert_eq!(pkt.len(), 20 + payload.len());
791
792 let layer = TcpLayer::at_offset(0);
793 let pkt_payload = layer.payload(&pkt);
794 assert_eq!(pkt_payload, payload);
795 }
796
797 #[test]
798 fn test_with_checksum() {
799 let src_ip = Ipv4Addr::new(192, 168, 1, 1);
800 let dst_ip = Ipv4Addr::new(192, 168, 1, 2);
801
802 let pkt = TcpBuilder::new()
803 .src_port(12345)
804 .dst_port(80)
805 .syn()
806 .ipv4_addrs(src_ip, dst_ip)
807 .build();
808
809 let layer = TcpLayer::at_offset(0);
810 let checksum = layer.checksum(&pkt).unwrap();
811 assert_ne!(checksum, 0);
812 }
813
814 #[test]
815 fn test_flags() {
816 let pkt = TcpBuilder::new().syn_ack().build();
817
818 let layer = TcpLayer::at_offset(0);
819 let flags = layer.flags(&pkt).unwrap();
820 assert!(flags.syn);
821 assert!(flags.ack);
822
823 let pkt = TcpBuilder::new().fin_ack().build();
824 let flags = layer.flags(&pkt).unwrap();
825 assert!(flags.fin);
826 assert!(flags.ack);
827 }
828
829 #[test]
830 fn test_from_bytes() {
831 let original = TcpBuilder::new()
832 .src_port(12345)
833 .dst_port(443)
834 .seq(0xDEADBEEF)
835 .ack_num(0xCAFEBABE)
836 .syn_ack()
837 .window(32768)
838 .mss(1460)
839 .build();
840
841 let rebuilt = TcpBuilder::from_bytes(&original)
842 .unwrap()
843 .auto_checksum(false)
844 .build();
845
846 assert_eq!(original.len(), rebuilt.len());
847
848 let layer = TcpLayer::at_offset(0);
849 assert_eq!(
850 layer.src_port(&original).unwrap(),
851 layer.src_port(&rebuilt).unwrap()
852 );
853 assert_eq!(
854 layer.dst_port(&original).unwrap(),
855 layer.dst_port(&rebuilt).unwrap()
856 );
857 assert_eq!(layer.seq(&original).unwrap(), layer.seq(&rebuilt).unwrap());
858 assert_eq!(layer.ack(&original).unwrap(), layer.ack(&rebuilt).unwrap());
859 }
860
861 #[test]
862 fn test_convenience_constructors() {
863 let syn = TcpBuilder::syn_packet().build();
864 let layer = TcpLayer::at_offset(0);
865 let flags = layer.flags(&syn).unwrap();
866 assert!(flags.syn);
867 assert!(!flags.ack);
868
869 let syn_ack = TcpBuilder::syn_ack_packet().build();
870 let flags = layer.flags(&syn_ack).unwrap();
871 assert!(flags.syn);
872 assert!(flags.ack);
873
874 let rst = TcpBuilder::rst_packet().build();
875 let flags = layer.flags(&rst).unwrap();
876 assert!(flags.rst);
877 }
878
879 #[test]
880 fn test_timestamp_option() {
881 let pkt = TcpBuilder::new()
882 .syn()
883 .mss(1460)
884 .timestamp(12345, 0)
885 .build();
886
887 let layer = TcpLayer::at_offset(0);
888 let opts = layer.options(&pkt).unwrap();
889
890 let ts = opts.timestamp().unwrap();
891 assert_eq!(ts.ts_val, 12345);
892 assert_eq!(ts.ts_ecr, 0);
893 }
894
895 #[test]
896 fn test_tfo_option() {
897 let cookie = vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];
898 let pkt = TcpBuilder::new().syn().tfo(Some(cookie.clone())).build();
899
900 let layer = TcpLayer::at_offset(0);
901 let opts = layer.options(&pkt).unwrap();
902
903 assert_eq!(opts.tfo_cookie(), Some(cookie.as_slice()));
904 }
905
906 #[test]
907 fn test_flags_str() {
908 let pkt = TcpBuilder::new().flags_str("SA").build();
909
910 let layer = TcpLayer::at_offset(0);
911 let flags = layer.flags(&pkt).unwrap();
912 assert!(flags.syn);
913 assert!(flags.ack);
914
915 let pkt = TcpBuilder::new().flags_str("FA").build();
916 let flags = layer.flags(&pkt).unwrap();
917 assert!(flags.fin);
918 assert!(flags.ack);
919 }
920}