tail-fin-common 0.7.5

Shared infrastructure for tail-fin: error types, page_fetch, cookies, CDP helpers
Documentation
use std::path::Path;

use base64::Engine as _;
use image::codecs::jpeg::JpegEncoder;
use image::imageops::FilterType;

use crate::TailFinError;

const DEFAULT_MAX_IMAGE_COUNT: usize = 4;
const DEFAULT_MAX_IMAGE_BYTES: usize = 6 * 1024 * 1024;
const DEFAULT_MAX_TOTAL_IMAGE_BYTES: usize = 12 * 1024 * 1024;
const COMPRESS_TRIGGER_BYTES: usize = 1_500_000;
const TARGET_COMPRESSED_BYTES: usize = 1_200_000;
const MAX_EDGE_PIXELS: u32 = 2048;

#[derive(Debug, Clone, Copy)]
pub struct AttachmentLimits {
    pub max_image_count: usize,
    pub max_image_bytes: usize,
    pub max_total_image_bytes: usize,
}

impl Default for AttachmentLimits {
    fn default() -> Self {
        Self {
            max_image_count: DEFAULT_MAX_IMAGE_COUNT,
            max_image_bytes: DEFAULT_MAX_IMAGE_BYTES,
            max_total_image_bytes: DEFAULT_MAX_TOTAL_IMAGE_BYTES,
        }
    }
}

