sie_sdk/
media.rs

1//! Turning caller-supplied images, audio and documents into the bytes the wire carries.
2//!
3//! Format detection is by magic bytes and fails closed: an undetectable image is an error,
4//! never a guess. Encoded inputs pass through untouched, so a JPEG the caller already has
5//! is never re-encoded.
6
7use std::io::Cursor;
8use std::path::Path;
9
10use crate::error::{Error, Result};
11
12const JPEG_QUALITY: u8 = 95;
13const MAX_FORMAT_LEN: usize = 32;
14
15/// Normalise a caller-declared media format token.
16pub(crate) fn canonical_format(value: &str) -> Result<String> {
17    let invalid = || {
18        Error::invalid(format!(
19            "Image format must be a short ASCII media-format token, got {value:?}"
20        ))
21    };
22    if value.is_empty() || value.len() > MAX_FORMAT_LEN {
23        return Err(invalid());
24    }
25    if !value
26        .bytes()
27        .all(|b| b.is_ascii_alphanumeric() || b == b'.' || b == b'+' || b == b'-')
28    {
29        return Err(invalid());
30    }
31    let lowered = value.to_ascii_lowercase();
32    Ok(match lowered.as_str() {
33        "jpg" | "jpe" => "jpeg".to_string(),
34        _ => lowered,
35    })
36}
37
38/// Identify an encoded image by its magic bytes.
39///
40/// Formats without a signature the SDK recognises (HEIC and AVIF among them) are rejected
41/// rather than mislabelled.
42pub(crate) fn detect_image_format(data: &[u8]) -> Result<&'static str> {
43    let format = if data.starts_with(&[0xff, 0xd8, 0xff]) {
44        "jpeg"
45    } else if data.starts_with(b"\x89PNG\r\n\x1a\n") {
46        "png"
47    } else if data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a") {
48        "gif"
49    } else if data.len() >= 12 && data.starts_with(b"RIFF") && &data[8..12] == b"WEBP" {
50        "webp"
51    } else if data.starts_with(b"BM") {
52        "bmp"
53    } else if data.starts_with(b"II*\x00")
54        || data.starts_with(b"MM\x00*")
55        || data.starts_with(b"II+\x00")
56        || data.starts_with(b"MM\x00+")
57    {
58        "tiff"
59    } else {
60        return Err(Error::invalid(
61            "Could not detect encoded image format from bytes",
62        ));
63    };
64    Ok(format)
65}
66
67/// Encode a decoded image as JPEG.
68///
69/// Greyscale stays greyscale; everything else becomes RGB, since JPEG has no alpha channel.
70pub(crate) fn encode_jpeg(image: &image::DynamicImage) -> Result<Vec<u8>> {
71    use image::ColorType;
72
73    let converted = match image.color() {
74        ColorType::L8 | ColorType::L16 => image::DynamicImage::ImageLuma8(image.to_luma8()),
75        _ => image::DynamicImage::ImageRgb8(image.to_rgb8()),
76    };
77    let mut buffer = Cursor::new(Vec::new());
78    let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buffer, JPEG_QUALITY);
79    encoder
80        .encode_image(&converted)
81        .map_err(|err| Error::invalid(format!("could not encode image as JPEG: {err}")))?;
82    Ok(buffer.into_inner())
83}
84
85fn suffix(path: &Path) -> Option<String> {
86    Some(path.extension()?.to_str()?.to_ascii_lowercase())
87}
88
89/// Audio container implied by a filename.
90pub(crate) fn infer_audio_format(path: &Path) -> Option<&'static str> {
91    Some(match suffix(path)?.as_str() {
92        "flac" => "flac",
93        "m4a" => "m4a",
94        "mp3" => "mp3",
95        "mp4" => "mp4",
96        "mpeg" => "mpeg",
97        "mpga" => "mpga",
98        "ogg" => "ogg",
99        "wav" => "wav",
100        "webm" => "webm",
101        _ => return None,
102    })
103}
104
105/// Document type implied by a filename.
106pub(crate) fn infer_document_format(path: &Path) -> Option<&'static str> {
107    Some(match suffix(path)?.as_str() {
108        "pdf" => "pdf",
109        "docx" => "docx",
110        "doc" => "doc",
111        "html" | "htm" | "xhtml" => "html",
112        "md" | "markdown" => "md",
113        "txt" => "txt",
114        "rtf" => "rtf",
115        "odt" => "odt",
116        "pptx" => "pptx",
117        "xlsx" => "xlsx",
118        "csv" => "csv",
119        _ => return None,
120    })
121}
122
123/// Video container implied by a filename.
124pub(crate) fn infer_video_format(path: &Path) -> Option<&'static str> {
125    Some(match suffix(path)?.as_str() {
126        "mp4" => "mp4",
127        "webm" => "webm",
128        "mkv" => "mkv",
129        "mov" => "mov",
130        "avi" => "avi",
131        _ => return None,
132    })
133}
134
135/// One PCM sample buffer, in the precision the caller has.
136#[derive(Debug, Clone, PartialEq)]
137pub enum Samples {
138    /// Normalised floats, nominally in `[-1.0, 1.0]`.
139    F32(Vec<f32>),
140    /// Signed 16-bit PCM, written through untouched.
141    I16(Vec<i16>),
142}
143
144impl Samples {
145    fn len(&self) -> usize {
146        match self {
147            Self::F32(values) => values.len(),
148            Self::I16(values) => values.len(),
149        }
150    }
151
152    /// Convert to signed 16-bit PCM.
153    ///
154    /// The scale is asymmetric because the i16 range is: `-1.0` maps to `-32768` and `1.0`
155    /// to `32767`, so neither extreme clips.
156    fn to_pcm16(&self) -> Result<Vec<i16>> {
157        match self {
158            Self::I16(values) => Ok(values.clone()),
159            Self::F32(values) => values
160                .iter()
161                .map(|sample| {
162                    if !sample.is_finite() {
163                        return Err(Error::invalid("Audio samples must all be finite"));
164                    }
165                    let clamped = f64::from(*sample).clamp(-1.0, 1.0);
166                    let scaled = clamped * if clamped < 0.0 { 32768.0 } else { 32767.0 };
167                    // numpy rounds half to even; matching it keeps encoded audio identical
168                    // to what the Python SDK produces for the same input.
169                    Ok(scaled.round_ties_even().clamp(-32768.0, 32767.0) as i16)
170                })
171                .collect(),
172        }
173    }
174}
175
176/// Wrap raw samples in a WAV container.
177///
178/// Samples are frame-interleaved: for stereo, `[l0, r0, l1, r1, ...]`.
179pub(crate) fn encode_wav(samples: &Samples, channels: u16, sample_rate: u32) -> Result<Vec<u8>> {
180    if sample_rate == 0 {
181        return Err(Error::invalid(
182            "Audio sample_rate must be a positive integer",
183        ));
184    }
185    if !(1..=2).contains(&channels) {
186        return Err(Error::invalid(format!(
187            "Audio must contain 1 or 2 channels, got {channels}"
188        )));
189    }
190    if samples.len() == 0 {
191        return Err(Error::invalid("Audio must not be empty"));
192    }
193    if !samples.len().is_multiple_of(usize::from(channels)) {
194        return Err(Error::invalid(format!(
195            "Audio has {} samples, which is not a whole number of {channels}-channel frames",
196            samples.len()
197        )));
198    }
199
200    let spec = hound::WavSpec {
201        channels,
202        sample_rate,
203        bits_per_sample: 16,
204        sample_format: hound::SampleFormat::Int,
205    };
206    let mut buffer = Cursor::new(Vec::new());
207    let mut writer = hound::WavWriter::new(&mut buffer, spec)
208        .map_err(|err| Error::invalid(format!("could not start a WAV stream: {err}")))?;
209    for sample in samples.to_pcm16()? {
210        writer
211            .write_sample(sample)
212            .map_err(|err| Error::invalid(format!("could not write a WAV sample: {err}")))?;
213    }
214    writer
215        .finalize()
216        .map_err(|err| Error::invalid(format!("could not finalize the WAV stream: {err}")))?;
217    Ok(buffer.into_inner())
218}
219
220#[cfg(test)]
221mod tests {
222    // These assertions are about exact values, so exact comparison is the point.
223    #![allow(clippy::float_cmp)]
224
225    use super::*;
226
227    #[test]
228    fn canonicalizes_format_aliases() {
229        assert_eq!(canonical_format("JPG").unwrap(), "jpeg");
230        assert_eq!(canonical_format("jpe").unwrap(), "jpeg");
231        assert_eq!(canonical_format("PNG").unwrap(), "png");
232        assert_eq!(canonical_format("image+xml").unwrap(), "image+xml");
233    }
234
235    #[test]
236    fn rejects_format_tokens_that_are_not_short_ascii() {
237        for value in ["", &"x".repeat(33), "png/", "png ", "påäng"] {
238            assert!(canonical_format(value).is_err(), "{value:?}");
239        }
240    }
241
242    #[test]
243    fn sniffs_every_supported_signature() {
244        let cases: [(&[u8], &str); 8] = [
245            (&[0xff, 0xd8, 0xff, 0xe0], "jpeg"),
246            (b"\x89PNG\r\n\x1a\n....", "png"),
247            (b"GIF87a...", "gif"),
248            (b"GIF89a...", "gif"),
249            (b"RIFF\0\0\0\0WEBPVP8 ", "webp"),
250            (b"BM......", "bmp"),
251            (b"II*\x00....", "tiff"),
252            (b"MM\x00+....", "tiff"),
253        ];
254        for (data, expected) in cases {
255            assert_eq!(detect_image_format(data).unwrap(), expected);
256        }
257    }
258
259    #[test]
260    fn sniffing_fails_closed() {
261        assert!(detect_image_format(b"").is_err());
262        assert!(detect_image_format(b"\x00\x00\x00\x20ftypheic").is_err());
263        // A truncated RIFF header is not a WebP.
264        assert!(detect_image_format(b"RIFF\0\0\0\0WEB").is_err());
265    }
266
267    #[test]
268    fn encodes_a_decoded_image_as_jpeg() {
269        let image = image::DynamicImage::ImageRgba8(image::RgbaImage::new(4, 4));
270        let bytes = encode_jpeg(&image).unwrap();
271        assert_eq!(detect_image_format(&bytes).unwrap(), "jpeg");
272
273        let grey = image::DynamicImage::ImageLuma8(image::GrayImage::new(4, 4));
274        assert_eq!(
275            detect_image_format(&encode_jpeg(&grey).unwrap()).unwrap(),
276            "jpeg"
277        );
278    }
279
280    #[test]
281    fn infers_formats_from_suffixes() {
282        assert_eq!(infer_audio_format(Path::new("/tmp/clip.WAV")), Some("wav"));
283        assert_eq!(infer_audio_format(Path::new("/tmp/clip.aiff")), None);
284        assert_eq!(
285            infer_document_format(Path::new("report.Markdown")),
286            Some("md")
287        );
288        assert_eq!(infer_document_format(Path::new("page.htm")), Some("html"));
289        assert_eq!(infer_document_format(Path::new("noext")), None);
290        assert_eq!(infer_video_format(Path::new("a.mp4")), Some("mp4"));
291    }
292
293    #[test]
294    fn float_samples_use_the_asymmetric_pcm_scale() {
295        let pcm = Samples::F32(vec![-1.0, 0.0, 1.0, 0.5]).to_pcm16().unwrap();
296        assert_eq!(pcm, vec![-32768, 0, 32767, 16384]);
297    }
298
299    #[test]
300    fn float_samples_are_clamped_and_must_be_finite() {
301        assert_eq!(
302            Samples::F32(vec![-2.0, 2.0]).to_pcm16().unwrap(),
303            vec![-32768, 32767]
304        );
305        assert!(Samples::F32(vec![f32::NAN]).to_pcm16().is_err());
306        assert!(Samples::F32(vec![f32::INFINITY]).to_pcm16().is_err());
307    }
308
309    #[test]
310    fn wav_round_trips_through_hound() {
311        let samples = Samples::I16(vec![1, -1, 100, -100]);
312        let wav = encode_wav(&samples, 2, 16_000).unwrap();
313        let mut reader = hound::WavReader::new(Cursor::new(wav)).unwrap();
314        let spec = reader.spec();
315        assert_eq!(spec.channels, 2);
316        assert_eq!(spec.sample_rate, 16_000);
317        assert_eq!(spec.bits_per_sample, 16);
318        let decoded: Vec<i16> = reader
319            .samples::<i16>()
320            .map(std::result::Result::unwrap)
321            .collect();
322        assert_eq!(decoded, vec![1, -1, 100, -100]);
323    }
324
325    #[test]
326    fn wav_rejects_impossible_geometry() {
327        assert!(encode_wav(&Samples::I16(vec![1]), 1, 0).is_err());
328        assert!(encode_wav(&Samples::I16(vec![1]), 3, 16_000).is_err());
329        assert!(encode_wav(&Samples::I16(Vec::new()), 1, 16_000).is_err());
330        // Three samples cannot be split into stereo frames.
331        assert!(encode_wav(&Samples::I16(vec![1, 2, 3]), 2, 16_000).is_err());
332    }
333}
sie_sdk/media.rs

sie_sdk/
media.rs
