1use cra::{ArcEntry, ArcError, ArcReader, ArcWriter};
6use image::{
7 codecs::{
8 jpeg::JpegEncoder,
9 png::{CompressionType, FilterType, PngEncoder},
10 webp::WebPEncoder,
11 },
12 ColorType, DynamicImage, ImageError, ImageReader,
13};
14use indicatif::{style::TemplateError, ProgressBar, ProgressStyle};
15use infer::image::is_jxl;
16use jxl_oxide::integration::JxlDecoder;
17use libavif_image::{is_avif, read as read_avif, save as save_avif, Error as AvifError};
18use rayon::prelude::{IntoParallelIterator, ParallelIterator};
19use sha2::{Digest, Sha256};
20use std::{
21 fmt::Display,
22 fs::{rename, File},
23 io::{self, Cursor, Read, Write},
24 net::TcpStream,
25 str::FromStr,
26 sync::{Arc, Mutex},
27 time::Duration,
28};
29use thiserror::Error;
30use zune_core::{bit_depth::BitDepth, colorspace::ColorSpace, options::EncoderOptions};
31use zune_jpegxl::{JxlEncodeErrors, JxlSimpleEncoder};
32
33#[derive(Error, Debug)]
35#[error(transparent)]
36pub enum ConvError {
37 ArcError(#[from] ArcError),
38 IoError(#[from] io::Error),
39 TemplateError(#[from] TemplateError),
40 AvifError(#[from] AvifError),
41 ImageError(#[from] ImageError),
42 #[error("{0:?}")]
43 JxlEncodeError(JxlEncodeErrors),
44 #[error("Invalid server response")]
45 InvalidResponse,
46 #[error("Hash mismatch")]
47 HashMismatch,
48}
49
50pub type ConvResult<T> = Result<T, ConvError>;
51
52#[derive(Clone, Copy, Debug)]
54pub enum Format {
55 Jpeg,
56 JpegXL,
57 Png,
58 Webp,
59 Avif,
60}
61
62impl Display for Format {
63 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64 f.write_str(match self {
65 Format::Jpeg => "jpg",
66 Format::JpegXL => "jxl",
67 Format::Png => "png",
68 Format::Webp => "webp",
69 Format::Avif => "avif",
70 })
71 }
72}
73
74impl FromStr for Format {
75 type Err = String;
76
77 fn from_str(s: &str) -> Result<Self, Self::Err> {
78 match s.to_ascii_lowercase().as_str() {
79 "avif" => Ok(Format::Avif),
80 "jpeg" | "jpg" => Ok(Format::Jpeg),
81 "jxl" => Ok(Format::JpegXL),
82 "webp" => Ok(Format::Webp),
83 "png" => Ok(Format::Png),
84 _ => Err(format!("Invalid format: {s}")),
85 }
86 }
87}
88
89#[derive(Clone, Copy, Debug)]
92pub struct Converter {
93 pub quality: u8,
94 pub speed: u8,
95 pub format: Format,
96 pub backup: bool,
97 pub quiet: bool,
98}
99
100impl Default for Converter {
101 fn default() -> Self {
102 Self {
103 quality: 30,
104 speed: 3,
105 format: Format::Avif,
106 backup: false,
107 quiet: false,
108 }
109 }
110}
111
112impl Converter {
113 pub fn convert_file(self, file: &str) -> ConvResult<()> {
115 let buf = {
116 let mut buf = Vec::new();
117 File::open(file)?.read_to_end(&mut buf)?;
118 buf
119 };
120 if !self.quiet {
121 println!("Converting {}...", file);
122 }
123 let data = self.convert(&buf, None)?;
124 if self.backup {
125 rename(file, format!("{}.bak", file))?;
126 }
127 File::create(file)?.write_all(&data)?;
128 Ok(())
129 }
130
131 pub fn convert_file_online(self, file: &str, stream: &mut TcpStream) -> ConvResult<()> {
133 let buf = {
134 let mut buf = Vec::new();
135 File::open(file)?.read_to_end(&mut buf)?;
136 buf
137 };
138 if !self.quiet {
139 println!("Converting {}...", file);
140 }
141 let data = self.convert_online(&buf, stream)?;
142 if self.backup {
143 rename(file, format!("{}.bak", file))?;
144 }
145 File::create(file)?.write_all(&data)?;
146 Ok(())
147 }
148
149 pub fn convert(
152 mut self,
153 buf: &[u8],
154 mut status_stream: Option<&mut TcpStream>,
155 ) -> ConvResult<Vec<u8>> {
156 self.speed = self.speed.clamp(0, 10);
157 self.quality = self.quality.clamp(0, 100);
158 let mut archive = ArcReader::new(buf)?;
159 let mut writer = ArcWriter::new(archive.format());
160 let file_count = archive
161 .by_ref()
162 .filter(|entry| !matches!(entry, ArcEntry::Directory(_)))
163 .count();
164 if let Some(ref mut stream) = status_stream {
165 stream.write_all(&(file_count as u32).to_be_bytes())?;
166 }
167 let mut bar = if self.quiet {
168 ProgressBar::hidden()
169 } else {
170 ProgressBar::new(file_count as u64)
171 };
172 bar.set_style(
173 ProgressStyle::default_bar()
174 .template("Convert [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})")?
175 .progress_chars("=>-"),
176 );
177 let status_stream = status_stream.map(|stream| Arc::new(Mutex::new(stream)));
178 let pb = Arc::new(Mutex::new(&mut bar));
179 writer.extend(
180 &archive
181 .entries()
182 .into_par_iter()
183 .map(|entry| {
184 Ok(match entry {
185 ArcEntry::File(name, data) => {
186 let data = self.convert_image(data)?;
187 if let Some(stream) = status_stream.clone() {
188 stream.lock().unwrap().write_all(b"plus")?
189 }
190 pb.clone().lock().unwrap().inc(1);
191 ArcEntry::File(
192 format!(
193 "{}.{}",
194 name.rsplit_once('.').unwrap_or((name, "")).0,
195 self.format
196 ),
197 data,
198 )
199 }
200 other => other.clone(),
201 })
202 })
203 .collect::<ConvResult<Vec<ArcEntry>>>()?,
204 );
205 bar.finish();
206 Ok(writer.archive()?)
207 }
208
209 pub fn convert_online(mut self, buf: &[u8], stream: &mut TcpStream) -> ConvResult<Vec<u8>> {
211 self.speed = self.speed.clamp(0, 10);
212 self.quality = self.quality.clamp(0, 100);
213 stream.set_nodelay(true)?;
214 stream.set_read_timeout(Some(Duration::from_secs(10)))?;
215 stream.write_all(b"comi")?;
216 {
217 let mut buf = [0; 4];
218 stream.read_exact(&mut buf)?;
219 if &buf != b"conv" {
220 return Err(ConvError::InvalidResponse);
221 }
222 }
223 let format = match self.format {
224 Format::Avif => b'A',
225 Format::Webp => b'W',
226 Format::Png => b'P',
227 Format::Jpeg => b'J',
228 Format::JpegXL => todo!(),
229 };
230 let mut left = buf.len();
231 {
232 let mut buf = [0; 8];
233 buf[0] = format;
234 buf[1] = self.speed;
235 buf[2] = self.quality;
236 buf[4..].copy_from_slice(&(left as u32).to_be_bytes());
237 stream.write_all(&buf)?;
238 }
239 let hash = {
240 let mut hasher = Sha256::new();
241 hasher.update(buf);
242 hasher.finalize()
243 };
244 stream.write_all(&hash)?;
245 let mut sent = 0;
246 let pb = if self.quiet {
247 ProgressBar::hidden()
248 } else {
249 ProgressBar::new(left as u64)
250 };
251 pb.set_style(
252 ProgressStyle::default_bar()
253 .template("Upload [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")?
254 .progress_chars("=>-"),
255 );
256 while left > 0 {
257 let size = left.min(1024 * 1024);
258 stream.write_all(&buf[sent..sent + size])?;
259 let mut buf = [0; 2];
260 stream.read_exact(&mut buf)?;
261 if &buf != b"ok" {
262 return Err(ConvError::InvalidResponse);
263 }
264 sent += size;
265 left = left.saturating_sub(size);
266 pb.inc(size as u64);
267 }
268 pb.finish();
269 let mut left = {
270 let mut buf = [0; 4];
271 stream.read_exact(&mut buf)?;
272 u32::from_be_bytes(buf)
273 };
274 let pb = if self.quiet {
275 ProgressBar::hidden()
276 } else {
277 ProgressBar::new(left as u64)
278 };
279 pb.set_style(
280 ProgressStyle::default_bar()
281 .template("Convert [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})")?
282 .progress_chars("=>-"),
283 );
284 while left > 0 {
285 let response = {
286 let mut buf = [0; 4];
287 stream.read_exact(&mut buf)?;
288 buf
289 };
290 if &response != b"plus" {
291 return Err(ConvError::InvalidResponse);
292 }
293 pb.inc(1);
294 left -= 1;
295 }
296 pb.finish();
297 let mut left = {
298 let mut buf = [0; 4];
299 stream.read_exact(&mut buf)?;
300 u32::from_be_bytes(buf)
301 };
302 let hash = {
303 let mut buf = [0; 32];
304 stream.read_exact(&mut buf)?;
305 buf
306 };
307 let pb = if self.quiet {
308 ProgressBar::hidden()
309 } else {
310 ProgressBar::new(left as u64)
311 };
312 pb.set_style(
313 ProgressStyle::default_bar()
314 .template("Download [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")?
315 .progress_chars("=>-"),
316 );
317 let mut data = Vec::new();
318 while left > 0 {
319 let mut buf = [0; 1024 * 1024];
320 let read = stream.read(&mut buf)?;
321 pb.inc(read as u64);
322 data.extend_from_slice(&buf[..read]);
323 left = left.saturating_sub(read as u32);
324 }
325 pb.finish();
326 let mut hasher = Sha256::new();
327 hasher.update(&data);
328 if hasher.finalize() != hash.into() {
329 return Err(ConvError::HashMismatch);
330 }
331 Ok(data)
332 }
333
334 fn convert_image(self, buf: &[u8]) -> ConvResult<Vec<u8>> {
335 let image = if is_avif(buf) {
336 read_avif(buf)?
337 } else if is_jxl(buf) {
338 DynamicImage::from_decoder(JxlDecoder::new(buf)?)?
339 } else {
340 ImageReader::new(Cursor::new(buf))
341 .with_guessed_format()?
342 .decode()?
343 };
344 let mut data = Vec::new();
345 match self.format {
346 Format::Avif => {
347 data = save_avif(&image)?.to_vec();
348 }
349 Format::Webp => image.write_with_encoder(WebPEncoder::new_lossless(&mut data))?,
350 Format::Png => image.write_with_encoder(PngEncoder::new_with_quality(
351 &mut data,
352 match self.speed.clamp(0, 2) {
353 0 => CompressionType::Fast,
354 1 => CompressionType::Default,
355 2 => CompressionType::Best,
356 _ => unreachable!(),
357 },
358 FilterType::Adaptive,
359 ))?,
360 Format::Jpeg => image
361 .into_rgb8()
362 .write_with_encoder(JpegEncoder::new_with_quality(&mut data, self.quality))?,
363 Format::JpegXL => {
364 let (color, depth) = image_to_zune_colot_type(&image);
365 data = JxlSimpleEncoder::new(
366 image.as_bytes(),
367 EncoderOptions::new(image.width() as _, image.height() as _, color, depth),
368 )
369 .encode()
370 .map_err(ConvError::JxlEncodeError)?;
371 }
372 }
373 Ok(data)
374 }
375}
376
377fn image_to_zune_colot_type(image: &DynamicImage) -> (ColorSpace, BitDepth) {
378 match image.color() {
379 ColorType::L8 => (ColorSpace::Luma, BitDepth::Eight),
380 ColorType::La16 => (ColorSpace::LumaA, BitDepth::Sixteen),
381 ColorType::Rgb16 => (ColorSpace::RGB, BitDepth::Sixteen),
382 ColorType::Rgba16 => (ColorSpace::RGBA, BitDepth::Sixteen),
383 ColorType::Rgb32F => (ColorSpace::RGB, BitDepth::Float32),
384 ColorType::Rgba32F => (ColorSpace::RGBA, BitDepth::Float32),
385 ColorType::Rgba8 => (ColorSpace::RGBA, BitDepth::Eight),
386 ColorType::L16 => (ColorSpace::Luma, BitDepth::Sixteen),
387 ColorType::La8 => (ColorSpace::LumaA, BitDepth::Eight),
388 ColorType::Rgb8 => (ColorSpace::RGB, BitDepth::Eight),
389 _ => unimplemented!(),
390 }
391}