Merge pull request #39 from osusec/dr/deploy-upload

S3 asset upload for deploy command

Author: detjensrobert

src/access_handlers/docker.rs

@@ -30,12 +30,12 @@ pub async fn check(profile_name: &str) -> Result<()> {
     debug!("will push test image to {}", test_image);
 
     // push alpine image with build credentials
-    check_build_credentials(&client, &test_image)
+    check_build_credentials(client, &test_image)
         .await
         .with_context(|| "Could not push images to registry (bad build credentials?)")?;
 
     // try pulling that image with cluster credentials
-    check_cluster_credentials(&client, &test_image)
+    check_cluster_credentials(client, &test_image)
         .await
         .with_context(|| "Could not pull images from registry (bad cluster credentials?)")?;
 

src/clients.rs

@@ -1,5 +1,7 @@
 // Builders for the various client structs for Docker/Kube etc.
 
+use std::sync::OnceLock;
+
 use anyhow::{anyhow, bail, Context, Error, Result};
 use bollard;
 use futures::TryFutureExt;
@@ -17,16 +19,27 @@ use crate::configparser::config;
 //
 // Docker stuff
 //
-pub async fn docker() -> Result<bollard::Docker> {
-    debug!("connecting to docker...");
-    let client = bollard::Docker::connect_with_defaults()?;
-    client
-        .ping()
-        .await
-        // truncate error chain with new error (returned error is way too verbose)
-        .map_err(|_| anyhow!("could not talk to Docker daemon (is DOCKER_HOST correct?)"))?;
 
-    Ok(client)
+static DOCKER_CLIENT: OnceLock<bollard::Docker> = OnceLock::new();
+
+/// Return existing or create new Docker client
+pub async fn docker() -> Result<&'static bollard::Docker> {
+    match DOCKER_CLIENT.get() {
+        Some(d) => Ok(d),
+        None => {
+            debug!("connecting to docker...");
+            let client = bollard::Docker::connect_with_defaults()?;
+            client
+                .ping()
+                .await
+                // truncate error chain with new error (returned error is way too verbose)
+                .map_err(|_| {
+                    anyhow!("could not talk to Docker daemon (is DOCKER_HOST correct?)")
+                })?;
+
+            Ok(DOCKER_CLIENT.get_or_init(|| client))
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -54,27 +67,38 @@ pub async fn engine_type() -> EngineType {
 // S3 stuff
 //
 
-/// create bucket client for passed profile config
-pub fn bucket_client(config: &config::S3Config) -> Result<Box<s3::Bucket>> {
-    trace!("creating bucket client");
-    // TODO: once_cell this so it reuses the same bucket?
-    let region = s3::Region::Custom {
-        region: config.region.clone(),
-        endpoint: config.endpoint.clone(),
-    };
-    let creds = s3::creds::Credentials::new(
-        Some(&config.access_key),
-        Some(&config.secret_key),
-        None,
-        None,
-        None,
-    )?;
-    let bucket = s3::Bucket::new(&config.bucket_name, region, creds)?.with_path_style();
-
-    Ok(bucket)
+// this does need to be a OnceLock instead of a LazyLock, even though how this
+// is used is more inline with a LazyLock. Lazy does not allow for passing
+// anything into the init function, and this needs a parameter to know what
+// profile to fetch creds for.
+static BUCKET_CLIENT: OnceLock<Box<s3::Bucket>> = OnceLock::new();
+
+/// return existing or create new bucket client for passed profile config
+pub fn bucket_client(config: &config::S3Config) -> Result<&s3::Bucket> {
+    match BUCKET_CLIENT.get() {
+        Some(b) => Ok(b),
+        None => {
+            trace!("creating bucket client");
+            let region = s3::Region::Custom {
+                region: config.region.clone(),
+                endpoint: config.endpoint.clone(),
+            };
+            let creds = s3::creds::Credentials::new(
+                Some(&config.access_key),
+                Some(&config.secret_key),
+                None,
+                None,
+                None,
+            )?;
+            let bucket = s3::Bucket::new(&config.bucket_name, region, creds)?.with_path_style();
+
+            Ok(BUCKET_CLIENT.get_or_init(|| bucket))
+        }
+    }
 }
 
 /// create public/anonymous bucket client for passed profile config
+// this does not need a oncelock and can be created on-demand, as this is not used in very many places
 pub fn bucket_client_anonymous(config: &config::S3Config) -> Result<Box<s3::Bucket>> {
     trace!("creating anon bucket client");
     // TODO: once_cell this so it reuses the same bucket?
@@ -92,6 +116,10 @@ pub fn bucket_client_anonymous(config: &config::S3Config) -> Result<Box<s3::Buck
 // Kubernetes stuff
 //
 
+// no OnceLock caching for K8S client. Some operations with the client require
+// their own owned `kube::Client`, so always returning a borrowed client from
+// the OnceLock would not work.
+
 /// Returns Kubernetes Client for selected profile
 pub async fn kube_client(profile: &config::ProfileConfig) -> Result<kube::Client> {
     debug!("building kube client");

src/commands/deploy.rs

@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use simplelog::*;
 use std::process::exit;
 
@@ -31,6 +32,11 @@ pub async fn run(profile_name: &str, no_build: &bool, _dry_run: &bool) {
         }
     };
 
+    trace!(
+        "got built results: {:#?}",
+        build_results.iter().map(|b| &b.1).collect_vec()
+    );
+
     // deploy needs to:
     // A) render kubernetes manifests
     //    - namespace, deployment, service, ingress
@@ -41,21 +47,18 @@ pub async fn run(profile_name: &str, no_build: &bool, _dry_run: &bool) {
     // C) update frontend with new state of challenges
 
     // A)
-    info!("deploying challenges...");
     if let Err(e) = deploy::kubernetes::deploy_challenges(profile_name, &build_results).await {
         error!("{e:?}");
         exit(1);
     }
 
     // B)
-    info!("deploying challenges...");
     if let Err(e) = deploy::s3::upload_assets(profile_name, &build_results).await {
         error!("{e:?}");
         exit(1);
     }
 
-    // A)
-    info!("deploying challenges...");
+    // C)
     if let Err(e) = deploy::frontend::update_frontend(profile_name, &build_results).await {
         error!("{e:?}");
         exit(1);

src/deploy/s3.rs

@@ -1,9 +1,12 @@
 use std::fs::File;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use anyhow::{anyhow, bail, Context, Error, Ok, Result};
+use futures::future::try_join_all;
 use itertools::Itertools;
+use s3::Bucket;
 use simplelog::*;
+use tokio;
 
 use crate::builder::BuildResult;
 use crate::clients::bucket_client;
@@ -16,11 +19,66 @@ use crate::utils::TryJoinAll;
 pub async fn upload_assets(
     profile_name: &str,
     build_results: &[(&ChallengeConfig, BuildResult)],
-) -> Result<Vec<String>> {
+) -> Result<Vec<BuildResult>> {
     let profile = get_profile_config(profile_name)?;
     let enabled_challenges = enabled_challenges(profile_name)?;
 
     let bucket = bucket_client(&profile.s3)?;
 
-    todo!();
+    info!("uploading assets...");
+
+    // upload all files for each challenge
+    build_results
+        .iter()
+        .map(|(chal, result)| async move {
+            // upload all files for a specific challenge
+
+            info!("  for chal {:?}...", chal.directory);
+
+            let uploaded = result
+                .assets
+                .iter()
+                .map(|asset_file| async move {
+                    upload_single_file(bucket, chal, asset_file)
+                        .await
+                        .with_context(|| format!("failed to upload file {asset_file:?}"))
+                })
+                .try_join_all()
+                .await
+                .with_context(|| {
+                    format!("failed to upload asset files for chal {:?}", chal.directory)
+                })?;
+
+            // return new BuildResult with assets as bucket path
+            Ok(BuildResult {
+                tags: result.tags.clone(),
+                assets: uploaded,
+            })
+        })
+        .try_join_all()
+        .await
+}
+
+async fn upload_single_file(
+    bucket: &Bucket,
+    chal: &ChallengeConfig,
+    file: &Path,
+) -> Result<PathBuf> {
+    // e.g. s3.example.domain/assets/misc/foo/stuff.zip
+    let path_in_bucket = format!(
+        "assets/{chal_slug}/{file}",
+        chal_slug = chal.directory.to_string_lossy(),
+        file = file.file_name().unwrap().to_string_lossy()
+    );
+
+    trace!("uploading {:?} to bucket path {:?}", file, &path_in_bucket);
+
+    // TODO: move to async/streaming to better handle large files and report progress
+    let mut asset_file = tokio::fs::File::open(file).await?;
+    let r = bucket
+        .put_object_stream(&mut asset_file, &path_in_bucket)
+        .await?;
+    trace!("uploaded {} bytes for file {:?}", r.uploaded_bytes(), file);
+
+    Ok(PathBuf::from(path_in_bucket))
 }