cloudflare-soos 2.3.1

//! Scans a progressive JPEG file to find which byte ranges of the file are critical for displaying it at key stages.
//! This knowledge can be used to serve JPEGs optimally over HTTP/2 connections. This library can generate cf-priority-change headers
//! compatible with [prioritization syntax used by Cloudflare](https://blog.cloudflare.com/parallel-streaming-of-progressive-images/).

use std::fmt::Display;

mod error;
pub use crate::error::*;

/// For advanced usage, low-level access to the basic JPEG structure
pub mod jpeg;

#[cfg(feature = "gif")]
mod gif;
mod png;

#[cfg(target_arch = "wasm32")]
use wasm_bindgen::prelude::*;

/// Key positions in a progressive image file
///
/// ```rust,no_run
/// # let input_file = vec![];
/// cloudflare_soos::Scans::from_file(&input_file)?.cf_priority_change_headers()?;
/// # Ok::<_, cloudflare_soos::Error>(())
/// ```
#[derive(Debug, Clone, Default, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Scans {
    /// Byte position where metadata ends.
    /// This many bytes are preceeding start of pixel data.
    /// It's usually <200 bytes, unless the image has color profiles or other bloat.
    pub metadata_end: Option<usize>,
    /// All metadata + minimum amount of data to make browsers render _anything_
    /// (in case they don't reserve space based on metadata)
    pub frame_render_start: Option<usize>,
    /// Byte position where the first (lowest-quality) progressive scan ends.
    /// This many bytes are needed to display anything on screen.
    /// It's usually 12-15% of the file size.
    pub first_scan_end: Option<usize>,
    /// Byte position where most of ok-quality pixel data ends.
    /// This many bytes are needed to display a good-enough image.
    /// It's usually about 50% of the file size.
    pub good_scan_end: Option<usize>,
    /// Size of the whole image file, in bytes. The size is only used for heuristics,
    /// so the value may be approximate, but it must be greater than all other
    /// positions set in this struct.
    pub file_size: usize,
}

#[cfg(target_arch = "wasm32")]
#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
pub fn cf_priority_change_header_wasm(image: &[u8]) -> Option<(String, String)> {
    Scans::from_file(image).and_then(|v| v.cf_priority_change_headers()).ok()
}

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
pub fn rfc9218_priority_change_headers_wasm(image: &[u8]) -> Option<String> {
    Scans::from_file(image).and_then(|v| v.rfc9218_priority_change_headers()).ok()
}

// minimum chunk size for h2 and rfc9218 (h3) prioritization changes
// there's no point changing priority for 20 bytes, H/2 frame takes half of that
const MIN_H2_CHUNK_SIZE: usize = 20;
// there's no point changing priority for 32 bytes, H/3 frame's headers takes 16 bytes
const MIN_H3_CHUNK_SIZE: usize = 32;

#[derive(Debug)]
enum Concurrency {
    ExclusiveSequential, // 0
    SharedSequential,    // 1
    Shared,              // n
}

impl Default for Concurrency {
    fn default() -> Self {
        Self::ExclusiveSequential
    }
}

impl Display for Concurrency {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let value = match self {
            Concurrency::ExclusiveSequential => "0",
            Concurrency::SharedSequential => "1",
            Concurrency::Shared => "n",
        };
        write!(f, "{value}")
    }
}

#[derive(Default)]
struct PriorityChanges {
    next_offset: usize,
    changes: Vec<PriorityChange>,
}

struct PriorityChange {
    offset: usize,
    http2_priority: u8,
    concurrency: Concurrency,
}

impl PriorityChange {
    fn get_http2_priority_chunk(&self) -> String {
        format!("{}:{}", self.offset, self.get_http2_priority())
    }

    #[inline(always)]
    fn get_http2_priority(&self) -> String {
        format!("{}/{}", self.http2_priority, self.concurrency)
    }

    fn get_http3_priority_chunk(&self) -> String {
        format!("{};{}", self.offset, self.get_http3_priority())
    }

