From 07a47b62aa5b9c4db552c58b71a925d42760b1c8 Mon Sep 17 00:00:00 2001 From: Sunshine Date: Fri, 16 Oct 2026 16:31:21 -0400 Subject: [PATCH] avoid copying large data --- CHANGELOG.md | 8 +++ README.md | 2 +- src/base64.rs | 60 +++++++++++++---- src/lib.rs | 156 ++++++++++++++++++++++++++++++++---------- src/main.rs | 30 ++++++--- src/percent.rs | 44 ++++++++++-- tests/large_data.rs | 161 ++++++++++++++++++++++++++++++++++++++++++++ tests/lib.rs | 2 +- 8 files changed, 399 insertions(+), 64 deletions(-) create mode 100644 tests/large_data.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index bff2141..df7c4b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ - `DataEncoding::Shortest`, which picks percent-encoding or base64, whichever results in a shorter data URL - `Escaping::Minimal`, which only escapes what has to be when percent-encoding, for data URLs that go inside double-quoted HTML attributes or CSS strings; this keeps text formats like SVG short and readable + - `DataUrl::write_to()`, which writes a data URL to any `std::io::Write` without building it as a string first + - `DataUrl::encoded_len()`, the exact length of the serialized data URL, e.g. for preallocating buffers ### Changed @@ -18,6 +20,12 @@ which use the new `DataEncoding` enum (`Percent`, `Base64` or `Shortest`) - Percent-encoding now also escapes `&`, so that HTML entities can't alter the data when a data URL is put into an HTML attribute by hand + - `set_data()` takes `impl Into>`: passing a `Vec` or `String` moves it in without copying, + while slices, arrays and string literals still work (`&Vec` needs `.as_slice()` or to be passed by value) + - Serializing writes in chunks instead of building the whole data URL in memory first, and parsing base64 + no longer makes intermediate copies of the data: for a 10 MB payload, peak memory use while parsing went + from about 40 MB to 10 MB (just the result), and parsing got about a third faster + - The CLI writes its output as it encodes, instead of building it in memory first ## 0.2.1 diff --git a/README.md b/README.md index 13f8270..cb2a0ab 100644 --- a/README.md +++ b/README.md @@ -75,7 +75,7 @@ assert_eq!(data_url.to_string(), "data:,Hello,%20World!"); let mut data_url = DataUrl::new(); data_url.set_media_type(Some("image/png"))?; data_url.set_data_encoding(DataEncoding::Base64); -data_url.set_data(&[0x89, 0x50, 0x4E, 0x47]); +data_url.set_data([0x89, 0x50, 0x4E, 0x47]); assert_eq!(data_url.to_string(), "data:image/png;base64,iVBORw=="); ``` diff --git a/src/base64.rs b/src/base64.rs index 31955cb..d4ec056 100644 --- a/src/base64.rs +++ b/src/base64.rs @@ -2,6 +2,26 @@ const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; +/// Writes base64 in chunks, so that large data doesn't need one big intermediate string. +pub(crate) fn write(writer: &mut W, data: &[u8]) -> std::fmt::Result { + // A multiple of 3, so that only the very last chunk can need padding + const CHUNK_SIZE: usize = 6 * 1024; + let mut buffer = String::with_capacity(CHUNK_SIZE / 3 * 4); + + for chunk in data.chunks(CHUNK_SIZE) { + buffer.clear(); + encode_into(&mut buffer, chunk); + writer.write_str(&buffer)?; + } + + Ok(()) +} + +/// Length of the base64 encoding of `length` bytes, padding included. +pub(crate) fn encoded_length(length: usize) -> usize { + length.div_ceil(3) * 4 +} + pub(crate) fn encode_into(output: &mut String, data: &[u8]) { output.reserve(data.len().div_ceil(3) * 4); @@ -40,29 +60,30 @@ fn decode_char(c: u8) -> Option { /// (): ASCII whitespace is ignored, /// padding is optional, and anything else outside the alphabet makes decoding fail. pub(crate) fn decode_forgiving(input: &[u8]) -> Option> { - let mut data: Vec = input - .iter() - .copied() - .filter(|c| !matches!(c, b'\t' | b'\n' | 0x0C | b'\r' | b' ')) - .collect(); - - if data.len() % 4 == 0 { - if data.ends_with(b"==") { - data.truncate(data.len() - 2); - } else if data.ends_with(b"=") { - data.truncate(data.len() - 1); + let is_not_whitespace = |c: &&u8| !matches!(c, b'\t' | b'\n' | 0x0C | b'\r' | b' '); + + // Work out the length without whitespace and padding first, instead of making a stripped copy + let mut length = input.iter().filter(is_not_whitespace).count(); + if length % 4 == 0 { + let mut from_the_end = input.iter().rev().filter(is_not_whitespace); + if from_the_end.next() == Some(&b'=') { + length -= 1; + if from_the_end.next() == Some(&b'=') { + length -= 1; + } } } - if data.len() % 4 == 1 { + if length % 4 == 1 { return None; } - let mut output: Vec = Vec::with_capacity(data.len() / 4 * 3 + 2); + let mut output: Vec = Vec::with_capacity(length / 4 * 3 + 2); let mut buffer: u32 = 0; let mut bits: u8 = 0; - for c in data { + // The padding that was accounted for comes last, so stopping after `length` characters skips it + for &c in input.iter().filter(is_not_whitespace).take(length) { buffer = buffer << 6 | decode_char(c)?; bits += 6; if bits == 24 { @@ -116,6 +137,17 @@ mod tests { assert_eq!(decode_forgiving(s.as_bytes()).unwrap(), data); } + #[test] + fn chunked_writer_matches_one_go_encoding() { + let data: Vec = (0..=255).cycle().take(6 * 1024 * 3 + 7).collect(); + let mut chunked = String::new(); + write(&mut chunked, &data).unwrap(); + let mut whole = String::new(); + encode_into(&mut whole, &data); + assert_eq!(chunked, whole); + assert_eq!(chunked.len(), encoded_length(data.len())); + } + #[test] fn rejects_url_safe_alphabet() { assert_eq!(decode_forgiving(b"-_8="), None); diff --git a/src/lib.rs b/src/lib.rs index d41bb89..6a08f95 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,7 +18,7 @@ //! let mut data_url = DataUrl::new(); //! data_url.set_media_type(Some("image/png"))?; //! data_url.set_data_encoding(DataEncoding::Base64); -//! data_url.set_data(&[0x89, 0x50, 0x4E, 0x47]); +//! data_url.set_data([0x89, 0x50, 0x4E, 0x47]); //! assert_eq!(data_url.to_string(), "data:image/png;base64,iVBORw=="); //! //! // For text, percent-encoding with minimal escaping is often much shorter than base64 @@ -47,6 +47,7 @@ pub use mime::MimeType; use encoding_rs::{Encoding, REPLACEMENT, UTF_16BE, UTF_16LE, WINDOWS_1252}; use std::borrow::Cow; use std::fmt; +use std::io; use std::str::FromStr; const DEFAULT_MEDIA_TYPE: &str = "text/plain"; @@ -212,15 +213,21 @@ impl DataUrl { ); let mut mime_type: &str = mime_type.trim_matches(|c: char| c.is_ascii_whitespace()); - let mut data: Vec = percent::decode(encoded_data.as_bytes()); + // Base64 rarely contains percent-encoded characters, in which case decoding it can be done in one go + let percent_decoded: Cow<[u8]> = if encoded_data.contains('%') { + Cow::Owned(percent::decode(encoded_data.as_bytes())) + } else { + Cow::Borrowed(encoded_data.as_bytes()) + }; - let data_encoding = match strip_base64_marker(mime_type) { + let (data, data_encoding) = match strip_base64_marker(mime_type) { Some(stripped) => { - data = base64::decode_forgiving(&data).ok_or(Error::InvalidBase64)?; mime_type = stripped; - DataEncoding::Base64 + let data = + base64::decode_forgiving(&percent_decoded).ok_or(Error::InvalidBase64)?; + (data, DataEncoding::Base64) } - None => DataEncoding::Percent, + None => (percent_decoded.into_owned(), DataEncoding::Percent), }; let (media_type, parameters) = if mime_type.starts_with(';') { @@ -440,8 +447,11 @@ impl DataUrl { } /// Sets the raw data, which will be serialized byte-for-byte. - pub fn set_data(&mut self, data: &[u8]) { - self.data = data.to_vec(); + /// + /// Passing a `Vec` (or a `String`) moves it into the data URL without copying; + /// slices, arrays and string literals get copied. + pub fn set_data(&mut self, data: impl Into>) { + self.data = data.into(); } /// The data decoded as text using [`encoding`](DataUrl::encoding). @@ -474,54 +484,130 @@ impl DataUrl { Ok(()) } - /// The fragment (the part after `#`), `None` if absent. - pub fn fragment(&self) -> Option<&str> { - self.fragment.as_deref() + /// The exact length of the serialized data URL, without serializing it, + /// e.g. for preallocating a buffer. Serialized data URLs are always ASCII, + /// so this is both the number of bytes and of characters. + pub fn encoded_len(&self) -> usize { + let should_encode = self.percent_encoding_set(); + let data_length = if self.uses_base64(should_encode) { + ";base64,".len() + base64::encoded_length(self.data.len()) + } else { + ",".len() + percent::encoded_length(&self.data, should_encode) + }; + + self.metadata().len() + data_length + self.encoded_fragment().map_or(0, |f| 1 + f.len()) } - /// Sets the fragment, or removes it with `None`. - pub fn set_fragment(&mut self, fragment: Option<&str>) { - self.fragment = fragment.map(str::to_string); + /// Writes the data URL to `writer` without building it as one string first, + /// which saves memory for large data (`to_string()` and `format!` work too). + /// + /// ``` + /// let mut data_url = dataurl::DataUrl::new(); + /// data_url.set_data("Hello!"); + /// + /// let mut output: Vec = Vec::new(); + /// data_url.write_to(&mut output)?; + /// assert_eq!(output, b"data:,Hello!"); + /// # Ok::<(), std::io::Error>(()) + /// ``` + pub fn write_to(&self, writer: W) -> io::Result<()> { + // Adapts io::Write to fmt::Write, keeping the actual I/O error + struct Adapter { + inner: W, + error: Option, + } + + impl fmt::Write for Adapter { + fn write_str(&mut self, s: &str) -> fmt::Result { + self.inner.write_all(s.as_bytes()).map_err(|e| { + self.error = Some(e); + fmt::Error + }) + } + } + + let mut adapter = Adapter { + inner: writer, + error: None, + }; + self.write_into(&mut adapter).map_err(|_| { + adapter + .error + .take() + .unwrap_or_else(|| io::Error::other("unable to write data URL")) + }) } -} -impl fmt::Display for DataUrl { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let mut output = String::from("data:"); + fn write_into(&self, writer: &mut W) -> fmt::Result { + writer.write_str(&self.metadata())?; + + let should_encode = self.percent_encoding_set(); + if self.uses_base64(should_encode) { + writer.write_str(";base64,")?; + base64::write(writer, &self.data)?; + } else { + writer.write_char(',')?; + percent::write_data(writer, &self.data, should_encode)?; + } + + if let Some(fragment) = self.encoded_fragment() { + writer.write_char('#')?; + writer.write_str(&fragment)?; + } + Ok(()) + } + + /// "data:" followed by the media type and its parameters + fn metadata(&self) -> String { + let mut output = String::from("data:"); if let Some(media_type) = &self.media_type { output.push_str(media_type); } mime::serialize_parameters(&mut output, &self.parameters); + output + } - let should_encode: fn(u8) -> bool = match self.escaping { + fn percent_encoding_set(&self) -> fn(u8) -> bool { + match self.escaping { Escaping::Full => percent::in_data_set, Escaping::Minimal => percent::in_minimal_data_set, - }; - let use_base64 = match self.data_encoding { + } + } + + fn uses_base64(&self, should_encode: fn(u8) -> bool) -> bool { + match self.data_encoding { DataEncoding::Percent => false, DataEncoding::Base64 => true, DataEncoding::Shortest => { - // ";base64" plus 4 characters for every 3 bytes, versus 1 or 3 characters per byte - let base64_length = 7 + self.data.len().div_ceil(3) * 4; + let base64_length = ";base64".len() + base64::encoded_length(self.data.len()); base64_length < percent::encoded_length(&self.data, should_encode) } - }; - - if use_base64 { - output.push_str(";base64,"); - base64::encode_into(&mut output, &self.data); - } else { - output.push(','); - percent::encode_data_into(&mut output, &self.data, should_encode); } + } - if let Some(fragment) = &self.fragment { - output.push('#'); + fn encoded_fragment(&self) -> Option { + self.fragment.as_ref().map(|fragment| { + let mut output = String::with_capacity(fragment.len()); percent::encode_into(&mut output, fragment.as_bytes(), percent::in_fragment_set); - } + output + }) + } + + /// The fragment (the part after `#`), `None` if absent. + pub fn fragment(&self) -> Option<&str> { + self.fragment.as_deref() + } + + /// Sets the fragment, or removes it with `None`. + pub fn set_fragment(&mut self, fragment: Option<&str>) { + self.fragment = fragment.map(str::to_string); + } +} - f.write_str(&output) +impl fmt::Display for DataUrl { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.write_into(f) } } diff --git a/src/main.rs b/src/main.rs index f024d4f..72c969b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -141,12 +141,21 @@ fn parse_args(args: impl Iterator) -> Options { options } -fn write_output(output_file: &Option, data: &[u8]) { +fn write_output( + output_file: &Option, + write: impl FnOnce(&mut dyn Write) -> io::Result<()>, +) { let result = match output_file { - Some(path) => fs::write(path, data).map_err(|e| (e, path.to_string_lossy())), + Some(path) => fs::File::create(path) + .and_then(|file| { + let mut writer = io::BufWriter::new(file); + write(&mut writer)?; + writer.flush() + }) + .map_err(|e| (e, path.to_string_lossy())), None => { let mut stdout = io::stdout().lock(); - match stdout.write_all(data).and_then(|_| stdout.flush()) { + match write(&mut stdout).and_then(|_| stdout.flush()) { // Being piped into something like `head` is not an error Err(e) if e.kind() == io::ErrorKind::BrokenPipe => Ok(()), result => result.map_err(|e| (e, "STDOUT".into())), @@ -194,9 +203,11 @@ fn main() { if options.output_file.is_none() && io::stdout().is_terminal() && !data_url.is_binary() { // Text printed into a terminal gets converted from the data URL's charset // TODO: check the terminal's encoding instead of assuming it's UTF-8 - write_output(&None, data_url.text().as_bytes()); + write_output(&None, |writer| writer.write_all(data_url.text().as_bytes())); } else { - write_output(&options.output_file, data_url.data()); + write_output(&options.output_file, |writer| { + writer.write_all(data_url.data()) + }); } } else { let mut data_url = DataUrl::new(); @@ -236,7 +247,7 @@ fn main() { )) }), // Files and STDIN are taken as-is - None => data_url.set_data(&input), + None => data_url.set_data(input), } } (None, Some(text)) => { @@ -245,13 +256,16 @@ fn main() { data_url.set_charset(Some("UTF-8")).unwrap(); } } - (None, None) => data_url.set_data(&input), + (None, None) => data_url.set_data(input), } if let Some(fragment) = &options.fragment { data_url.set_fragment(Some(fragment)); } - write_output(&options.output_file, format!("{data_url}\n").as_bytes()); + write_output(&options.output_file, |writer| { + data_url.write_to(&mut *writer)?; + writer.write_all(b"\n") + }); } } diff --git a/src/percent.rs b/src/percent.rs index 94b58be..1f84722 100644 --- a/src/percent.rs +++ b/src/percent.rs @@ -21,18 +21,36 @@ fn trailing_spaces_start(data: &[u8]) -> usize { data.len() - data.iter().rev().take_while(|&&b| b == b' ').count() } -/// Percent-encodes the data portion of a data URL. -pub(crate) fn encode_data_into(output: &mut String, data: &[u8], should_encode: fn(u8) -> bool) { +/// Percent-encodes the data portion of a data URL, writing it out in chunks, +/// so that large data doesn't need one big intermediate string. +pub(crate) fn write_data( + writer: &mut W, + data: &[u8], + should_encode: fn(u8) -> bool, +) -> std::fmt::Result { + const CHUNK_SIZE: usize = 8 * 1024; let trailing_spaces_start = trailing_spaces_start(data); - output.reserve(data.len()); + let mut buffer = String::with_capacity(CHUNK_SIZE + 3); for (i, &byte) in data.iter().enumerate() { if should_encode(byte) || i >= trailing_spaces_start { - push_encoded(output, byte); + push_encoded(&mut buffer, byte); } else { - output.push(byte as char); + buffer.push(byte as char); + } + if buffer.len() >= CHUNK_SIZE { + writer.write_str(&buffer)?; + buffer.clear(); } } + + writer.write_str(&buffer) +} + +/// Percent-encodes the data portion of a data URL. +#[cfg(test)] +pub(crate) fn encode_data_into(output: &mut String, data: &[u8], should_encode: fn(u8) -> bool) { + write_data(output, data, should_encode).unwrap(); } /// How long the data portion would be once percent-encoded, without encoding it. @@ -116,6 +134,22 @@ mod tests { assert_eq!(encoded_length(b"a b ", in_minimal_data_set), s.len()); } + #[test] + fn chunk_boundaries_do_not_affect_output() { + // Trailing spaces spanning a chunk boundary, and escapes landing right on it + for length in [8 * 1024 - 1, 8 * 1024, 8 * 1024 + 1, 3 * 8 * 1024 + 5] { + for filler in *b"a% " { + let mut data = vec![filler; length]; + data.extend_from_slice(b"x "); + let mut s = String::new(); + encode_data_into(&mut s, &data, in_minimal_data_set); + assert_eq!(s.len(), encoded_length(&data, in_minimal_data_set)); + assert_eq!(decode(s.as_bytes()), data); + assert!(s.ends_with("x%20%20%20")); + } + } + } + #[test] fn decodes_valid_and_keeps_malformed_sequences() { assert_eq!(decode(b"%41%4a%zz%4%"), b"AJ%zz%4%"); diff --git a/tests/large_data.rs b/tests/large_data.rs new file mode 100644 index 0000000..906b57b --- /dev/null +++ b/tests/large_data.rs @@ -0,0 +1,161 @@ +use dataurl::{DataEncoding, DataUrl, Error, Escaping}; +use std::io; + +const LIZARD_PNG: &[u8] = include_bytes!("_data_/lizard.png"); + +/// Deterministic pseudo-random bytes (xorshift), so tests don't need a dependency for randomness +fn pseudo_random_bytes(length: usize, mut seed: u64) -> Vec { + (0..length) + .map(|_| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + seed as u8 + }) + .collect() +} + +fn all_settings() -> Vec<(DataEncoding, Escaping)> { + let mut settings = Vec::new(); + for data_encoding in [ + DataEncoding::Percent, + DataEncoding::Base64, + DataEncoding::Shortest, + ] { + for escaping in [Escaping::Full, Escaping::Minimal] { + settings.push((data_encoding, escaping)); + } + } + settings +} + +#[test] +fn set_data_accepts_owned_and_borrowed_data() { + let mut data_url = DataUrl::new(); + + data_url.set_data(vec![1, 2, 3]); + assert_eq!(data_url.data(), [1, 2, 3]); + data_url.set_data(&[4, 5][..]); + assert_eq!(data_url.data(), [4, 5]); + data_url.set_data([6, 7, 8]); + assert_eq!(data_url.data(), [6, 7, 8]); + data_url.set_data("text"); + assert_eq!(data_url.data(), b"text"); + data_url.set_data(String::from("owned text")); + assert_eq!(data_url.data(), b"owned text"); +} + +#[test] +fn set_data_moves_vectors_without_copying() { + let data = pseudo_random_bytes(1024, 1); + let pointer = data.as_ptr(); + + let mut data_url = DataUrl::new(); + data_url.set_data(data); + assert_eq!(data_url.data().as_ptr(), pointer); + let data = data_url.into_data(); + assert_eq!(data.as_ptr(), pointer); +} + +#[test] +fn encoded_len_is_exact() -> Result<(), Error> { + let datas: Vec> = vec![ + vec![], + b" ".to_vec(), + b"trailing ".to_vec(), + LIZARD_PNG.to_vec(), + (0..=255).collect(), + pseudo_random_bytes(10_000, 2), + ]; + + for data in &datas { + for (data_encoding, escaping) in all_settings() { + for fragment in [None, Some(""), Some("some fragment")] { + let mut data_url = DataUrl::new(); + data_url.set_media_type(Some("image/svg+xml"))?; + data_url.set_parameter("name", Some("a \"b\""))?; + data_url.set_data_encoding(data_encoding); + data_url.set_escaping(escaping); + data_url.set_fragment(fragment); + data_url.set_data(data.as_slice()); + + assert_eq!( + data_url.encoded_len(), + data_url.to_string().len(), + "{data_encoding:?} {escaping:?} {fragment:?}" + ); + } + } + } + Ok(()) +} + +#[test] +fn write_to_matches_display() -> Result<(), Error> { + for (data_encoding, escaping) in all_settings() { + let mut data_url = DataUrl::new(); + data_url.set_charset(Some("utf-8"))?; + data_url.set_data_encoding(data_encoding); + data_url.set_escaping(escaping); + data_url.set_fragment(Some("frag")); + data_url.set_data(pseudo_random_bytes(100_000, 3)); + + let mut written: Vec = Vec::new(); + data_url.write_to(&mut written).unwrap(); + assert_eq!(String::from_utf8(written).unwrap(), data_url.to_string()); + } + Ok(()) +} + +#[test] +fn write_to_reports_io_errors() { + struct FailingWriter; + + impl io::Write for FailingWriter { + fn write(&mut self, _: &[u8]) -> io::Result { + Err(io::Error::new(io::ErrorKind::BrokenPipe, "closed")) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + let mut data_url = DataUrl::new(); + data_url.set_data("x"); + let error = data_url.write_to(FailingWriter).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); +} + +#[test] +fn large_data_round_trips() -> Result<(), Error> { + // Big enough to span many of the serializer's chunks + let data = pseudo_random_bytes(1024 * 1024 + 1, 4); + + for (data_encoding, escaping) in all_settings() { + let mut data_url = DataUrl::new(); + data_url.set_media_type(Some("application/octet-stream"))?; + data_url.set_data_encoding(data_encoding); + data_url.set_escaping(escaping); + data_url.set_data(data.as_slice()); + + let serialized = data_url.to_string(); + assert_eq!(serialized.len(), data_url.encoded_len()); + assert_eq!(DataUrl::parse(&serialized)?.into_data(), data); + } + Ok(()) +} + +#[test] +fn base64_with_and_without_percent_encoded_characters() -> Result<(), Error> { + // Both decoding paths: straight from the input, and after percent-decoding + assert_eq!(DataUrl::parse("data:;base64,SGVsbG8=")?.data(), b"Hello"); + assert_eq!(DataUrl::parse("data:;base64,SGVsbG8%3D")?.data(), b"Hello"); + assert_eq!(DataUrl::parse("data:;base64,SGVs%20bG8")?.data(), b"Hello"); + assert_eq!( + DataUrl::parse("data:;base64,%ZZ"), + Err(Error::InvalidBase64), + "malformed escapes are kept as is, and aren't valid base64" + ); + Ok(()) +} diff --git a/tests/lib.rs b/tests/lib.rs index e523a75..6dba9e5 100644 --- a/tests/lib.rs +++ b/tests/lib.rs @@ -182,7 +182,7 @@ mod serialize { let mut data_url = DataUrl::new(); data_url.set_media_type(Some("image/png")).unwrap(); data_url.set_data_encoding(DataEncoding::Base64); - data_url.set_data(&[0xFB, 0xFF]); + data_url.set_data([0xFB, 0xFF]); assert_eq!(data_url.to_string(), "data:image/png;base64,+/8="); }