Merge pull request #5200 from arlosi/invalidate

Add CloudFront invalidation for index files

Author: jtgeibel

.env.sample

@@ -41,6 +41,12 @@ export TEST_DATABASE_URL=
 # not needed if the S3 bucket is in US standard
 # export S3_INDEX_REGION=
 
+# Credentials for invalidating cached files on CloudFront. You can leave these
+# commented out if you're not using CloudFront caching for the index files.
+# export CLOUDFRONT_DISTRIBUTION=
+# export CLOUDFRONT_ACCESS_KEY=
+# export CLOUDFRONT_SECRET_KEY=
+
 # Upstream location of the registry index. Background jobs will push to
 # this URL. The default points to a local index for development.
 # Run `./script/init-local-index.sh` to initialize this repo.

Cargo.lock

@@ -23,6 +23,7 @@ path = "src/tests/all.rs"
 
 [dependencies]
 anyhow = "=1.0.65"
+aws-sigv4 = "=0.48.0"
 base64 = "=0.13.0"
 cargo-registry-index = { path = "cargo-registry-index" }
 cargo-registry-markdown = { path = "cargo-registry-markdown" }

Cargo.toml

@@ -6,6 +6,7 @@ use swirl::PerformError;
 
 use crate::db::{DieselPool, DieselPooledConn, PoolError};
 use crate::uploaders::Uploader;
