diff --git a/packages/os/Cargo.toml b/packages/os/Cargo.toml index 54b0c76a3de..6e122888410 100644 --- a/packages/os/Cargo.toml +++ b/packages/os/Cargo.toml @@ -19,7 +19,8 @@ source-groups = [ "webpki-roots-shim", "logdog", "models", - "imdsclient" + "imdsclient", + "retry-read", ] [lib] diff --git a/sources/Cargo.lock b/sources/Cargo.lock index 8b659496afd..5714ffc5003 100644 --- a/sources/Cargo.lock +++ b/sources/Cargo.lock @@ -962,6 +962,7 @@ dependencies = [ "imdsclient", "lazy_static", "log", + "retry-read", "serde", "serde-xml-rs", "serde_json", @@ -2521,6 +2522,13 @@ dependencies = [ "winreg", ] +[[package]] +name = "retry-read" +version = "0.1.0" +dependencies = [ + "cargo-readme", +] + [[package]] name = "ring" version = "0.16.20" diff --git a/sources/Cargo.toml b/sources/Cargo.toml index 2bd5c9c1e58..b7a3b71dcff 100644 --- a/sources/Cargo.toml +++ b/sources/Cargo.toml @@ -47,6 +47,8 @@ members = [ "parse-datetime", + "retry-read", + "updater/block-party", "updater/signpost", "updater/update_metadata", diff --git a/sources/api/early-boot-config/Cargo.toml b/sources/api/early-boot-config/Cargo.toml index 0bce2b7330f..50dff4ca2f0 100644 --- a/sources/api/early-boot-config/Cargo.toml +++ b/sources/api/early-boot-config/Cargo.toml @@ -18,6 +18,7 @@ flate2 = { version = "1.0", default-features = false, features = ["rust_backend" http = "0.2" imdsclient = { path = "../../imdsclient", version = "0.1.0" } log = "0.4" +retry-read = { path = "../../retry-read", version = "0.1.0" } serde = { version = "1.0", features = ["derive"] } serde_json = "1" serde_plain = "1.0" diff --git a/sources/api/early-boot-config/src/compression.rs b/sources/api/early-boot-config/src/compression.rs index 2c0ee052b4d..817647cdbcc 100644 --- a/sources/api/early-boot-config/src/compression.rs +++ b/sources/api/early-boot-config/src/compression.rs @@ -3,8 +3,9 @@ //! Currently gzip compression is supported. use flate2::read::GzDecoder; +use retry_read::RetryRead; use std::fs::File; -use std::io::{BufReader, Chain, Cursor, ErrorKind, Read, Result, Take}; +use std::io::{BufReader, Chain, Cursor, Read, Result, Take}; use std::path::Path; /// "File magic" that indicates file type is stored in a few bytes at the start at the start of the @@ -120,42 +121,6 @@ impl Read for OptionalCompressionReader { // =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= -/// This trait represents a `Read` operation where we want to retry after standard interruptions -/// (unlike `read()`) but also need to know the number of bytes we read (unlike `read_exact()`). -trait RetryRead { - fn retry_read(&mut self, buf: &mut [u8]) -> Result; -} - -impl RetryRead for R { - // This implementation is based on stdlib Read::read_exact, but hitting EOF isn't a failure, we - // just want to return the number of bytes we could read. - fn retry_read(&mut self, mut buf: &mut [u8]) -> Result { - let mut count = 0; - - // Read until we have no more space in the output buffer - while !buf.is_empty() { - match self.read(buf) { - // No bytes left, done - Ok(0) => break, - // Read n bytes, slide ahead n in the output buffer and read more - Ok(n) => { - count += n; - let tmp = buf; - buf = &mut tmp[n..]; - } - // Retry on interrupt - Err(e) if e.kind() == ErrorKind::Interrupted => {} - // Other failures are fatal - Err(e) => return Err(e), - } - } - - Ok(count) - } -} - -// =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= =^..^= - #[cfg(test)] mod test { use super::*; diff --git a/sources/retry-read/Cargo.toml b/sources/retry-read/Cargo.toml new file mode 100644 index 00000000000..dbf4faf57d1 --- /dev/null +++ b/sources/retry-read/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "retry-read" +version = "0.1.0" +authors = ["Tom Kirchner "] +license = "Apache-2.0 OR MIT" +edition = "2018" +publish = false +# Don't rebuild crate just because of changes to README. +exclude = ["README.md"] + +[build-dependencies] +cargo-readme = "3.1" diff --git a/sources/retry-read/README.md b/sources/retry-read/README.md new file mode 100644 index 00000000000..2ec99fefa4a --- /dev/null +++ b/sources/retry-read/README.md @@ -0,0 +1,12 @@ +# retry-read + +Current version: 0.1.0 + +This library provides a `RetryRead` trait with a `retry_read` function that's available for any +`Read` type. `retry_read` retries after standard interruptions (unlike `read`) but also +returns the number of bytes read (unlike `read_exact`), and without needing to read to the end +of the input (unlike `read_to_end` and `read_to_string`). + +## Colophon + +This text was generated from `README.tpl` using [cargo-readme](https://crates.io/crates/cargo-readme), and includes the rustdoc from `src/lib.rs`. \ No newline at end of file diff --git a/sources/retry-read/README.tpl b/sources/retry-read/README.tpl new file mode 100644 index 00000000000..91fb62910c8 --- /dev/null +++ b/sources/retry-read/README.tpl @@ -0,0 +1,9 @@ +# {{crate}} + +Current version: {{version}} + +{{readme}} + +## Colophon + +This text was generated from `README.tpl` using [cargo-readme](https://crates.io/crates/cargo-readme), and includes the rustdoc from `src/lib.rs`. diff --git a/sources/retry-read/build.rs b/sources/retry-read/build.rs new file mode 100644 index 00000000000..86c7bbc026e --- /dev/null +++ b/sources/retry-read/build.rs @@ -0,0 +1,32 @@ +// Automatically generate README.md from rustdoc. + +use std::env; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; + +fn main() { + // Check for environment variable "SKIP_README". If it is set, + // skip README generation + if env::var_os("SKIP_README").is_some() { + return; + } + + let mut source = File::open("src/lib.rs").unwrap(); + let mut template = File::open("README.tpl").unwrap(); + + let content = cargo_readme::generate_readme( + &PathBuf::from("."), // root + &mut source, // source + Some(&mut template), // template + // The "add x" arguments don't apply when using a template. + true, // add title + false, // add badges + false, // add license + true, // indent headings + ) + .unwrap(); + + let mut readme = File::create("README.md").unwrap(); + readme.write_all(content.as_bytes()).unwrap(); +} diff --git a/sources/retry-read/src/lib.rs b/sources/retry-read/src/lib.rs new file mode 100644 index 00000000000..9b1f4add766 --- /dev/null +++ b/sources/retry-read/src/lib.rs @@ -0,0 +1,114 @@ +//! This library provides a `RetryRead` trait with a `retry_read` function that's available for any +//! `Read` type. `retry_read` retries after standard interruptions (unlike `read`) but also +//! returns the number of bytes read (unlike `read_exact`), and without needing to read to the end +//! of the input (unlike `read_to_end` and `read_to_string`). + +use std::io::{ErrorKind, Read, Result}; + +/// Provides a way to retry standard read operations while also returning the number of bytes read. +pub trait RetryRead { + fn retry_read(&mut self, buf: &mut [u8]) -> Result; +} + +impl RetryRead for R { + // This implementation is based on stdlib Read::read_exact, but hitting EOF isn't a failure, we + // just want to return the number of bytes we could read. + /// Like `Read::read` but retries on ErrorKind::Interrupted, returning the number of bytes read. + fn retry_read(&mut self, mut buf: &mut [u8]) -> Result { + let mut count = 0; + + // Read until we have no more space in the output buffer + while !buf.is_empty() { + match self.read(buf) { + // No bytes left, done + Ok(0) => break, + // Read n bytes, slide ahead n in the output buffer and read more + Ok(n) => { + count += n; + let tmp = buf; + buf = &mut tmp[n..]; + } + // Retry on interrupt + Err(e) if e.kind() == ErrorKind::Interrupted => {} + // Other failures are fatal + Err(e) => return Err(e), + } + } + + Ok(count) + } +} + +#[cfg(test)] +mod test { + use super::{ErrorKind, Read, Result, RetryRead}; + use std::io::{Error, Write}; + + // Helper method for simple test cases, confirming we read the full given slice. + fn test(data: &[u8]) { + let mut output = vec![0; data.len()]; + let count = (&data[..]).retry_read(&mut output).unwrap(); + assert_eq!(count, data.len()); + assert_eq!(&data[..], &output); + } + + #[test] + fn zero_read() { + test(&[]); + } + + #[test] + fn small_read() { + test(&[0, 1, 2, 3, 42]); + } + + #[test] + fn large_read() { + test(&[42; 9999]); + } + + // Confirm we retry reads when interrupted. + #[test] + fn retried_read() { + let mut reader = InterruptedReader::new(5); + let mut output = vec![0; 5]; + let count = reader.retry_read(&mut output).unwrap(); + assert_eq!(count, 5); + assert_eq!(output, vec![42, 42, 42, 42, 42]); + } + + // Helper that implements Read, eventually returning the requested number of bytes, but returns + // ErrorKind::Interrupted every other call. + struct InterruptedReader { + requested_reads: u64, + finished_reads: u64, + interrupt: bool, + } + + impl InterruptedReader { + fn new(requested_reads: u64) -> Self { + Self { + requested_reads, + finished_reads: 0, + interrupt: false, + } + } + } + + impl Read for InterruptedReader { + fn read(&mut self, mut buf: &mut [u8]) -> Result { + if self.finished_reads > self.requested_reads { + return Ok(0); + } + + if self.interrupt { + self.interrupt = false; + Err(Error::new(ErrorKind::Interrupted, "you asked for it")) + } else { + self.interrupt = true; + self.finished_reads += 1; + buf.write(&[42]) + } + } + } +}