Skip to main content

pingap_imageoptim/
lib.rs

1// Copyright 2024-2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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
44/// The largest original that is converted, in bytes. A larger one is
45/// passed on as it is.
46const MAX_IMAGE_SIZE: usize = 20 * 1024 * 1024;
47
48/// The formats an image can be converted to.
49const OUTPUT_FORMATS: &[&str] = &["avif", "webp", "jpeg", "png"];
50
51type Result<T, E = Error> = std::result::Result<T, E>;
52
53struct ImageOptimizer {
54    /// The format of the original, `png` or `jpeg`.
55    image_type: String,
56    png_quality: u8,
57    jpeg_quality: u8,
58    avif_quality: u8,
59    avif_speed: u8,
60    webp_quality: u8,
61    /// The format to convert to, one of `OUTPUT_FORMATS`.
62    format_type: String,
63    buffer: BytesMut,
64    /// The original turned out too large to convert: what was held back
65    /// has been released, and the rest goes through as it comes.
66    passthrough: bool,
67}
68
69/// Runs `f`, a long computation, without holding up the other connections
70/// of this worker thread.
71///
72/// Decoding and encoding an image takes from tens of milliseconds to
73/// seconds, and the body filter it happens in is not async. On the
74/// work-stealing runtime `block_in_place` moves the worker's other tasks to
75/// another thread first. A runtime without work stealing has no such
76/// thread to move them to, and there `f` simply runs.
77fn 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
99/// What goes out for an image that was converted: the conversion, or the
100/// original where it was written again in the format it came in and is
101/// no smaller for it. An icon, a png that has a palette already, a jpeg
102/// of a lower quality than the one configured - what is small has little
103/// to give, and a palette alone is a kilobyte.
104fn 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            // Without a `Content-Length` the size is only known once it
127            // has been exceeded. The body used to be collected whatever
128            // its size.
129            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        // No body at all, the answer to a HEAD request for one.
142        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            // The original is what there is to send. A failure used to
151            // leave the body empty: a 200 with no image in it.
152            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    /// The name this instance keeps its body handler under, see
171    /// `new_body_handler_id`.
172    handler_id: String,
173    /// A unique identifier for this plugin instance.
174    /// Used for internal tracking and debugging purposes.
175    hash_value: String,
176    support_types: HashSet<String>,
177    /// The formats to convert to, in order of preference.
178    output_types: Vec<String>,
179    /// The media type of each of `output_types`, in the same order.
180    output_mimes: Vec<String>,
181    /// The cache key components a request can have: each selection of
182    /// `output_mimes`, sorted, as `Accept` may name any of them.
183    key_alternatives: Arc<Vec<Vec<String>>>,
184    png_quality: u8,
185    jpeg_quality: u8,
186    avif_quality: u8,
187    avif_speed: u8,
188}
189
190/// Every selection of `mimes`, the empty one first, each in sorted order.
191fn 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        // A format that is not known can not be produced. It used to be
221        // taken as it was: the response was announced as `image/<format>`
222        // and its body encoded as png.
223        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        // Unset, or 0, is the default. Anything else outside the range is
241        // an error: cast to a byte as it came, 300 was a quality of 44.
242        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    /// Returns a unique identifier for this plugin instance
285    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        // Ahead of the plugins that answer in the request step. The cache
295        // plugin answers a `PURGE` there, and has to know by then that the
296        // formats are a part of the key and which ones there are: listed
297        // after it, this plugin had not run, and the purge removed the
298        // entry of a request that accepts no image format and no other.
299        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        // As in `selections`: a format listed twice is one format.
315        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        // A partial or an empty response is not an image to convert, and
329        // neither are the bytes of a compressed one: decoding those can
330        // only fail, after the header has promised the new format.
331        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        // The media type without its parameters.
347        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        // The size the upstream announces. It used to be the size of the
364        // buffer allocated right here, whatever it said. Too large an
365        // image is left alone, headers included.
366        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        // The first of the configured formats the client accepts, else the
377        // format the image already has. This is the format's name; the
378        // media type, with `image/` in front, used to be kept in its place,
379        // so no conversion ever matched it and everything came out as png,
380        // announced as avif or webp.
381        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        // Remove content-length since we're modifying the body
389        upstream_response.remove_header(&http::header::CONTENT_LENGTH);
390        // Ranges of the new image are not ranges of the upstream's. The
391        // `ETag` stays the upstream's for the cache to revalidate with,
392        // and is weakened on the way out, in `handle_response`.
393        upstream_response.remove_header(&http::header::ACCEPT_RANGES);
394        // Switch to chunked transfer encoding
395        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                // only support lossless
414                webp_quality: 100,
415                format_type,
416                buffer: BytesMut::with_capacity(capacity),
417                passthrough: false,
418            }),
419        );
420        Ok(ResponsePluginResult::Modified)
421    }
422    /// The upstream's strong `ETag` names the image it sent, and what
423    /// goes out is that image encoded anew: the validator is made a weak
424    /// one. On the way to the client, for a response from the cache as
425    /// well, and not on the response that is stored: the cache
426    /// revalidates with the validator as the upstream gave it.
427    ///
428    /// Every image of a kind this plugin reads or writes gets it, also
429    /// one that was left as it came (too large, say): which it was is not
430    /// known here, and a weak validator is never wrong.
431    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    /// The validator of an image this plugin may have encoded anew is a
521    /// weak one for the client, whatever the cache holds.
522    #[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        // Not an image of this plugin's.
544        assert_eq!("\"v1\"", sent("image/svg+xml").await);
545        assert_eq!("\"v1\"", sent("text/html").await);
546    }
547
548    /// Regression: the levels were cast to a byte before they were
549    /// checked, so 300 passed as 44 instead of being out of range.
550    #[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        // Unset or 0 is the default, the ends of the range are in it.
568        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        // not accept value
604        {
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            // No format accepted: nothing in the key, and the place of
628            // the formats noted with all that could stand there.
629            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        // accept avif
645        {
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        // accept avif, webp
675        {
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        // no content type
735        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        // content type is not image
745        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        // response image png
758        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        // The first configured format the client accepts, as a media type.
778        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    /// Regression: an image that is no smaller for being written again in
807    /// its own format goes out as it came. It used to go out larger.
808    #[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        // the same format: only what is smaller replaces it
815        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        // another format is what was asked for, whatever its size
819        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    /// Which responses are converted, and what the header says afterwards.
852    #[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        // Regression: no configured format accepted, so the image keeps its
885        // own. The header used to say `png`, which is not a media type.
886        assert_eq!(
887            (true, text("image/png"), None),
888            run("image/*", 200, &png).await
889        );
890        // Parameters after the media type do not hide it.
891        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        // Left alone, header and all: too large to convert,
902        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        // part of an image,
916        assert_eq!(
917            (false, text("image/png"), text("100")),
918            run("image/webp", 206, &png).await
919        );
920        // not a format that is converted,
921        assert_eq!(
922            (false, text("image/gif"), None),
923            run("image/webp", 200, &[("content-type", "image/gif")]).await
924        );
925        // compressed, so not the bytes of an image.
926        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    /// The body that goes out for each target format, and when the
942    /// conversion can not be done.
943    #[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        // Regression: every target used to come out as png, whatever the
960        // header announced.
961        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            // in two chunks: nothing goes out before the last one
969            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        // Nothing came, as for a HEAD request: nothing is made up.
983        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        // A png that gets smaller is replaced: grain no png packs well,
989        // of a few thousand colours.
990        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        // Another format is what was asked for, whatever its size.
1019        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        // Regression: not an image after all. The body used to be empty.
1030        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        // Regression: no `Content-Length`, and more than is converted. It
1036        // goes out as it came, and is not collected any further.
1037        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    /// On the work-stealing runtime the conversion goes through
1052    /// `block_in_place`, which is only allowed there.
1053    #[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}