+use crate::worker::cloudfront::CloudFront;
 use cargo_registry_index::Repository;
 
 impl<'a> swirl::db::BorrowedConnection<'a> for DieselPool {
@@ -24,6 +25,7 @@ pub struct Environment {
     index: Arc<Mutex<Repository>>,
     pub uploader: Uploader,
     http_client: AssertUnwindSafe<Client>,
+    cloudfront: Option<CloudFront>,
 }
 
 impl Clone for Environment {
@@ -32,24 +34,37 @@ impl Clone for Environment {
             index: self.index.clone(),
             uploader: self.uploader.clone(),
             http_client: AssertUnwindSafe(self.http_client.0.clone()),
+            cloudfront: self.cloudfront.clone(),
         }
     }
 }
 
 impl Environment {
-    pub fn new(index: Repository, uploader: Uploader, http_client: Client) -> Self {
-        Self::new_shared(Arc::new(Mutex::new(index)), uploader, http_client)
+    pub fn new(
+        index: Repository,
+        uploader: Uploader,
+        http_client: Client,
+        cloudfront: Option<CloudFront>,
+    ) -> Self {
+        Self::new_shared(
+            Arc::new(Mutex::new(index)),
+            uploader,
+            http_client,
+            cloudfront,
+        )
     }
 
     pub fn new_shared(
         index: Arc<Mutex<Repository>>,
         uploader: Uploader,
         http_client: Client,
+        cloudfront: Option<CloudFront>,
     ) -> Self {
         Self {
             index,
             uploader,
             http_client: AssertUnwindSafe(http_client),
+            cloudfront,
         }
     }
 
@@ -63,4 +78,8 @@ impl Environment {
     pub(crate) fn http_client(&self) -> &Client {
         &self.http_client
     }
+
+    pub(crate) fn cloudfront(&self) -> Option<&CloudFront> {
+        self.cloudfront.as_ref()
+    }
 }

src/background_jobs.rs

@@ -13,6 +13,7 @@
 #![warn(clippy::all, rust_2018_idioms)]
 
 use cargo_registry::config;
+use cargo_registry::worker::cloudfront::CloudFront;
 use cargo_registry::{background_jobs::*, db};
 use cargo_registry_index::{Repository, RepositoryConfig};
 use diesel::r2d2;
@@ -52,12 +53,19 @@ fn main() {
     ));
     println!("Index cloned");
 
+    let cloudfront = CloudFront::from_environment();
+
     let build_runner = || {
         let client = Client::builder()
             .timeout(Duration::from_secs(45))
             .build()
             .expect("Couldn't build client");
-        let environment = Environment::new_shared(repository.clone(), uploader.clone(), client);
+        let environment = Environment::new_shared(
+            repository.clone(),
+            uploader.clone(),
+            client,
+            cloudfront.clone(),
+        );
         let db_config = r2d2::Pool::builder().min_idle(Some(0));
         swirl::Runner::builder(environment)
             .connection_pool_builder(&db_url, db_config)

src/bin/background-worker.rs

@@ -236,6 +236,7 @@ impl TestAppBuilder {
                 index,
                 app.config.uploader().clone(),
                 app.http_client().clone(),
+                None,
             );
 
             Some(

src/tests/util/test_app.rs

@@ -0,0 +1,75 @@
+use std::time::SystemTime;
+
+use aws_sigv4::{
+    http_request::{self, SignableRequest, SigningSettings},
+    SigningParams,
+};
+use reqwest::blocking::Client;
+
+#[derive(Clone)]
+pub struct CloudFront {
+    distribution_id: String,
+    access_key: String,
+    secret_key: String,
+}
+
+impl CloudFront {
+    pub fn from_environment() -> Option<Self> {
+        let distribution_id = dotenv::var("CLOUDFRONT_DISTRIBUTION").ok()?;
+        let access_key =
+            dotenv::var("CLOUDFRONT_ACCESS_KEY").expect("missing CLOUDFRONT_ACCESS_KEY");
+        let secret_key =
+            dotenv::var("CLOUDFRONT_SECRET_KEY").expect("missing CLOUDFRONT_SECRET_KEY");
+        Some(Self {
+            distribution_id,
+            access_key,
+            secret_key,
+        })
+    }
+
+    /// Invalidate a file on CloudFront
+    ///
+    /// `path` is the path to the file to invalidate, such as `config.json`, or `re/ge/regex`
+    pub fn invalidate(&self, client: &Client, path: &str) -> anyhow::Result<()> {
+        let path = path.trim_start_matches('/');
+        let url = format!(
+            "https://cloudfront.amazonaws.com/2020-05-31/distribution/{}/invalidation",
+            self.distribution_id
+        );
+        let now = chrono::offset::Utc::now().timestamp_micros();
+        let body = format!(
+            r#"
+<?xml version="1.0" encoding="UTF-8"?>
+<InvalidationBatch xmlns="http://cloudfront.amazonaws.com/doc/2020-05-31/">
+    <CallerReference>{now}</CallerReference>
+    <Paths>
+        <Items>
+            <Path>/{path}</Path>
+        </Items>
+        <Quantity>1</Quantity>
+    </Paths>
+</InvalidationBatch>
+"#
+        );
+
+        let request = http::Request::post(&url).body(&body)?;
+        let request = SignableRequest::from(&request);
+        let params = SigningParams::builder()
+            .access_key(&self.access_key)
+            .secret_key(&self.secret_key)
+            .region("us-east-1") // cloudfront is a regionless service, use the default region for signing.
+            .service_name("cloudfront")
+            .settings(SigningSettings::default())
+            .time(SystemTime::now())
+            .build()
+            .unwrap(); // all required fields are set
+        let (mut signature_headers, _) = http_request::sign(request, &params).unwrap().into_parts();
+        client
+            .post(url)
+            .headers(signature_headers.take_headers().unwrap_or_default())
+            .body(body)
+            .send()?
+            .error_for_status()?;
+        Ok(())
+    }
+}

src/worker/cloudfront.rs

@@ -1,7 +1,7 @@
 use crate::background_jobs::Environment;
 use crate::schema;
 use anyhow::Context;
-use cargo_registry_index::Crate;
+use cargo_registry_index::{Crate, Repository};
 use chrono::Utc;
 use diesel::prelude::*;
 use std::fs::{self, OpenOptions};
@@ -44,6 +44,14 @@ pub fn update_crate_index(env: &Environment, crate_name: String) -> Result<(), P
     env.uploader
         .sync_index(env.http_client(), &crate_name, contents)?;
 
+    if let Some(cloudfront) = env.cloudfront() {
+        trace!(?crate_name, "invalidate CloudFront");
+        cloudfront.invalidate(
+            env.http_client(),
+            &Repository::relative_index_file_for_url(&crate_name),
+        )?;
+    }
+
     Ok(())
 }
 

src/worker/git.rs

@@ -3,6 +3,7 @@
 //! the daily database maintenance, but also operations like rendering READMEs
 //! and uploading them to S3.
 
+pub mod cloudfront;
 mod daily_db_maintenance;
 pub mod dump_db;
 mod git;