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::pin::Pin; 15 use core::task::{Context, Poll}; 16 use std::collections::VecDeque; 17 18 use crate::body::mime::common::{data_copy, SizeResult}; 19 use crate::body::mime::EncodeHeaders; 20 use crate::body::{async_impl, sync_impl, MimeMulti, MimePartEncoder, TokenStatus, XPart}; 21 use crate::headers::Headers; 22 23 /// `MimeMultiEncoder` can get a [`MimeMulti`] to encode into data that can be 24 /// transmitted. 25 /// 26 /// [`MimeMulti`]: MimeMulti 27 /// 28 /// # Examples 29 /// 30 /// ``` 31 /// use ylong_http::body::{sync_impl, MimeMulti, MimeMultiEncoder, MimePart}; 32 /// 33 /// let multi1 = MimeMulti::builder() 34 /// .set_content_type(b"multipart/mixed", b"abc".to_vec()) 35 /// .set_boundary(b"abc".to_vec()) 36 /// .add_part( 37 /// MimePart::builder() 38 /// .header("key1", "value1") 39 /// .body_from_reader("111111".as_bytes()) 40 /// .build() 41 /// .unwrap(), 42 /// ) 43 /// .add_part( 44 /// MimePart::builder() 45 /// .header("key2", "value2") 46 /// .body_from_reader("22222".as_bytes()) 47 /// .build() 48 /// .unwrap(), 49 /// ) 50 /// .build() 51 /// .unwrap(); 52 /// 53 /// let multi = MimeMulti::builder() 54 /// .set_boundary(b"abcde".to_vec()) 55 /// .add_part( 56 /// MimePart::builder() 57 /// .header("key3", "value3") 58 /// .body_from_reader("33333".as_bytes()) 59 /// .build() 60 /// .unwrap(), 61 /// ) 62 /// .add_multi(multi1) 63 /// .add_part( 64 /// MimePart::builder() 65 /// .header("key4", "value4") 66 /// .body_from_async_reader("444444".as_bytes()) // nothing in sync 67 /// .build() 68 /// .unwrap(), 69 /// ) 70 /// .build() 71 /// .unwrap(); 72 /// 73 /// let mut multi_encoder = MimeMultiEncoder::from_multi(multi); 74 /// let mut buf = vec![0u8; 30]; 75 /// let mut v_size = vec![]; 76 /// let mut v_str = vec![]; 77 /// 78 /// loop { 79 /// let len = sync_impl::Body::data(&mut multi_encoder, &mut buf).unwrap(); 80 /// if len == 0 { 81 /// break; 82 /// } 83 /// v_size.push(len); 84 /// v_str.extend_from_slice(&buf[..len]); 85 /// } 86 /// assert_eq!(v_size, vec![30, 30, 30, 30, 30, 30, 13]); 87 /// assert_eq!( 88 /// v_str, 89 /// b"--abcde\r\nkey3:value3\r\n\r\n33333\r\n--abcde\r\n\ 90 /// content-type:multipart/mixed; boundary=abc\r\n\r\n--abc\r\n\ 91 /// key1:value1\r\n\r\n111111\r\n--abc\r\nkey2:value2\r\n\r\n22222\r\n\ 92 /// --abc--\r\n\r\n--abcde\r\nkey4:value4\r\n\r\n\r\n--abcde--\r\n" 93 /// ); 94 /// ``` 95 #[derive(Debug)] 96 pub struct MimeMultiEncoder<'a> { 97 stages: VecDeque<MultiStage<'a>>, 98 // default is not part, don't encode headers 99 headers: Option<Headers>, 100 src_idx: usize, 101 src: Vec<u8>, 102 } 103 104 impl<'a> MimeMultiEncoder<'a> { 105 /// Creates a `MimeMultiEncoder` by a [`MimeMulti`]. 106 /// 107 /// [`MimeMulti`]: MimeMulti 108 /// 109 /// # Examples 110 /// 111 /// ``` 112 /// use ylong_http::body::{MimeMulti, MimeMultiEncoder}; 113 /// 114 /// let multi = MimeMulti::builder().build().unwrap(); 115 /// let multi_encoder = MimeMultiEncoder::from_multi(multi); 116 /// ``` from_multi(multi: MimeMulti<'a>) -> Self117 pub fn from_multi(multi: MimeMulti<'a>) -> Self { 118 let mut encoder = Self::new(); 119 encoder.headers = Some(multi.headers); 120 encoder.push_list_to_stages(multi.boundary, multi.list); 121 encoder 122 } 123 new() -> Self124 fn new() -> Self { 125 MimeMultiEncoder { 126 stages: VecDeque::new(), 127 headers: None, 128 src: vec![], 129 src_idx: 0, 130 } 131 } 132 push_list_to_stages(&mut self, boundary: Vec<u8>, list: Vec<XPart<'a>>)133 fn push_list_to_stages(&mut self, boundary: Vec<u8>, list: Vec<XPart<'a>>) { 134 // dash-boundary := "--" boundary 135 self.stages.push_back(MultiStage::Dash); 136 self.stages 137 .push_back(MultiStage::Boundary(boundary.clone())); 138 self.stages.push_back(MultiStage::Crlf); 139 140 if list.is_empty() { 141 // close-delimiter CRLF "--" boundary "--" CRLF 142 self.push_end_boundary(boundary); 143 return; 144 } 145 146 let len = list.len(); 147 148 for (idx, xpart) in list.into_iter().enumerate() { 149 match xpart { 150 XPart::Part(part) => { 151 let part_encoder = MimePartEncoder::from_part(part); 152 self.stages.push_back(MultiStage::MimePart(part_encoder)); 153 } 154 XPart::Multi(multi) => self.push_multi_as_part(multi), 155 } 156 if idx == len - 1 { 157 self.push_end_boundary(boundary.clone()); 158 } else { 159 self.push_middle_boundary(boundary.clone()); 160 } 161 } 162 } 163 push_multi_as_part(&mut self, multi: MimeMulti<'a>)164 fn push_multi_as_part(&mut self, multi: MimeMulti<'a>) { 165 let headers_encoder = EncodeHeaders::new(multi.headers); 166 self.stages.push_back(MultiStage::Headers(headers_encoder)); 167 self.stages.push_back(MultiStage::Crlf); 168 self.push_list_to_stages(multi.boundary, multi.list); 169 } 170 171 // "--" boundary CRLF push_first_boundary(&mut self, boundary: Vec<u8>)172 fn push_first_boundary(&mut self, boundary: Vec<u8>) { 173 // dash-boundary := "--" boundary 174 self.stages.push_back(MultiStage::Dash); 175 self.stages.push_back(MultiStage::Boundary(boundary)); 176 // end CRLF 177 self.stages.push_back(MultiStage::Crlf); 178 } 179 180 // CRLF "--" boundary CRLF push_middle_boundary(&mut self, boundary: Vec<u8>)181 fn push_middle_boundary(&mut self, boundary: Vec<u8>) { 182 // delimiter := CRLF dash-boundary 183 self.stages.push_back(MultiStage::Crlf); 184 self.push_first_boundary(boundary); 185 } 186 187 // CRLF "--" boundary "--" CRLF push_end_boundary(&mut self, boundary: Vec<u8>)188 fn push_end_boundary(&mut self, boundary: Vec<u8>) { 189 // close-delimiter: = CRLF "--" boundary "--" 190 self.stages.push_back(MultiStage::Crlf); 191 self.stages.push_back(MultiStage::Dash); 192 self.stages.push_back(MultiStage::Boundary(boundary)); 193 self.stages.push_back(MultiStage::Dash); 194 // end CRLF 195 self.stages.push_back(MultiStage::Crlf); 196 } 197 init_src(&mut self)198 fn init_src(&mut self) { 199 self.src = vec![]; 200 self.src_idx = 0; 201 } 202 temp_src(&mut self, stage: &MultiStage)203 fn temp_src(&mut self, stage: &MultiStage) { 204 match stage { 205 MultiStage::Crlf => self.src = b"\r\n".to_vec(), 206 MultiStage::Dash => self.src = b"--".to_vec(), 207 _ => {} 208 } 209 } 210 211 //"\r\n" crlf_encode(&mut self, dst: &mut [u8]) -> SizeResult212 fn crlf_encode(&mut self, dst: &mut [u8]) -> SizeResult { 213 match data_copy(&self.src, &mut self.src_idx, dst)? { 214 TokenStatus::Partial(size) => { 215 self.stages.push_front(MultiStage::Crlf); 216 Ok(size) 217 } 218 TokenStatus::Complete(size) => { 219 self.init_src(); 220 Ok(size) 221 } 222 } 223 } 224 225 // "--" dash_encode(&mut self, dst: &mut [u8]) -> SizeResult226 fn dash_encode(&mut self, dst: &mut [u8]) -> SizeResult { 227 match data_copy(&self.src, &mut self.src_idx, dst)? { 228 TokenStatus::Partial(size) => { 229 self.stages.push_front(MultiStage::Dash); 230 Ok(size) 231 } 232 TokenStatus::Complete(size) => { 233 self.init_src(); 234 Ok(size) 235 } 236 } 237 } 238 boundary_encode(&mut self, dst: &mut [u8], boundary: Vec<u8>) -> SizeResult239 fn boundary_encode(&mut self, dst: &mut [u8], boundary: Vec<u8>) -> SizeResult { 240 match data_copy(&boundary, &mut self.src_idx, dst)? { 241 TokenStatus::Partial(size) => { 242 self.stages.push_front(MultiStage::Boundary(boundary)); 243 Ok(size) 244 } 245 TokenStatus::Complete(size) => { 246 self.init_src(); 247 Ok(size) 248 } 249 } 250 } 251 headers_encode(&mut self, dst: &mut [u8], mut headers_encoder: EncodeHeaders) -> SizeResult252 fn headers_encode(&mut self, dst: &mut [u8], mut headers_encoder: EncodeHeaders) -> SizeResult { 253 match headers_encoder.encode(dst)? { 254 TokenStatus::Partial(size) => { 255 self.stages.push_front(MultiStage::Headers(headers_encoder)); 256 Ok(size) 257 } 258 TokenStatus::Complete(size) => { 259 self.init_src(); 260 Ok(size) 261 } 262 } 263 } 264 sync_mimepart_encode( &mut self, dst: &mut [u8], mut part_encoder: MimePartEncoder<'a>, ) -> SizeResult265 fn sync_mimepart_encode( 266 &mut self, 267 dst: &mut [u8], 268 mut part_encoder: MimePartEncoder<'a>, 269 ) -> SizeResult { 270 let size = sync_impl::Body::data(&mut part_encoder, dst)?; 271 if size != 0 { 272 self.stages.push_front(MultiStage::MimePart(part_encoder)); 273 } else { 274 self.init_src(); 275 } 276 Ok(size) 277 } 278 } 279 280 impl sync_impl::Body for MimeMultiEncoder<'_> { 281 type Error = std::io::Error; 282 data(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error>283 fn data(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> { 284 let mut count = 0; 285 while count != buf.len() { 286 let stage = match self.stages.pop_front() { 287 Some(stage) => stage, 288 None => break, 289 }; 290 self.temp_src(&stage); 291 let size = match stage { 292 MultiStage::Crlf => self.crlf_encode(&mut buf[count..]), 293 MultiStage::Dash => self.dash_encode(&mut buf[count..]), 294 MultiStage::Boundary(boundary) => self.boundary_encode(&mut buf[count..], boundary), 295 MultiStage::Headers(headers_encoder) => { 296 self.headers_encode(&mut buf[count..], headers_encoder) 297 } 298 MultiStage::MimePart(part_encoder) => { 299 self.sync_mimepart_encode(&mut buf[count..], part_encoder) 300 } 301 }?; 302 count += size; 303 } 304 Ok(count) 305 } 306 } 307 308 impl async_impl::Body for MimeMultiEncoder<'_> { 309 type Error = std::io::Error; 310 poll_data( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8], ) -> Poll<Result<usize, Self::Error>>311 fn poll_data( 312 mut self: Pin<&mut Self>, 313 cx: &mut Context<'_>, 314 buf: &mut [u8], 315 ) -> Poll<Result<usize, Self::Error>> { 316 let mut count = 0; 317 while count != buf.len() { 318 let stage = match self.stages.pop_front() { 319 Some(stage) => stage, 320 None => break, 321 }; 322 self.temp_src(&stage); 323 let poll_size: Poll<SizeResult> = match stage { 324 MultiStage::Crlf => Poll::Ready(self.crlf_encode(&mut buf[count..])), 325 MultiStage::Dash => Poll::Ready(self.dash_encode(&mut buf[count..])), 326 MultiStage::Boundary(boundary) => { 327 Poll::Ready(self.boundary_encode(&mut buf[count..], boundary)) 328 } 329 MultiStage::Headers(headers_encoder) => { 330 Poll::Ready(self.headers_encode(&mut buf[count..], headers_encoder)) 331 } 332 MultiStage::MimePart(mut part_encoder) => { 333 let poll_result = Pin::new(&mut part_encoder).poll_data(cx, &mut buf[count..]); 334 if let Poll::Ready(Ok(0)) = poll_result { 335 // complete async read body 336 self.init_src(); 337 } else if let Poll::Ready(Ok(_)) = poll_result { 338 self.stages.push_front(MultiStage::MimePart(part_encoder)); 339 } 340 poll_result 341 } 342 }; 343 344 match poll_size { 345 Poll::Ready(Ok(size)) => count += size, 346 Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), 347 Poll::Pending => return Poll::Pending, 348 } 349 } 350 Poll::Ready(Ok(count)) 351 } 352 } 353 354 #[derive(Debug)] 355 enum MultiStage<'a> { 356 // \r\n 357 Crlf, 358 // -- 359 Dash, 360 // boundary 361 Boundary(Vec<u8>), 362 // headers 363 Headers(EncodeHeaders), 364 // part 365 MimePart(MimePartEncoder<'a>), 366 } 367 368 #[cfg(test)] 369 mod ut_mime_multi_encoder { 370 use crate::body::{async_impl, sync_impl, MimeMulti, MimeMultiEncoder, MimePart}; 371 372 // TODO: Uses a commonly macro like `part_encode_compare` to replace this one. 373 macro_rules! encode_and_compare { 374 ( $buf: expr, $encode: block, $v_size: expr, $v_str: expr) => { 375 let mut v_size = vec![]; 376 let mut v_str = vec![]; 377 378 loop { 379 let len = $encode; 380 if len == 0 { 381 break; 382 } 383 v_size.push(len); 384 v_str.extend_from_slice(&$buf[..len]); 385 } 386 assert_eq!(v_size, $v_size); 387 assert_eq!(v_str, $v_str); 388 }; 389 } 390 391 /// UT test cases for `syn::Body::data` of `MimeMultiEncoder`. 392 /// 393 /// # Brief 394 /// 1. Creates a `MimeMulti`. 395 /// 2. Builds a `MimeMultiEncoder` by `from_multi` and encodes. 396 /// 3. Checks whether the result is correct. 397 #[test] ut_mime_multi_encoder_from_multi()398 fn ut_mime_multi_encoder_from_multi() { 399 let multi = MimeMulti::builder().build().unwrap(); 400 let mut multi_encoder = MimeMultiEncoder::from_multi(multi); 401 let mut buf = vec![0u8; 10]; 402 encode_and_compare!( 403 buf, 404 { sync_impl::Body::data(&mut multi_encoder, &mut buf).unwrap() }, 405 vec![10, 4], 406 b"---\r\n\r\n-----\r\n" 407 ); 408 } 409 410 /// UT test cases for `syn::Body::data` of `MimeMultiEncoder`. 411 /// 412 /// # Brief 413 /// 1. Creates a `MimeMulti`. 414 /// 2. Sets headers. 415 /// 3. Builds a `MimeMultiEncoder` by `from_multi` and encodes. 416 /// 4. Checks whether the result is correct. 417 #[test] ut_mime_multi_encoder_data_one_part()418 fn ut_mime_multi_encoder_data_one_part() { 419 let part = MimePart::builder() 420 .header("key1", "value1") 421 .body_from_reader("9876543210\r\n".as_bytes()) 422 .build() 423 .unwrap(); 424 let multi = MimeMulti::builder().add_part(part).build().unwrap(); 425 let mut multi_encoder = MimeMultiEncoder::from_multi(multi); 426 let mut buf = vec![0u8; 10]; 427 encode_and_compare!( 428 buf, 429 { sync_impl::Body::data(&mut multi_encoder, &mut buf).unwrap() }, 430 vec![10, 10, 10, 10, 1], 431 b"---\r\nkey1:value1\r\n\r\n9876543210\r\n\r\n-----\r\n" 432 ); 433 } 434 435 /// UT test cases for `syn::Body::data` of `MimeMultiEncoder`. 436 /// 437 /// # Brief 438 /// 1. Creates a `MimeMulti`. 439 /// 2. Creates several `MimePart`, sets headers, sets body, builds. 440 /// 3. Adds `MimePart` to `MimeMulti` by `add_part`. 441 /// 4. Builds a `MimeMultiEncoder` by `from_multi` and encodes 442 /// synchronously. 443 /// 5. Checks whether the result is correct. 444 #[test] ut_mime_multi_encoder_data_many_parts()445 fn ut_mime_multi_encoder_data_many_parts() { 446 let multi = MimeMulti::builder() 447 .add_part( 448 MimePart::builder() 449 .header("key1", "value1") 450 .body_from_reader("98765432\r\n".as_bytes()) 451 .build() 452 .unwrap(), 453 ) 454 .add_part( 455 MimePart::builder() 456 .header("key2", "value2") 457 .body_from_reader("22222".as_bytes()) 458 .build() 459 .unwrap(), 460 ) 461 .build() 462 .unwrap(); 463 let mut multi_encoder = MimeMultiEncoder::from_multi(multi); 464 let mut buf = vec![0u8; 10]; 465 encode_and_compare!( 466 buf, 467 { sync_impl::Body::data(&mut multi_encoder, &mut buf).unwrap() }, 468 vec![10, 10, 10, 10, 10, 10, 6], 469 b"---\r\nkey1:value1\r\n\r\n98765432\r\n\r\n---\r\nkey2:value2\r\n\r\n22222\r\n-----\r\n" 470 ); 471 } 472 473 /// UT test cases for `syn::Body::data` of `MimeMultiEncoder`. 474 /// 475 /// # Brief 476 /// 1. Creates a `MimeMulti`. 477 /// 2. Creates several `MimePart`, sets headers, sets body, builds. 478 /// 3. Creates a main `MimeMulti`, adds parts. 479 /// 4. Builds a `MimeMultiEncoder` by `from_multi` and encodes 480 /// synchronously. 481 /// 5. Checks whether the result is correct. 482 #[test] ut_mime_multi_encoder_data_many_parts_nesting()483 fn ut_mime_multi_encoder_data_many_parts_nesting() { 484 let multi1 = MimeMulti::builder() 485 .set_content_type(b"multipart/mixed", b"abc".to_vec()) 486 .add_part( 487 MimePart::builder() 488 .header("key1", "value1") 489 .body_from_reader("111111".as_bytes()) 490 .build() 491 .unwrap(), 492 ) 493 .add_part( 494 MimePart::builder() 495 .header("key2", "value2") 496 .body_from_reader("22222".as_bytes()) 497 .build() 498 .unwrap(), 499 ) 500 .build() 501 .unwrap(); 502 503 let multi = MimeMulti::builder() 504 .set_boundary(b"abcde".to_vec()) 505 .add_part( 506 MimePart::builder() 507 .header("key3", "value3") 508 .body_from_reader("33333".as_bytes()) 509 .build() 510 .unwrap(), 511 ) 512 .add_multi(multi1) 513 .add_part( 514 MimePart::builder() 515 .header("key4", "value4") 516 .body_from_async_reader("444444".as_bytes()) // nothing in sync 517 .build() 518 .unwrap(), 519 ) 520 .build() 521 .unwrap(); 522 523 let mut multi_encoder = MimeMultiEncoder::from_multi(multi); 524 let mut buf = vec![0u8; 30]; 525 encode_and_compare!( 526 buf, 527 { sync_impl::Body::data(&mut multi_encoder, &mut buf).unwrap() }, 528 vec![30, 30, 30, 30, 30, 30, 13], 529 b"--abcde\r\nkey3:value3\r\n\r\n33333\r\n--abcde\r\n\ 530 content-type:multipart/mixed; boundary=abc\r\n\r\n--abc\r\n\ 531 key1:value1\r\n\r\n111111\r\n--abc\r\nkey2:value2\r\n\r\n22222\r\n\ 532 --abc--\r\n\r\n--abcde\r\nkey4:value4\r\n\r\n\r\n--abcde--\r\n" 533 ); 534 } 535 536 /// UT test cases for `asyn::Body::data` of `MimeMultiEncoder`. 537 /// 538 /// # Brief 539 /// 1. Creates a `MimeMulti`. 540 /// 2. Creates several `MimePart`, sets headers, sets body, builds. 541 /// 3. Creates a main `MimeMulti`, adds parts. 542 /// 4. Builds a `MimeMultiEncoder` by `from_multi` and encodes 543 /// asynchronously. 544 /// 5. Checks whether the result is correct. 545 #[cfg(feature = "ylong_base")] 546 #[test] ut_mime_multi_encoder_data_many_parts_nesting_then_async_data()547 fn ut_mime_multi_encoder_data_many_parts_nesting_then_async_data() { 548 let handle = ylong_runtime::spawn(async move { 549 mime_multi_encoder_data_many_parts_nesting_then_async_data().await; 550 }); 551 ylong_runtime::block_on(handle).unwrap(); 552 } 553 554 #[cfg(feature = "ylong_base")] mime_multi_encoder_data_many_parts_nesting_then_async_data()555 async fn mime_multi_encoder_data_many_parts_nesting_then_async_data() { 556 let multi1 = MimeMulti::builder() 557 .set_content_type(b"multipart/mixed", b"abc".to_vec()) 558 .add_part( 559 MimePart::builder() 560 .header("key1", "value1") 561 .body_from_async_reader("111111".as_bytes()) 562 .build() 563 .unwrap(), 564 ) 565 .add_part( 566 MimePart::builder() 567 .header("key2", "value2") 568 .body_from_async_reader("22222".as_bytes()) 569 .build() 570 .unwrap(), 571 ) 572 .build() 573 .unwrap(); 574 575 let multi = MimeMulti::builder() 576 .set_boundary(b"abcde".to_vec()) 577 .add_part( 578 MimePart::builder() 579 .header("key3", "value3") 580 .body_from_bytes("33333".as_bytes()) 581 .build() 582 .unwrap(), 583 ) 584 .add_multi(multi1) 585 .add_part( 586 MimePart::builder() 587 .header("key4", "value4") 588 .body_from_reader("444444".as_bytes()) // nothing in async 589 .build() 590 .unwrap(), 591 ) 592 .build() 593 .unwrap(); 594 595 let mut multi_encoder = MimeMultiEncoder::from_multi(multi); 596 let mut buf = vec![0u8; 30]; 597 encode_and_compare!( 598 buf, 599 { 600 async_impl::Body::data(&mut multi_encoder, &mut buf) 601 .await 602 .unwrap() 603 }, 604 vec![30, 30, 30, 30, 30, 30, 13], 605 b"--abcde\r\nkey3:value3\r\n\r\n33333\r\n--abcde\r\n\ 606 content-type:multipart/mixed; boundary=abc\r\n\r\n--abc\r\n\ 607 key1:value1\r\n\r\n111111\r\n--abc\r\nkey2:value2\r\n\r\n22222\r\n\ 608 --abc--\r\n\r\n--abcde\r\nkey4:value4\r\n\r\n\r\n--abcde--\r\n" 609 ); 610 } 611 } 612