Explicit proxy support (#337)

* Add proxy support * Improve clap definition and check * Include missing protocol * Improve error
2023-03-14 07:56:57 -07:00 · 2023-03-14 07:56:57 -07:00 · c13b08987b
parent a977370e74
commit c13b08987b
5 changed files with 75 additions and 18 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -1313,6 +1313,7 @@ dependencies = [
 "serde_urlencoded",
 "tokio",
 "tokio-rustls",
+ "tokio-socks",
 "tokio-util",
 "tower-service",
 "url",
@ -1823,6 +1824,18 @@ dependencies = [
 "webpki",
 ]

+[[package]]
+name = "tokio-socks"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "51165dfa029d2a65969413a6cc96f354b86b464498702f174a4efa13608fd8c0"
+dependencies = [
+ "either",
+ "futures-util",
+ "thiserror",
+ "tokio",
+]
+
 [[package]]
 name = "tokio-util"
 version = "0.7.7"
--- a/Cargo.toml
+++ b/Cargo.toml
@ -33,7 +33,7 @@ eyre = { version = "0.6.8", default-features = false, features = [ "track-caller
 glob = { version = "0.3.0", default-features = false }
 nix = { version = "0.26.0", default-features = false, features = ["user", "fs", "process", "term"] }
 owo-colors = { version = "3.5.0", default-features = false, features = [ "supports-colors" ] }
-reqwest = { version = "0.11.11", default-features = false, features = ["rustls-tls-native-roots", "stream"] }
+reqwest = { version = "0.11.11", default-features = false, features = ["rustls-tls-native-roots", "stream", "socks"] }
 serde = { version = "1.0.144", default-features = false, features = [ "std", "derive" ] }
 serde_json = { version = "1.0.85", default-features = false, features = [ "std" ] }
 serde_with = { version = "2.0.1", default-features = false, features = [ "std", "macros" ] }
--- a/src/action/base/fetch_and_unpack_nix.rs
+++ b/src/action/base/fetch_and_unpack_nix.rs
@ -13,11 +13,16 @@ Fetch a URL to the given path
 pub struct FetchAndUnpackNix {
    url: Url,
    dest: PathBuf,
+    proxy: Option<Url>,
 }

 impl FetchAndUnpackNix {
    #[tracing::instrument(level = "debug", skip_all)]
-    pub async fn plan(url: Url, dest: PathBuf) -> Result<StatefulAction<Self>, ActionError> {
+    pub async fn plan(
+        url: Url,
+        dest: PathBuf,
+        proxy: Option<Url>,
+    ) -> Result<StatefulAction<Self>, ActionError> {
        // TODO(@hoverbear): Check URL exists?
        // TODO(@hoverbear): Check tempdir exists

@ -30,7 +35,18 @@ impl FetchAndUnpackNix {
            },
        };

-        Ok(Self { url, dest }.into())
+        if let Some(proxy) = &proxy {
+            match proxy.scheme() {
+                "https" | "http" | "socks5" => (),
+                _ => {
+                    return Err(ActionError::Custom(Box::new(
+                        FetchUrlError::UnknownProxyScheme,
+                    )))
+                },
+            };
+        }
+
+        Ok(Self { url, dest, proxy }.into())
    }
 }

@ -45,12 +61,17 @@ impl Action for FetchAndUnpackNix {
    }

    fn tracing_span(&self) -> Span {
-        span!(
+        let span = span!(
            tracing::Level::DEBUG,
            "fetch_and_unpack_nix",
            url = tracing::field::display(&self.url),
+            proxy = tracing::field::Empty,
            dest = tracing::field::display(self.dest.display()),
-        )
+        );
+        if let Some(proxy) = &self.proxy {
+            span.record("proxy", tracing::field::display(&proxy));
+        }
+        span
    }

    fn execute_description(&self) -> Vec<ActionDescription> {
@ -59,11 +80,24 @@ impl Action for FetchAndUnpackNix {

    #[tracing::instrument(level = "debug", skip_all)]
    async fn execute(&mut self) -> Result<(), ActionError> {
-        let Self { url, dest } = self;
-
-        let bytes = match url.scheme() {
+        let bytes = match self.url.scheme() {
            "https" | "http" => {
-                let res = reqwest::get(url.clone())
+                let mut buildable_client = reqwest::Client::builder();
+                if let Some(proxy) = &self.proxy {
+                    buildable_client =
+                        buildable_client.proxy(reqwest::Proxy::all(proxy.clone()).map_err(|e| {
+                            ActionError::Custom(Box::new(FetchUrlError::Reqwest(e)))
+                        })?)
+                }
+                let client = buildable_client
+                    .build()
+                    .map_err(|e| ActionError::Custom(Box::new(FetchUrlError::Reqwest(e))))?;
+                let req = client
+                    .get(self.url.clone())
+                    .build()
+                    .map_err(|e| ActionError::Custom(Box::new(FetchUrlError::Reqwest(e))))?;
+                let res = client
+                    .execute(req)
                    .await
                    .map_err(|e| ActionError::Custom(Box::new(FetchUrlError::Reqwest(e))))?;
                res.bytes()
@ -71,9 +105,9 @@ impl Action for FetchAndUnpackNix {
                    .map_err(|e| ActionError::Custom(Box::new(FetchUrlError::Reqwest(e))))?
            },
            "file" => {
-                let buf = tokio::fs::read(url.path())
+                let buf = tokio::fs::read(self.url.path())
                    .await
-                    .map_err(|e| ActionError::Read(PathBuf::from(url.path()), e))?;
+                    .map_err(|e| ActionError::Read(PathBuf::from(self.url.path()), e))?;
                Bytes::from(buf)
            },
            _ => {
@ -85,7 +119,7 @@ impl Action for FetchAndUnpackNix {

        // TODO(@Hoverbear): Pick directory
        tracing::trace!("Unpacking tar.xz");
-        let dest_clone = dest.clone();
+        let dest_clone = self.dest.clone();

        let decoder = xz2::read::XzDecoder::new(bytes.reader());
        let mut archive = tar::Archive::new(decoder);
@ -102,8 +136,6 @@ impl Action for FetchAndUnpackNix {

    #[tracing::instrument(level = "debug", skip_all)]
    async fn revert(&mut self) -> Result<(), ActionError> {
-        let Self { url: _, dest: _ } = self;
-
        Ok(())
    }
 }
@ -119,6 +151,8 @@ pub enum FetchUrlError {
    ),
    #[error("Unarchiving error")]
    Unarchive(#[source] std::io::Error),
-    #[error("Unknown url scheme")]
+    #[error("Unknown url scheme, `file://`, `https://` and `http://` supported")]
    UnknownUrlScheme,
+    #[error("Unknown url scheme, `https://`, `socks5://`, and `http://` supported")]
+    UnknownProxyScheme,
 }
--- a/src/action/common/provision_nix.rs
+++ b/src/action/common/provision_nix.rs
@ -24,9 +24,12 @@ pub struct ProvisionNix {
 impl ProvisionNix {
    #[tracing::instrument(level = "debug", skip_all)]
    pub async fn plan(settings: &CommonSettings) -> Result<StatefulAction<Self>, ActionError> {
-        let fetch_nix =
-            FetchAndUnpackNix::plan(settings.nix_package_url.clone(), PathBuf::from(SCRATCH_DIR))
-                .await?;
+        let fetch_nix = FetchAndUnpackNix::plan(
+            settings.nix_package_url.clone(),
+            PathBuf::from(SCRATCH_DIR),
+            settings.proxy.clone(),
+        )
+        .await?;
        let create_users_and_group = CreateUsersAndGroups::plan(settings.clone())
            .await
            .map_err(|e| ActionError::Child(CreateUsersAndGroups::action_tag(), Box::new(e)))?;
--- a/src/settings.rs
+++ b/src/settings.rs
@ -172,6 +172,10 @@ pub struct CommonSettings {
    )]
    pub(crate) nix_package_url: Url,

+    /// The proxy to use (if any), valid proxy bases are `https://$URL`, `http://$URL` and `socks5://$URL`
+    #[cfg_attr(feature = "cli", clap(long, env = "NIX_INSTALLER_PROXY"))]
+    pub(crate) proxy: Option<Url>,
+
    /// Extra configuration lines for `/etc/nix.conf`
    #[cfg_attr(feature = "cli", clap(long, action = ArgAction::Set, num_args = 0.., value_delimiter = ',', env = "NIX_INSTALLER_EXTRA_CONF", global = true))]
    pub extra_conf: Vec<String>,
@ -272,6 +276,7 @@ impl CommonSettings {
            nix_build_user_prefix: nix_build_user_prefix.to_string(),
            nix_build_user_id_base,
            nix_package_url: url.parse()?,
+            proxy: Default::default(),
            extra_conf: Default::default(),
            force: false,
            #[cfg(feature = "diagnostics")]
@ -291,6 +296,7 @@ impl CommonSettings {
            nix_build_user_prefix,
            nix_build_user_id_base,
            nix_package_url,
+            proxy,
            extra_conf,
            force,
            #[cfg(feature = "diagnostics")]
@ -326,6 +332,7 @@ impl CommonSettings {
            "nix_package_url".into(),
            serde_json::to_value(nix_package_url)?,
        );
+        map.insert("proxy".into(), serde_json::to_value(proxy)?);
        map.insert("extra_conf".into(), serde_json::to_value(extra_conf)?);
        map.insert("force".into(), serde_json::to_value(force)?);