From 05a5e9cca8c53000852867c3f74d5923ad3ef1e8 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Sun, 8 Jan 2023 00:57:22 -0700
Subject: [PATCH] server: Set a custom header if the cache is public

This can be used as a signal to improve caching. Only done for
the Nix Binary Cache APIs for now.
---
 attic/src/api/binary_cache.rs  |  7 +++++++
 attic/src/api/mod.rs           |  1 +
 server/src/api/binary_cache.rs |  6 ++++++
 server/src/lib.rs              | 15 ++++++++++++++-
 server/src/middleware.rs       | 24 ++++++++++++++++++++++--
 5 files changed, 50 insertions(+), 3 deletions(-)
 create mode 100644 attic/src/api/binary_cache.rs

diff --git a/attic/src/api/binary_cache.rs b/attic/src/api/binary_cache.rs
new file mode 100644
index 0000000..02a0a8e
--- /dev/null
+++ b/attic/src/api/binary_cache.rs
@@ -0,0 +1,7 @@
+//! Nix Binary Cache server.
+//!
+//! This module contains Attic-specific extensions to the
+//! Nix Binary Cache API.
+
+/// Header indicating a cache's visibility.
+pub const ATTIC_CACHE_VISIBILITY: &str = "X-Attic-Cache-Visibility";
diff --git a/attic/src/api/mod.rs b/attic/src/api/mod.rs
index a3a6d96..ebb4a7d 100644
--- a/attic/src/api/mod.rs
+++ b/attic/src/api/mod.rs
@@ -1 +1,2 @@
+pub mod binary_cache;
 pub mod v1;
diff --git a/server/src/api/binary_cache.rs b/server/src/api/binary_cache.rs
index 02ed28b..ec4bd0a 100644
--- a/server/src/api/binary_cache.rs
+++ b/server/src/api/binary_cache.rs
@@ -85,6 +85,8 @@ async fn get_nix_cache_info(
         })
         .await?;
 
+    req_state.set_public_cache(cache.is_public);
+
     let info = NixCacheInfo {
         want_mass_query: true,
         store_dir: cache.store_dir.into(),
@@ -137,6 +139,8 @@ async fn get_store_path_info(
         .get_permission_for_cache(&cache_name, cache.is_public);
     permission.require_pull()?;
 
+    req_state.set_public_cache(cache.is_public);
+
     let mut narinfo = object.to_nar_info(&nar)?;
 
     if narinfo.signature().is_none() {
@@ -189,6 +193,8 @@ async fn get_nar(
         .get_permission_for_cache(&cache_name, cache.is_public);
     permission.require_pull()?;
 
+    req_state.set_public_cache(cache.is_public);
+
     database.bump_object_last_accessed(object.id).await?;
 
     let remote_file = nar.remote_file.0;
diff --git a/server/src/lib.rs b/server/src/lib.rs
index 546e3fb..2f7c1e6 100644
--- a/server/src/lib.rs
+++ b/server/src/lib.rs
@@ -26,6 +26,7 @@ pub mod oobe;
 mod storage;
 
 use std::net::SocketAddr;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -46,7 +47,7 @@ use attic::cache::CacheName;
 use config::{Config, StorageConfig};
 use database::migration::{Migrator, MigratorTrait};
 use error::{ErrorKind, ServerError, ServerResult};
-use middleware::{init_request_state, restrict_host};
+use middleware::{init_request_state, restrict_host, set_visibility_header};
 use storage::{LocalBackend, S3Backend, StorageBackend};
 
 type State = Arc<StateInner>;
@@ -79,6 +80,12 @@ struct RequestStateInner {
 
     /// Whether the client claims the connection is HTTPS or not.
     client_claims_https: bool,
+
+    /// Whether the cache the client's interacting with is public.
+    ///
+    /// This is purely informational and used to add the `X-Attic-Cache-Visibility`.
+    /// header in responses.
+    public_cache: AtomicBool,
 }
 
 impl StateInner {
@@ -167,6 +174,11 @@ impl RequestStateInner {
     fn substituter_endpoint(&self, cache: CacheName) -> ServerResult<String> {
         Ok(format!("{}{}", self.api_endpoint()?, cache.as_str()))
     }
+
+    /// Indicates whether the cache the client is interacting with is public.
+    fn set_public_cache(&self, public: bool) {
+        self.public_cache.store(public, Ordering::Relaxed);
+    }
 }
 
 /// The fallback route.
@@ -192,6 +204,7 @@ pub async fn run_api_server(cli_listen: Option<SocketAddr>, config: Config) -> R
         .fallback(fallback)
         // middlewares
         .layer(axum::middleware::from_fn(apply_auth))
+        .layer(axum::middleware::from_fn(set_visibility_header))
         .layer(axum::middleware::from_fn(init_request_state))
         .layer(axum::middleware::from_fn(restrict_host))
         .layer(Extension(state.clone()))
diff --git a/server/src/middleware.rs b/server/src/middleware.rs
index 3009882..ccaabeb 100644
--- a/server/src/middleware.rs
+++ b/server/src/middleware.rs
@@ -1,15 +1,17 @@
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use anyhow::anyhow;
 use axum::{
     extract::{Extension, Host},
-    http::Request,
+    http::{HeaderValue, Request},
     middleware::Next,
     response::Response,
 };
 
-use super::{AuthState, RequestStateInner, State};
+use super::{AuthState, RequestState, RequestStateInner, State};
 use crate::error::{ErrorKind, ServerResult};
+use attic::api::binary_cache::ATTIC_CACHE_VISIBILITY;
 
 /// Initializes per-request state.
 pub async fn init_request_state<B>(
@@ -31,6 +33,7 @@ pub async fn init_request_state<B>(
         api_endpoint: state.config.api_endpoint.to_owned(),
         host,
         client_claims_https,
+        public_cache: AtomicBool::new(false),
     });
 
     req.extensions_mut().insert(req_state);
@@ -55,3 +58,20 @@ pub async fn restrict_host<B>(
 
     Ok(next.run(req).await)
 }
+
+/// Sets the `X-Attic-Cache-Visibility` header in responses.
+pub(crate) async fn set_visibility_header<B>(
+    Extension(req_state): Extension<RequestState>,
+    req: Request<B>,
+    next: Next<B>,
+) -> ServerResult<Response> {
+    let mut response = next.run(req).await;
+
+    if req_state.public_cache.load(Ordering::Relaxed) {
+        response
+            .headers_mut()
+            .append(ATTIC_CACHE_VISIBILITY, HeaderValue::from_static("public"));
+    }
+
+    Ok(response)
+}