Extract retry-read to new crate

packages/os/Cargo.toml

@@ -19,7 +19,8 @@ source-groups = [
     "webpki-roots-shim",
     "logdog",
     "models",
-    "imdsclient"
+    "imdsclient",
+    "retry-read",
 ]
 
 [lib]

sources/Cargo.toml

@@ -47,6 +47,8 @@ members = [
 
     "parse-datetime",
 
+    "retry-read",
+
     "updater/block-party",
     "updater/signpost",
     "updater/update_metadata",

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"

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<R: Read> Read for OptionalCompressionReader<R> {
 
 // =^..^=   =^..^=   =^..^=   =^..^=   =^..^=   =^..^=   =^..^=   =^..^=   =^..^=   =^..^=   =^..^=
 
-/// 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<R> {
-    fn retry_read(&mut self, buf: &mut [u8]) -> Result<usize>;
-}
-
-impl<R: Read> RetryRead<R> 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<usize> {
-        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::*;

sources/retry-read/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "retry-read"
+version = "0.1.0"
+authors = ["Tom Kirchner <tjk@amazon.com>"]
+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"

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`.

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`.

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();
+}

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<R> {
+    fn retry_read(&mut self, buf: &mut [u8]) -> Result<usize>;
+}
+
+impl<R: Read> RetryRead<R> 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<usize> {
+        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<usize> {
+            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])
+            }
+        }
+    }
+}