impl AttachmentLimits {
    pub fn with_overrides(
        max_image_mb: Option<usize>,
        max_total_image_mb: Option<usize>,
    ) -> Result<Self, TailFinError> {
        let mut v = Self::default();
        if let Some(mb) = max_image_mb {
            if mb == 0 {
                return Err(TailFinError::Api("max_image_mb must be >= 1".into()));
            }
            v.max_image_bytes = mb.saturating_mul(1024 * 1024);
            v.max_total_image_bytes = v.max_image_bytes.saturating_mul(v.max_image_count);
        }
        if let Some(total_mb) = max_total_image_mb {
            if total_mb == 0 {
                return Err(TailFinError::Api("max_total_image_mb must be >= 1".into()));
            }
            v.max_total_image_bytes = total_mb.saturating_mul(1024 * 1024);
        }
        if v.max_total_image_bytes < v.max_image_bytes {
            return Err(TailFinError::Api(
                "max_total_image_mb must be >= max_image_mb".into(),
            ));
        }
        Ok(v)
    }
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ImageAttachmentDebug {
    pub original_path: String,
    pub sent_name: String,
    pub original_bytes: usize,
    pub sent_bytes: usize,
    pub mime: String,
    pub compressed: bool,
}

#[derive(Debug, Clone)]
pub struct PreparedImageAttachment {
    pub data_url: String,
    pub debug: ImageAttachmentDebug,
}

pub fn prepare_image_data_urls(
    paths: &[String],
    limits: AttachmentLimits,
) -> Result<Vec<PreparedImageAttachment>, TailFinError> {
    if paths.len() > limits.max_image_count {
        return Err(TailFinError::Api(format!(
            "too many images: {} (limit {})",
            paths.len(),
            limits.max_image_count
        )));
    }

    let mut out = Vec::with_capacity(paths.len());
    let mut total = 0usize;
    for p in paths {
        let prepared = prepare_one_image(p, limits)?;
        total += prepared.debug.sent_bytes;
        if total > limits.max_total_image_bytes {
            return Err(TailFinError::Api(format!(
                "total image payload too large: {} (limit {})",
                human_bytes(total),
                human_bytes(limits.max_total_image_bytes)
            )));
        }
        out.push(prepared);
    }
    Ok(out)
}

fn prepare_one_image(
    path: &str,
    limits: AttachmentLimits,
) -> Result<PreparedImageAttachment, TailFinError> {
    let src =
        std::fs::read(path).map_err(|e| TailFinError::Io(format!("read image '{}': {e}", path)))?;

    let original_bytes = src.len();
    let mut sent = src;
    let mut mime = mime_from_path(path).to_string();
    let mut sent_name = Path::new(path)
        .file_name()
        .and_then(|s| s.to_str())
        .unwrap_or("image")
        .to_string();
    let mut compressed = false;

    if original_bytes > COMPRESS_TRIGGER_BYTES {
        if let Ok(reencoded) = compress_image_to_jpeg(&sent) {
            if reencoded.len() < sent.len() {
                sent = reencoded;
                mime = "image/jpeg".to_string();
                let stem = Path::new(path)
                    .file_stem()
                    .and_then(|s| s.to_str())
                    .unwrap_or("image");
                sent_name = format!("{stem}.jpg");
                compressed = true;
            }
        }
    }

    if sent.len() > limits.max_image_bytes {
        return Err(TailFinError::Api(format!(
            "image too large: '{}' is {} (limit {}). Try a smaller image.",
            path,
            human_bytes(sent.len()),
            human_bytes(limits.max_image_bytes)
        )));
    }

    let b64 = base64::engine::general_purpose::STANDARD.encode(&sent);
    let data_url = format!("data:{mime};base64,{b64}");
    Ok(PreparedImageAttachment {
        data_url,
        debug: ImageAttachmentDebug {
            original_path: path.to_string(),
            sent_name,
            original_bytes,
            sent_bytes: sent.len(),
            mime,
            compressed,
        },
    })
}

fn compress_image_to_jpeg(src: &[u8]) -> Result<Vec<u8>, TailFinError> {
    let decoded = image::load_from_memory(src)
        .map_err(|e| TailFinError::Api(format!("decode image for compression: {e}")))?;

    let resized = if decoded.width() > MAX_EDGE_PIXELS || decoded.height() > MAX_EDGE_PIXELS {
        decoded.resize(MAX_EDGE_PIXELS, MAX_EDGE_PIXELS, FilterType::Lanczos3)
    } else {
        decoded
    };

    let rgb = resized.to_rgb8();
    let (w, h) = rgb.dimensions();

    let mut best = Vec::new();
    for quality in [82u8, 72, 62, 52] {
        let mut buf = Vec::new();
        let mut enc = JpegEncoder::new_with_quality(&mut buf, quality);
        enc.encode(&rgb, w, h, image::ColorType::Rgb8.into())
            .map_err(|e| TailFinError::Api(format!("encode jpeg for compression: {e}")))?;

        if best.is_empty() || buf.len() < best.len() {
            best = buf;
        }
        if best.len() <= TARGET_COMPRESSED_BYTES {
            break;
        }
    }

    if best.is_empty() {
        return Err(TailFinError::Api("image compression failed".into()));
    }
    Ok(best)
}

pub fn human_bytes(n: usize) -> String {
    const KB: f64 = 1024.0;
    const MB: f64 = KB * 1024.0;
    let v = n as f64;
    if v >= MB {
        format!("{:.2} MB", v / MB)
    } else if v >= KB {
        format!("{:.1} KB", v / KB)
    } else {
        format!("{} B", n)
    }
}

pub fn mime_from_path(path: &str) -> &'static str {
    let ext = Path::new(path)
        .extension()
        .and_then(|s| s.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();
    match ext.as_str() {
        "png" => "image/png",
        "jpg" | "jpeg" => "image/jpeg",
        "webp" => "image/webp",
        "gif" => "image/gif",
        "bmp" => "image/bmp",
        _ => "application/octet-stream",
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn limits_with_overrides_validate_bounds() {
        assert!(AttachmentLimits::with_overrides(Some(0), None).is_err());
        assert!(AttachmentLimits::with_overrides(None, Some(0)).is_err());
    }

    #[test]
    fn prepare_image_data_urls_rejects_too_many_images() {
        let limits = AttachmentLimits::default();
        let paths: Vec<String> = (0..(limits.max_image_count + 1))
            .map(|i| format!("/tmp/fake-{i}.png"))
            .collect();
        let err = prepare_image_data_urls(&paths, limits).unwrap_err();
        assert!(err.to_string().contains("too many images"));
    }

    #[test]
    fn human_bytes_formats_units() {
        assert_eq!(human_bytes(512), "512 B");
        assert!(human_bytes(2048).contains("KB"));
        assert!(human_bytes(2 * 1024 * 1024).contains("MB"));
    }
}