photos/src/routes/zip.rs

//! Streaming ZIP downloads, hand-written for spec compliance.
//!
//! Every entry is Stored (raws/jpegs don't compress) with its exact size and
//! CRC-32 in the local file header — no data descriptors — so the archives
//! work with strict *streaming* extractors (Java ZipInputStream, bsdtar from
//! a pipe), not just central-directory readers. Sizes are known up front,
//! which also makes the total byte length deterministic: the response carries
//! a real Content-Length, so browsers show progress and flag truncated
//! downloads as failed.
//!
//! Each file is spooled from S3 to an anonymous temp file to compute its CRC
//! before its header is written; the next file spools while the current one
//! streams out, so S3 latency doesn't stall the download.

use std::collections::HashSet;

use axum::body::{Body, Bytes};
use axum::extract::{Form, Path, State};
use axum::http::{header, StatusCode};
use axum::response::Response;
use axum_extra::extract::cookie::SignedCookieJar;
use chrono::{DateTime, Datelike, Timelike, Utc};
use serde::Deserialize;
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
use tokio::task::JoinHandle;
use uuid::Uuid;

use crate::error::{ApiError, ApiResult};
use crate::models::{Photo, PhotoStatus};
use crate::routes::client::authorize_share;
use crate::s3;
use crate::state::AppState;

const U32_SENTINEL: u64 = 0xFFFF_FFFF;

/// Sent as a plain form POST (not fetch) so the browser streams the download
/// natively; `ids` is a comma-separated list, empty = every ready photo.
#[derive(Deserialize)]
pub struct ZipRequest {
    #[serde(default)]
    ids: String,
}

/// Strict: a present-but-malformed id is a 400, never silently reinterpreted
/// as the download-everything sentinel.
fn parse_ids(raw: &str) -> Result<Vec<Uuid>, ApiError> {
    raw.split(',')
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .map(|s| {
            Uuid::parse_str(s).map_err(|_| ApiError::bad_request(format!("invalid photo id: {s}")))
        })
        .collect()
}

async fn ready_photos(
    state: &AppState,
    album_id: Uuid,
    ids: &[Uuid],
) -> Result<Vec<Photo>, sqlx::Error> {
    sqlx::query_as(
        "select * from photos
         where album_id = $1 and status = $3
           and (cardinality($2::uuid[]) = 0 or id = any($2))
         order by coalesce(taken_at, created_at), filename",
    )
    .bind(album_id)
    .bind(ids)
    .bind(PhotoStatus::Ready.as_str())
    .fetch_all(&state.db)
    .await
}

async fn album_name(state: &AppState, album_id: Uuid) -> Result<String, ApiError> {
    let name: Option<(String,)> = sqlx::query_as("select name from albums where id = $1")
        .bind(album_id)
        .fetch_optional(&state.db)
        .await?;
    Ok(name.ok_or_else(ApiError::not_found)?.0)
}

pub async fn album_zip(
    State(state): State<AppState>,
    Path(album_id): Path<Uuid>,
    Form(request): Form<ZipRequest>,
) -> ApiResult<Response> {
    let name = album_name(&state, album_id).await?;
    let photos = ready_photos(&state, album_id, &parse_ids(&request.ids)?).await?;
    stream_zip(state, photos, &name)
}

pub async fn share_zip(
    State(state): State<AppState>,
    Path(token): Path<String>,
    jar: SignedCookieJar,
    Form(request): Form<ZipRequest>,
) -> ApiResult<Response> {
    let share = authorize_share(&state, &jar, &token, true).await?;
    let name = album_name(&state, share.album_id).await?;
    let photos = ready_photos(&state, share.album_id, &parse_ids(&request.ids)?).await?;
    stream_zip(state, photos, &name)
}

