add new option to fetch data from remote URL into string

Add new option --payload-url to fetch data from the given URL into
String. If both --input-xml and --payload-url are given, --input-xml
would be taken.

src/bin/download_sysext.rs

@@ -5,7 +5,7 @@ use std::fs::File;
 use std::fs;
 use std::io;
 use std::io::{Read, Seek, SeekFrom};
-use std::io::BufReader;
+use std::io::{BufReader, BufWriter};
 
 #[macro_use]
 extern crate log;
@@ -14,6 +14,7 @@ use anyhow::{Context, Result, bail};
 use globset::{Glob, GlobSet, GlobSetBuilder};
 use hard_xml::XmlRead;
 use argh::FromArgs;
+use reqwest::Client;
 use reqwest::redirect::Policy;
 use url::Url;
 
@@ -147,7 +148,7 @@ impl<'a> Package<'a> {
  let path = into_dir.join(&*self.name);
  let mut file = File::create(path.clone()).context(format!("failed to create path ({:?})", path.display()))?;
 
- let res = match ue_rs::download_and_hash(&client, self.url.clone(), &mut file).await {
+ let res = match fetch_url_to_bufwriter(BufWriter::new(&mut file), self.url.clone(), &client).await {
  Ok(ok) => ok,
  Err(err) => {
  error!("Downloading failed with error {}", err);
@@ -282,6 +283,44 @@ fn get_pkgs_to_download<'a>(resp: &'a omaha::Response, glob_set: &GlobSet)
  Ok(to_download)
 }
 
+// Read data from remote URL into BufWriter.
+async fn fetch_url_to_bufwriter<'a, U, W>(writer: BufWriter<W>, url: U, client: &'a Client) -> Result<ue_rs::DownloadResult<io::BufWriter<W>>>
+where
+ U: reqwest::IntoUrl + From<U> + std::clone::Clone + std::fmt::Debug,
+ W: io::Write,
+ Url: From<U>,
+{
+ let res = match ue_rs::download_and_hash(&client, url.clone().into(), writer).await {
+ Ok(ok) => ok,
+ Err(err) => {
+ bail!("unable to download data(url {:?}), err {:?}", url, err);
+ }
+ };
+
+ Ok(res)
+}
+
+// Read data from remote URL into String.
+async fn fetch_url_to_string<U>(path: &Path, url: U, client: &Client) -> Result<String>
+where
+ U: reqwest::IntoUrl + From<U> + std::clone::Clone + std::fmt::Debug,
+ Url: From<U>,
+{
+ let mut file = File::create(path).context(format!("failed to create path ({:?})", path.display()))?;
+
+ let res = match ue_rs::download_and_hash(&client, url.clone().into(), &mut file).await {
+ Ok(ok) => ok,
+ Err(err) => {
+ bail!("unable to download data(url {:?}), err {:?}", url, err);
+ }
+ };
+
+ let mut outtext: String = Default::default();
+ res.data.read_to_string(&mut outtext).context(format!("failed to read data from URL {:?}", url))?;
+
+ Ok(outtext)
+}
+
 #[derive(FromArgs, Debug)]
 /// Parse an update-engine Omaha XML response to extract sysext images, then download and verify
 /// their signatures.
@@ -292,7 +331,11 @@ struct Args {
 
  /// path to the Omaha XML file, or - to read from stdin
  #[argh(option, short = 'i')]
- input_xml: String,
+ input_xml: Option<String>,
+
+ /// URL to fetch remote update payload
+ #[argh(option, short = 'u')]
+ payload_url: Option<String>,
 
  /// path to the public key file
  #[argh(option, short = 'p')]
@@ -325,14 +368,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
  let glob_set = args.image_match_glob_set()?;
 
- let response_text = match &*args.input_xml {
- "-" => io::read_to_string(io::stdin())?,
- path => {
- let file = File::open(path)?;
- io::read_to_string(file)?
- }
- };
-
  let output_dir = Path::new(&*args.output_dir);
  if !output_dir.try_exists()? {
  return Err(format!("output directory `{}` does not exist", args.output_dir).into());
@@ -343,6 +378,28 @@ async fn main() -> Result<(), Box<dyn Error>> {
  fs::create_dir_all(&unverified_dir)?;
  fs::create_dir_all(&temp_dir)?;
 
+ // The default policy of reqwest Client supports max 10 attempts on HTTP redirect.
+ let client = Client::builder().redirect(Policy::default()).build()?;
+
+ // If input_xml exists, simply read it.
+ // If not, try to read from payload_url.
+ let response_text = match args.input_xml {
+ Some(name) => {
+ if name == "-".to_string() {
+ io::read_to_string(io::stdin())?
+ } else {
+ let file = File::open(name)?;
+ io::read_to_string(file)?
+ }
+ }
+ None => match args.payload_url {
+ Some(url) => fetch_url_to_string(&temp_dir, Url::try_from(url.as_str()).unwrap(), &client).await?,
+ None => return Err(format!("Either --input-xml or --payload-url must be given.").into()),
+ },
+ };
+
+ debug!("response_text: {:?}", response_text);
+
  ////
  // parse response
  ////
@@ -356,8 +413,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
  ////
  // download
  ////
- // The default policy of reqwest Client supports max 10 attempts on HTTP redirect.
- let client = reqwest::Client::builder().redirect(Policy::default()).build()?;
 
  for pkg in pkgs_to_dl.iter_mut() {
  pkg.check_download(&unverified_dir)?;

src/lib.rs

@@ -1,4 +1,5 @@
 mod download;
 pub use download::download_and_hash;
+pub use download::DownloadResult;
 
 pub mod request;