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    load_image, optimize_avif, optimize_jpeg, optimize_png, optimize_webp,
17};
18use async_trait::async_trait;
19use bytes::{Bytes, BytesMut};
20use ctor::ctor;
21use pingap_config::PluginConf;
22use pingap_core::HTTP_HEADER_TRANSFER_CHUNKED;
23use pingap_core::{
24    Ctx, Plugin, PluginStep, RequestPluginResult, ResponsePluginResult,
25};
26use pingap_core::{ModifyResponseBody, ResponseBodyPluginResult};
27use pingap_plugin::{
28    Error, get_hash_key, get_int_conf, get_plugin_factory, get_str_conf,
29};
30use pingora::http::ResponseHeader;
31use pingora::proxy::Session;
32use std::borrow::Cow;
33use std::collections::HashSet;
34use std::convert::TryFrom;
35use std::sync::Arc;
36use tracing::debug;
37
38mod optimizer;
39
40const PLUGIN_ID: &str = "_image_optimize_";
41
42type Result<T, E = Error> = std::result::Result<T, E>;
43
44struct ImageOptimizer {
45    image_type: String,
46    png_quality: u8,
47    jpeg_quality: u8,
48    avif_quality: u8,
49    avif_speed: u8,
50    webp_quality: u8,
51    format_type: String,
52    buffer: BytesMut,
53}
54
55impl ModifyResponseBody for ImageOptimizer {
56    fn handle(
57        &mut self,
58        _session: &Session,
59        body: &mut Option<bytes::Bytes>,
60        end_of_stream: bool,
61    ) -> pingora::Result<()> {
62        if let Some(data) = body {
63            self.buffer.extend(&data[..]);
64            data.clear();
65        }
66        if !end_of_stream {
67            return Ok(());
68        }
69        if let Ok(info) = load_image(&self.buffer, &self.image_type) {
70            let result = match self.format_type.as_str() {
71                "jpeg" => optimize_jpeg(&info, self.jpeg_quality),
72                "avif" => {
73                    optimize_avif(&info, self.avif_quality, self.avif_speed)
74                },
75                "webp" => optimize_webp(&info, self.webp_quality),
76                _ => optimize_png(&info, self.png_quality),
77            };
78            if let Ok(data) = result {
79                *body = Some(Bytes::from(data));
80            }
81        }
82        Ok(())
83    }
84    fn name(&self) -> String {
85        "image_optimization".to_string()
86    }
87}
88
89pub struct ImageOptim {
90    /// A unique identifier for this plugin instance.
91    /// Used for internal tracking and debugging purposes.
92    hash_value: String,
93    support_types: HashSet<String>,
94    output_mimes: Vec<String>,
95    png_quality: u8,
96    jpeg_quality: u8,
97    avif_quality: u8,
98    avif_speed: u8,
99}
100
101impl TryFrom<&PluginConf> for ImageOptim {
102    type Error = Error;
103    fn try_from(value: &PluginConf) -> Result<Self> {
104        debug!(params = value.to_string(), "new image optimizer plugin");
105        let hash_value = get_hash_key(value);
106
107        let output_types: Vec<String> = get_str_conf(value, "output_types")
108            .split(',')
109            .map(|s| s.trim())
110            .filter(|s| !s.is_empty())
111            .map(|s| s.to_string())
112            .collect();
113
114        let output_mimes = output_types
115            .iter()
116            .map(|format| format!("image/{}", format))
117            .collect();
118
119        let mut png_quality = get_int_conf(value, "png_quality") as u8;
120        if png_quality == 0 || png_quality > 100 {
121            png_quality = 90;
122        }
123        let mut jpeg_quality = get_int_conf(value, "jpeg_quality") as u8;
124        if jpeg_quality == 0 || jpeg_quality > 100 {
125            jpeg_quality = 80;
126        }
127        let mut avif_quality = get_int_conf(value, "avif_quality") as u8;
128        if avif_quality == 0 || avif_quality > 100 {
129            avif_quality = 75;
130        }
131        let mut avif_speed = get_int_conf(value, "avif_speed") as u8;
132        if avif_speed == 0 || avif_speed > 10 {
133            avif_speed = 3;
134        }
135        Ok(Self {
136            hash_value,
137            support_types: HashSet::from([
138                "jpeg".to_string(),
139                "png".to_string(),
140            ]),
141            output_mimes,
142            png_quality,
143            jpeg_quality,
144            avif_quality,
145            avif_speed,
146        })
147    }
148}
149
150impl ImageOptim {
151    pub fn new(params: &PluginConf) -> Result<Self> {
152        Self::try_from(params)
153    }
154}
155
156#[async_trait]
157impl Plugin for ImageOptim {
158    /// Returns a unique identifier for this plugin instance
159    fn config_key(&self) -> Cow<'_, str> {
160        Cow::Borrowed(&self.hash_value)
161    }
162    async fn handle_request(
163        &self,
164        step: PluginStep,
165        session: &mut Session,
166        ctx: &mut Ctx,
167    ) -> pingora::Result<RequestPluginResult> {
168        if step != PluginStep::Request {
169            return Ok(RequestPluginResult::Skipped);
170        }
171
172        if let Some(accept) = session.get_header(http::header::ACCEPT) {
173            if let Ok(accept_str) = accept.to_str() {
174                let mut accept_images: Vec<_> = self
175                    .output_mimes
176                    .iter()
177                    .filter(|mime| accept_str.contains(*mime))
178                    .cloned()
179                    .collect();
180
181                if !accept_images.is_empty() {
182                    accept_images.sort();
183                    ctx.extend_cache_keys(accept_images);
184                }
185            }
186        }
187        Ok(RequestPluginResult::Continue)
188    }
189    fn handle_upstream_response(
190        &self,
191        session: &mut Session,
192        ctx: &mut Ctx,
193        upstream_response: &mut ResponseHeader,
194    ) -> pingora::Result<ResponsePluginResult> {
195        let content_type = if let Some(value) =
196            upstream_response.headers.get(http::header::CONTENT_TYPE)
197        {
198            value.to_str().unwrap_or_default()
199        } else {
200            return Ok(ResponsePluginResult::Unchanged);
201        };
202
203        let Some(image_type) = content_type.strip_prefix("image/") else {
204            return Ok(ResponsePluginResult::Unchanged);
205        };
206
207        if !self.support_types.contains(image_type) {
208            return Ok(ResponsePluginResult::Unchanged);
209        }
210
211        let Some(accept) = session.get_header(http::header::ACCEPT) else {
212            return Ok(ResponsePluginResult::Unchanged);
213        };
214        let Ok(accept_str) = accept.to_str() else {
215            return Ok(ResponsePluginResult::Unchanged);
216        };
217
218        let image_type = image_type.to_string();
219        let mut format_type = image_type.clone();
220        for item in self.output_mimes.iter() {
221            if accept_str.contains(item.as_str()) {
222                format_type = item.clone();
223                break;
224            }
225        }
226        let mut capacity = 8192;
227        // if the image is not changed, it will be optimized again, so we need to remove the content-length
228        // Remove content-length since we're modifying the body
229        if let Some(value) =
230            upstream_response.remove_header(&http::header::CONTENT_LENGTH)
231        {
232            if let Ok(size) =
233                value.to_str().unwrap_or_default().parse::<usize>()
234            {
235                capacity = size;
236            }
237        }
238        // Switch to chunked transfer encoding
239        let _ = upstream_response.insert_header(
240            http::header::TRANSFER_ENCODING,
241            HTTP_HEADER_TRANSFER_CHUNKED.1.clone(),
242        );
243        let _ = upstream_response
244            .insert_header(http::header::CONTENT_TYPE, &format_type);
245
246        ctx.add_modify_body_handler(
247            PLUGIN_ID,
248            Box::new(ImageOptimizer {
249                image_type,
250                png_quality: self.png_quality,
251                jpeg_quality: self.jpeg_quality,
252                avif_quality: self.avif_quality,
253                avif_speed: self.avif_speed,
254                // only support lossless
255                webp_quality: 100,
256                format_type,
257                buffer: BytesMut::with_capacity(capacity),
258            }),
259        );
260        Ok(ResponsePluginResult::Modified)
261    }
262    fn handle_upstream_response_body(
263        &self,
264        session: &mut Session,
265        ctx: &mut Ctx,
266        body: &mut Option<bytes::Bytes>,
267        end_of_stream: bool,
268    ) -> pingora::Result<ResponseBodyPluginResult> {
269        if let Some(modifier) = ctx.get_modify_body_handler(PLUGIN_ID) {
270            modifier.handle(session, body, end_of_stream)?;
271            let result = if end_of_stream {
272                ResponseBodyPluginResult::FullyReplaced
273            } else {
274                ResponseBodyPluginResult::PartialReplaced
275            };
276            Ok(result)
277        } else {
278            Ok(ResponseBodyPluginResult::Unchanged)
279        }
280    }
281}
282
283#[ctor]
284fn init() {
285    get_plugin_factory().register("image_optim", |params| {
286        Ok(Arc::new(ImageOptim::new(params)?))
287    });
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    use pingap_config::PluginConf;
294    use pingap_core::{Ctx, Plugin};
295    use pingora::modules::http::HttpModules;
296    use pretty_assertions::assert_eq;
297    use tokio_test::io::Builder;
298
299    #[test]
300    fn test_new_image_optimize() {
301        let optim = ImageOptim::try_from(
302            &toml::from_str::<PluginConf>(
303                r###"
304avif_quality = 75
305avif_speed = 3
306category = "image_optim"
307jpeg_quality = 80
308output_types = "avif,webp"
309png_quality = 90
310"###,
311            )
312            .unwrap(),
313        )
314        .unwrap();
315
316        assert_eq!(
317            HashSet::from(["jpeg".to_string(), "png".to_string(),]),
318            optim.support_types
319        );
320        assert_eq!(
321            vec!["image/avif".to_string(), "image/webp".to_string()],
322            optim.output_mimes
323        );
324        assert_eq!(90, optim.png_quality);
325        assert_eq!(80, optim.jpeg_quality);
326        assert_eq!(75, optim.avif_quality);
327        assert_eq!(3, optim.avif_speed);
328    }
329
330    #[tokio::test]
331    async fn test_image_optimize_handle_request() {
332        let optim = ImageOptim::try_from(
333            &toml::from_str::<PluginConf>(
334                r###"
335avif_quality = 75
336avif_speed = 3
337category = "image_optim"
338jpeg_quality = 80
339output_types = "avif,webp"
340png_quality = 90
341"###,
342            )
343            .unwrap(),
344        )
345        .unwrap();
346        // not accept value
347        {
348            let headers = [""].join("\r\n");
349            let input_header = format!(
350                "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
351            );
352            let mock_io = Builder::new().read(input_header.as_bytes()).build();
353            let mut session = Session::new_h1_with_modules(
354                Box::new(mock_io),
355                &HttpModules::new(),
356            );
357            session.read_request().await.unwrap();
358            let mut ctx = Ctx::default();
359
360            let result = optim
361                .handle_request(PluginStep::Request, &mut session, &mut ctx)
362                .await
363                .unwrap();
364
365            assert_eq!(true, RequestPluginResult::Continue == result);
366            assert_eq!(true, ctx.cache.is_none());
367        }
368
369        // accept avif
370        {
371            let headers = ["Accept: image/avif"].join("\r\n");
372            let input_header = format!(
373                "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
374            );
375            let mock_io = Builder::new().read(input_header.as_bytes()).build();
376            let mut session = Session::new_h1_with_modules(
377                Box::new(mock_io),
378                &HttpModules::new(),
379            );
380            session.read_request().await.unwrap();
381            let mut ctx = Ctx::default();
382
383            let result = optim
384                .handle_request(PluginStep::Request, &mut session, &mut ctx)
385                .await
386                .unwrap();
387
388            assert_eq!(true, RequestPluginResult::Continue == result);
389            assert_eq!(
390                vec!["image/avif".to_string()],
391                ctx.cache.unwrap().keys.unwrap()
392            );
393        }
394
395        // accept avif, webp
396        {
397            let headers = ["Accept: image/webp, image/avif"].join("\r\n");
398            let input_header = format!(
399                "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
400            );
401            let mock_io = Builder::new().read(input_header.as_bytes()).build();
402            let mut session = Session::new_h1_with_modules(
403                Box::new(mock_io),
404                &HttpModules::new(),
405            );
406            session.read_request().await.unwrap();
407            let mut ctx = Ctx::default();
408
409            let result = optim
410                .handle_request(PluginStep::Request, &mut session, &mut ctx)
411                .await
412                .unwrap();
413
414            assert_eq!(true, RequestPluginResult::Continue == result);
415            assert_eq!(
416                vec!["image/avif".to_string(), "image/webp".to_string()],
417                ctx.cache.unwrap().keys.unwrap()
418            );
419        }
420    }
421
422    #[tokio::test]
423    async fn test_image_optimize_handle_upstream_response() {
424        let optim = ImageOptim::try_from(
425            &toml::from_str::<PluginConf>(
426                r###"
427avif_quality = 75
428avif_speed = 3
429category = "image_optim"
430jpeg_quality = 80
431output_types = "avif,webp"
432png_quality = 90
433"###,
434            )
435            .unwrap(),
436        )
437        .unwrap();
438
439        let headers = ["Accept: image/webp, image/avif"].join("\r\n");
440        let input_header =
441            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
442        let mock_io = Builder::new().read(input_header.as_bytes()).build();
443        let mut session = Session::new_h1_with_modules(
444            Box::new(mock_io),
445            &HttpModules::new(),
446        );
447        session.read_request().await.unwrap();
448        let mut ctx = Ctx::default();
449        let mut upstream_response = ResponseHeader::build(200, None).unwrap();
450
451        // no content type
452        let result = optim
453            .handle_upstream_response(
454                &mut session,
455                &mut ctx,
456                &mut upstream_response,
457            )
458            .unwrap();
459        assert_eq!(true, ResponsePluginResult::Unchanged == result);
460
461        // content type is not image
462        upstream_response
463            .append_header("content-type", "application/json")
464            .unwrap();
465        let result = optim
466            .handle_upstream_response(
467                &mut session,
468                &mut ctx,
469                &mut upstream_response,
470            )
471            .unwrap();
472        assert_eq!(true, ResponsePluginResult::Unchanged == result);
473
474        // response image png
475        upstream_response
476            .insert_header("content-type", "image/png")
477            .unwrap();
478        let result = optim
479            .handle_upstream_response(
480                &mut session,
481                &mut ctx,
482                &mut upstream_response,
483            )
484            .unwrap();
485        assert_eq!(
486            "chunked",
487            upstream_response.headers.get("transfer-encoding").unwrap()
488        );
489        assert_eq!(true, ResponsePluginResult::Modified == result);
490        assert_eq!(true, ctx.get_modify_body_handler(PLUGIN_ID).is_some());
491    }
492}
pingap_imageoptim/lib.rs

pingap_imageoptim/
lib.rs
