1use crate::optimizer::{
16 ImageError, load_image, optimize_avif, optimize_jpeg, optimize_png,
17 optimize_webp,
18};
19use async_trait::async_trait;
20use bytes::{Bytes, BytesMut};
21use ctor::ctor;
22use pingap_config::PluginConf;
23use pingap_core::HTTP_HEADER_TRANSFER_CHUNKED;
24use pingap_core::{
25 Ctx, Plugin, PluginStep, RequestPluginResult, ResponsePluginResult,
26};
27use pingap_core::{ModifyResponseBody, ResponseBodyPluginResult};
28use pingap_plugin::{
29 Error, get_hash_key, get_int_conf, get_plugin_factory, get_str_conf,
30};
31use pingora::http::ResponseHeader;
32use pingora::proxy::Session;
33use std::borrow::Cow;
34use std::collections::HashSet;
35use std::convert::TryFrom;
36use std::sync::Arc;
37use tracing::{debug, warn};
38
39mod optimizer;
40mod quantize;
41
42const PLUGIN_ID: &str = "_image_optimize_";
43
44const MAX_IMAGE_SIZE: usize = 20 * 1024 * 1024;
47
48const OUTPUT_FORMATS: &[&str] = &["avif", "webp", "jpeg", "png"];
50
51type Result<T, E = Error> = std::result::Result<T, E>;
52
53struct ImageOptimizer {
54 image_type: String,
56 png_quality: u8,
57 jpeg_quality: u8,
58 avif_quality: u8,
59 avif_speed: u8,
60 webp_quality: u8,
61 format_type: String,
63 buffer: BytesMut,
64 passthrough: bool,
67}
68
69fn run_blocking<T>(f: impl FnOnce() -> T) -> T {
78 use tokio::runtime::{Handle, RuntimeFlavor};
79 match Handle::try_current() {
80 Ok(handle) if handle.runtime_flavor() == RuntimeFlavor::MultiThread => {
81 tokio::task::block_in_place(f)
82 },
83 _ => f(),
84 }
85}
86
87impl ImageOptimizer {
88 fn optimize(&self, data: &[u8]) -> Result<Vec<u8>, ImageError> {
89 let info = load_image(data, &self.image_type)?;
90 match self.format_type.as_str() {
91 "jpeg" => optimize_jpeg(&info, self.jpeg_quality),
92 "avif" => optimize_avif(&info, self.avif_quality, self.avif_speed),
93 "webp" => optimize_webp(&info, self.webp_quality),
94 _ => optimize_png(info, self.png_quality),
95 }
96 }
97}
98
99fn worth_sending(
105 original: Bytes,
106 converted: Vec<u8>,
107 same_format: bool,
108) -> Bytes {
109 if same_format && converted.len() >= original.len() {
110 return original;
111 }
112 Bytes::from(converted)
113}
114
115impl ModifyResponseBody for ImageOptimizer {
116 fn handle(
117 &mut self,
118 _session: &Session,
119 body: &mut Option<bytes::Bytes>,
120 end_of_stream: bool,
121 ) -> pingora::Result<()> {
122 if self.passthrough {
123 return Ok(());
124 }
125 if let Some(data) = body {
126 if self.buffer.len() + data.len() > MAX_IMAGE_SIZE {
130 self.passthrough = true;
131 self.buffer.extend_from_slice(data);
132 *data = self.buffer.split().freeze();
133 return Ok(());
134 }
135 self.buffer.extend_from_slice(data);
136 data.clear();
137 }
138 if !end_of_stream {
139 return Ok(());
140 }
141 if self.buffer.is_empty() {
143 return Ok(());
144 }
145 let original = self.buffer.split().freeze();
146 let optimized = run_blocking(|| self.optimize(&original));
147 let same_format = self.format_type == self.image_type;
148 *body = Some(match optimized {
149 Ok(data) => worth_sending(original, data, same_format),
150 Err(e) => {
153 warn!(
154 error = %e,
155 image_type = self.image_type,
156 format_type = self.format_type,
157 "optimize image fail, the original is sent"
158 );
159 original
160 },
161 });
162 Ok(())
163 }
164 fn name(&self) -> &str {
165 "image_optimization"
166 }
167}
168
169pub struct ImageOptim {
170 handler_id: String,
173 hash_value: String,
176 support_types: HashSet<String>,
177 output_types: Vec<String>,
179 output_mimes: Vec<String>,
181 key_alternatives: Arc<Vec<Vec<String>>>,
184 png_quality: u8,
185 jpeg_quality: u8,
186 avif_quality: u8,
187 avif_speed: u8,
188}
189
190fn selections(mimes: &[String]) -> Vec<Vec<String>> {
192 let mut sorted = mimes.to_vec();
193 sorted.sort();
194 sorted.dedup();
195 (0..1usize << sorted.len())
196 .map(|picked| {
197 sorted
198 .iter()
199 .enumerate()
200 .filter(|(index, _)| picked & (1 << index) != 0)
201 .map(|(_, mime)| mime.clone())
202 .collect()
203 })
204 .collect()
205}
206
207impl TryFrom<&PluginConf> for ImageOptim {
208 type Error = Error;
209 fn try_from(value: &PluginConf) -> Result<Self> {
210 debug!(params = value.to_string(), "new image optimizer plugin");
211 let hash_value = get_hash_key(value);
212
213 let output_types: Vec<String> = get_str_conf(value, "output_types")
214 .split(',')
215 .map(|s| s.trim())
216 .filter(|s| !s.is_empty())
217 .map(|s| s.to_string())
218 .collect();
219
220 if let Some(format) = output_types
224 .iter()
225 .find(|format| !OUTPUT_FORMATS.contains(&format.as_str()))
226 {
227 return Err(Error::Invalid {
228 category: "image_optim".to_string(),
229 message: format!(
230 "output type {format} is not supported, expected one of {}",
231 OUTPUT_FORMATS.join(", ")
232 ),
233 });
234 }
235 let output_mimes: Vec<String> = output_types
236 .iter()
237 .map(|format| format!("image/{}", format))
238 .collect();
239
240 let level = |key: &str, max: i64, default: u8| -> Result<u8, Error> {
243 match get_int_conf(value, key) {
244 0 => Ok(default),
245 level if (1..=max).contains(&level) => Ok(level as u8),
246 level => Err(Error::Invalid {
247 category: "image_optim".to_string(),
248 message: format!(
249 "{key} should be between 1 and {max}, got {level}"
250 ),
251 }),
252 }
253 };
254 let png_quality = level("png_quality", 100, 90)?;
255 let jpeg_quality = level("jpeg_quality", 100, 80)?;
256 let avif_quality = level("avif_quality", 100, 75)?;
257 let avif_speed = level("avif_speed", 10, 3)?;
258 Ok(Self {
259 hash_value,
260 handler_id: pingap_plugin::new_body_handler_id(PLUGIN_ID),
261 support_types: HashSet::from([
262 "jpeg".to_string(),
263 "png".to_string(),
264 ]),
265 output_types,
266 key_alternatives: Arc::new(selections(&output_mimes)),
267 output_mimes,
268 png_quality,
269 jpeg_quality,
270 avif_quality,
271 avif_speed,
272 })
273 }
274}
275
276impl ImageOptim {
277 pub fn new(params: &PluginConf) -> Result<Self> {
278 Self::try_from(params)
279 }
280}
281
282#[async_trait]
283impl Plugin for ImageOptim {
284 fn config_key(&self) -> Cow<'_, str> {
286 Cow::Borrowed(&self.hash_value)
287 }
288 async fn handle_request(
289 &self,
290 step: PluginStep,
291 session: &mut Session,
292 ctx: &mut Ctx,
293 ) -> pingora::Result<RequestPluginResult> {
294 if step != PluginStep::EarlyRequest {
300 return Ok(RequestPluginResult::Skipped);
301 }
302
303 let accept = session
304 .get_header(http::header::ACCEPT)
305 .and_then(|accept| accept.to_str().ok())
306 .unwrap_or_default();
307 let mut accept_images: Vec<_> = self
308 .output_mimes
309 .iter()
310 .filter(|mime| accept.contains(*mime))
311 .cloned()
312 .collect();
313 accept_images.sort();
314 accept_images.dedup();
316 ctx.push_cache_key_variant(
317 accept_images,
318 self.key_alternatives.clone(),
319 );
320 Ok(RequestPluginResult::Continue)
321 }
322 fn handle_upstream_response(
323 &self,
324 session: &mut Session,
325 ctx: &mut Ctx,
326 upstream_response: &mut ResponseHeader,
327 ) -> pingora::Result<ResponsePluginResult> {
328 if upstream_response.status != http::StatusCode::OK
332 || upstream_response
333 .headers
334 .contains_key(http::header::CONTENT_ENCODING)
335 {
336 return Ok(ResponsePluginResult::Unchanged);
337 }
338 let content_type = if let Some(value) =
339 upstream_response.headers.get(http::header::CONTENT_TYPE)
340 {
341 value.to_str().unwrap_or_default()
342 } else {
343 return Ok(ResponsePluginResult::Unchanged);
344 };
345
346 let media_type = content_type.split(';').next().unwrap_or_default();
348 let Some(image_type) = media_type.trim().strip_prefix("image/") else {
349 return Ok(ResponsePluginResult::Unchanged);
350 };
351
352 if !self.support_types.contains(image_type) {
353 return Ok(ResponsePluginResult::Unchanged);
354 }
355
356 let Some(accept) = session.get_header(http::header::ACCEPT) else {
357 return Ok(ResponsePluginResult::Unchanged);
358 };
359 let Ok(accept_str) = accept.to_str() else {
360 return Ok(ResponsePluginResult::Unchanged);
361 };
362
363 let content_length = upstream_response
367 .headers
368 .get(http::header::CONTENT_LENGTH)
369 .and_then(|value| value.to_str().ok())
370 .and_then(|value| value.parse::<usize>().ok());
371 if content_length.is_some_and(|size| size > MAX_IMAGE_SIZE) {
372 return Ok(ResponsePluginResult::Unchanged);
373 }
374
375 let image_type = image_type.to_string();
376 let format_type = self
382 .output_types
383 .iter()
384 .zip(self.output_mimes.iter())
385 .find(|(_, mime)| accept_str.contains(mime.as_str()))
386 .map(|(format, _)| format.clone())
387 .unwrap_or_else(|| image_type.clone());
388 upstream_response.remove_header(&http::header::CONTENT_LENGTH);
390 upstream_response.remove_header(&http::header::ACCEPT_RANGES);
394 let _ = upstream_response.insert_header(
396 http::header::TRANSFER_ENCODING,
397 HTTP_HEADER_TRANSFER_CHUNKED.1.clone(),
398 );
399 let _ = upstream_response.insert_header(
400 http::header::CONTENT_TYPE,
401 format!("image/{format_type}"),
402 );
403 let capacity = content_length.unwrap_or(8192);
404
405 ctx.add_modify_body_handler(
406 &self.handler_id,
407 Box::new(ImageOptimizer {
408 image_type,
409 png_quality: self.png_quality,
410 jpeg_quality: self.jpeg_quality,
411 avif_quality: self.avif_quality,
412 avif_speed: self.avif_speed,
413 webp_quality: 100,
415 format_type,
416 buffer: BytesMut::with_capacity(capacity),
417 passthrough: false,
418 }),
419 );
420 Ok(ResponsePluginResult::Modified)
421 }
422 async fn handle_response(
432 &self,
433 _session: &mut Session,
434 _ctx: &mut Ctx,
435 upstream_response: &mut ResponseHeader,
436 ) -> pingora::Result<ResponsePluginResult> {
437 let ours = upstream_response
438 .headers
439 .get(http::header::CONTENT_TYPE)
440 .and_then(|value| value.to_str().ok())
441 .and_then(|value| value.split(';').next())
442 .and_then(|value| value.trim().strip_prefix("image/"))
443 .is_some_and(|kind| {
444 self.support_types.contains(kind)
445 || self.output_types.iter().any(|output| output == kind)
446 });
447 if ours && pingap_plugin::weaken_etag(upstream_response) {
448 return Ok(ResponsePluginResult::Modified);
449 }
450 Ok(ResponsePluginResult::Unchanged)
451 }
452 fn handle_upstream_response_body(
453 &self,
454 session: &mut Session,
455 ctx: &mut Ctx,
456 body: &mut Option<bytes::Bytes>,
457 end_of_stream: bool,
458 ) -> pingora::Result<ResponseBodyPluginResult> {
459 if let Some(modifier) = ctx.get_modify_body_handler(&self.handler_id) {
460 modifier.handle(session, body, end_of_stream)?;
461 let result = if end_of_stream {
462 ResponseBodyPluginResult::FullyReplaced
463 } else {
464 ResponseBodyPluginResult::PartialReplaced
465 };
466 Ok(result)
467 } else {
468 Ok(ResponseBodyPluginResult::Unchanged)
469 }
470 }
471}
472
473#[ctor(unsafe)]
474fn init() {
475 get_plugin_factory().register("image_optim", |params| {
476 Ok(Arc::new(ImageOptim::new(params)?))
477 });
478}
479
480#[cfg(test)]
481mod tests {
482 use super::*;
483 use pingap_config::PluginConf;
484 use pingap_core::{Ctx, Plugin};
485 use pingora::modules::http::HttpModules;
486 use pretty_assertions::assert_eq;
487 use tokio_test::io::Builder;
488
489 #[test]
490 fn test_new_image_optimize() {
491 let optim = ImageOptim::try_from(
492 &toml::from_str::<PluginConf>(
493 r###"
494avif_quality = 75
495avif_speed = 3
496category = "image_optim"
497jpeg_quality = 80
498output_types = "avif,webp"
499png_quality = 90
500"###,
501 )
502 .unwrap(),
503 )
504 .unwrap();
505
506 assert_eq!(
507 HashSet::from(["jpeg".to_string(), "png".to_string(),]),
508 optim.support_types
509 );
510 assert_eq!(
511 vec!["image/avif".to_string(), "image/webp".to_string()],
512 optim.output_mimes
513 );
514 assert_eq!(90, optim.png_quality);
515 assert_eq!(80, optim.jpeg_quality);
516 assert_eq!(75, optim.avif_quality);
517 assert_eq!(3, optim.avif_speed);
518 }
519
520 #[tokio::test]
523 async fn test_image_validator_is_weak_on_the_way_out() {
524 let optim = new_optim("avif,webp");
525 let sent = async |content_type: &str| {
526 let mut session = new_session("").await;
527 let mut resp = ResponseHeader::build(200, None).unwrap();
528 resp.insert_header("Content-Type", content_type).unwrap();
529 resp.insert_header("ETag", "\"v1\"").unwrap();
530 optim
531 .handle_response(&mut session, &mut Ctx::default(), &mut resp)
532 .await
533 .unwrap();
534 resp.headers
535 .get("ETag")
536 .unwrap()
537 .to_str()
538 .unwrap()
539 .to_string()
540 };
541 assert_eq!("W/\"v1\"", sent("image/webp").await);
542 assert_eq!("W/\"v1\"", sent("image/jpeg; q=1").await);
543 assert_eq!("\"v1\"", sent("image/svg+xml").await);
545 assert_eq!("\"v1\"", sent("text/html").await);
546 }
547
548 #[test]
551 fn test_levels_out_of_range_are_rejected() {
552 let build = |conf: &str| {
553 ImageOptim::try_from(&toml::from_str::<PluginConf>(conf).unwrap())
554 };
555 for conf in [
556 "png_quality = 300",
557 "png_quality = 101",
558 "jpeg_quality = 256",
559 "avif_quality = -1",
560 "avif_speed = 11",
561 "avif_speed = 266",
562 ] {
563 let message =
564 build(conf).err().map(|e| e.to_string()).unwrap_or_default();
565 assert_eq!(true, message.contains("should be between"), "{conf}");
566 }
567 let optim = build("jpeg_quality = 0").unwrap();
569 assert_eq!(
570 (90, 80, 75, 3),
571 (
572 optim.png_quality,
573 optim.jpeg_quality,
574 optim.avif_quality,
575 optim.avif_speed
576 )
577 );
578 let optim =
579 build("png_quality = 1\njpeg_quality = 100\navif_speed = 10")
580 .unwrap();
581 assert_eq!(
582 (1, 100, 10),
583 (optim.png_quality, optim.jpeg_quality, optim.avif_speed)
584 );
585 }
586
587 #[tokio::test]
588 async fn test_image_optimize_handle_request() {
589 let optim = ImageOptim::try_from(
590 &toml::from_str::<PluginConf>(
591 r###"
592avif_quality = 75
593avif_speed = 3
594category = "image_optim"
595jpeg_quality = 80
596output_types = "avif,webp"
597png_quality = 90
598"###,
599 )
600 .unwrap(),
601 )
602 .unwrap();
603 {
605 let headers = [""].join("\r\n");
606 let input_header = format!(
607 "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
608 );
609 let mock_io = Builder::new().read(input_header.as_bytes()).build();
610 let mut session = Session::new_h1_with_modules(
611 Box::new(mock_io),
612 &HttpModules::new(),
613 );
614 session.read_request().await.unwrap();
615 let mut ctx = Ctx::default();
616
617 let result = optim
618 .handle_request(
619 PluginStep::EarlyRequest,
620 &mut session,
621 &mut ctx,
622 )
623 .await
624 .unwrap();
625
626 assert_eq!(true, RequestPluginResult::Continue == result);
627 let cache = ctx.cache.unwrap();
630 assert_eq!(Some(vec![]), cache.keys);
631 let variants = cache.key_variants.unwrap();
632 assert_eq!((0, 0), (variants[0].at, variants[0].len));
633 assert_eq!(
634 vec![
635 vec![],
636 vec!["image/avif".to_string()],
637 vec!["image/webp".to_string()],
638 vec!["image/avif".to_string(), "image/webp".to_string()],
639 ],
640 *variants[0].alternatives
641 );
642 }
643
644 {
646 let headers = ["Accept: image/avif"].join("\r\n");
647 let input_header = format!(
648 "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
649 );
650 let mock_io = Builder::new().read(input_header.as_bytes()).build();
651 let mut session = Session::new_h1_with_modules(
652 Box::new(mock_io),
653 &HttpModules::new(),
654 );
655 session.read_request().await.unwrap();
656 let mut ctx = Ctx::default();
657
658 let result = optim
659 .handle_request(
660 PluginStep::EarlyRequest,
661 &mut session,
662 &mut ctx,
663 )
664 .await
665 .unwrap();
666
667 assert_eq!(true, RequestPluginResult::Continue == result);
668 assert_eq!(
669 vec!["image/avif".to_string()],
670 ctx.cache.unwrap().keys.unwrap()
671 );
672 }
673
674 {
676 let headers = ["Accept: image/webp, image/avif"].join("\r\n");
677 let input_header = format!(
678 "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
679 );
680 let mock_io = Builder::new().read(input_header.as_bytes()).build();
681 let mut session = Session::new_h1_with_modules(
682 Box::new(mock_io),
683 &HttpModules::new(),
684 );
685 session.read_request().await.unwrap();
686 let mut ctx = Ctx::default();
687
688 let result = optim
689 .handle_request(
690 PluginStep::EarlyRequest,
691 &mut session,
692 &mut ctx,
693 )
694 .await
695 .unwrap();
696
697 assert_eq!(true, RequestPluginResult::Continue == result);
698 assert_eq!(
699 vec!["image/avif".to_string(), "image/webp".to_string()],
700 ctx.cache.unwrap().keys.unwrap()
701 );
702 }
703 }
704
705 #[tokio::test]
706 async fn test_image_optimize_handle_upstream_response() {
707 let optim = ImageOptim::try_from(
708 &toml::from_str::<PluginConf>(
709 r###"
710avif_quality = 75
711avif_speed = 3
712category = "image_optim"
713jpeg_quality = 80
714output_types = "avif,webp"
715png_quality = 90
716"###,
717 )
718 .unwrap(),
719 )
720 .unwrap();
721
722 let headers = ["Accept: image/webp, image/avif"].join("\r\n");
723 let input_header =
724 format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
725 let mock_io = Builder::new().read(input_header.as_bytes()).build();
726 let mut session = Session::new_h1_with_modules(
727 Box::new(mock_io),
728 &HttpModules::new(),
729 );
730 session.read_request().await.unwrap();
731 let mut ctx = Ctx::default();
732 let mut upstream_response = ResponseHeader::build(200, None).unwrap();
733
734 let result = optim
736 .handle_upstream_response(
737 &mut session,
738 &mut ctx,
739 &mut upstream_response,
740 )
741 .unwrap();
742 assert_eq!(true, ResponsePluginResult::Unchanged == result);
743
744 upstream_response
746 .append_header("content-type", "application/json")
747 .unwrap();
748 let result = optim
749 .handle_upstream_response(
750 &mut session,
751 &mut ctx,
752 &mut upstream_response,
753 )
754 .unwrap();
755 assert_eq!(true, ResponsePluginResult::Unchanged == result);
756
757 upstream_response
759 .insert_header("content-type", "image/png")
760 .unwrap();
761 let result = optim
762 .handle_upstream_response(
763 &mut session,
764 &mut ctx,
765 &mut upstream_response,
766 )
767 .unwrap();
768 assert_eq!(
769 "chunked",
770 upstream_response.headers.get("transfer-encoding").unwrap()
771 );
772 assert_eq!(true, ResponsePluginResult::Modified == result);
773 assert_eq!(
774 true,
775 ctx.get_modify_body_handler(&optim.handler_id).is_some()
776 );
777 assert_eq!(
779 "image/avif",
780 upstream_response.headers.get("content-type").unwrap()
781 );
782 }
783
784 fn new_optim(output_types: &str) -> ImageOptim {
785 ImageOptim::try_from(
786 &toml::from_str::<PluginConf>(&format!(
787 "category = \"image_optim\"\noutput_types = \"{output_types}\""
788 ))
789 .unwrap(),
790 )
791 .unwrap()
792 }
793
794 async fn new_session(accept: &str) -> Session {
795 let input_header =
796 format!("GET /a.png HTTP/1.1\r\nAccept: {accept}\r\n\r\n");
797 let mock_io = Builder::new().read(input_header.as_bytes()).build();
798 let mut session = Session::new_h1_with_modules(
799 Box::new(mock_io),
800 &HttpModules::new(),
801 );
802 session.read_request().await.unwrap();
803 session
804 }
805
806 #[test]
809 fn test_original_is_sent_when_nothing_is_gained() {
810 let original = Bytes::from_static(b"12345678");
811 let run = |converted: &[u8], same_format| {
812 worth_sending(original.clone(), converted.to_vec(), same_format)
813 };
814 assert_eq!(Bytes::from_static(b"1234567"), run(b"1234567", true));
816 assert_eq!(original, run(b"abcdefgh", true));
817 assert_eq!(original, run(b"abcdefghi", true));
818 assert_eq!(Bytes::from_static(b"abcdefghi"), run(b"abcdefghi", false));
820 }
821
822 fn sample_png() -> Vec<u8> {
823 use image::ImageEncoder;
824 let img = image::RgbaImage::from_fn(8, 6, |x, y| {
825 image::Rgba([(x * 30) as u8, (y * 40) as u8, 0, 255])
826 });
827 let mut png = Vec::new();
828 image::codecs::png::PngEncoder::new(&mut png)
829 .write_image(img.as_raw(), 8, 6, image::ExtendedColorType::Rgba8)
830 .unwrap();
831 png
832 }
833
834 #[test]
835 fn test_unknown_output_type_is_rejected() {
836 let err = ImageOptim::try_from(
837 &toml::from_str::<PluginConf>(
838 "category = \"image_optim\"\noutput_types = \"avif,bmp\"",
839 )
840 .unwrap(),
841 )
842 .err()
843 .unwrap()
844 .to_string();
845 assert_eq!(
846 "Plugin image_optim invalid, message: output type bmp is not supported, expected one of avif, webp, jpeg, png",
847 err
848 );
849 }
850
851 #[tokio::test]
853 async fn test_image_optimize_response_header() {
854 let optim = new_optim("avif,webp");
855 let run =
856 async |accept: &str, status: u16, headers: &[(&str, &str)]| {
857 let mut session = new_session(accept).await;
858 let mut ctx = Ctx::default();
859 let mut resp = ResponseHeader::build(status, None).unwrap();
860 for (name, value) in headers {
861 resp.insert_header(name.to_string(), *value).unwrap();
862 }
863 let result = optim
864 .handle_upstream_response(&mut session, &mut ctx, &mut resp)
865 .unwrap();
866 let header = |name: &str| {
867 resp.headers
868 .get(name)
869 .map(|value| value.to_str().unwrap().to_string())
870 };
871 (
872 result == ResponsePluginResult::Modified,
873 header("content-type"),
874 header("content-length"),
875 )
876 };
877 let text = |value: &str| Some(value.to_string());
878 let png = [("content-type", "image/png"), ("content-length", "100")];
879
880 assert_eq!(
881 (true, text("image/webp"), None),
882 run("image/webp", 200, &png).await
883 );
884 assert_eq!(
887 (true, text("image/png"), None),
888 run("image/*", 200, &png).await
889 );
890 assert_eq!(
892 (true, text("image/avif"), None),
893 run(
894 "image/avif",
895 200,
896 &[("content-type", "image/jpeg; charset=binary")]
897 )
898 .await
899 );
900
901 let too_large = (MAX_IMAGE_SIZE + 1).to_string();
903 assert_eq!(
904 (false, text("image/png"), Some(too_large.clone())),
905 run(
906 "image/webp",
907 200,
908 &[
909 ("content-type", "image/png"),
910 ("content-length", &too_large)
911 ]
912 )
913 .await
914 );
915 assert_eq!(
917 (false, text("image/png"), text("100")),
918 run("image/webp", 206, &png).await
919 );
920 assert_eq!(
922 (false, text("image/gif"), None),
923 run("image/webp", 200, &[("content-type", "image/gif")]).await
924 );
925 assert_eq!(
927 (false, text("image/png"), text("100")),
928 run(
929 "image/webp",
930 200,
931 &[
932 ("content-type", "image/png"),
933 ("content-length", "100"),
934 ("content-encoding", "gzip")
935 ]
936 )
937 .await
938 );
939 }
940
941 #[tokio::test]
944 async fn test_image_optimize_body() {
945 let session = new_session("image/webp").await;
946 let new_optimizer = |format_type: &str| ImageOptimizer {
947 image_type: "png".to_string(),
948 png_quality: 90,
949 jpeg_quality: 80,
950 avif_quality: 75,
951 avif_speed: 10,
952 webp_quality: 100,
953 format_type: format_type.to_string(),
954 buffer: BytesMut::new(),
955 passthrough: false,
956 };
957 let png = sample_png();
958
959 for (format_type, expected) in [
962 ("webp", image::ImageFormat::WebP),
963 ("avif", image::ImageFormat::Avif),
964 ("jpeg", image::ImageFormat::Jpeg),
965 ("png", image::ImageFormat::Png),
966 ] {
967 let mut optimizer = new_optimizer(format_type);
968 let (head, tail) = png.split_at(png.len() / 2);
970 let mut body = Some(Bytes::copy_from_slice(head));
971 optimizer.handle(&session, &mut body, false).unwrap();
972 assert_eq!(Some(Bytes::new()), body, "{format_type}");
973 let mut body = Some(Bytes::copy_from_slice(tail));
974 optimizer.handle(&session, &mut body, true).unwrap();
975 assert_eq!(
976 Some(expected),
977 image::guess_format(&body.unwrap()).ok(),
978 "{format_type}"
979 );
980 }
981
982 let mut optimizer = new_optimizer("webp");
984 let mut body = None;
985 optimizer.handle(&session, &mut body, true).unwrap();
986 assert_eq!(None, body);
987
988 let grainy = image::RgbaImage::from_fn(120, 80, |x, y| {
991 let grain = (x * 7919 + y * 104_729) % 61;
992 image::Rgba([
993 (x + grain) as u8,
994 (y * 2 + grain) as u8,
995 (90 + grain) as u8,
996 255,
997 ])
998 });
999 let mut large = Vec::new();
1000 use image::ImageEncoder;
1001 image::codecs::png::PngEncoder::new(&mut large)
1002 .write_image(
1003 grainy.as_raw(),
1004 120,
1005 80,
1006 image::ExtendedColorType::Rgba8,
1007 )
1008 .unwrap();
1009 let mut optimizer = new_optimizer("png");
1010 let mut body = Some(Bytes::from(large.clone()));
1011 optimizer.handle(&session, &mut body, true).unwrap();
1012 let small = body.unwrap();
1013 assert_eq!(true, small.len() < large.len(), "{}", small.len());
1014 assert_eq!(
1015 Some(image::ImageFormat::Png),
1016 image::guess_format(&small).ok()
1017 );
1018 let mut optimizer = new_optimizer("webp");
1020 let mut body = Some(small.clone());
1021 optimizer.handle(&session, &mut body, true).unwrap();
1022 let webp = body.unwrap();
1023 assert_eq!(
1024 Some(image::ImageFormat::WebP),
1025 image::guess_format(&webp).ok()
1026 );
1027 assert_eq!(true, webp.len() > small.len());
1028
1029 let mut optimizer = new_optimizer("webp");
1031 let mut body = Some(Bytes::from_static(b"not an image"));
1032 optimizer.handle(&session, &mut body, true).unwrap();
1033 assert_eq!(Some(Bytes::from_static(b"not an image")), body);
1034
1035 let mut optimizer = new_optimizer("webp");
1038 let chunk = Bytes::from(vec![1u8; MAX_IMAGE_SIZE / 2 + 1]);
1039 let mut body = Some(chunk.clone());
1040 optimizer.handle(&session, &mut body, false).unwrap();
1041 assert_eq!(Some(Bytes::new()), body);
1042 let mut body = Some(chunk.clone());
1043 optimizer.handle(&session, &mut body, false).unwrap();
1044 assert_eq!(Some(2 * chunk.len()), body.map(|data| data.len()));
1045 assert_eq!(0, optimizer.buffer.len());
1046 let mut body = Some(Bytes::from_static(b"tail"));
1047 optimizer.handle(&session, &mut body, true).unwrap();
1048 assert_eq!(Some(Bytes::from_static(b"tail")), body);
1049 }
1050
1051 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1054 async fn test_image_optimize_on_multi_thread_runtime() {
1055 let png = sample_png();
1056 let output = tokio::spawn(async move {
1057 let session = new_session("image/webp").await;
1058 let mut optimizer = ImageOptimizer {
1059 image_type: "png".to_string(),
1060 png_quality: 90,
1061 jpeg_quality: 80,
1062 avif_quality: 75,
1063 avif_speed: 10,
1064 webp_quality: 100,
1065 format_type: "webp".to_string(),
1066 buffer: BytesMut::new(),
1067 passthrough: false,
1068 };
1069 let mut body = Some(Bytes::from(png));
1070 optimizer.handle(&session, &mut body, true).unwrap();
1071 body.unwrap()
1072 })
1073 .await
1074 .unwrap();
1075 assert_eq!(
1076 Some(image::ImageFormat::WebP),
1077 image::guess_format(&output).ok()
1078 );
1079 }
1080}