1pub mod error;
2
3pub use error::*;
4
5use std::collections::{BTreeMap, HashMap, VecDeque};
6use std::path::PathBuf;
7use std::str::from_utf8;
8use std::sync::{Arc, Mutex, OnceLock};
9
10use base64::Engine;
11use base64::engine::general_purpose::STANDARD as BASE64;
12use graphitepdf_svg::{SvgNode, try_parse_svg};
13
14const DEFAULT_CACHE_LIMIT: usize = 30;
15const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];
16
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
18pub enum ImageFormat {
19 Jpeg,
20 Png,
21 Svg,
22}
23
24impl ImageFormat {
25 fn from_str(value: &str) -> Option<Self> {
26 match value.trim().to_ascii_lowercase().as_str() {
27 "jpg" | "jpeg" => Some(Self::Jpeg),
28 "png" => Some(Self::Png),
29 "svg" | "svg+xml" => Some(Self::Svg),
30 _ => None,
31 }
32 }
33
34 pub const fn as_str(self) -> &'static str {
35 match self {
36 Self::Jpeg => "jpeg",
37 Self::Png => "png",
38 Self::Svg => "svg",
39 }
40 }
41}
42
43#[derive(Clone, Debug, PartialEq, Eq)]
44pub struct DataImageSource {
45 pub data: Vec<u8>,
46 pub format: ImageFormat,
47}
48
49impl DataImageSource {
50 pub fn new(data: impl Into<Vec<u8>>, format: ImageFormat) -> Self {
51 Self {
52 data: data.into(),
53 format,
54 }
55 }
56}
57
58#[derive(Clone, Debug, PartialEq, Eq)]
59pub struct LocalImageSource {
60 pub path: PathBuf,
61 pub format: Option<ImageFormat>,
62}
63
64impl LocalImageSource {
65 pub fn new(path: impl Into<PathBuf>) -> Self {
66 Self {
67 path: path.into(),
68 format: None,
69 }
70 }
71
72 pub fn with_format(mut self, format: ImageFormat) -> Self {
73 self.format = Some(format);
74 self
75 }
76}
77
78#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
79pub enum RemoteMethod {
80 #[default]
81 Get,
82 Head,
83 Post,
84 Put,
85 Delete,
86 Patch,
87}
88
89impl RemoteMethod {
90 fn as_reqwest_method(self) -> reqwest::Method {
91 match self {
92 Self::Get => reqwest::Method::GET,
93 Self::Head => reqwest::Method::HEAD,
94 Self::Post => reqwest::Method::POST,
95 Self::Put => reqwest::Method::PUT,
96 Self::Delete => reqwest::Method::DELETE,
97 Self::Patch => reqwest::Method::PATCH,
98 }
99 }
100}
101
102#[derive(Clone, Copy, Debug, PartialEq, Eq)]
103pub enum RemoteCredentials {
104 Omit,
105 SameOrigin,
106 Include,
107}
108
109#[derive(Clone, Debug, PartialEq, Eq)]
110pub struct RemoteImageSource {
111 pub uri: String,
112 pub method: RemoteMethod,
113 pub headers: BTreeMap<String, String>,
114 pub format: Option<ImageFormat>,
115 pub body: Option<Vec<u8>>,
116 pub credentials: Option<RemoteCredentials>,
117}
118
119impl RemoteImageSource {
120 pub fn new(uri: impl Into<String>) -> Self {
121 Self {
122 uri: uri.into(),
123 method: RemoteMethod::Get,
124 headers: BTreeMap::new(),
125 format: None,
126 body: None,
127 credentials: None,
128 }
129 }
130
131 pub fn with_method(mut self, method: RemoteMethod) -> Self {
132 self.method = method;
133 self
134 }
135
136 pub fn with_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
137 self.headers.insert(name.into(), value.into());
138 self
139 }
140
141 pub fn with_body(mut self, body: impl Into<Vec<u8>>) -> Self {
142 self.body = Some(body.into());
143 self
144 }
145
146 pub fn with_format(mut self, format: ImageFormat) -> Self {
147 self.format = Some(format);
148 self
149 }
150
151 pub fn with_credentials(mut self, credentials: RemoteCredentials) -> Self {
152 self.credentials = Some(credentials);
153 self
154 }
155}
156
157#[derive(Clone, Debug, PartialEq, Eq)]
158pub struct DataUriImageSource {
159 pub uri: String,
160}
161
162impl DataUriImageSource {
163 pub fn new(uri: impl Into<String>) -> Self {
164 Self { uri: uri.into() }
165 }
166}
167
168#[derive(Clone, Debug, PartialEq, Eq)]
169pub enum ImageSource {
170 Bytes(Vec<u8>),
171 Data(DataImageSource),
172 Local(LocalImageSource),
173 Remote(RemoteImageSource),
174 DataUri(DataUriImageSource),
175}
176
177impl ImageSource {
178 fn cache_key(&self) -> Option<String> {
179 match self {
180 Self::Bytes(_) => None,
181 Self::Data(source) => Some(BASE64.encode(&source.data)),
182 Self::Local(source) => Some(source.path.to_string_lossy().into_owned()),
183 Self::Remote(source) => Some(source.uri.clone()),
184 Self::DataUri(source) => Some(source.uri.clone()),
185 }
186 }
187}
188
189impl From<Vec<u8>> for ImageSource {
190 fn from(value: Vec<u8>) -> Self {
191 Self::Bytes(value)
192 }
193}
194
195impl From<&[u8]> for ImageSource {
196 fn from(value: &[u8]) -> Self {
197 Self::Bytes(value.to_vec())
198 }
199}
200
201impl From<DataImageSource> for ImageSource {
202 fn from(value: DataImageSource) -> Self {
203 Self::Data(value)
204 }
205}
206
207impl From<LocalImageSource> for ImageSource {
208 fn from(value: LocalImageSource) -> Self {
209 Self::Local(value)
210 }
211}
212
213impl From<RemoteImageSource> for ImageSource {
214 fn from(value: RemoteImageSource) -> Self {
215 Self::Remote(value)
216 }
217}
218
219impl From<DataUriImageSource> for ImageSource {
220 fn from(value: DataUriImageSource) -> Self {
221 Self::DataUri(value)
222 }
223}
224
225#[derive(Clone, Debug, PartialEq, Eq)]
226pub struct RasterImage {
227 pub width: u32,
228 pub height: u32,
229 pub data: Vec<u8>,
230 pub format: ImageFormat,
231 pub key: Option<String>,
232}
233
234#[derive(Clone, Debug, PartialEq)]
235pub struct SvgImage {
236 pub width: f32,
237 pub height: f32,
238 pub data: SvgNode,
239 pub raw_data: Vec<u8>,
240 pub key: Option<String>,
241}
242
243#[derive(Clone, Debug, PartialEq)]
244pub enum Image {
245 Raster(RasterImage),
246 Svg(SvgImage),
247}
248
249pub type ImageAsset = Image;
250
251impl Image {
252 pub fn format(&self) -> ImageFormat {
253 match self {
254 Self::Raster(image) => image.format,
255 Self::Svg(_) => ImageFormat::Svg,
256 }
257 }
258
259 pub fn width(&self) -> f32 {
260 match self {
261 Self::Raster(image) => image.width as f32,
262 Self::Svg(image) => image.width,
263 }
264 }
265
266 pub fn height(&self) -> f32 {
267 match self {
268 Self::Raster(image) => image.height as f32,
269 Self::Svg(image) => image.height,
270 }
271 }
272
273 pub fn key(&self) -> Option<&str> {
274 match self {
275 Self::Raster(image) => image.key.as_deref(),
276 Self::Svg(image) => image.key.as_deref(),
277 }
278 }
279
280 fn set_key(&mut self, key: String) {
281 match self {
282 Self::Raster(image) => image.key = Some(key),
283 Self::Svg(image) => image.key = Some(key),
284 }
285 }
286}
287
288#[derive(Debug)]
289pub struct ImageCache {
290 limit: usize,
291 state: Mutex<CacheState>,
292}
293
294#[derive(Debug, Default)]
295struct CacheState {
296 entries: HashMap<String, Arc<Image>>,
297 order: VecDeque<String>,
298}
299
300impl ImageCache {
301 pub fn new(limit: usize) -> Self {
302 Self {
303 limit: limit.max(1),
304 state: Mutex::new(CacheState::default()),
305 }
306 }
307
308 pub fn get(&self, key: &str) -> Option<Arc<Image>> {
309 let mut state = self.state.lock().expect("image cache mutex poisoned");
310 let value = state.entries.get(key).cloned();
311
312 if value.is_some() {
313 touch_key(&mut state.order, key);
314 }
315
316 value
317 }
318
319 pub fn set(&self, key: impl Into<String>, value: Arc<Image>) {
320 let key = key.into();
321 let mut state = self.state.lock().expect("image cache mutex poisoned");
322
323 state.entries.insert(key.clone(), value);
324 touch_key(&mut state.order, &key);
325
326 while state.entries.len() > self.limit {
327 if let Some(oldest) = state.order.pop_front() {
328 if oldest != key {
329 state.entries.remove(&oldest);
330 }
331 } else {
332 break;
333 }
334 }
335 }
336
337 pub fn reset(&self) {
338 let mut state = self.state.lock().expect("image cache mutex poisoned");
339 state.entries.clear();
340 state.order.clear();
341 }
342
343 pub fn len(&self) -> usize {
344 let state = self.state.lock().expect("image cache mutex poisoned");
345 state.entries.len()
346 }
347
348 pub fn is_empty(&self) -> bool {
349 self.len() == 0
350 }
351}
352
353fn touch_key(order: &mut VecDeque<String>, key: &str) {
354 if let Some(index) = order.iter().position(|existing| existing == key) {
355 order.remove(index);
356 }
357
358 order.push_back(key.to_string());
359}
360
361fn global_image_cache() -> &'static ImageCache {
362 static CACHE: OnceLock<ImageCache> = OnceLock::new();
363 CACHE.get_or_init(|| ImageCache::new(DEFAULT_CACHE_LIMIT))
364}
365
366#[derive(Clone, Debug, PartialEq, Eq)]
367pub struct ResolveImageOptions {
368 pub cache: bool,
369}
370
371impl Default for ResolveImageOptions {
372 fn default() -> Self {
373 Self { cache: true }
374 }
375}
376
377pub async fn resolve_image(source: impl Into<ImageSource>) -> Result<Arc<Image>> {
378 resolve_image_with_options(source, ResolveImageOptions::default()).await
379}
380
381pub async fn resolve_image_with_options(
382 source: impl Into<ImageSource>,
383 options: ResolveImageOptions,
384) -> Result<Arc<Image>> {
385 resolve_image_with_cache(source.into(), &options, global_image_cache()).await
386}
387
388async fn resolve_image_with_cache(
389 source: ImageSource,
390 options: &ResolveImageOptions,
391 cache: &ImageCache,
392) -> Result<Arc<Image>> {
393 let cache_key = source.cache_key();
394
395 if options.cache
396 && let Some(ref key) = cache_key
397 && let Some(image) = cache.get(key)
398 {
399 return Ok(image);
400 }
401
402 let mut image = match source {
403 ImageSource::Bytes(bytes) => resolve_bytes_image(bytes, None)?,
404 ImageSource::Data(source) => resolve_data_image(source)?,
405 ImageSource::Local(source) => resolve_local_image(source).await?,
406 ImageSource::Remote(source) => resolve_remote_image(source).await?,
407 ImageSource::DataUri(source) => resolve_data_uri_image(source)?,
408 };
409
410 if let Some(key) = cache_key {
411 image.set_key(key.clone());
412
413 let image = Arc::new(image);
414 if options.cache {
415 cache.set(key, Arc::clone(&image));
416 }
417
418 return Ok(image);
419 }
420
421 Ok(Arc::new(image))
422}
423
424fn resolve_data_image(source: DataImageSource) -> Result<Image> {
425 parse_image(source.data, source.format)
426}
427
428async fn resolve_local_image(source: LocalImageSource) -> Result<Image> {
429 let bytes = tokio::fs::read(source.path).await?;
430 resolve_bytes_image(bytes, source.format)
431}
432
433async fn resolve_remote_image(source: RemoteImageSource) -> Result<Image> {
434 let mut request =
435 reqwest::Client::new().request(source.method.as_reqwest_method(), &source.uri);
436
437 for (name, value) in source.headers {
438 request = request.header(name, value);
439 }
440
441 if let Some(body) = source.body {
442 request = request.body(body);
443 }
444
445 let bytes = request
446 .send()
447 .await?
448 .error_for_status()?
449 .bytes()
450 .await?
451 .to_vec();
452
453 resolve_bytes_image(bytes, source.format)
454}
455
456fn resolve_data_uri_image(source: DataUriImageSource) -> Result<Image> {
457 let payload =
458 source
459 .uri
460 .strip_prefix("data:image/")
461 .ok_or_else(|| Error::InvalidImageData {
462 message: format!("invalid image data URI: {}", source.uri),
463 })?;
464
465 let (metadata, encoded) = payload
466 .split_once(',')
467 .ok_or_else(|| Error::InvalidImageData {
468 message: format!("invalid image data URI: {}", source.uri),
469 })?;
470 let (format, encoding) = metadata
471 .split_once(';')
472 .ok_or_else(|| Error::InvalidImageData {
473 message: format!("invalid image data URI metadata: {metadata}"),
474 })?;
475
476 if !encoding.eq_ignore_ascii_case("base64") {
477 return Err(Error::InvalidImageData {
478 message: format!("unsupported image data URI encoding: {encoding}"),
479 });
480 }
481
482 let format = ImageFormat::from_str(format).ok_or_else(|| Error::UnsupportedImageFormat {
483 format: format.to_string(),
484 })?;
485 let data = BASE64.decode(encoded)?;
486
487 parse_image(data, format)
488}
489
490fn resolve_bytes_image(bytes: Vec<u8>, declared_format: Option<ImageFormat>) -> Result<Image> {
491 let format = sniff_image_format(&bytes)
492 .or(declared_format)
493 .ok_or_else(|| Error::InvalidImageData {
494 message: "unable to determine image format from bytes".to_string(),
495 })?;
496
497 parse_image(bytes, format)
498}
499
500fn parse_image(bytes: Vec<u8>, format: ImageFormat) -> Result<Image> {
501 match format {
502 ImageFormat::Png => parse_png(bytes).map(Image::Raster),
503 ImageFormat::Jpeg => parse_jpeg(bytes).map(Image::Raster),
504 ImageFormat::Svg => parse_svg(bytes).map(Image::Svg),
505 }
506}
507
508fn parse_png(data: Vec<u8>) -> Result<RasterImage> {
509 if !is_png(&data) {
510 return Err(Error::InvalidImageData {
511 message: "PNG signature not found".to_string(),
512 });
513 }
514
515 if data.len() < 24 || &data[12..16] != b"IHDR" {
516 return Err(Error::InvalidImageData {
517 message: "PNG missing IHDR chunk".to_string(),
518 });
519 }
520
521 let width = read_be_u32(&data[16..20])?;
522 let height = read_be_u32(&data[20..24])?;
523
524 Ok(RasterImage {
525 width,
526 height,
527 data,
528 format: ImageFormat::Png,
529 key: None,
530 })
531}
532
533fn parse_jpeg(data: Vec<u8>) -> Result<RasterImage> {
534 if !is_jpeg(&data) {
535 return Err(Error::InvalidImageData {
536 message: "SOI not found in JPEG".to_string(),
537 });
538 }
539
540 let mut offset = 2;
541 let mut width = None;
542 let mut height = None;
543 let mut orientation = None;
544
545 while offset + 1 < data.len() {
546 if data[offset] != 0xFF {
547 offset += 1;
548 continue;
549 }
550
551 while offset < data.len() && data[offset] == 0xFF {
552 offset += 1;
553 }
554
555 if offset >= data.len() {
556 break;
557 }
558
559 let marker = data[offset];
560 offset += 1;
561
562 if marker == 0xD9 || marker == 0xDA {
563 break;
564 }
565
566 if matches!(marker, 0x01 | 0xD0..=0xD7) {
567 continue;
568 }
569
570 if offset + 2 > data.len() {
571 break;
572 }
573
574 let segment_length = read_be_u16(&data[offset..offset + 2])? as usize;
575 if segment_length < 2 || offset + segment_length > data.len() {
576 return Err(Error::InvalidImageData {
577 message: "JPEG segment exceeds input length".to_string(),
578 });
579 }
580
581 let segment_start = offset + 2;
582 let segment_end = offset + segment_length;
583 let segment = &data[segment_start..segment_end];
584
585 if marker == 0xE1 {
586 orientation = parse_exif_orientation(segment)?;
587 } else if is_start_of_frame(marker) {
588 if segment.len() < 5 {
589 return Err(Error::InvalidImageData {
590 message: "JPEG SOF segment too short".to_string(),
591 });
592 }
593
594 height = Some(read_be_u16(&segment[1..3])? as u32);
595 width = Some(read_be_u16(&segment[3..5])? as u32);
596 }
597
598 offset = segment_end;
599 }
600
601 let mut width = width.ok_or_else(|| Error::InvalidImageData {
602 message: "JPEG dimensions not found".to_string(),
603 })?;
604 let mut height = height.ok_or_else(|| Error::InvalidImageData {
605 message: "JPEG dimensions not found".to_string(),
606 })?;
607
608 if matches!(orientation, Some(5..=8)) {
609 std::mem::swap(&mut width, &mut height);
610 }
611
612 Ok(RasterImage {
613 width,
614 height,
615 data,
616 format: ImageFormat::Jpeg,
617 key: None,
618 })
619}
620
621fn parse_svg(data: Vec<u8>) -> Result<SvgImage> {
622 if !is_svg(&data) {
623 return Err(Error::InvalidImageData {
624 message: "SVG signature not found".to_string(),
625 });
626 }
627
628 let svg_string = from_utf8(strip_utf8_bom(&data))?;
629 let parsed = try_parse_svg(svg_string)?;
630 let view_box = parsed
631 .props
632 .get("viewBox")
633 .and_then(|value| parse_view_box(value));
634 let width = parsed
635 .props
636 .get("width")
637 .and_then(|value| parse_svg_dimension(value))
638 .or_else(|| view_box.map(|view_box| view_box.width))
639 .unwrap_or(0.0);
640 let height = parsed
641 .props
642 .get("height")
643 .and_then(|value| parse_svg_dimension(value))
644 .or_else(|| view_box.map(|view_box| view_box.height))
645 .unwrap_or(0.0);
646
647 Ok(SvgImage {
648 width,
649 height,
650 data: parsed,
651 raw_data: data,
652 key: None,
653 })
654}
655
656fn sniff_image_format(data: &[u8]) -> Option<ImageFormat> {
657 if is_jpeg(data) {
658 Some(ImageFormat::Jpeg)
659 } else if is_png(data) {
660 Some(ImageFormat::Png)
661 } else if is_svg(data) {
662 Some(ImageFormat::Svg)
663 } else {
664 None
665 }
666}
667
668fn is_png(data: &[u8]) -> bool {
669 data.starts_with(&PNG_SIGNATURE)
670}
671
672fn is_jpeg(data: &[u8]) -> bool {
673 data.len() >= 2 && data[0] == 0xFF && data[1] == 0xD8
674}
675
676fn is_svg(data: &[u8]) -> bool {
677 let Ok(text) = from_utf8(strip_utf8_bom(data)) else {
678 return false;
679 };
680
681 let trimmed = text.trim_start();
682 trimmed.starts_with("<?xml") || trimmed.starts_with("<svg")
683}
684
685fn strip_utf8_bom(data: &[u8]) -> &[u8] {
686 data.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(data)
687}
688
689fn is_start_of_frame(marker: u8) -> bool {
690 matches!(
691 marker,
692 0xC0 | 0xC1 | 0xC2 | 0xC3 | 0xC5 | 0xC6 | 0xC7 | 0xC9 | 0xCA | 0xCB | 0xCD | 0xCE | 0xCF
693 )
694}
695
696fn parse_exif_orientation(segment: &[u8]) -> Result<Option<u16>> {
697 if !segment.starts_with(b"Exif\0\0") {
698 return Ok(None);
699 }
700
701 let tiff = &segment[6..];
702 if tiff.len() < 8 {
703 return Ok(None);
704 }
705
706 let big_endian = match &tiff[..2] {
707 b"MM" => true,
708 b"II" => false,
709 _ => return Ok(None),
710 };
711
712 let ifd_offset = read_endian_u32(&tiff[4..8], big_endian)? as usize;
713 if ifd_offset + 2 > tiff.len() {
714 return Ok(None);
715 }
716
717 let entry_count = read_endian_u16(&tiff[ifd_offset..ifd_offset + 2], big_endian)? as usize;
718 let mut entry_offset = ifd_offset + 2;
719
720 for _ in 0..entry_count {
721 if entry_offset + 12 > tiff.len() {
722 return Ok(None);
723 }
724
725 let entry = &tiff[entry_offset..entry_offset + 12];
726 let tag = read_endian_u16(&entry[0..2], big_endian)?;
727 let field_type = read_endian_u16(&entry[2..4], big_endian)?;
728 let count = read_endian_u32(&entry[4..8], big_endian)?;
729
730 if tag == 0x0112 && field_type == 3 && count >= 1 {
731 let value = if big_endian {
732 read_be_u16(&entry[8..10])?
733 } else {
734 read_le_u16(&entry[8..10])?
735 };
736
737 return Ok(Some(value));
738 }
739
740 entry_offset += 12;
741 }
742
743 Ok(None)
744}
745
746fn read_be_u16(bytes: &[u8]) -> Result<u16> {
747 let array = bytes.try_into().map_err(|_| Error::InvalidImageData {
748 message: "expected a 2-byte big-endian integer".to_string(),
749 })?;
750 Ok(u16::from_be_bytes(array))
751}
752
753fn read_be_u32(bytes: &[u8]) -> Result<u32> {
754 let array = bytes.try_into().map_err(|_| Error::InvalidImageData {
755 message: "expected a 4-byte big-endian integer".to_string(),
756 })?;
757 Ok(u32::from_be_bytes(array))
758}
759
760fn read_le_u16(bytes: &[u8]) -> Result<u16> {
761 let array = bytes.try_into().map_err(|_| Error::InvalidImageData {
762 message: "expected a 2-byte little-endian integer".to_string(),
763 })?;
764 Ok(u16::from_le_bytes(array))
765}
766
767fn read_endian_u16(bytes: &[u8], big_endian: bool) -> Result<u16> {
768 if big_endian {
769 read_be_u16(bytes)
770 } else {
771 read_le_u16(bytes)
772 }
773}
774
775fn read_endian_u32(bytes: &[u8], big_endian: bool) -> Result<u32> {
776 let array: [u8; 4] = bytes.try_into().map_err(|_| Error::InvalidImageData {
777 message: "expected a 4-byte endian integer".to_string(),
778 })?;
779
780 Ok(if big_endian {
781 u32::from_be_bytes(array)
782 } else {
783 u32::from_le_bytes(array)
784 })
785}
786
787#[derive(Clone, Copy, Debug)]
788struct ViewBox {
789 width: f32,
790 height: f32,
791}
792
793fn parse_view_box(value: &str) -> Option<ViewBox> {
794 let parts: Vec<_> = value
795 .split(|character: char| character.is_ascii_whitespace() || character == ',')
796 .filter(|part| !part.is_empty())
797 .collect();
798
799 if parts.len() != 4 {
800 return None;
801 }
802
803 let width = parts[2].parse::<f32>().ok()?;
804 let height = parts[3].parse::<f32>().ok()?;
805
806 Some(ViewBox { width, height })
807}
808
809fn parse_svg_dimension(value: &str) -> Option<f32> {
810 let value = value.trim();
811
812 for (suffix, multiplier) in [
813 ("px", 72.0 / 96.0),
814 ("pt", 1.0),
815 ("in", 72.0),
816 ("cm", 72.0 / 2.54),
817 ("mm", 72.0 / 25.4),
818 ] {
819 if let Some(number) = value.strip_suffix(suffix) {
820 return number
821 .trim()
822 .parse::<f32>()
823 .ok()
824 .map(|value| value * multiplier);
825 }
826 }
827
828 value.parse::<f32>().ok()
829}
830
831#[cfg(test)]
832#[allow(clippy::await_holding_lock)]
833mod tests {
834 use super::*;
835
836 use std::fs;
837 use std::time::{SystemTime, UNIX_EPOCH};
838
839 use tokio::io::{AsyncReadExt, AsyncWriteExt};
840 use tokio::net::TcpListener;
841
842 static TEST_MUTEX: OnceLock<Mutex<()>> = OnceLock::new();
843
844 const PNG_1X1: &[u8] = &[
845 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6,
846 0, 0, 0, 31, 21, 196, 137, 0, 0, 0, 13, 73, 68, 65, 84, 120, 218, 99, 248, 207, 192, 0, 0,
847 3, 1, 1, 0, 201, 254, 146, 239, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130,
848 ];
849
850 fn test_guard() -> std::sync::MutexGuard<'static, ()> {
851 TEST_MUTEX
852 .get_or_init(|| Mutex::new(()))
853 .lock()
854 .expect("test mutex poisoned")
855 }
856
857 fn reset_cache() {
858 global_image_cache().reset();
859 }
860
861 fn unique_temp_path(extension: &str) -> PathBuf {
862 let timestamp = SystemTime::now()
863 .duration_since(UNIX_EPOCH)
864 .expect("system clock should be after unix epoch")
865 .as_nanos();
866 std::env::temp_dir().join(format!("graphitepdf-image-{timestamp}.{extension}"))
867 }
868
869 fn png_data_uri() -> String {
870 format!("data:image/png;base64,{}", BASE64.encode(PNG_1X1))
871 }
872
873 fn svg_bytes() -> Vec<u8> {
874 br#"<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 320 180"></svg>"#.to_vec()
875 }
876
877 fn jpeg_with_orientation(width: u16, height: u16, orientation: u16) -> Vec<u8> {
878 let exif_payload = [
879 b'E',
880 b'x',
881 b'i',
882 b'f',
883 0,
884 0,
885 b'M',
886 b'M',
887 0,
888 42,
889 0,
890 0,
891 0,
892 8,
893 0,
894 1,
895 0x01,
896 0x12,
897 0,
898 3,
899 0,
900 0,
901 0,
902 1,
903 (orientation >> 8) as u8,
904 orientation as u8,
905 0,
906 0,
907 0,
908 0,
909 0,
910 0,
911 ];
912 let app1_length = (exif_payload.len() + 2) as u16;
913
914 let mut bytes = vec![0xFF, 0xD8, 0xFF, 0xE1];
915 bytes.extend_from_slice(&app1_length.to_be_bytes());
916 bytes.extend_from_slice(&exif_payload);
917 bytes.extend_from_slice(&[
918 0xFF,
919 0xC0,
920 0x00,
921 0x11,
922 0x08,
923 (height >> 8) as u8,
924 height as u8,
925 (width >> 8) as u8,
926 width as u8,
927 0x03,
928 0x01,
929 0x11,
930 0x00,
931 0x02,
932 0x11,
933 0x01,
934 0x03,
935 0x11,
936 0x01,
937 0xFF,
938 0xD9,
939 ]);
940 bytes
941 }
942
943 async fn serve_once(response_body: Vec<u8>) -> String {
944 let listener = TcpListener::bind("127.0.0.1:0")
945 .await
946 .expect("listener should bind");
947 let address = listener
948 .local_addr()
949 .expect("listener should have local address");
950
951 tokio::spawn(async move {
952 let (mut stream, _) = listener
953 .accept()
954 .await
955 .expect("connection should be accepted");
956 let mut request = vec![0_u8; 2048];
957 let _ = stream.read(&mut request).await;
958
959 let response = format!(
960 "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: image/png\r\nConnection: close\r\n\r\n",
961 response_body.len()
962 );
963
964 stream
965 .write_all(response.as_bytes())
966 .await
967 .expect("headers should be written");
968 stream
969 .write_all(&response_body)
970 .await
971 .expect("body should be written");
972 });
973
974 format!("http://{address}/image.png")
975 }
976
977 #[tokio::test]
978 async fn resolves_png_bytes_without_using_cache() {
979 let _guard = test_guard();
980 reset_cache();
981
982 let first = resolve_image(PNG_1X1.to_vec())
983 .await
984 .expect("first byte image should resolve");
985 let second = resolve_image(PNG_1X1.to_vec())
986 .await
987 .expect("second byte image should resolve");
988
989 assert!(
990 matches!(&*first, Image::Raster(image) if image.width == 1 && image.height == 1 && image.format == ImageFormat::Png)
991 );
992 assert!(!Arc::ptr_eq(&first, &second));
993 assert_eq!(global_image_cache().len(), 0);
994 }
995
996 #[tokio::test]
997 async fn resolves_data_source_and_reuses_cached_result() {
998 let _guard = test_guard();
999 reset_cache();
1000
1001 let source = DataImageSource::new(PNG_1X1, ImageFormat::Png);
1002
1003 let first = resolve_image(source.clone())
1004 .await
1005 .expect("data source should resolve");
1006 let second = resolve_image(source)
1007 .await
1008 .expect("cached data source should resolve");
1009
1010 assert!(Arc::ptr_eq(&first, &second));
1011 assert_eq!(first.key(), Some(BASE64.encode(PNG_1X1).as_str()));
1012 assert_eq!(global_image_cache().len(), 1);
1013 }
1014
1015 #[tokio::test]
1016 async fn resolves_png_from_data_uri() {
1017 let _guard = test_guard();
1018 reset_cache();
1019
1020 let image = resolve_image(DataUriImageSource::new(png_data_uri()))
1021 .await
1022 .expect("data URI image should resolve");
1023
1024 assert!(
1025 matches!(&*image, Image::Raster(raster) if raster.width == 1 && raster.height == 1)
1026 );
1027 assert_eq!(image.format(), ImageFormat::Png);
1028 assert_eq!(global_image_cache().len(), 1);
1029 }
1030
1031 #[tokio::test]
1032 async fn resolves_png_from_local_file() {
1033 let _guard = test_guard();
1034 reset_cache();
1035
1036 let path = unique_temp_path("png");
1037 fs::write(&path, PNG_1X1).expect("temp PNG should be written");
1038
1039 let result = resolve_image(LocalImageSource::new(&path)).await;
1040 let _ = fs::remove_file(&path);
1041
1042 let image = result.expect("local image should resolve");
1043 assert!(
1044 matches!(&*image, Image::Raster(raster) if raster.width == 1 && raster.height == 1)
1045 );
1046 assert_eq!(image.key(), Some(path.to_string_lossy().as_ref()));
1047 }
1048
1049 #[tokio::test]
1050 async fn resolves_png_from_remote_url() {
1051 let _guard = test_guard();
1052 reset_cache();
1053
1054 let uri = serve_once(PNG_1X1.to_vec()).await;
1055 let image = resolve_image(RemoteImageSource::new(uri.clone()))
1056 .await
1057 .expect("remote image should resolve");
1058
1059 assert!(
1060 matches!(&*image, Image::Raster(raster) if raster.width == 1 && raster.height == 1)
1061 );
1062 assert_eq!(image.key(), Some(uri.as_str()));
1063 }
1064
1065 #[tokio::test]
1066 async fn parses_svg_dimensions_from_view_box() {
1067 let _guard = test_guard();
1068 reset_cache();
1069
1070 let image = resolve_image(svg_bytes())
1071 .await
1072 .expect("SVG bytes should resolve");
1073
1074 assert!(matches!(&*image, Image::Svg(svg) if svg.width == 320.0 && svg.height == 180.0));
1075 assert_eq!(image.format(), ImageFormat::Svg);
1076 }
1077
1078 #[tokio::test]
1079 async fn parses_jpeg_and_applies_exif_orientation() {
1080 let _guard = test_guard();
1081 reset_cache();
1082
1083 let jpeg = jpeg_with_orientation(3, 2, 6);
1084 let image = resolve_image(jpeg).await.expect("JPEG should resolve");
1085
1086 assert!(
1087 matches!(&*image, Image::Raster(raster) if raster.width == 2 && raster.height == 3 && raster.format == ImageFormat::Jpeg)
1088 );
1089 }
1090
1091 #[tokio::test]
1092 async fn supports_disabling_cache() {
1093 let _guard = test_guard();
1094 reset_cache();
1095
1096 let source = DataImageSource::new(PNG_1X1, ImageFormat::Png);
1097 let options = ResolveImageOptions { cache: false };
1098
1099 let first = resolve_image_with_options(source.clone(), options.clone())
1100 .await
1101 .expect("uncached image should resolve");
1102 let second = resolve_image_with_options(source, options)
1103 .await
1104 .expect("second uncached image should resolve");
1105
1106 assert!(!Arc::ptr_eq(&first, &second));
1107 assert_eq!(global_image_cache().len(), 0);
1108 }
1109
1110 #[tokio::test]
1111 async fn evicts_least_recently_used_entries() {
1112 let _guard = test_guard();
1113
1114 let cache = ImageCache::new(2);
1115 let first = Arc::new(Image::Raster(RasterImage {
1116 width: 1,
1117 height: 1,
1118 data: PNG_1X1.to_vec(),
1119 format: ImageFormat::Png,
1120 key: Some("first".to_string()),
1121 }));
1122 let second = Arc::new(Image::Raster(RasterImage {
1123 width: 1,
1124 height: 1,
1125 data: PNG_1X1.to_vec(),
1126 format: ImageFormat::Png,
1127 key: Some("second".to_string()),
1128 }));
1129 let third = Arc::new(Image::Raster(RasterImage {
1130 width: 1,
1131 height: 1,
1132 data: PNG_1X1.to_vec(),
1133 format: ImageFormat::Png,
1134 key: Some("third".to_string()),
1135 }));
1136
1137 cache.set("first", Arc::clone(&first));
1138 cache.set("second", Arc::clone(&second));
1139 let cached_first = cache.get("first").expect("first entry should be present");
1140 assert!(Arc::ptr_eq(&cached_first, &first));
1141
1142 cache.set("third", Arc::clone(&third));
1143
1144 assert!(cache.get("first").is_some());
1145 assert!(cache.get("third").is_some());
1146 assert!(cache.get("second").is_none());
1147 }
1148
1149 #[test]
1150 fn parses_svg_dimensions_with_supported_units() {
1151 assert_eq!(parse_svg_dimension("96px"), Some(72.0));
1152 assert_eq!(parse_svg_dimension("2in"), Some(144.0));
1153 assert_eq!(parse_svg_dimension("10"), Some(10.0));
1154 }
1155
1156 #[test]
1157 fn reports_invalid_data_uris() {
1158 let error =
1159 resolve_data_uri_image(DataUriImageSource::new("data:text/plain;base64,SGVsbG8="))
1160 .expect_err("non-image data URI should be rejected");
1161
1162 assert!(matches!(error, Error::InvalidImageData { .. }));
1163 }
1164}