    #[inline(always)]
    fn get_http3_priority(&self) -> String {
        // the lower urgency level the more priority it has over other frames
        let urgency = match self.http2_priority {
            0..=9 => 6,
            10..=19 => 5,
            20..=29 => 4,
            30..=39 => 3,
            40..=49 => 2,
            50..=59 => 1,
            60..=63 => 0,
            _ => 7, // not supported
        };

        let (urgency, incremental) = match self.concurrency {
            Concurrency::ExclusiveSequential => (urgency, false),
            // Shared concurrency reduces priority (urgency is a reversed measure)
            Concurrency::SharedSequential => (urgency + 1, false),
            Concurrency::Shared => (urgency + 1, true),
        };

        if incremental {
            format!("u={urgency};i")
        } else {
            format!("u={urgency};i=?0")
        }
    }
}

impl PriorityChanges {
    fn add(&mut self, offset: usize, http2_priority: u8, concurrency: Concurrency) {
        self.changes.push(PriorityChange {
            offset: self.next_offset,
            http2_priority,
            concurrency,
        });

        self.next_offset = offset;
    }

    fn has_changes(&self) -> bool {
        self.changes.len() > 1
    }

    fn get_http2_priority(&self) -> Result<String> {
        self.changes
            .first()
            .map(PriorityChange::get_http2_priority)
            .ok_or(Error::Format("no priority header"))
    }

    fn get_http2_priority_changes(&self) -> Result<String> {
        let mut min_pos = 0;
        let change = self
            .changes
            .iter()
            // return changes w/o first priority
            .skip(1)
            .filter(|&p| {
                // filter out small chunks
                let keep = p.offset > min_pos;
                min_pos = p.offset + MIN_H2_CHUNK_SIZE;
                keep
            })
            .map(PriorityChange::get_http2_priority_chunk)
            .collect::<Vec<String>>()
            .join(",");

        if change.is_empty() {
            return Err(Error::Format("Can't find useful scans"));
        }

        Ok(change)
    }

    fn get_http3_priority(&self) -> Result<String> {
        self.changes
            .first()
            .map(PriorityChange::get_http3_priority)
            .ok_or(Error::Format("no priority header"))
    }

    fn get_rfc9218_priority_changes(&self) -> Result<String> {
        let mut min_pos = 0;
        let change = self
            .changes
            .iter()
            // returns changes w/o first priority
            .skip(1)
            .filter(|&p| {
                // filter out small chunks
                let keep = p.offset > min_pos;
                min_pos = p.offset + MIN_H3_CHUNK_SIZE;
                keep
            })
            .map(PriorityChange::get_http3_priority_chunk)
            .collect::<Vec<String>>()
            .join(" ");

        if change.is_empty() {
            return Err(Error::Format("Can't find useful scans"));
        }

        Ok(format!("cf-chb=({change})"))
    }
}

impl Scans {
    /// Analyze an image file to find byte ranges of its metadata and progressive scans
    pub fn from_file(input_file: &[u8]) -> Result<Self> {
        match input_file {
            [0xff, ..] => crate::jpeg::scans(input_file),
            #[cfg(feature = "gif")]
            [b'G', ..] => crate::gif::scans(input_file),
            [0x89, ..] => crate::png::scans(input_file),
            _ => Err(Error::Unsupported),
        }
    }

    /// Assumes the HTTP2 priorities are:
    ///
    ///  * 50 = critical js/css, fonts + image metadata
    ///  * 30 = regular js/css, followed by other images + JPEG DC
    ///  * 20 = low-priority image bodies
    ///  * 10 = idle
    ///
    /// Returns `cf-priority` and `cf-priority-change` header values, respectively.
    pub fn cf_priority_change_headers(&self) -> Result<(String, String)> {
        let priority_changes = self.get_priority_changes()?;
        Ok((
            priority_changes.get_http2_priority()?,
            priority_changes.get_http2_priority_changes()?,
        ))
    }

    /// <https://www.rfc-editor.org/rfc/rfc9218.html>
    pub fn rfc9218_priority_change_headers(&self) -> Result<String> {
        let priority_changes = self.get_priority_changes()?;
        Ok(format!("{},{}",
            priority_changes.get_http3_priority()?,
            priority_changes.get_rfc9218_priority_changes()?,
        ))
    }