/// Dedupe case-insensitively — archives are extracted onto case-insensitive
/// filesystems (macOS/Windows), where "DSC1.JPG" and "dsc1.jpg" would collide.
fn unique_entry_name(used: &mut HashSet<String>, filename: &str) -> String {
    if used.insert(filename.to_lowercase()) {
        return filename.to_string();
    }
    let (stem, ext) = match filename.rsplit_once('.') {
        Some((stem, ext)) if !stem.is_empty() => (stem, format!(".{ext}")),
        _ => (filename, String::new()),
    };
    for n in 2.. {
        let candidate = format!("{stem} ({n}){ext}");
        if used.insert(candidate.to_lowercase()) {
            return candidate;
        }
    }
    unreachable!()
}

/// MS-DOS timestamp (2-second resolution, no timezone; years 1980+ only —
/// callers clamp earlier dates).
fn dos_datetime(t: DateTime<Utc>) -> (u16, u16) {
    let time = ((t.hour() as u16) << 11) | ((t.minute() as u16) << 5) | (t.second() as u16 / 2);
    let date = (((t.year() - 1980) as u16) << 9) | ((t.month() as u16) << 5) | (t.day() as u16);
    (time, date)
}

struct Entry {
    photo_id: Uuid,
    s3_key: String,
    name: Vec<u8>,
    size: u64,
    offset: u64,
    dos_time: u16,
    dos_date: u16,
    /// Stored at upload/processing time; photos from before hashes existed
    /// have None and take the slower spool path (which backfills it).
    crc: Option<u32>,
}

struct ZipPlan {
    entries: Vec<Entry>,
    cd_offset: u64,
    cd_size: u64,
    zip64_eocd: bool,
    total_len: u64,
}

fn plan_zip(photos: &[Photo]) -> Result<ZipPlan, ApiError> {
    let mut used_names = HashSet::new();
    let mut entries = Vec::with_capacity(photos.len());
    let mut offset: u64 = 0;
    for photo in photos {
        let size = photo.size_bytes as u64;
        if size >= U32_SENTINEL {
            return Err(ApiError::bad_request(format!(
                "{} is too large for a zip download",
                photo.filename
            )));
        }
        let name = unique_entry_name(&mut used_names, &photo.filename).into_bytes();
        // DOS timestamps can't represent pre-1980 dates (cameras with unset
        // clocks); fall back to the upload time.
        let timestamp = photo
            .taken_at
            .filter(|t| t.year() >= 1980)
            .unwrap_or(photo.created_at);
        let (dos_time, dos_date) = dos_datetime(timestamp);
        let entry_offset = offset;
        offset += 30 + name.len() as u64 + size;
        entries.push(Entry {
            photo_id: photo.id,
            s3_key: s3::original_key(photo.id, &photo.filename),
            name,
            size,
            offset: entry_offset,
            dos_time,
            dos_date,
            crc: photo.crc32.map(|v| v as u32),
        });
    }
    let cd_offset = offset;
    let cd_size: u64 = entries
        .iter()
        .map(|e| 46 + e.name.len() as u64 + if e.offset >= U32_SENTINEL { 12 } else { 0 })
        .sum();
    let zip64_eocd = entries.len() >= 0xFFFF
        || cd_size >= U32_SENTINEL
        || cd_offset >= U32_SENTINEL;
    let total_len = cd_offset + cd_size + 22 + if zip64_eocd { 56 + 20 } else { 0 };
    Ok(ZipPlan {
        entries,
        cd_offset,
        cd_size,
        zip64_eocd,
        total_len,
    })
}

