From 136a3d43d60059f03b28235cf66f9919c2eac89d Mon Sep 17 00:00:00 2001
From: Luc Perkins <lucperkins@gmail.com>
Date: Thu, 16 May 2024 14:29:20 -0300
Subject: [PATCH] Periodically fetch JWT only in GHA

---
 magic-nix-cache/src/env.rs      | 12 ++++++++++-
 magic-nix-cache/src/flakehub.rs | 38 +++++++++++++++++++--------------
 magic-nix-cache/src/main.rs     | 19 +++++++++++++++--
 3 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/magic-nix-cache/src/env.rs b/magic-nix-cache/src/env.rs
index 2af3313..464221c 100644
--- a/magic-nix-cache/src/env.rs
+++ b/magic-nix-cache/src/env.rs
@@ -7,6 +7,16 @@ pub enum Environment {
     Other,
 }
 
+impl Environment {
+    pub fn is_github_actions(&self) -> bool {
+        matches!(self, Self::GitHubActions)
+    }
+
+    pub fn is_gitlab_ci(&self) -> bool {
+        matches!(self, Self::GitLabCI)
+    }
+}
+
 impl ToString for Environment {
     fn to_string(&self) -> String {
         use Environment::*;
@@ -14,7 +24,7 @@ impl ToString for Environment {
         String::from(match self {
             GitHubActions => "GitHub Actions",
             GitLabCI => "GitLab CI",
-            _ => "unspecified",
+            _ => "an unspecified environment",
         })
     }
 }
diff --git a/magic-nix-cache/src/flakehub.rs b/magic-nix-cache/src/flakehub.rs
index 1dbd79a..b3d3a6f 100644
--- a/magic-nix-cache/src/flakehub.rs
+++ b/magic-nix-cache/src/flakehub.rs
@@ -1,3 +1,4 @@
+use crate::env::Environment;
 use crate::error::{Error, Result};
 use attic::cache::CacheName;
 use attic::nix_store::{NixStore, StorePath};
@@ -27,6 +28,7 @@ pub struct State {
 }
 
 pub async fn init_cache(
+    environment: Environment,
     flakehub_api_server: &Url,
     flakehub_api_server_netrc: &Path,
     flakehub_cache_server: &Url,
@@ -102,22 +104,26 @@ pub async fn init_cache(
     let api_inner = ApiClient::from_server_config(server_config)?;
     let api = Arc::new(RwLock::new(api_inner));
 
-    // NOTE(cole-h): This is a workaround -- at the time of writing, GitHub Actions JWTs are only
-    // valid for 5 minutes after being issued. FlakeHub uses these JWTs for authentication, which
-    // means that after those 5 minutes have passed and the token is expired, FlakeHub (and by
-    // extension FlakeHub Cache) will no longer allow requests using this token. However, GitHub
-    // gives us a way to repeatedly request new tokens, so we utilize that and refresh the token
-    // every 2 minutes (less than half of the lifetime of the token).
-    let netrc_path_clone = flakehub_api_server_netrc.to_path_buf();
-    let initial_github_jwt_clone = flakehub_password.clone();
-    let flakehub_cache_server_clone = flakehub_cache_server.to_string();
-    let api_clone = api.clone();
-    tokio::task::spawn(refresh_github_actions_jwt_worker(
-        netrc_path_clone,
-        initial_github_jwt_clone,
-        flakehub_cache_server_clone,
-        api_clone,
-    ));
+    // Periodically refresh JWT in GitHub Actions environment
+    if environment.is_github_actions() {
+        // NOTE(cole-h): This is a workaround -- at the time of writing, GitHub Actions JWTs are only
+        // valid for 5 minutes after being issued. FlakeHub uses these JWTs for authentication, which
+        // means that after those 5 minutes have passed and the token is expired, FlakeHub (and by
+        // extension FlakeHub Cache) will no longer allow requests using this token. However, GitHub
+        // gives us a way to repeatedly request new tokens, so we utilize that and refresh the token
+        // every 2 minutes (less than half of the lifetime of the token).
+        let netrc_path_clone = flakehub_api_server_netrc.to_path_buf();
+        let initial_github_jwt_clone = flakehub_password.clone();
+        let flakehub_cache_server_clone = flakehub_cache_server.to_string();
+        let api_clone = api.clone();
+
+        tokio::task::spawn(refresh_github_actions_jwt_worker(
+            netrc_path_clone,
+            initial_github_jwt_clone,
+            flakehub_cache_server_clone,
+            api_clone,
+        ));
+    }
 
     // Get the cache UUID for this project.
     let cache_name = {
diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs
index 037cb68..cd90c12 100644
--- a/magic-nix-cache/src/main.rs
+++ b/magic-nix-cache/src/main.rs
@@ -32,9 +32,11 @@ use ::attic::nix_store::NixStore;
 use anyhow::{anyhow, Context, Result};
 use axum::{extract::Extension, routing::get, Router};
 use clap::Parser;
+use env::Environment;
 use tempfile::NamedTempFile;
 use tokio::process::Command;
 use tokio::sync::{oneshot, Mutex, RwLock};
+use tracing_subscriber::field::debug;
 use tracing_subscriber::filter::EnvFilter;
 
 use gha_cache::Credentials;
@@ -115,6 +117,18 @@ struct Args {
     startup_notification_url: Option<reqwest::Url>,
 }
 
+impl Args {
+    fn validate(&self, environment: Environment) -> Result<(), error::Error> {
+        if environment.is_gitlab_ci() && self.use_gha_cache {
+            return Err(Error::Config(String::from(
+                "the --use-gha-cache flag should not be applied in GitLab CI",
+            )));
+        }
+
+        Ok(())
+    }
+}
+
 /// The global server state.
 struct StateInner {
     /// State for uploading to the GHA cache.
@@ -143,8 +157,9 @@ async fn main_cli() -> Result<()> {
     init_logging();
 
     let args = Args::parse();
-
-    let _environment = determine_environment();
+    let environment = determine_environment();
+    tracing::debug!("Running in {environment}");
+    args.validate(environment)?;
 
     let metrics = Arc::new(telemetry::TelemetryReport::new());