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) -> &str {
85        "image_optimization"
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            && let Ok(accept_str) = accept.to_str()
174        {
175            let mut accept_images: Vec<_> = self
176                .output_mimes
177                .iter()
178                .filter(|mime| accept_str.contains(*mime))
179                .cloned()
180                .collect();
181
182            if !accept_images.is_empty() {
183                accept_images.sort();
184                ctx.extend_cache_keys(accept_images);
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            && let Ok(size) =
232                value.to_str().unwrap_or_default().parse::<usize>()
233        {
234            capacity = size;
235        }
236        // Switch to chunked transfer encoding
237        let _ = upstream_response.insert_header(
238            http::header::TRANSFER_ENCODING,
239            HTTP_HEADER_TRANSFER_CHUNKED.1.clone(),
240        );
241        let _ = upstream_response
242            .insert_header(http::header::CONTENT_TYPE, &format_type);
243
244        ctx.add_modify_body_handler(
245            PLUGIN_ID,
246            Box::new(ImageOptimizer {
247                image_type,
248                png_quality: self.png_quality,
249                jpeg_quality: self.jpeg_quality,
250                avif_quality: self.avif_quality,
251                avif_speed: self.avif_speed,
252                // only support lossless
253                webp_quality: 100,
254                format_type,
255                buffer: BytesMut::with_capacity(capacity),
256            }),
257        );
258        Ok(ResponsePluginResult::Modified)
259    }
260    fn handle_upstream_response_body(
261        &self,
262        session: &mut Session,
263        ctx: &mut Ctx,
264        body: &mut Option<bytes::Bytes>,
265        end_of_stream: bool,
266    ) -> pingora::Result<ResponseBodyPluginResult> {
267        if let Some(modifier) = ctx.get_modify_body_handler(PLUGIN_ID) {
268            modifier.handle(session, body, end_of_stream)?;
269            let result = if end_of_stream {
270                ResponseBodyPluginResult::FullyReplaced
271            } else {
272                ResponseBodyPluginResult::PartialReplaced
273            };
274            Ok(result)
275        } else {
276            Ok(ResponseBodyPluginResult::Unchanged)
277        }
278    }
279}
280
281#[ctor(unsafe)]
282fn init() {
283    get_plugin_factory().register("image_optim", |params| {
284        Ok(Arc::new(ImageOptim::new(params)?))
285    });
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291    use pingap_config::PluginConf;
292    use pingap_core::{Ctx, Plugin};
293    use pingora::modules::http::HttpModules;
294    use pretty_assertions::assert_eq;
295    use tokio_test::io::Builder;
296
297    #[test]
298    fn test_new_image_optimize() {
299        let optim = ImageOptim::try_from(
300            &toml::from_str::<PluginConf>(
301                r###"
302avif_quality = 75
303avif_speed = 3
304category = "image_optim"
305jpeg_quality = 80
306output_types = "avif,webp"
307png_quality = 90
308"###,
309            )
310            .unwrap(),
311        )
312        .unwrap();
313
314        assert_eq!(
315            HashSet::from(["jpeg".to_string(), "png".to_string(),]),
316            optim.support_types
317        );
318        assert_eq!(
319            vec!["image/avif".to_string(), "image/webp".to_string()],
320            optim.output_mimes
321        );
322        assert_eq!(90, optim.png_quality);
323        assert_eq!(80, optim.jpeg_quality);
324        assert_eq!(75, optim.avif_quality);
325        assert_eq!(3, optim.avif_speed);
326    }
327
328    #[tokio::test]
329    async fn test_image_optimize_handle_request() {
330        let optim = ImageOptim::try_from(
331            &toml::from_str::<PluginConf>(
332                r###"
333avif_quality = 75
334avif_speed = 3
335category = "image_optim"
336jpeg_quality = 80
337output_types = "avif,webp"
338png_quality = 90
339"###,
340            )
341            .unwrap(),
342        )
343        .unwrap();
344        // not accept value
345        {
346            let headers = [""].join("\r\n");
347            let input_header = format!(
348                "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
349            );
350            let mock_io = Builder::new().read(input_header.as_bytes()).build();
351            let mut session = Session::new_h1_with_modules(
352                Box::new(mock_io),
353                &HttpModules::new(),
354            );
355            session.read_request().await.unwrap();
356            let mut ctx = Ctx::default();
357
358            let result = optim
359                .handle_request(PluginStep::Request, &mut session, &mut ctx)
360                .await
361                .unwrap();
362
363            assert_eq!(true, RequestPluginResult::Continue == result);
364            assert_eq!(true, ctx.cache.is_none());
365        }
366
367        // accept avif
368        {
369            let headers = ["Accept: image/avif"].join("\r\n");
370            let input_header = format!(
371                "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
372            );
373            let mock_io = Builder::new().read(input_header.as_bytes()).build();
374            let mut session = Session::new_h1_with_modules(
375                Box::new(mock_io),
376                &HttpModules::new(),
377            );
378            session.read_request().await.unwrap();
379            let mut ctx = Ctx::default();
380
381            let result = optim
382                .handle_request(PluginStep::Request, &mut session, &mut ctx)
383                .await
384                .unwrap();
385
386            assert_eq!(true, RequestPluginResult::Continue == result);
387            assert_eq!(
388                vec!["image/avif".to_string()],
389                ctx.cache.unwrap().keys.unwrap()
390            );
391        }
392
393        // accept avif, webp
394        {
395            let headers = ["Accept: image/webp, image/avif"].join("\r\n");
396            let input_header = format!(
397                "GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n"
398            );
399            let mock_io = Builder::new().read(input_header.as_bytes()).build();
400            let mut session = Session::new_h1_with_modules(
401                Box::new(mock_io),
402                &HttpModules::new(),
403            );
404            session.read_request().await.unwrap();
405            let mut ctx = Ctx::default();
406
407            let result = optim
408                .handle_request(PluginStep::Request, &mut session, &mut ctx)
409                .await
410                .unwrap();
411
412            assert_eq!(true, RequestPluginResult::Continue == result);
413            assert_eq!(
414                vec!["image/avif".to_string(), "image/webp".to_string()],
415                ctx.cache.unwrap().keys.unwrap()
416            );
417        }
418    }
419
420    #[tokio::test]
421    async fn test_image_optimize_handle_upstream_response() {
422        let optim = ImageOptim::try_from(
423            &toml::from_str::<PluginConf>(
424                r###"
425avif_quality = 75
426avif_speed = 3
427category = "image_optim"
428jpeg_quality = 80
429output_types = "avif,webp"
430png_quality = 90
431"###,
432            )
433            .unwrap(),
434        )
435        .unwrap();
436
437        let headers = ["Accept: image/webp, image/avif"].join("\r\n");
438        let input_header =
439            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
440        let mock_io = Builder::new().read(input_header.as_bytes()).build();
441        let mut session = Session::new_h1_with_modules(
442            Box::new(mock_io),
443            &HttpModules::new(),
444        );
445        session.read_request().await.unwrap();
446        let mut ctx = Ctx::default();
447        let mut upstream_response = ResponseHeader::build(200, None).unwrap();
448
449        // no content type
450        let result = optim
451            .handle_upstream_response(
452                &mut session,
453                &mut ctx,
454                &mut upstream_response,
455            )
456            .unwrap();
457        assert_eq!(true, ResponsePluginResult::Unchanged == result);
458
459        // content type is not image
460        upstream_response
461            .append_header("content-type", "application/json")
462            .unwrap();
463        let result = optim
464            .handle_upstream_response(
465                &mut session,
466                &mut ctx,
467                &mut upstream_response,
468            )
469            .unwrap();
470        assert_eq!(true, ResponsePluginResult::Unchanged == result);
471
472        // response image png
473        upstream_response
474            .insert_header("content-type", "image/png")
475            .unwrap();
476        let result = optim
477            .handle_upstream_response(
478                &mut session,
479                &mut ctx,
480                &mut upstream_response,
481            )
482            .unwrap();
483        assert_eq!(
484            "chunked",
485            upstream_response.headers.get("transfer-encoding").unwrap()
486        );
487        assert_eq!(true, ResponsePluginResult::Modified == result);
488        assert_eq!(true, ctx.get_modify_body_handler(PLUGIN_ID).is_some());
489    }
490}
pingap_imageoptim/lib.rs

pingap_imageoptim/
lib.rs