fn stream_zip(state: AppState, photos: Vec<Photo>, album_name: &str) -> ApiResult<Response> {
    if photos.is_empty() {
        return Err(ApiError::bad_request("no downloadable photos selected"));
    }
    let plan = plan_zip(&photos)?;
    let permit = state.zip_permits.clone().try_acquire_owned().map_err(|_| {
        ApiError(
            StatusCode::SERVICE_UNAVAILABLE,
            "too many downloads in progress — try again in a moment".into(),
        )
    })?;

    let zip_name: String = album_name
        .trim()
        .chars()
        .map(|c| if c.is_ascii_alphanumeric() || matches!(c, ' ' | '-' | '_') { c } else { '_' })
        .take(80)
        .collect();
    let zip_name = if zip_name.is_empty() { "photos".to_string() } else { zip_name };
    let total_len = plan.total_len;

    let (writer, reader) = tokio::io::duplex(256 * 1024);
    let write_task = tokio::spawn(async move {
        let result = write_zip(&state, plan, writer).await;
        drop(permit);
        result
    });

    // Forward bytes to the response; when the writer finishes, surface any
    // zip error as a stream error so hyper ABORTS the connection — combined
    // with the exact Content-Length, browsers report truncation as a failed
    // download instead of keeping a silently corrupt file.
    struct Pump {
        reader: tokio::io::DuplexStream,
        task: Option<JoinHandle<anyhow::Result<()>>>,
    }
    let pump = Pump {
        reader,
        task: Some(write_task),
    };
    let stream = futures::stream::unfold(pump, |mut pump| async move {
        pump.task.as_ref()?; // stream is over after a terminal item
        let mut buf = vec![0u8; 64 * 1024];
        match pump.reader.read(&mut buf).await {
            Ok(0) => {
                let task = pump.task.take()?;
                match task.await {
                    Ok(Ok(())) => None,
                    Ok(Err(e)) => Some((
                        Err(std::io::Error::other(format!("zip stream failed: {e:#}"))),
                        pump,
                    )),
                    Err(e) => Some((
                        Err(std::io::Error::other(format!("zip task panicked: {e}"))),
                        pump,
                    )),
                }
            }
            Ok(n) => {
                buf.truncate(n);
                Some((Ok(Bytes::from(buf)), pump))
            }
            Err(e) => {
                pump.task = None;
                Some((Err(e), pump))
            }
        }
    });

    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "application/zip")
        .header(header::CONTENT_LENGTH, total_len)
        .header(
            header::CONTENT_DISPOSITION,
            format!("attachment; filename=\"{zip_name}.zip\""),
        )
        .body(Body::from_stream(stream))
        .map_err(|e| anyhow::anyhow!("building response: {e}").into())
}

enum Fetched {
    /// CRC already known — the body streams straight into the response.
    Direct(Box<aws_sdk_s3::operation::get_object::GetObjectOutput>),
    /// Pre-hash photo: spooled to a temp file to compute the CRC first.
    Spooled(tokio::fs::File, u32),
}

/// Start fetching an original. With a known CRC this only opens the S3
/// response (the body is consumed later, straight into the zip stream);
/// otherwise the object is spooled to an anonymous temp file to compute the
/// CRC, verifying the byte count the zip plan promised.
fn fetch_entry(
    state: &AppState,
    key: String,
    expected_size: u64,
    crc_known: bool,
) -> JoinHandle<anyhow::Result<Fetched>> {
    let state = state.clone();
    tokio::spawn(async move {
        let object = state
            .s3
            .get_object()
            .bucket(&state.config.s3_bucket)
            .key(&key)
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("fetching {key}: {e}"))?;
        if crc_known {
            return Ok(Fetched::Direct(Box::new(object)));
        }
        let mut file = tokio::fs::File::from_std(tempfile::tempfile()?);
        let mut reader = object.body.into_async_read();
        let mut hasher = crc32fast::Hasher::new();
        let mut written: u64 = 0;
        let mut buf = vec![0u8; 128 * 1024];
        loop {
            let n = reader.read(&mut buf).await?;
            if n == 0 {
                break;
            }
            hasher.update(&buf[..n]);
            file.write_all(&buf[..n]).await?;
            written += n as u64;
        }
        anyhow::ensure!(
            written == expected_size,
            "{key} is {written} bytes in s3 but {expected_size} in the database"
        );
        file.flush().await?;
        file.seek(std::io::SeekFrom::Start(0)).await?;
        Ok(Fetched::Spooled(file, hasher.finalize()))
    })
}

