Merge pull request #542 from amazonlinux/updog-whats

updog: Add support for query parameters in TUF requests

workspaces/updater/updog/Cargo.toml

@@ -26,6 +26,7 @@ tough = { version = "0.1.0", features = ["http"] }
 update_metadata = { path = "../update_metadata" }
 structopt = "0.3"
 migrator = { path = "../../api/migration/migrator" }
+url = "2.1.0"
 
 [dev-dependencies]
 tempfile = "3.1.0"

workspaces/updater/updog/src/error.rs

@@ -274,6 +274,12 @@ pub(crate) enum Error {
         source: std::io::Error,
     },
 
+    #[snafu(display("2Borrow2Fast"))]
+    TransportBorrow {
+        backtrace: Backtrace,
+        source: std::cell::BorrowMutError,
+    },
+
     #[snafu(display("Failed to serialize update information: {}", source))]
     UpdateSerialize {
         source: serde_json::Error,

workspaces/updater/updog/src/main.rs

@@ -2,8 +2,10 @@
 #![warn(clippy::pedantic)]
 
 mod error;
+mod transport;
 
 use crate::error::Result;
+use crate::transport::{HttpQueryRepo, HttpQueryTransport};
 use chrono::{DateTime, Utc};
 use data_store_version::Version as DataVersion;
 use semver::Version as SemVer;
@@ -17,11 +19,9 @@ use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
 use std::process;
 use std::str::FromStr;
-use tough::{HttpTransport, Limits, Repository, Settings};
+use tough::{Limits, Repository, Settings};
 use update_metadata::{Manifest, Update};
 
-type HttpRepo<'a> = Repository<'a, HttpTransport>;
-
 #[cfg(target_arch = "x86_64")]
 const TARGET_ARCH: &str = "x86_64";
 #[cfg(target_arch = "aarch64")]
@@ -96,7 +96,10 @@ fn load_config() -> Result<Config> {
     Ok(config)
 }
 
-fn load_repository<'a>(transport: &'a HttpTransport, config: &'a Config) -> Result<HttpRepo<'a>> {
+fn load_repository<'a>(
+    transport: &'a HttpQueryTransport,
+    config: &'a Config,
+) -> Result<HttpQueryRepo<'a>> {
     fs::create_dir_all("/var/lib/thar/updog").context(error::CreateMetadataCache)?;
     Repository::load(
         transport,
@@ -118,7 +121,7 @@ fn load_repository<'a>(transport: &'a HttpTransport, config: &'a Config) -> Resu
     .context(error::Metadata)
 }
 
-fn load_manifest(repository: &HttpRepo<'_>) -> Result<Manifest> {
+fn load_manifest(repository: &HttpQueryRepo<'_>) -> Result<Manifest> {
     let target = "manifest.json";
     serde_json::from_reader(
         repository
@@ -201,7 +204,7 @@ fn update_required<'a>(
 }
 
 fn write_target_to_disk<P: AsRef<Path>>(
-    repository: &HttpRepo<'_>,
+    repository: &HttpQueryRepo<'_>,
     target: &str,
     disk_path: P,
 ) -> Result<()> {
@@ -261,7 +264,7 @@ fn migration_targets(
 /// storage. All intermediate migrations between the current version and the
 /// target version must be retrieved.
 fn retrieve_migrations(
-    repository: &HttpRepo<'_>,
+    repository: &HttpQueryRepo<'_>,
     manifest: &Manifest,
     update: &Update,
 ) -> Result<()> {
@@ -309,7 +312,7 @@ fn retrieve_migrations(
     Ok(())
 }
 
-fn update_image(update: &Update, repository: &HttpRepo<'_>) -> Result<()> {
+fn update_image(update: &Update, repository: &HttpQueryRepo<'_>) -> Result<()> {
     let mut gpt_state = State::load().context(error::PartitionTableRead)?;
     gpt_state.clear_inactive();
     // Write out the clearing of the inactive partition immediately, because we're about to
@@ -444,10 +447,14 @@ fn main_inner() -> Result<()> {
         serde_plain::from_str::<Command>(&arguments.subcommand).unwrap_or_else(|_| usage());
 
     let config = load_config()?;
-    let transport = HttpTransport::new();
+    let (current_version, flavor) = running_version()?;
+    let transport = HttpQueryTransport::new();
+    transport
+        .queries_get_mut()
+        .context(error::TransportBorrow)?
+        .push((String::from("version"), current_version.to_string()));
     let repository = load_repository(&transport, &config)?;
     let manifest = load_manifest(&repository)?;
-    let (current_version, flavor) = running_version().unwrap();
 
     match command {
         Command::CheckUpdate | Command::Whats => {

workspaces/updater/updog/src/transport.rs

@@ -0,0 +1,45 @@
+use std::cell::{BorrowMutError, RefCell};
+use tough::{HttpTransport, Repository, Transport};
+use url::Url;
+
+#[derive(Debug)]
+#[allow(clippy::module_name_repetitions)]
+pub struct HttpQueryTransport {
+    pub inner: HttpTransport,
+    parameters: RefCell<Vec<(String, String)>>,
+}
+
+impl HttpQueryTransport {
+    pub fn new() -> Self {
+        Self {
+            inner: HttpTransport::new(),
+            parameters: RefCell::new(vec![]),
+        }
+    }
+
+    /// Try to borrow a mutable reference to parameters; returns an error if
+    /// a borrow is already active
+    pub fn queries_get_mut(
+        &self,
+    ) -> Result<std::cell::RefMut<'_, Vec<(String, String)>>, BorrowMutError> {
+        self.parameters.try_borrow_mut()
+    }
+
+    fn set_query_string(&self, mut url: Url) -> Url {
+        for (key, val) in self.parameters.borrow().iter() {
+            url.query_pairs_mut().append_pair(&key, &val);
+        }
+        url
+    }
+}
+
+pub type HttpQueryRepo<'a> = Repository<'a, HttpQueryTransport>;
+
+impl Transport for HttpQueryTransport {
+    type Stream = reqwest::Response;
+    type Error = reqwest::Error;
+
+    fn fetch(&self, url: Url) -> Result<Self::Stream, Self::Error> {
+        self.inner.fetch(self.set_query_string(url))
+    }
+}