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;
20use ctor::ctor;
21use pingap_config::PluginConf;
22use pingap_core::HttpResponse;
23use pingap_core::ModifyResponseBody;
24use pingap_core::HTTP_HEADER_TRANSFER_CHUNKED;
25use pingap_core::{Ctx, Plugin, PluginStep};
26use pingap_plugin::{
27    get_hash_key, get_int_conf, get_plugin_factory, get_str_conf, Error,
28};
29use pingora::http::ResponseHeader;
30use pingora::proxy::Session;
31use std::sync::Arc;
32use tracing::debug;
33
34mod optimizer;
35
36type Result<T, E = Error> = std::result::Result<T, E>;
37
38struct ImageOptimizer {
39    image_type: String,
40    png_quality: u8,
41    jpeg_quality: u8,
42    avif_quality: u8,
43    avif_speed: u8,
44    webp_quality: u8,
45    format_type: String,
46}
47
48impl ModifyResponseBody for ImageOptimizer {
49    fn handle(&self, data: Bytes) -> pingora::Result<Bytes> {
50        if let Ok(info) = load_image(&data, &self.image_type) {
51            let result = match self.format_type.as_str() {
52                "jpeg" => optimize_jpeg(&info, self.jpeg_quality),
53                "avif" => {
54                    optimize_avif(&info, self.avif_quality, self.avif_speed)
55                },
56                "webp" => optimize_webp(&info, self.webp_quality),
57                _ => optimize_png(&info, self.png_quality),
58            };
59            if let Ok(data) = result {
60                return Ok(Bytes::from(data));
61            }
62        }
63        Ok(data)
64    }
65    fn name(&self) -> String {
66        "image_optimization".to_string()
67    }
68}
69
70pub struct ImageOptim {
71    /// A unique identifier for this plugin instance.
72    /// Used for internal tracking and debugging purposes.
73    hash_value: String,
74    support_types: Vec<String>,
75    plugin_step: PluginStep,
76    output_types: Vec<String>,
77    png_quality: u8,
78    jpeg_quality: u8,
79    avif_quality: u8,
80    avif_speed: u8,
81}
82
83impl TryFrom<&PluginConf> for ImageOptim {
84    type Error = Error;
85    fn try_from(value: &PluginConf) -> Result<Self> {
86        debug!(params = value.to_string(), "new image optimizer plugin");
87        let hash_value = get_hash_key(value);
88        let mut output_types = vec![];
89        for item in get_str_conf(value, "output_types").split(",") {
90            let item = item.trim();
91            if item.is_empty() {
92                continue;
93            }
94            output_types.push(item.to_string());
95        }
96        let mut png_quality = get_int_conf(value, "png_quality") as u8;
97        if png_quality == 0 || png_quality > 100 {
98            png_quality = 90;
99        }
100        let mut jpeg_quality = get_int_conf(value, "jpeg_quality") as u8;
101        if jpeg_quality == 0 || jpeg_quality > 100 {
102            jpeg_quality = 80;
103        }
104        let mut avif_quality = get_int_conf(value, "avif_quality") as u8;
105        if avif_quality == 0 || avif_quality > 100 {
106            avif_quality = 75;
107        }
108        let mut avif_speed = get_int_conf(value, "avif_speed") as u8;
109        if avif_speed == 0 || avif_speed > 10 {
110            avif_speed = 3;
111        }
112        Ok(Self {
113            hash_value,
114            support_types: vec!["jpeg".to_string(), "png".to_string()],
115            output_types,
116            plugin_step: PluginStep::UpstreamResponse,
117            png_quality,
118            jpeg_quality,
119            avif_quality,
120            avif_speed,
121        })
122    }
123}
124
125impl ImageOptim {
126    pub fn new(params: &PluginConf) -> Result<Self> {
127        Self::try_from(params)
128    }
129}
130
131#[async_trait]
132impl Plugin for ImageOptim {
133    /// Returns a unique identifier for this plugin instance
134    fn hash_key(&self) -> String {
135        self.hash_value.clone()
136    }
137    async fn handle_request(
138        &self,
139        step: PluginStep,
140        session: &mut Session,
141        ctx: &mut Ctx,
142    ) -> pingora::Result<(bool, Option<HttpResponse>)> {
143        if step != PluginStep::Request {
144            return Ok((false, None));
145        }
146        // set cache key with accept image type
147        let mut accept_images = Vec::with_capacity(2);
148        if let Some(accept) = session.get_header(http::header::ACCEPT) {
149            if let Ok(accept) = accept.to_str() {
150                for item in self.output_types.iter() {
151                    let key = format!("image/{item}");
152                    if accept.contains(key.as_str()) {
153                        accept_images.push(key);
154                    }
155                }
156                accept_images.sort();
157            }
158        }
159        if !accept_images.is_empty() {
160            ctx.extend_cache_keys(accept_images);
161        }
162        Ok((false, None))
163    }
164    fn handle_upstream_response(
165        &self,
166        step: PluginStep,
167        session: &mut Session,
168        ctx: &mut Ctx,
169        upstream_response: &mut ResponseHeader,
170    ) -> pingora::Result<bool> {
171        // Skip if not at the correct plugin step
172        if self.plugin_step != step {
173            return Ok(false);
174        }
175        let Some(content_type) =
176            upstream_response.headers.get(http::header::CONTENT_TYPE)
177        else {
178            return Ok(false);
179        };
180        let Ok(content_type) = content_type.to_str() else {
181            return Ok(false);
182        };
183        let Some((content_type, image_type)) = content_type.split_once("/")
184        else {
185            return Ok(false);
186        };
187        if content_type != "image" {
188            return Ok(false);
189        }
190        let image_type = image_type.to_string();
191        if !self.support_types.contains(&image_type) {
192            return Ok(false);
193        }
194
195        let mut format_type = image_type.clone();
196        if let Some(accept) = session.get_header(http::header::ACCEPT) {
197            if let Ok(accept) = accept.to_str() {
198                for item in self.output_types.iter() {
199                    if accept.contains(format!("image/{item}").as_str()) {
200                        format_type = item.to_string();
201                        break;
202                    }
203                }
204            }
205        }
206        if format_type.is_empty() {
207            return Ok(false);
208        }
209        // Remove content-length since we're modifying the body
210        upstream_response.remove_header(&http::header::CONTENT_LENGTH);
211        // Switch to chunked transfer encoding
212        let _ = upstream_response.insert_header(
213            http::header::TRANSFER_ENCODING,
214            HTTP_HEADER_TRANSFER_CHUNKED.1.clone(),
215        );
216        let _ = upstream_response.insert_header(
217            http::header::CONTENT_TYPE,
218            format!("image/{format_type}").as_str(),
219        );
220
221        ctx.modify_upstream_response_body = Some(Box::new(ImageOptimizer {
222            image_type,
223            png_quality: self.png_quality,
224            jpeg_quality: self.jpeg_quality,
225            avif_quality: self.avif_quality,
226            avif_speed: self.avif_speed,
227            // only support lossless
228            webp_quality: 100,
229            format_type,
230        }));
231        Ok(true)
232    }
233}
234
235#[ctor]
236fn init() {
237    get_plugin_factory().register("image_optim", |params| {
238        Ok(Arc::new(ImageOptim::new(params)?))
239    });
240}
pingap_imageoptim/lib.rs

pingap_imageoptim/
lib.rs
