1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 use core::mem::take;
15 
16 use crate::body::mime::common::{get_crlf_contain, trim_back_lwsp_if_end_with_lf, trim_front_lwsp};
17 use crate::body::mime::decode::BoundaryTag;
18 use crate::body::mime::{DecodeHeaders, MimePart, PartStatus};
19 use crate::body::TokenStatus;
20 use crate::error::{ErrorKind, HttpError};
21 use crate::headers::Headers;
22 
23 #[derive(Debug, PartialEq)]
24 pub(crate) struct MimePartDecoder {
25     // Encode stage now
26     stage: PartStatus,
27     headers_decoder: DecodeHeaders,
28     part: MimePart<'static>,
29     // boundary
30     boundary: Vec<u8>,
31     tag: BoundaryTag,
32     // marks last "\n"
33     src_idx: usize,
34 }
35 
36 impl MimePartDecoder {
new() -> Self37     pub(crate) fn new() -> Self {
38         MimePartDecoder {
39             stage: PartStatus::Start,
40             headers_decoder: DecodeHeaders::new(),
41             part: MimePart::new(),
42             boundary: b"-".to_vec(),
43             tag: BoundaryTag::Init,
44             src_idx: 0,
45         }
46     }
47 
init(&mut self)48     pub(crate) fn init(&mut self) {
49         self.stage = PartStatus::Start;
50         self.headers_decoder = DecodeHeaders::new();
51 
52         self.part = MimePart::new();
53         self.boundary = b"-".to_vec();
54         self.tag = BoundaryTag::Init;
55         self.src_idx = 0;
56     }
57 
58     // by MimeMultiDecode
init_from_multi(&mut self, headers: Headers, boundary: Vec<u8>)59     pub(crate) fn init_from_multi(&mut self, headers: Headers, boundary: Vec<u8>) {
60         self.init();
61         self.part.set_headers(headers);
62         self.boundary = boundary;
63         self.stage = PartStatus::Body;
64     }
65 
set_boundary(&mut self, boundary: Vec<u8>)66     pub(crate) fn set_boundary(&mut self, boundary: Vec<u8>) {
67         self.boundary = boundary;
68     }
69 
decode<'a>( &mut self, buf: &'a [u8], ) -> Result<(TokenStatus<MimePart<'static>, ()>, &'a [u8]), HttpError>70     pub(crate) fn decode<'a>(
71         &mut self,
72         buf: &'a [u8],
73     ) -> Result<(TokenStatus<MimePart<'static>, ()>, &'a [u8]), HttpError> {
74         // Option
75         if buf.is_empty() {
76             return Err(ErrorKind::InvalidInput.into());
77         }
78 
79         let mut results = TokenStatus::Partial(());
80         let mut remains = buf;
81         loop {
82             let rest = match self.stage {
83                 PartStatus::Start => self.start_decode(remains),
84                 PartStatus::Headers => self.headers_decode(remains),
85                 // not use
86                 PartStatus::Crlf => Ok(remains),
87                 PartStatus::Body => self.body_decode(remains),
88                 PartStatus::End => {
89                     results = TokenStatus::Complete(take(&mut self.part));
90                     break;
91                 }
92             }?;
93             remains = rest;
94             if remains.is_empty() && self.stage != PartStatus::End {
95                 break;
96             }
97         }
98         Ok((results, remains))
99     }
100 
start_decode<'a>(&mut self, buf: &'a [u8]) -> Result<&'a [u8], HttpError>101     fn start_decode<'a>(&mut self, buf: &'a [u8]) -> Result<&'a [u8], HttpError> {
102         let buf = trim_front_lwsp(buf);
103         self.stage = PartStatus::Headers;
104         Ok(buf)
105     }
106 
107     // Headers + Crlf
headers_decode<'a>(&mut self, buf: &'a [u8]) -> Result<&'a [u8], HttpError>108     fn headers_decode<'a>(&mut self, buf: &'a [u8]) -> Result<&'a [u8], HttpError> {
109         let (elem, rest) = self.headers_decoder.decode(buf)?;
110         if let TokenStatus::Complete(headers) = elem {
111             self.part.set_headers(headers);
112             self.stage = PartStatus::Body;
113         }
114 
115         Ok(rest)
116     }
117 
118     // Writes to the part body and checks boundary.
body_decode<'a>(&mut self, buf: &'a [u8]) -> Result<&'a [u8], HttpError>119     fn body_decode<'a>(&mut self, buf: &'a [u8]) -> Result<&'a [u8], HttpError> {
120         match get_crlf_contain(buf) {
121             TokenStatus::Partial(unparsed) => {
122                 self.part.body_extend_from_slice(unparsed);
123                 Ok(&[])
124             }
125             TokenStatus::Complete((src, unparsed)) => {
126                 // clone in this.
127                 self.part.body_extend_from_slice(src);
128                 let body = match self.part.body() {
129                     Some(vec) => vec,
130                     None => return Err(ErrorKind::InvalidInput.into()),
131                 };
132                 let line = &body[self.src_idx..];
133                 let trim_line = trim_back_lwsp_if_end_with_lf(line);
134 
135                 // checks whether is middle boundary
136                 let mut deriv_boundary = b"--".to_vec();
137                 deriv_boundary.extend_from_slice(&self.boundary);
138                 if trim_line == deriv_boundary {
139                     self.tag = BoundaryTag::Middle;
140                     self.when_is_end();
141                     return Ok(unparsed);
142                 }
143                 // checks whether is end boundary
144                 deriv_boundary.extend_from_slice(b"--");
145                 if trim_line == deriv_boundary {
146                     self.tag = BoundaryTag::End;
147                     self.when_is_end();
148                     return Ok(unparsed);
149                 }
150 
151                 // is not end
152                 self.src_idx = body.len();
153                 Ok(unparsed)
154             }
155         }
156     }
157 
when_is_end(&mut self)158     fn when_is_end(&mut self) {
159         self.stage = PartStatus::End;
160         self.part.body_truncate(self.src_idx);
161         self.part.body_trim_crlf_once();
162     }
163 
164     // Checks whether is the last part of multi.
is_last_part(&self) -> bool165     pub(crate) fn is_last_part(&self) -> bool {
166         self.tag == BoundaryTag::End
167     }
168 }
169 
170 #[cfg(test)]
171 mod ut_mime_part_decoder {
172     use crate::body::mime::{MimePart, MimePartDecoder};
173     use crate::body::TokenStatus;
174 
175     /// UT test cases for `MimePartDecoder::decode`.
176     ///
177     /// # Brief
178     /// 1. Creates a `MimePartDecoder` by `MimePartDecoder::new`.
179     /// 2. Uses `MimePartDecoder::decode` to decode `MimePart`.
180     /// 3. Creates a `MimePart` and sets the same parameters to compare.
181     /// 4. Checks whether the result is correct.
182     #[test]
ut_mime_part_decoder_new_by_crlf()183     fn ut_mime_part_decoder_new_by_crlf() {
184         let part1 = MimePart::builder()
185             .body_from_owned(b"abcd".to_vec())
186             .build()
187             .unwrap();
188 
189         // common
190         let buf = b"\r\nabcd\r\n--abc\r\nabcd";
191         let mut decoder = MimePartDecoder::new();
192         decoder.set_boundary(b"abc".to_vec());
193         let (elem, rest) = decoder.decode(buf).unwrap();
194         assert!(elem.is_complete());
195 
196         if let TokenStatus::Complete(part) = elem {
197             assert_eq!(part, part1);
198         }
199         assert_eq!(rest, b"abcd");
200 
201         // has LWSP
202         let buf = b"   \r\nabcd\r\n--abc\r\nabcd";
203         let mut decoder = MimePartDecoder::new();
204         decoder.set_boundary(b"abc".to_vec());
205         let (elem, rest) = decoder.decode(buf).unwrap();
206         assert!(elem.is_complete());
207 
208         if let TokenStatus::Complete(part) = elem {
209             assert_eq!(part, part1);
210         }
211         assert_eq!(rest, b"abcd");
212 
213         // has RFC LWSP
214         let buf = b"\r\nabcd\r\n--abc    \r\nabcd";
215         let mut decoder = MimePartDecoder::new();
216         decoder.set_boundary(b"abc".to_vec());
217         let (elem, rest) = decoder.decode(buf).unwrap();
218         assert!(elem.is_complete());
219 
220         if let TokenStatus::Complete(part) = elem {
221             assert_eq!(part, part1);
222         }
223         assert_eq!(rest, b"abcd");
224     }
225 
226     /// UT test cases for `MimePartDecoder::decode`.
227     ///
228     /// # Brief
229     /// 1. Creates a `MimePartDecoder` by `MimePartDecoder::new`.
230     /// 2. Uses `MimePartDecoder::decode` to decode `MimePart`.
231     /// 3. Creates a `MimePart` and sets the same parameters to compare.
232     /// 4. Checks whether the result is correct.
233     #[test]
ut_mime_part_decoder_new_by_lf()234     fn ut_mime_part_decoder_new_by_lf() {
235         let part1 = MimePart::builder()
236             .body_from_owned(b"abcd".to_vec())
237             .build()
238             .unwrap();
239 
240         // common
241         let buf = b"\nabcd\n--abc\nabcd";
242         let mut decoder = MimePartDecoder::new();
243         decoder.set_boundary(b"abc".to_vec());
244         let (elem, rest) = decoder.decode(buf).unwrap();
245         assert!(elem.is_complete());
246 
247         if let TokenStatus::Complete(part) = elem {
248             assert_eq!(part, part1);
249         }
250         assert_eq!(rest, b"abcd");
251 
252         // has LWSP
253         let buf = b"    \nabcd\n--abc\nabcd";
254         let mut decoder = MimePartDecoder::new();
255         decoder.set_boundary(b"abc".to_vec());
256         let (elem, rest) = decoder.decode(buf).unwrap();
257         assert!(elem.is_complete());
258 
259         if let TokenStatus::Complete(part) = elem {
260             assert_eq!(part, part1);
261         }
262         assert_eq!(rest, b"abcd");
263 
264         // has RFC LWSP
265         let buf = b"    \nabcd\n--abc   \nabcd";
266         let mut decoder = MimePartDecoder::new();
267         decoder.set_boundary(b"abc".to_vec());
268         let (elem, rest) = decoder.decode(buf).unwrap();
269         assert!(elem.is_complete());
270 
271         if let TokenStatus::Complete(part) = elem {
272             assert_eq!(part, part1);
273         }
274         assert_eq!(rest, b"abcd");
275     }
276 
277     /// UT test cases for `MimePartDecoder::decode`.
278     ///
279     /// # Brief
280     /// 1. Creates a `MimePartDecoder` by `MimePartDecoder::new`.
281     /// 2. Uses `MimePartDecoder::decode` to decode `MimePart`.
282     /// 3. Creates a `MimePart` and sets the same parameters to compare.
283     /// 4. Checks whether the result is correct.
284     #[test]
ut_mime_part_decoder_decode()285     fn ut_mime_part_decoder_decode() {
286         let buf = b"    name1:   value1\r\n    name2:    value2\r\n\r\nabcd\r\n--abc  \r\nabcd";
287         let mut decoder = MimePartDecoder::new();
288         decoder.set_boundary(b"abc".to_vec());
289         let (elem, rest) = decoder.decode(buf).unwrap();
290         assert!(elem.is_complete());
291         let part1 = MimePart::builder()
292             .header("name1", "value1")
293             .header("name2", "value2")
294             .body_from_owned(b"abcd".to_vec())
295             .build()
296             .unwrap();
297         if let TokenStatus::Complete(part) = elem {
298             assert_eq!(part, part1);
299         }
300         assert_eq!(rest, b"abcd");
301     }
302 
303     /// UT test cases for `MimePartDecoder::decode`.
304     ///
305     /// # Brief
306     /// 1. Creates a `MimePartDecoder` by `MimePartDecoder::new`.
307     /// 2. Uses `MimePartDecoder::decode` to decode `MimePart`.
308     /// 3. Creates a `MimePart` and sets the same parameters to compare.
309     /// 4. Checks whether the result is correct.
310     #[test]
ut_mime_part_decoder_no_headers_no_body()311     fn ut_mime_part_decoder_no_headers_no_body() {
312         let buf = b"\r\n--abc\r\nabcd";
313         let mut decoder = MimePartDecoder::new();
314         decoder.set_boundary(b"abc".to_vec());
315         let (elem, rest) = decoder.decode(buf).unwrap();
316         assert!(elem.is_complete());
317         let part1 = MimePart::builder().build().unwrap();
318         if let TokenStatus::Complete(part) = elem {
319             assert_eq!(part, part1);
320         }
321         assert_eq!(rest, b"abcd");
322     }
323 
324     /// UT test cases for `MimePartDecoder::decode`.
325     ///
326     /// # Brief
327     /// 1. Creates a `MimePartDecoder` by `MimePartDecoder::new`.
328     /// 2. Uses `MimePartDecoder::decode` to decode `MimePart`.
329     /// 3. The decode bytes are divided into several executions.
330     /// 4. Creates a `MimePart` and sets the same parameters to compare.
331     /// 5. Checks whether the result is correct.
332     #[test]
ut_decode_headers_decode_times()333     fn ut_decode_headers_decode_times() {
334         let buf = b"name1:value1\r\nname2:value2\r\n\r\nabcd\r\n--abc\r\nabcd";
335         let mut decoder = MimePartDecoder::new();
336         decoder.set_boundary(b"abc".to_vec());
337         // nam
338         let (elem, rest) = decoder.decode(&buf[0..3]).unwrap();
339         assert!(!elem.is_complete());
340         assert_eq!(rest, b"");
341 
342         // e1:value1\r\nname2:value2\r\n\r\n
343         let (elem, rest) = decoder.decode(&buf[3..30]).unwrap();
344         assert!(!elem.is_complete());
345         assert_eq!(rest, b"");
346 
347         // abcd\r\n--abc\r\nabcd
348         let (elem, rest) = decoder.decode(&buf[30..]).unwrap();
349         assert!(elem.is_complete());
350         let part1 = MimePart::builder()
351             .header("name1", "value1")
352             .header("name2", "value2")
353             .body_from_owned(b"abcd".to_vec())
354             .build()
355             .unwrap();
356         if let TokenStatus::Complete(part) = elem {
357             assert_eq!(part, part1);
358         }
359         assert_eq!(rest, b"abcd");
360     }
361 }
362