    fn get_priority_changes(&self) -> Result<PriorityChanges> {
        let mut metadata_end = self.metadata_end.unwrap_or(0);
        let is_progressive = self
            .first_scan_end
            .map_or(false, |len| len < self.file_size / 2);

        // if there's a fat color profile or Adobe Garbage,
        // then sending "just" the metadata is not fast, and shouldn't be prioritized.
        let fat_metadata_limit = ((self.file_size / 8) + 180).min(2000);
        let fat_frame_start_limit = ((self.file_size / 8) + 500).min(8000);

        // if the metadata can cheaply include first frame/scan info, then send both at once (reduce H/2 framing overhead)
        let rendered_anything = self
            .frame_render_start
            .or(self.first_scan_end)
            .or(self.good_scan_end);
        if let Some(rendered_anything) = rendered_anything {
            if rendered_anything < fat_metadata_limit
                && rendered_anything < metadata_end + metadata_end / 8 + 100
            {
                metadata_end = rendered_anything;
            }
        }

        // let mut chunks = Vec::with_capacity(8);
        let mut priority_changes = PriorityChanges::default();

        // This is important, because it decides when the whole image starts sending
        if is_progressive && metadata_end < fat_metadata_limit {
            // "50/0" fast and worth accelerating
            priority_changes.add(metadata_end, 50, Concurrency::ExclusiveSequential);
        } else if self.file_size < 1200
            || is_progressive
            || rendered_anything.map_or(false, |n| n < 2000)
        {
            // "30/1" at least we can show something quick
            priority_changes.add(metadata_end, 30, Concurrency::SharedSequential);
        } else {
            // "21/n" lost cause: baseline image with bloated metadata
            priority_changes.add(metadata_end, 21, Concurrency::Shared);
        };

        // Browsers don't always reserve space based on metadata availability alone,
        // so here's trying again, with more data
        if let Some(frame_render_start) = self.frame_render_start {
            if frame_render_start > metadata_end
                && self
                    .first_scan_end
                    .map_or(true, |dc| frame_render_start < dc)
            {
                if frame_render_start < fat_frame_start_limit {
                    if frame_render_start < 1000 {
                        priority_changes.add(frame_render_start, 50, Concurrency::SharedSequential);
                    } else {
                        priority_changes.add(frame_render_start, 40, Concurrency::SharedSequential);
                    }
                } else {
                    priority_changes.add(frame_render_start, 30, Concurrency::Shared);
                }
            }
        }

        if let Some(first_scan_end) = self.first_scan_end {
            // small DC can be downloaded in one go, to save on re-rendering
            priority_changes.add(
                first_scan_end,
                30,
                if first_scan_end < 25000 {
                    Concurrency::ExclusiveSequential
                } else {
                    Concurrency::SharedSequential
                },
            );
        }

        if let Some(good_scan_end) = self.good_scan_end {
            priority_changes.add(
                good_scan_end,
                20,
                if good_scan_end < 100_000 {
                    Concurrency::SharedSequential
                } else {
                    Concurrency::Shared
                },
            );
        }

        let rendered_already = self.first_scan_end.is_some() || self.good_scan_end.is_some();
        let bytes_left = self
            .file_size
            .saturating_sub(self.good_scan_end.or(self.first_scan_end).unwrap_or(0));
        let is_big = bytes_left > 80_000;
        let is_tiny = bytes_left < 1_000;

        let (priority, concurrency) = if rendered_already {
            // if it's on screen, it's not urgent to send anything more
            (
                10,
                if is_big {
                    Concurrency::Shared
                } else {
                    Concurrency::SharedSequential
                },
            )
        } else if is_tiny {
            (30, Concurrency::SharedSequential)
        } else if is_big {
            (20, Concurrency::Shared)
        } else {
            (20, Concurrency::SharedSequential)
        };

        priority_changes.add(self.file_size, priority, concurrency);
        if !priority_changes.has_changes() {
            return Err(Error::Format("Can't find useful scans"));
        }
        Ok(priority_changes)
    }
}

#[cfg(test)]
fn s(a: &str, b: &str) -> (String, String) {
    (a.into(), b.into())
}

