miltr_common/modifications/
headers.rs1use std::borrow::Cow;
4
5use bytes::{BufMut, BytesMut};
6
7use crate::commands::Header;
8use crate::decoding::Parsable;
9use crate::encoding::Writable;
10use crate::error::STAGE_DECODING;
11use crate::{NotEnoughData, ProtocolError};
12use miltr_utils::ByteParsing;
13
14#[derive(Debug, Clone)]
16pub struct AddHeader {
17 header: Header,
18}
19
20impl AddHeader {
21 const CODE: u8 = b'h';
22
23 #[must_use]
25 pub fn new(name: &[u8], value: &[u8]) -> Self {
26 Self {
27 header: Header::new(name, value),
28 }
29 }
30
31 #[must_use]
33 pub fn name(&self) -> Cow<str> {
34 self.header.name()
35 }
36
37 #[must_use]
39 pub fn value(&self) -> Cow<str> {
40 self.header.value()
41 }
42}
43
44impl Parsable for AddHeader {
45 const CODE: u8 = Self::CODE;
46
47 fn parse(buffer: BytesMut) -> Result<Self, ProtocolError> {
48 let header = Header::parse(buffer)?;
49
50 Ok(Self { header })
51 }
52}
53
54impl Writable for AddHeader {
55 fn write(&self, buffer: &mut BytesMut) {
56 self.header.write(buffer);
57 }
58
59 fn len(&self) -> usize {
60 self.header.len()
61 }
62
63 fn code(&self) -> u8 {
64 Self::CODE
65 }
66 fn is_empty(&self) -> bool {
67 self.header.is_empty()
68 }
69}
70
71#[derive(Debug, Clone)]
73pub struct ChangeHeader {
74 index: u32,
79
80 header: Header,
81}
82
83impl ChangeHeader {
84 const CODE: u8 = b'm';
85
86 #[must_use]
88 pub fn new(index: u32, name: &[u8], value: &[u8]) -> Self {
89 Self {
90 index,
91 header: Header::new(name, value),
92 }
93 }
94
95 #[must_use]
97 pub fn name(&self) -> Cow<str> {
98 self.header.name()
99 }
100
101 #[must_use]
103 pub fn value(&self) -> Cow<str> {
104 self.header.value()
105 }
106
107 #[must_use]
112 pub fn index(&self) -> u32 {
113 self.index
114 }
115}
116
117impl Parsable for ChangeHeader {
118 const CODE: u8 = Self::CODE;
119
120 fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
121 let Some(index) = buffer.safe_get_u32() else {
122 return Err(NotEnoughData::new(
123 STAGE_DECODING,
124 "ChangeHeader",
125 "Index byte missing",
126 1,
127 0,
128 buffer,
129 )
130 .into());
131 };
132 let header = Header::parse(buffer)?;
133
134 Ok(Self { index, header })
135 }
136}
137
138impl Writable for ChangeHeader {
139 fn write(&self, buffer: &mut BytesMut) {
147 let index = u32::to_be_bytes(self.index);
148 buffer.put_slice(&index);
149 self.header.write(buffer);
150 }
151
152 fn len(&self) -> usize {
153 4 + self.header.len()
154 }
155
156 fn code(&self) -> u8 {
157 Self::CODE
158 }
159 fn is_empty(&self) -> bool {
160 self.header.is_empty()
161 }
162}
163
164#[derive(Debug, Clone)]
166pub struct InsertHeader {
167 index: u32,
168 header: Header,
169}
170
171impl InsertHeader {
172 const CODE: u8 = b'i';
173
174 #[must_use]
176 pub fn new(index: u32, name: &[u8], value: &[u8]) -> Self {
177 Self {
178 index,
179 header: Header::new(name, value),
180 }
181 }
182
183 #[must_use]
185 pub fn name(&self) -> Cow<str> {
186 self.header.name()
187 }
188
189 #[must_use]
191 pub fn value(&self) -> Cow<str> {
192 self.header.value()
193 }
194
195 #[must_use]
197 pub fn index(&self) -> u32 {
198 self.index
199 }
200}
201
202impl Parsable for InsertHeader {
203 const CODE: u8 = Self::CODE;
204
205 fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
206 let Some(index) = buffer.safe_get_u32() else {
207 return Err(NotEnoughData::new(
208 STAGE_DECODING,
209 "InsertHeader",
210 "Index byte missing",
211 1,
212 0,
213 buffer,
214 )
215 .into());
216 };
217 let header = Header::parse(buffer)?;
218
219 Ok(Self { index, header })
220 }
221}
222
223impl Writable for InsertHeader {
224 fn write(&self, buffer: &mut BytesMut) {
226 let index = u32::to_be_bytes(self.index);
227 buffer.put_slice(&index);
228 self.header.write(buffer);
229 }
230
231 fn len(&self) -> usize {
232 4 + self.header.len()
233 }
234
235 fn code(&self) -> u8 {
236 Self::CODE
237 }
238 fn is_empty(&self) -> bool {
239 self.header.is_empty()
240 }
241}
242
243#[cfg(test)]
244mod test {
245
246 use super::*;
247 use pretty_assertions::assert_eq;
248 use rstest::rstest;
249
250 #[test]
251 fn test_add_header() {
252 let mut buffer = BytesMut::from("h");
253 let add_header = AddHeader {
254 header: Header::new(b"name", b"value"),
255 };
256
257 add_header.write(&mut buffer);
258 assert_eq!(buffer, BytesMut::from("hname\0value\0"));
259 }
260
261 #[rstest]
262 #[case((1, String::from("name"), String::from("value")), BytesMut::from("m\0\0\0\x01name\0value\0"))]
263 #[case((0, String::from("name"), String::from("value")), BytesMut::from("m\0\0\0\0name\0value\0"))]
264 #[case((2, String::from("name"), String::from("\0")), BytesMut::from("m\0\0\0\x02name\0\0\0"))]
265 fn test_change_header(#[case] input: (u32, String, String), #[case] expected: BytesMut) {
266 let mut buffer = BytesMut::from("m");
267
268 let change_header = ChangeHeader {
269 index: input.0,
270 header: Header::new(input.1.as_bytes(), input.2.as_bytes()),
271 };
272
273 change_header.write(&mut buffer);
274
275 assert_eq!(buffer, expected);
276 }
277 #[rstest]
278 #[case((1, String::from("name"), String::from("value")), BytesMut::from("i\0\0\0\x01name\0value\0"))]
279 #[case((0, String::from("name"), String::from("value")), BytesMut::from("i\0\0\0\0name\0value\0"))]
280 #[case((2, String::from("name"), String::from("\0")), BytesMut::from("i\0\0\0\x02name\0\0\0"))]
281 fn test_insert_header(#[case] input: (u32, String, String), #[case] expected: BytesMut) {
282 let mut buffer = BytesMut::from("i");
283
284 let change_header = ChangeHeader {
285 index: input.0,
286 header: Header::new(input.1.as_bytes(), input.2.as_bytes()),
287 };
288
289 change_header.write(&mut buffer);
290
291 assert_eq!(buffer, expected);
292 }
293}