async fn write_zip(
    state: &AppState,
    plan: ZipPlan,
    mut out: tokio::io::DuplexStream,
) -> anyhow::Result<()> {
    // UTF-8 filename flag; no data descriptors (bit 3 unset).
    const FLAGS: u16 = 0x0800;
    let mut crcs = Vec::with_capacity(plan.entries.len());

    // Prefetch: start fetching the next object while streaming the current one.
    let mut pending: Option<JoinHandle<anyhow::Result<Fetched>>> = None;
    for (i, entry) in plan.entries.iter().enumerate() {
        let current = match pending.take() {
            Some(handle) => handle,
            None => fetch_entry(state, entry.s3_key.clone(), entry.size, entry.crc.is_some()),
        };
        if let Some(next) = plan.entries.get(i + 1) {
            pending = Some(fetch_entry(
                state,
                next.s3_key.clone(),
                next.size,
                next.crc.is_some(),
            ));
        }
        let fetched = current
            .await
            .map_err(|e| anyhow::anyhow!("fetch task failed: {e}"))?
            .map_err(|e| anyhow::anyhow!("photo {}: {e:#}", entry.photo_id))?;
        let crc = match &fetched {
            Fetched::Direct(_) => entry.crc.expect("direct fetch implies known crc"),
            Fetched::Spooled(_, crc) => *crc,
        };
        crcs.push(crc);

        let mut lfh = Vec::with_capacity(30 + entry.name.len());
        lfh.extend_from_slice(&0x04034b50u32.to_le_bytes());
        lfh.extend_from_slice(&20u16.to_le_bytes()); // version needed
        lfh.extend_from_slice(&FLAGS.to_le_bytes());
        lfh.extend_from_slice(&0u16.to_le_bytes()); // method: stored
        lfh.extend_from_slice(&entry.dos_time.to_le_bytes());
        lfh.extend_from_slice(&entry.dos_date.to_le_bytes());
        lfh.extend_from_slice(&crc.to_le_bytes());
        lfh.extend_from_slice(&(entry.size as u32).to_le_bytes()); // compressed
        lfh.extend_from_slice(&(entry.size as u32).to_le_bytes()); // uncompressed
        lfh.extend_from_slice(&(entry.name.len() as u16).to_le_bytes());
        lfh.extend_from_slice(&0u16.to_le_bytes()); // extra len
        lfh.extend_from_slice(&entry.name);
        out.write_all(&lfh).await?;
        match fetched {
            Fetched::Direct(object) => {
                let mut reader = object.body.into_async_read();
                let copied = tokio::io::copy(&mut reader, &mut out).await?;
                anyhow::ensure!(
                    copied == entry.size,
                    "{} is {copied} bytes in s3 but {} in the database",
                    entry.s3_key,
                    entry.size
                );
            }
            Fetched::Spooled(mut file, crc) => {
                tokio::io::copy(&mut file, &mut out).await?;
                // Self-heal: store the freshly computed crc so the next
                // download of this photo streams directly.
                let _ = sqlx::query("update photos set crc32 = coalesce(crc32, $2) where id = $1")
                    .bind(entry.photo_id)
                    .bind(i64::from(crc))
                    .execute(&state.db)
                    .await;
            }
        }
    }

    // Central directory.
    for (entry, crc) in plan.entries.iter().zip(&crcs) {
        let zip64_offset = entry.offset >= U32_SENTINEL;
        let mut cdh = Vec::with_capacity(46 + entry.name.len() + 12);
        cdh.extend_from_slice(&0x02014b50u32.to_le_bytes());
        cdh.extend_from_slice(&0x031Eu16.to_le_bytes()); // made by: unix, 3.0
        cdh.extend_from_slice(&(if zip64_offset { 45u16 } else { 20u16 }).to_le_bytes());
        cdh.extend_from_slice(&FLAGS.to_le_bytes());
        cdh.extend_from_slice(&0u16.to_le_bytes()); // method: stored
        cdh.extend_from_slice(&entry.dos_time.to_le_bytes());
        cdh.extend_from_slice(&entry.dos_date.to_le_bytes());
        cdh.extend_from_slice(&crc.to_le_bytes());
        cdh.extend_from_slice(&(entry.size as u32).to_le_bytes());
        cdh.extend_from_slice(&(entry.size as u32).to_le_bytes());
        cdh.extend_from_slice(&(entry.name.len() as u16).to_le_bytes());
        cdh.extend_from_slice(&(if zip64_offset { 12u16 } else { 0u16 }).to_le_bytes()); // extra len
        cdh.extend_from_slice(&0u16.to_le_bytes()); // comment len
        cdh.extend_from_slice(&0u16.to_le_bytes()); // disk number
        cdh.extend_from_slice(&0u16.to_le_bytes()); // internal attrs
        cdh.extend_from_slice(&(0o100644u32 << 16).to_le_bytes()); // unix -rw-r--r--
        let offset32 = if zip64_offset { U32_SENTINEL as u32 } else { entry.offset as u32 };
        cdh.extend_from_slice(&offset32.to_le_bytes());
        cdh.extend_from_slice(&entry.name);
        if zip64_offset {
            cdh.extend_from_slice(&0x0001u16.to_le_bytes()); // zip64 extra field
            cdh.extend_from_slice(&8u16.to_le_bytes());
            cdh.extend_from_slice(&entry.offset.to_le_bytes());
        }
        out.write_all(&cdh).await?;
    }

    // End of central directory (zip64 variants only when values overflow).
    let mut tail = Vec::with_capacity(98);
    if plan.zip64_eocd {
        let entries = plan.entries.len() as u64;
        tail.extend_from_slice(&0x06064b50u32.to_le_bytes());
        tail.extend_from_slice(&44u64.to_le_bytes()); // record size
        tail.extend_from_slice(&0x031Eu16.to_le_bytes()); // made by
        tail.extend_from_slice(&45u16.to_le_bytes()); // version needed
        tail.extend_from_slice(&0u32.to_le_bytes()); // this disk
        tail.extend_from_slice(&0u32.to_le_bytes()); // cd disk
        tail.extend_from_slice(&entries.to_le_bytes());
        tail.extend_from_slice(&entries.to_le_bytes());
        tail.extend_from_slice(&plan.cd_size.to_le_bytes());
        tail.extend_from_slice(&plan.cd_offset.to_le_bytes());
        // zip64 EOCD locator
        tail.extend_from_slice(&0x07064b50u32.to_le_bytes());
        tail.extend_from_slice(&0u32.to_le_bytes());
        tail.extend_from_slice(&(plan.cd_offset + plan.cd_size).to_le_bytes());
        tail.extend_from_slice(&1u32.to_le_bytes());
    }
    let clamp16 = |v: u64| -> u16 { v.min(0xFFFF) as u16 };
    let clamp32 = |v: u64| -> u32 { v.min(U32_SENTINEL) as u32 };
    tail.extend_from_slice(&0x06054b50u32.to_le_bytes());
    tail.extend_from_slice(&0u16.to_le_bytes()); // this disk
    tail.extend_from_slice(&0u16.to_le_bytes()); // cd disk
    tail.extend_from_slice(&clamp16(plan.entries.len() as u64).to_le_bytes());
    tail.extend_from_slice(&clamp16(plan.entries.len() as u64).to_le_bytes());
    tail.extend_from_slice(&clamp32(plan.cd_size).to_le_bytes());
    tail.extend_from_slice(&clamp32(plan.cd_offset).to_le_bytes());
    tail.extend_from_slice(&0u16.to_le_bytes()); // comment len
    out.write_all(&tail).await?;
    out.shutdown().await?;
    Ok(())
}