1use crate::protocol::v2::{
4 Addresses, LENGTH, MINIMUM_LENGTH, MINIMUM_TLV_LENGTH, PROTOCOL_PREFIX, Protocol, Type,
5 TypeLengthValue, TypeLengthValues,
6};
7use std::io::{self, Write};
8
9#[derive(Debug, Default)]
25pub struct Writer {
26 bytes: Vec<u8>,
27}
28
29#[derive(Debug)]
54pub struct Builder {
55 header: Option<Vec<u8>>,
56 version_command: u8,
57 address_family_protocol: u8,
58 addresses: Addresses,
59 length: Option<u16>,
60 additional_capacity: usize,
61 with_crc32c: bool,
65}
66
67impl Writer {
68 #[must_use]
71 pub fn finish(self) -> Vec<u8> {
72 self.bytes
73 }
74}
75
76impl From<Vec<u8>> for Writer {
77 fn from(bytes: Vec<u8>) -> Self {
78 Self { bytes }
79 }
80}
81
82impl Write for Writer {
83 fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
84 if self.bytes.len() > (u16::MAX as usize) + MINIMUM_LENGTH {
85 Err(io::ErrorKind::WriteZero.into())
86 } else {
87 self.bytes.extend_from_slice(buffer);
88 Ok(buffer.len())
89 }
90 }
91
92 fn flush(&mut self) -> io::Result<()> {
93 Ok(())
94 }
95}
96
97pub trait WriteToHeader {
99 fn write_to(&self, writer: &mut Writer) -> io::Result<usize>;
103
104 fn to_bytes(&self) -> io::Result<Vec<u8>> {
106 let mut writer = Writer::default();
107
108 self.write_to(&mut writer)?;
109
110 Ok(writer.finish())
111 }
112}
113
114impl WriteToHeader for Addresses {
115 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
116 match self {
117 Self::Unspecified => (),
118 Self::IPv4(a) => {
119 writer.write_all(a.source_address.octets().as_slice())?;
120 writer.write_all(a.destination_address.octets().as_slice())?;
121 writer.write_all(a.source_port.to_be_bytes().as_slice())?;
122 writer.write_all(a.destination_port.to_be_bytes().as_slice())?;
123 }
124 Self::IPv6(a) => {
125 writer.write_all(a.source_address.octets().as_slice())?;
126 writer.write_all(a.destination_address.octets().as_slice())?;
127 writer.write_all(a.source_port.to_be_bytes().as_slice())?;
128 writer.write_all(a.destination_port.to_be_bytes().as_slice())?;
129 }
130 Self::Unix(a) => {
131 writer.write_all(a.source.as_slice())?;
132 writer.write_all(a.destination.as_slice())?;
133 }
134 };
135
136 Ok(self.len())
137 }
138}
139
140impl WriteToHeader for TypeLengthValue<'_> {
141 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
142 if self.value.len() > u16::MAX as usize {
143 return Err(io::ErrorKind::WriteZero.into());
144 }
145
146 writer.write_all([u8::from(self.kind)].as_slice())?;
147 writer.write_all((self.value.len() as u16).to_be_bytes().as_slice())?;
148 writer.write_all(self.value.as_ref())?;
149
150 Ok(MINIMUM_TLV_LENGTH + self.value.len())
151 }
152}
153
154impl<T: Copy + Into<u8>> WriteToHeader for (T, &[u8]) {
155 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
156 let kind = self.0.into();
157 let value = self.1;
158
159 if value.len() > u16::MAX as usize {
160 return Err(io::ErrorKind::WriteZero.into());
161 }
162
163 writer.write_all([kind].as_slice())?;
164 writer.write_all((value.len() as u16).to_be_bytes().as_slice())?;
165 writer.write_all(value)?;
166
167 Ok(MINIMUM_TLV_LENGTH + value.len())
168 }
169}
170
171impl WriteToHeader for TypeLengthValues<'_> {
172 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
173 let bytes = self.as_bytes();
174
175 writer.write_all(bytes)?;
176
177 Ok(bytes.len())
178 }
179}
180
181impl WriteToHeader for [u8] {
182 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
183 let slice = self;
184
185 if slice.len() > u16::MAX as usize {
186 return Err(io::ErrorKind::WriteZero.into());
187 }
188
189 writer.write_all(slice)?;
190
191 Ok(slice.len())
192 }
193}
194
195impl<T: ?Sized + WriteToHeader> WriteToHeader for &T {
196 #[inline(always)]
197 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
198 (*self).write_to(writer)
199 }
200}
201
202impl WriteToHeader for Type {
203 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
204 writer.write([(*self).into()].as_slice())
205 }
206}
207
208macro_rules! impl_write_to_header {
209 ($t:ident) => {
210 impl WriteToHeader for $t {
211 fn write_to(&self, writer: &mut Writer) -> io::Result<usize> {
212 let bytes = self.to_be_bytes();
213
214 writer.write_all(bytes.as_slice())?;
215
216 Ok(bytes.len())
217 }
218 }
219 };
220}
221
222impl_write_to_header!(u8);
223impl_write_to_header!(u16);
224impl_write_to_header!(u32);
225impl_write_to_header!(u64);
226impl_write_to_header!(u128);
227impl_write_to_header!(usize);
228
229impl_write_to_header!(i8);
230impl_write_to_header!(i16);
231impl_write_to_header!(i32);
232impl_write_to_header!(i64);
233impl_write_to_header!(i128);
234impl_write_to_header!(isize);
235
236impl Builder {
237 #[must_use]
241 pub const fn new(version_command: u8, address_family_protocol: u8) -> Self {
242 Self {
243 header: None,
244 version_command,
245 address_family_protocol,
246 addresses: Addresses::Unspecified,
247 length: None,
248 additional_capacity: 0,
249 with_crc32c: false,
250 }
251 }
252
253 pub fn with_addresses<T: Into<Addresses>>(
257 version_command: u8,
258 protocol: Protocol,
259 addresses: T,
260 ) -> Self {
261 let addresses = addresses.into();
262
263 Self {
264 header: None,
265 version_command,
266 address_family_protocol: addresses.address_family() | protocol,
267 addresses,
268 length: None,
269 additional_capacity: 0,
270 with_crc32c: false,
271 }
272 }
273
274 rama_utils::macros::generate_set_and_with! {
275 pub fn reserve_capacity(mut self, capacity: usize) -> Self {
279 match self.header {
280 None => self.additional_capacity += capacity,
281 Some(ref mut header) => header.reserve(capacity),
282 }
283 self
284 }
285 }
286
287 rama_utils::macros::generate_set_and_with! {
288 pub fn length(mut self, length: impl Into<Option<u16>>) -> Self {
291 self.length = length.into();
292 self
293 }
294 }
295
296 pub fn write_payloads<T, I, II>(mut self, payloads: II) -> io::Result<Self>
299 where
300 T: WriteToHeader,
301 I: Iterator<Item = T>,
302 II: IntoIterator<IntoIter = I, Item = T>,
303 {
304 self.write_header()?;
305
306 let mut writer = Writer::from(self.header.take().unwrap_or_default());
307
308 for item in payloads {
309 item.write_to(&mut writer)?;
310 }
311
312 self.header = Some(writer.finish());
313
314 Ok(self)
315 }
316
317 pub fn write_payload<T: WriteToHeader>(mut self, payload: T) -> io::Result<Self> {
320 self.write_header()?;
321 self.write_internal(payload)?;
322
323 Ok(self)
324 }
325
326 pub fn write_tlv(self, kind: impl Into<Type>, value: &[u8]) -> io::Result<Self> {
331 self.write_payload(TypeLengthValue::new(kind, value))
332 }
333
334 rama_utils::macros::generate_set_and_with! {
335 pub fn crc32c(mut self, enabled: bool) -> Self {
345 self.with_crc32c = enabled;
346 self
347 }
348 }
349
350 fn write_internal<T: WriteToHeader>(&mut self, payload: T) -> io::Result<()> {
352 let mut writer = Writer::from(self.header.take().unwrap_or_default());
353
354 payload.write_to(&mut writer)?;
355
356 self.header = Some(writer.finish());
357
358 Ok(())
359 }
360
361 fn write_header(&mut self) -> io::Result<()> {
364 if self.header.is_some() {
365 return Ok(());
366 }
367
368 let mut header =
369 Vec::with_capacity(MINIMUM_LENGTH + self.addresses.len() + self.additional_capacity);
370
371 let length = self.length.unwrap_or_default();
372
373 header.extend_from_slice(PROTOCOL_PREFIX);
374 header.push(self.version_command);
375 header.push(self.address_family_protocol);
376 header.extend_from_slice(length.to_be_bytes().as_slice());
377
378 let mut writer = Writer::from(header);
379
380 self.addresses.write_to(&mut writer)?;
381 self.header = Some(writer.finish());
382
383 Ok(())
384 }
385
386 pub fn build(mut self) -> io::Result<Vec<u8>> {
395 self.write_header()?;
396
397 let crc_value_offset = if self.with_crc32c {
403 self = self.write_tlv(Type::CRC32C, &[0u8; 4])?;
404 let header_len = self.header.as_ref().map(Vec::len).unwrap_or_default();
405 Some(header_len - 4)
407 } else {
408 None
409 };
410
411 let mut header = self.header.take().unwrap_or_default();
412
413 if self.length.is_none() {
417 let Ok(payload_length) = u16::try_from(header[MINIMUM_LENGTH..].len()) else {
418 return Err(io::ErrorKind::WriteZero.into());
419 };
420 let length = payload_length.to_be_bytes();
421 header[LENGTH..LENGTH + length.len()].copy_from_slice(length.as_slice());
422 }
423
424 if let Some(value_offset) = crc_value_offset {
425 if value_offset + 4 > header.len() {
429 return Err(io::Error::new(
430 io::ErrorKind::InvalidData,
431 "CRC32C value offset out of bounds",
432 ));
433 }
434 let mut hasher = super::Crc32cHasher::new();
435 hasher.update(&header);
436 let crc = hasher.finalize();
437 header[value_offset..value_offset + 4].copy_from_slice(&crc.to_be_bytes());
438 }
439
440 Ok(header)
441 }
442}
443
444#[cfg(test)]
445mod tests {
446 use super::*;
447 use crate::protocol::v2::{AddressFamily, Command, IPv4, IPv6, Protocol, Type, Unix, Version};
448
449 #[test]
450 fn build_length_too_small() {
451 let mut expected = Vec::from(PROTOCOL_PREFIX);
452 expected.extend([0x21, 0x12, 0, 1, 0, 0, 0, 1]);
453
454 let actual = Builder::new(
455 Version::Two | Command::Proxy,
456 AddressFamily::IPv4 | Protocol::Datagram,
457 )
458 .with_length(1)
459 .write_payload(1u32)
460 .unwrap()
461 .build()
462 .unwrap();
463
464 assert_eq!(actual, expected);
465 }
466
467 #[test]
468 fn build_payload_too_long() {
469 let error = Builder::new(
470 Version::Two | Command::Proxy,
471 AddressFamily::IPv4 | Protocol::Datagram,
472 )
473 .write_payload(vec![0u8; (u16::MAX as usize) + 1].as_slice())
474 .unwrap_err();
475
476 assert_eq!(error.kind(), io::ErrorKind::WriteZero);
477 }
478
479 #[test]
480 fn build_no_payload() {
481 let mut expected = Vec::from(PROTOCOL_PREFIX);
482 expected.extend([0x21, 0x01, 0, 0]);
483
484 let header = Builder::new(
485 Version::Two | Command::Proxy,
486 AddressFamily::Unspecified | Protocol::Stream,
487 )
488 .build()
489 .unwrap();
490
491 assert_eq!(header, expected);
492 }
493
494 #[test]
495 fn build_arbitrary_payload() {
496 let mut expected = Vec::from(PROTOCOL_PREFIX);
497 expected.extend([0x21, 0x01, 0, 1, 42]);
498
499 let header = Builder::new(
500 Version::Two | Command::Proxy,
501 AddressFamily::Unspecified | Protocol::Stream,
502 )
503 .write_payload(42u8)
504 .unwrap()
505 .build()
506 .unwrap();
507
508 assert_eq!(header, expected);
509 }
510
511 #[test]
512 fn build_ipv4() {
513 let mut expected = Vec::from(PROTOCOL_PREFIX);
514 expected.extend([
515 0x21, 0x12, 0, 12, 127, 0, 0, 1, 192, 168, 1, 1, 0, 80, 1, 187,
516 ]);
517
518 let addresses: Addresses = IPv4::new([127, 0, 0, 1], [192, 168, 1, 1], 80, 443).into();
519 let header = Builder::new(
520 Version::Two | Command::Proxy,
521 AddressFamily::IPv4 | Protocol::Datagram,
522 )
523 .with_length(addresses.len() as u16)
524 .write_payload(addresses)
525 .unwrap()
526 .build()
527 .unwrap();
528
529 assert_eq!(header, expected);
530 }
531
532 #[test]
533 fn build_ipv6() {
534 let source_address = [
535 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
536 0xFF, 0xF2,
537 ];
538 let destination_address = [
539 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
540 0xFF, 0xF1,
541 ];
542 let mut expected = Vec::from(PROTOCOL_PREFIX);
543 expected.extend([0x20, 0x20, 0, 36]);
544 expected.extend(source_address);
545 expected.extend(destination_address);
546 expected.extend([0, 80, 1, 187]);
547
548 let header = Builder::with_addresses(
549 Version::Two | Command::Local,
550 Protocol::Unspecified,
551 IPv6::new(source_address, destination_address, 80, 443),
552 )
553 .build()
554 .unwrap();
555
556 assert_eq!(header, expected);
557 }
558
559 #[test]
560 fn build_unix() {
561 let source_address = [0xFFu8; 108];
562 let destination_address = [0xAAu8; 108];
563
564 let addresses: Addresses = Unix::new(source_address, destination_address).into();
565 let mut expected = Vec::from(PROTOCOL_PREFIX);
566 expected.extend([0x20, 0x31, 0, 216]);
567 expected.extend(source_address);
568 expected.extend(destination_address);
569
570 let header = Builder::new(
571 Version::Two | Command::Local,
572 AddressFamily::Unix | Protocol::Stream,
573 )
574 .with_reserve_capacity(addresses.len())
575 .write_payload(addresses)
576 .unwrap()
577 .build()
578 .unwrap();
579
580 assert_eq!(header, expected);
581 }
582
583 #[test]
584 fn build_ipv4_with_tlv() {
585 let mut expected = Vec::from(PROTOCOL_PREFIX);
586 expected.extend([
587 0x21, 0x12, 0, 17, 127, 0, 0, 1, 192, 168, 1, 1, 0, 80, 1, 187, 4, 0, 2, 0, 42,
588 ]);
589
590 let addresses: Addresses = IPv4::new([127, 0, 0, 1], [192, 168, 1, 1], 80, 443).into();
591 let header =
592 Builder::with_addresses(Version::Two | Command::Proxy, Protocol::Datagram, addresses)
593 .with_reserve_capacity(5)
594 .write_tlv(Type::NoOp, [0, 42].as_slice())
595 .unwrap()
596 .build()
597 .unwrap();
598
599 assert_eq!(header, expected);
600 }
601
602 #[test]
603 fn build_ipv4_with_nested_tlv() {
604 let mut expected = Vec::from(PROTOCOL_PREFIX);
605 expected.extend([
606 0x21, 0x12, 0, 20, 127, 0, 0, 1, 192, 168, 1, 1, 0, 80, 1, 187, 0x20, 0, 5, 0, 0, 0, 0,
607 0,
608 ]);
609
610 let addresses: Addresses = IPv4::new([127, 0, 0, 1], [192, 168, 1, 1], 80, 443).into();
611 let header = Builder::new(
612 Version::Two | Command::Proxy,
613 AddressFamily::IPv4 | Protocol::Datagram,
614 )
615 .write_payload(addresses)
616 .unwrap()
617 .write_payload(Type::SSL)
618 .unwrap()
619 .write_payload(5u16)
620 .unwrap()
621 .write_payload([0u8; 5].as_slice())
622 .unwrap()
623 .build()
624 .unwrap();
625
626 assert_eq!(header, expected);
627 }
628
629 #[test]
630 fn build_ipv6_with_tlvs() {
631 let source_address = [
632 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
633 0xFF, 0xF2,
634 ];
635 let destination_address = [
636 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
637 0xFF, 0xF1,
638 ];
639 let addresses: Addresses = IPv6::new(source_address, destination_address, 80, 443).into();
640 let mut expected = Vec::from(PROTOCOL_PREFIX);
641 expected.extend([0x20, 0x20, 0, 48]);
642 expected.extend(source_address);
643 expected.extend(destination_address);
644 expected.extend([0, 80, 1, 187]);
645 expected.extend([4, 0, 1, 0, 4, 0, 1, 0, 4, 0, 1, 42]);
646
647 let header = Builder::new(
648 Version::Two | Command::Local,
649 AddressFamily::IPv6 | Protocol::Unspecified,
650 )
651 .write_payload(addresses)
652 .unwrap()
653 .write_payloads([
654 (Type::NoOp, [0].as_slice()),
655 (Type::NoOp, [0].as_slice()),
656 (Type::NoOp, [42].as_slice()),
657 ])
658 .unwrap()
659 .build()
660 .unwrap();
661
662 assert_eq!(header, expected);
663 }
664
665 #[test]
666 fn build_unix_with_tlv() {
667 let source_address = [0xFFu8; 108];
668 let destination_address = [0xAAu8; 108];
669
670 let addresses: Addresses = Unix::new(source_address, destination_address).into();
671 let mut expected = Vec::from(PROTOCOL_PREFIX);
672 expected.extend([0x20, 0x31, 0, 216]);
673 expected.extend(source_address);
674 expected.extend(destination_address);
675 expected.extend([0x20, 0, 0]);
676
677 let header = Builder::new(
678 Version::Two | Command::Local,
679 AddressFamily::Unix | Protocol::Stream,
680 )
681 .with_length(216)
682 .write_payload(addresses)
683 .unwrap()
684 .write_tlv(Type::SSL, &[])
685 .unwrap()
686 .build()
687 .unwrap();
688
689 assert_eq!(header, expected);
690 }
691}