kroki_rs/utils/
mod.rs

1use anyhow::{Context, Result};
2use flate2::read::ZlibDecoder;
3use std::io::Read;
4
5pub mod font_manager;
6pub mod image_converter;
7
8/// Decodes a Base64URL-encoded + Zlib/Deflate-compressed string.
9///
10/// Attempts Zlib (RFC 1950) first, then falls back to raw Deflate (RFC 1951).
11pub fn decode(encoded: &str) -> Result<String> {
12    use base64::{
13        alphabet,
14        engine::{self, general_purpose},
15        Engine as _,
16    };
17
18    // Use a permissive engine that handles both padded and unpadded Base64URL strings (common in Kroki)
19    let engine = engine::GeneralPurpose::new(
20        &alphabet::URL_SAFE,
21        general_purpose::NO_PAD.with_decode_padding_mode(engine::DecodePaddingMode::Indifferent),
22    );
23
24    let decoded_bytes = engine
25        .decode(encoded)
26        .context("Base64URL decode failed — input is not valid Base64URL")?;
27
28    // Try Zlib first (RFC 1950)
29    let mut decoder = ZlibDecoder::new(&decoded_bytes[..]);
30    let mut s = String::new();
31    match decoder.read_to_string(&mut s) {
32        Ok(_) => return Ok(s),
33        Err(e) => {
34            tracing::debug!("Zlib decode failed ({}), trying raw Deflate", e);
35        }
36    }
37
38    // Fallback to raw Deflate (RFC 1951) — often used by pako
39    let mut decoder = flate2::read::DeflateDecoder::new(&decoded_bytes[..]);
40    let mut s = String::new();
41    match decoder.read_to_string(&mut s) {
42        Ok(_) => Ok(s),
43        Err(e) => {
44            // Distinguish decompression failure from UTF-8 encoding issues
45            if e.kind() == std::io::ErrorKind::InvalidData {
46                anyhow::bail!(
47                    "Decompression failed — input is not valid Zlib or Deflate compressed data: {}",
48                    e
49                )
50            } else {
51                anyhow::bail!("Decompressed data is not valid UTF-8 text: {}", e)
52            }
53        }
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn test_decode_graphviz() {
63        // "digraph G { Hello -> World }" encoded with Python's zlib.compress + base64.urlsafe_b64encode
64        let encoded = "eJxLyUwvSizIUHBXqFbwSM3JyVfQtVMIzy_KSVGoBQCJQglG";
65        let result = decode(encoded);
66        assert!(result.is_ok(), "decode failed: {:?}", result.err());
67        let decoded = result.unwrap();
68        assert_eq!(decoded, "digraph G { Hello -> World }");
69    }
70
71    #[test]
72    fn test_decode_invalid_base64() {
73        let result = decode("not-valid-base64!!!");
74        assert!(result.is_err());
75        let err = result.unwrap_err().to_string();
76        assert!(
77            err.contains("Base64URL"),
78            "Expected Base64 error, got: {}",
79            err
80        );
81    }
82
83    #[test]
84    fn test_decode_not_compressed() {
85        use base64::{engine::general_purpose, Engine as _};
86        // Valid base64 but not compressed data
87        let encoded = general_purpose::URL_SAFE.encode("hello world");
88        let result = decode(&encoded);
89        // This should fail because "hello world" is not zlib/deflate compressed
90        assert!(result.is_err());
91    }
92}