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::convert::TryFrom;
15 use core::pin::Pin;
16 use core::task::{Context, Poll};
17 use std::io::Read;
18 
19 use crate::body::async_impl::Body;
20 use crate::body::mime::common::{data_copy, SizeResult, TokenStatus};
21 use crate::body::mime::{EncodeHeaders, MixFrom, PartStatus};
22 use crate::body::{async_impl, sync_impl, MimePart};
23 use crate::error::{ErrorKind, HttpError};
24 use crate::{AsyncRead, ReadBuf};
25 
26 /// `MimePartEncoder` can get a [`MimePart`] to encode into data that can be
27 /// transmitted.
28 ///
29 /// [`MimePart`]: MimePart
30 ///
31 /// # Examples
32 ///
33 /// ```
34 /// use ylong_http::body::{sync_impl, MimePart, MimePartEncoder};
35 ///
36 /// let part = MimePart::builder()
37 ///     .header("accept", "text/html")
38 ///     .body_from_reader("01234567890123456789".as_bytes())
39 ///     .body_from_bytes(b"9876543210\r\n")
40 ///     .build()
41 ///     .unwrap();
42 /// let mut part_encoder = MimePartEncoder::from_part(part);
43 /// let mut buf = vec![0u8; 10];
44 /// let mut v_size = vec![];
45 /// let mut v_str = vec![];
46 ///
47 /// loop {
48 ///     let len = sync_impl::Body::data(&mut part_encoder, &mut buf).unwrap();
49 ///     if len == 0 {
50 ///         break;
51 ///     }
52 ///     v_size.push(len);
53 ///     v_str.extend_from_slice(&buf[..len]);
54 /// }
55 /// assert_eq!(v_size, vec![10, 10, 10, 2]);
56 /// assert_eq!(v_str, b"accept:text/html\r\n\r\n9876543210\r\n");
57 /// ```
58 #[derive(Debug)]
59 pub struct MimePartEncoder<'a> {
60     // Encode stage now
61     stage: PartStatus,
62     headers_encode: EncodeHeaders,
63     body: Option<MixFrom<'a>>,
64     // src which is need to encode
65     src: Vec<u8>,
66     // index of src
67     src_idx: usize,
68 }
69 
70 impl<'a> MimePartEncoder<'a> {
71     /// Creates a `MimePartEncoder` by a [`MimePart`].
72     ///
73     /// [`MimePart`]: MimePart
74     ///
75     /// # Examples
76     ///
77     /// ```
78     /// use ylong_http::body::{MimePart, MimePartEncoder};
79     ///
80     /// let part = MimePart::builder().build().unwrap();
81     /// let part_encoder = MimePartEncoder::from_part(part);
82     /// ```
from_partnull83     pub fn from_part(part: MimePart<'a>) -> Self {
84         let body = (!part.body.is_empty()).then_some(part.body);
85 
86         MimePartEncoder {
87             stage: PartStatus::Start,
88             headers_encode: EncodeHeaders::new(part.headers),
89             src: vec![],
90             src_idx: 0,
91             body,
92         }
93     }
94 
check_nextnull95     fn check_next(&mut self) {
96         match self.stage {
97             PartStatus::Start => {
98                 self.stage = PartStatus::Headers;
99                 // init
100                 self.src_idx = 0;
101                 self.src = vec![];
102             }
103             PartStatus::Headers => {
104                 self.stage = PartStatus::Crlf;
105                 if self.body.is_some() {
106                     // has body, so has Crlf
107                     self.src = b"\r\n".to_vec();
108                 } else {
109                     self.src = vec![];
110                 }
111                 self.src_idx = 0;
112             }
113             PartStatus::Crlf => {
114                 self.stage = PartStatus::Body;
115                 // Just record index
116                 self.src_idx = 0;
117             }
118             PartStatus::Body => {
119                 self.stage = PartStatus::End;
120                 self.src_idx = 0;
121             }
122             PartStatus::End => {
123                 // init
124                 self.src_idx = 0;
125                 self.src = vec![];
126             }
127         }
128     }
129 
start_encodenull130     fn start_encode(&mut self) -> SizeResult {
131         self.check_next();
132         Ok(0)
133     }
134 
135     // use EncodeHeaders.encode
headers_encodenull136     fn headers_encode(&mut self, dst: &mut [u8]) -> SizeResult {
137         match self.headers_encode.encode(dst)? {
138             TokenStatus::Partial(size) => Ok(size),
139             TokenStatus::Complete(size) => {
140                 self.check_next();
141                 Ok(size)
142             }
143         }
144     }
145 
crlf_encodenull146     fn crlf_encode(&mut self, dst: &mut [u8]) -> SizeResult {
147         match data_copy(&self.src, &mut self.src_idx, dst)? {
148             TokenStatus::Partial(size) => Ok(size),
149             TokenStatus::Complete(size) => {
150                 self.check_next();
151                 Ok(size)
152             }
153         }
154     }
155 
156     // Uses `syn::Body::data` of `MixFrom`.
body_sync_encodenull157     fn body_sync_encode(&mut self, dst: &mut [u8]) -> SizeResult {
158         if let Some(body) = &mut self.body {
159             let size = sync_impl::Body::data(body, dst)?;
160             if size == 0 {
161                 self.check_next();
162             }
163             return Ok(size);
164         }
165         self.check_next();
166         Ok(0)
167     }
168 }
169 
170 impl sync_impl::Body for MimePartEncoder<'_> {
171     type Error = std::io::Error;
172 
datanull173     fn data(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
174         let mut count = 0;
175         while count != buf.len() {
176             let encode_size = match self.stage {
177                 PartStatus::Start => self.start_encode(),
178                 PartStatus::Headers => self.headers_encode(&mut buf[count..]),
179                 PartStatus::Crlf => self.crlf_encode(&mut buf[count..]),
180                 PartStatus::Body => self.body_sync_encode(&mut buf[count..]),
181                 PartStatus::End => return Ok(count),
182             };
183             count += encode_size?;
184         }
185         Ok(count)
186     }
187 }
188 
189 impl async_impl::Body for MimePartEncoder<'_> {
190     type Error = std::io::Error;
191 
poll_datanull192     fn poll_data(
193         mut self: Pin<&mut Self>,
194         cx: &mut Context<'_>,
195         buf: &mut [u8],
196     ) -> Poll<Result<usize, Self::Error>> {
197         let mut count = 0;
198         while count != buf.len() {
199             let encode_size: Poll<SizeResult> = match self.stage {
200                 PartStatus::Start => Poll::Ready(self.start_encode()),
201                 PartStatus::Headers => Poll::Ready(self.headers_encode(&mut buf[count..])),
202                 PartStatus::Crlf => Poll::Ready(self.crlf_encode(&mut buf[count..])),
203                 PartStatus::Body => self.poll_mime_body(cx, &mut buf[count..]),
204                 PartStatus::End => return Poll::Ready(Ok(count)),
205             };
206 
207             match encode_size {
208                 Poll::Ready(Ok(size)) => count += size,
209                 Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
210                 Poll::Pending => return Poll::Pending,
211             }
212         }
213         Poll::Ready(Ok(count))
214     }
215 }
216 
217 impl MimePartEncoder<'_> {
poll_mime_bodynull218     fn poll_mime_body(
219         &mut self,
220         cx: &mut Context<'_>,
221         buf: &mut [u8],
222     ) -> Poll<Result<usize, std::io::Error>> {
223         match self.body {
224             Some(ref mut body) => {
225                 let poll_result = Pin::new(body).poll_data(cx, buf);
226                 if let Poll::Ready(Ok(0)) = poll_result {
227                     // complete async read body
228                     self.check_next();
229                 };
230                 poll_result
231             }
232             _ => {
233                 self.check_next();
234                 Poll::Ready(Ok(0))
235             }
236         }
237     }
238 }
239 
240 #[cfg(test)]
241 mod ut_mime_part_encoder {
242     use crate::body::{async_impl, sync_impl, MimePart, MimePartEncoder};
243     use crate::headers::Headers;
244 
245     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
246     ///
247     /// # Brief
248     /// 1. Creates a `MimePart`.
249     /// 2. Sets body by `body_from_owned`.
250     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
251     /// 4. Checks whether the result is correct.
252     #[test]
ut_mime_part_encoder_body_from_ownednull253     fn ut_mime_part_encoder_body_from_owned() {
254         part_encode_compare! (
255             MimePart: {
256                 BodyOwned: b"123456".to_vec(),
257             },
258             ResultData: b"\r\n123456",
259             Sync,
260         );
261     }
262 
263     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
264     ///
265     /// # Brief
266     /// 1. Creates a `MimePart`.
267     /// 2. Sets body by `body_from_bytes`.
268     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
269     /// 4. Checks whether the result is correct.
270     #[test]
ut_mime_part_encoder_body_from_bytesnull271     fn ut_mime_part_encoder_body_from_bytes() {
272         part_encode_compare! (
273             MimePart: {
274                 BodySlice: "123456".as_bytes(),
275             },
276             BufSize: 5,
277             ResultData: b"\r\n123456",
278             Sync,
279         );
280     }
281 
282     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
283     ///
284     /// # Brief
285     /// 1. Creates a `MimePart`.
286     /// 2. Sets body by `body_from_reader`.
287     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
288     /// 4. Checks whether the result is correct.
289     #[test]
ut_mime_part_encoder_body_from_readernull290     fn ut_mime_part_encoder_body_from_reader() {
291         part_encode_compare! (
292             MimePart: {
293                 BodyReader: "123456".as_bytes(),
294             },
295             BufSize: 5,
296             ResultData: b"\r\n123456",
297             Sync,
298         );
299     }
300 
301     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
302     ///
303     /// # Brief
304     /// 1. Creates a `MimePart`.
305     /// 2. Sets headers and sets body.
306     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
307     /// 4. Checks whether the result is correct.
308     #[test]
ut_mime_part_encoder_data_commonnull309     fn ut_mime_part_encoder_data_common() {
310         part_encode_compare! (
311             MimePart: {
312                 Header: "accept", "text/html",
313                 BodyReader: "9876543210\r\n".as_bytes(),
314             },
315             BufSize: 10,
316             ResultSize: vec![10, 10, 10, 2],
317             ResultData: b"accept:text/html\r\n\r\n9876543210\r\n",
318             Sync,
319         );
320     }
321 
322     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
323     ///
324     /// # Brief
325     /// 1. Creates a `MimePart`.
326     /// 2. Doesn't set headers and only sets body.
327     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
328     /// 4. Checks whether the result is correct.
329     #[test]
ut_mime_part_encoder_data_noheadersnull330     fn ut_mime_part_encoder_data_noheaders() {
331         part_encode_compare! (
332             MimePart: {
333                 BodySlice: b"0123456789--0123",
334             },
335             BufSize: 10,
336             ResultSize: vec![10, 8],
337             ResultData: b"\r\n0123456789--0123",
338             Sync,
339         );
340     }
341 
342     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
343     ///
344     /// # Brief
345     /// 1. Creates a `MimePart`.
346     /// 2. Doesn't set body and only sets headers.
347     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
348     /// 4. Checks whether the result is correct.
349     #[test]
ut_mime_part_encoder_data_nobodynull350     fn ut_mime_part_encoder_data_nobody() {
351         // no body no CRLF
352         part_encode_compare! (
353             MimePart: {
354                 Header: "key", "\"value\"",
355             },
356             ResultData: b"key:\"value\"\r\n",
357             Sync,
358         );
359 
360         part_encode_compare! (
361             MimePart: {
362                 Header: "key1", "value1",
363                 Header: "key2", "value2",
364             },
365             BufSize: 10,
366             ResultSize: vec![10, 10, 6],
367             Sync,
368         );
369     }
370 
371     /// UT test cases for `syn::Body::data` of `MimePartEncoder`.
372     ///
373     /// # Brief
374     /// 1. Creates a `MimePart`.
375     /// 2. Doesn't set headers and doesn't set headers.
376     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes.
377     /// 4. Checks whether the result is correct.
378     #[test]
ut_mime_part_encoder_data_noheaders_nobodynull379     fn ut_mime_part_encoder_data_noheaders_nobody() {
380         part_encode_compare! (
381             MimePart: {
382             },
383             ResultData: b"",
384             Sync,
385         );
386     }
387 
388     /// UT test cases for `asyn::Body::data` of `MimePartEncoder`.
389     ///
390     /// # Brief
391     /// 1. Creates a `MimePart`.
392     /// 2. Sets body by `body_from_owned`.
393     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes asynchronously.
394     /// 4. Checks whether the result is correct.
395     #[cfg(feature = "ylong_base")]
396     #[test]
ut_mime_part_encoder_body_from_owned_then_async_datanull397     fn ut_mime_part_encoder_body_from_owned_then_async_data() {
398         let handle = ylong_runtime::spawn(async move {
399             mime_part_encoder_body_from_owned_then_async_data().await;
400         });
401         ylong_runtime::block_on(handle).unwrap();
402     }
403 
404     #[cfg(feature = "ylong_base")]
405     async fn mime_part_encoder_body_from_owned_then_async_data() {
406         part_encode_compare! (
407             MimePart: {
408                 BodyOwned: b"123456".to_vec(),
409             },
410             BufSize: 5,
411             ResultSize: vec![5, 3],
412             ResultData: b"\r\n123456",
413             Async,
414         );
415     }
416 
417     /// UT test cases for `asyn::Body::data` of `MimePartEncoder`.
418     ///
419     /// # Brief
420     /// 1. Creates a `MimePart`.
421     /// 2. Sets body by `body_from_bytes`.
422     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes asynchronously.
423     /// 4. Checks whether the result is correct.
424     #[cfg(feature = "ylong_base")]
425     #[test]
ut_mime_part_encoder_body_from_bytes_then_async_datanull426     fn ut_mime_part_encoder_body_from_bytes_then_async_data() {
427         let handle = ylong_runtime::spawn(async move {
428             mime_part_encoder_body_from_bytes_then_async_data().await;
429         });
430         ylong_runtime::block_on(handle).unwrap();
431     }
432 
433     #[cfg(feature = "ylong_base")]
434     async fn mime_part_encoder_body_from_bytes_then_async_data() {
435         part_encode_compare! (
436             MimePart: {
437                 BodySlice: b"123456",
438             },
439             BufSize: 5,
440             ResultSize: vec![5, 3],
441             ResultData: b"\r\n123456",
442             Async,
443         );
444     }
445 
446     /// UT test cases for `asyn::Body::data` of `MimePartEncoder`.
447     ///
448     /// # Brief
449     /// 1. Creates a `MimePart`.
450     /// 2. Sets headers and sets body.
451     /// 3. Builds a `MimePartEncoder` by `from_part` and encodes asynchronously.
452     /// 4. Checks whether the result is correct.
453     #[cfg(feature = "ylong_base")]
454     #[test]
ut_mime_part_encoder_common_then_async_datanull455     fn ut_mime_part_encoder_common_then_async_data() {
456         let handle = ylong_runtime::spawn(async move {
457             mime_part_encoder_common_then_async_data().await;
458         });
459         ylong_runtime::block_on(handle).unwrap();
460     }
461 
462     #[cfg(feature = "ylong_base")]
463     async fn mime_part_encoder_common_then_async_data() {
464         part_encode_compare! (
465             MimePart: {
466                 Header: "accept", "text/html",
467                 BodySlice: b"9876543210\r\n",
468             },
469             BufSize: 10,
470             ResultSize: vec![10, 10, 10, 2],
471             ResultData: b"accept:text/html\r\n\r\n9876543210\r\n",
472             Async,
473         );
474     }
475 }
476