Skip to main content

lafere/packet/
bytes.rs

1use super::{BodyBytes, BodyBytesMut};
2
3use bytes::{Bytes, BytesMut};
4
5/// A trait that allows efficient allocation if encryption is
6/// used or not.
7pub trait PacketBytes: std::fmt::Debug {
8	/// Creates a new Bytes instance.
9	///
10	/// It must always have header len available and initialized
11	fn new(header_len: usize) -> Self;
12
13	/// Returns the header Bytes.
14	fn header(&self) -> Bytes<'_>;
15
16	/// Returns the header mutably.
17	fn header_mut(&mut self) -> BytesMut<'_>;
18
19	/// Returns the full header mutably.
20	///
21	/// ## Note
22	/// This should only be used to fill the buffer from a reader, in any other
23	/// case you should use `header_mut`.
24	fn full_header_mut(&mut self) -> BytesMut<'_>;
25
26	/// Returns the body.
27	fn body(&self) -> BodyBytes<'_>;
28
29	/// Returns the body mutably.
30	fn body_mut(&mut self) -> BodyBytesMut<'_>;
31
32	/// Returns the full body mutably.
33	///
34	/// ## Note
35	/// This should only be used to fill the buffer from a reader, in any other
36	/// case you should use `body_mut`.
37	fn full_body_mut(&mut self) -> BytesMut<'_>;
38}
39
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct PlainBytes {
42	bytes: Vec<u8>,
43	header_len: usize,
44}
45
46impl PacketBytes for PlainBytes {
47	fn new(header_len: usize) -> Self {
48		Self {
49			bytes: vec![0; header_len],
50			header_len,
51		}
52	}
53
54	fn header(&self) -> Bytes<'_> {
55		self.bytes[..self.header_len].into()
56	}
57
58	fn header_mut(&mut self) -> BytesMut<'_> {
59		(&mut self.bytes[..self.header_len]).into()
60	}
61
62	fn full_header_mut(&mut self) -> BytesMut<'_> {
63		self.header_mut()
64	}
65
66	fn body(&self) -> BodyBytes<'_> {
67		BodyBytes::new(&self.bytes[self.header_len..])
68	}
69
70	fn body_mut(&mut self) -> BodyBytesMut<'_> {
71		BodyBytesMut::new(self.header_len, &mut self.bytes)
72	}
73
74	fn full_body_mut(&mut self) -> BytesMut<'_> {
75		(&mut self.bytes[self.header_len..]).into()
76	}
77}
78
79impl PlainBytes {
80	#[cfg_attr(not(feature = "connection"), allow(dead_code))]
81	pub(crate) fn as_slice(&self) -> &[u8] {
82		&*self.bytes
83	}
84}
85
86#[cfg(test)]
87pub(super) mod tests {
88	use super::*;
89	use bytes::{BytesOwned, BytesRead, BytesWrite};
90
91	pub fn test_gen_msg<B: PacketBytes>() {
92		let header = [10u8; 30];
93		let mut bytes = B::new(header.len());
94		bytes.header_mut().write(&header);
95		assert_eq!(bytes.body().len(), 0);
96
97		{
98			let mut body = bytes.body_mut();
99
100			// now i can't use bytes
101			// but body
102
103			body.write_u32(10u32);
104			body.write_u32(20u32);
105		}
106
107		// assert_eq!(bytes.as_slice().len(), 16 + 30 + 16 + 4 + 4);
108
109		assert_eq!(bytes.header().as_slice(), header);
110
111		let mut body = BytesOwned::new();
112		body.write_u32(10);
113		body.write_u32(20);
114		assert_eq!(bytes.body().as_slice(), body.as_slice());
115		assert_eq!(bytes.body().len(), body.len());
116	}
117
118	#[test]
119	fn plain_bytes() {
120		test_gen_msg::<PlainBytes>();
121	}
122}