#[test]
fn test_baseline() {
    {
        let scans = Scans {
            metadata_end: None,
            frame_render_start: None,
            first_scan_end: None,
            good_scan_end: None,
            file_size: 100_000,
        };

        let res = scans.cf_priority_change_headers();
        assert!(res.is_err(), "expected error, h/2: {res:?}");

        let res = scans.rfc9218_priority_change_headers();
        assert!(res.is_err(), "expected error, h/3: {res:?}");
    }

    {
        // regular baseline image
        let scans = Scans {
            metadata_end: Some(101),
            frame_render_start: Some(181),
            first_scan_end: None,
            good_scan_end: None,
            file_size: 100_000,
        };

        assert_eq!(
            s("30/1", "181:20/n"),
            scans.cf_priority_change_headers().unwrap(),
            "regular baseline image, h/2"
        );

        assert_eq!(
            "u=4;i=?0,cf-chb=(181;u=5;i)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "regular baseline image, h/3"
        );
    }

    {
        // tiny image
        let scans = Scans {
            metadata_end: Some(101),
            frame_render_start: None,
            first_scan_end: None,
            good_scan_end: None,
            file_size: 1000,
        };

        assert_eq!(
            s("30/1", "101:20/1"),
            scans.cf_priority_change_headers().unwrap(),
            "tiny image, h/2"
        );
        assert_eq!(
            "u=4;i=?0,cf-chb=(101;u=5;i=?0)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "tiny image, h/3"
        );
    }

    {
        // fat metadata
        let scans = Scans {
            metadata_end: Some(9999),
            frame_render_start: None,
            first_scan_end: None,
            good_scan_end: None,
            file_size: 100_000,
        };

        assert_eq!(
            s("21/n", "9999:20/n"),
            scans.cf_priority_change_headers().unwrap(),
            "fat metadata, h/2"
        );

        assert_eq!(
            "u=5;i,cf-chb=(9999;u=5;i)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "fat metadata, h/3"
        );
    }

    {
        // relatively fat metadata
        let scans = Scans {
            metadata_end: Some(1000),
            frame_render_start: None,
            first_scan_end: None,
            good_scan_end: None,
            file_size: 3000,
        };
        assert_eq!(
            s("21/n", "1000:20/1"),
            scans.cf_priority_change_headers().unwrap(),
            "relatively fat metadata, h/2"
        );

        assert_eq!(
            "u=5;i,cf-chb=(1000;u=5;i=?0)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "relatively fat metadata, h/3"
        );
    }
}

#[test]
fn test_progressive() {
    {
        let scans = Scans {
            metadata_end: Some(1_000),
            frame_render_start: None,
            first_scan_end: Some(10_000),
            good_scan_end: Some(100_000),
            file_size: 200_000,
        };
        assert_eq!(
            s("50/0", "1000:30/0,10000:20/n,100000:10/n"),
            scans.cf_priority_change_headers().unwrap(),
            "scan, h/2"
        );

        assert_eq!(
            "u=1;i=?0,cf-chb=(1000;u=3;i=?0 10000;u=5;i 100000;u=6;i)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "scan, h/3"
        );
    }

    {
        // fat metadata
        let scans = Scans {
            metadata_end: Some(4_000),
            frame_render_start: None,
            first_scan_end: Some(10_000),
            good_scan_end: Some(100_000),
            file_size: 200_000,
        };

        assert_eq!(
            s("30/1", "4000:30/0,10000:20/n,100000:10/n"),
            scans.cf_priority_change_headers().unwrap(),
            "fat metadata, h/2"
        );

        assert_eq!(
            "u=4;i=?0,cf-chb=(4000;u=3;i=?0 10000;u=5;i 100000;u=6;i)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "fat metadata, h/3"
        );
    }

    {
        // fat DC
        let scans = Scans {
            metadata_end: Some(1_000),
            frame_render_start: None,
            first_scan_end: Some(50_000),
            good_scan_end: Some(100_000),
            file_size: 200_000,
        };

        assert_eq!(
            s("50/0", "1000:30/1,50000:20/n,100000:10/n"),
            scans.cf_priority_change_headers().unwrap(),
            "fat DC, h/2"
        );

        assert_eq!(
            "u=1;i=?0,cf-chb=(1000;u=4;i=?0 50000;u=5;i 100000;u=6;i)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "fat DC, h/3"
        );
    }

    {
        // small good scan
        let scans = Scans {
            metadata_end: Some(1_000),
            frame_render_start: None,
            first_scan_end: Some(10_000),
            good_scan_end: Some(11_000),
            file_size: 200_000,
        };

        assert_eq!(
            s("50/0", "1000:30/0,10000:20/1,11000:10/n"),
            scans.cf_priority_change_headers().unwrap(),
            "small good scan, h/2"
        );

        assert_eq!(
            "u=1;i=?0,cf-chb=(1000;u=3;i=?0 10000;u=5;i=?0 11000;u=6;i)".to_string(),
            scans.rfc9218_priority_change_headers().unwrap(),
            "small good scan, h/3"
        );
    }
}