Add list-upstream-caches command

Author: pascalj

flake.nix

@@ -130,6 +130,7 @@
 
             snarf.succeed('export SNARF_CLIENT_TOKEN=$(snarf -s snarfd:9000 create-token); \
               snarf -s snarfd:9000 add-upstream-cache --token $SNARF_CLIENT_TOKEN "https://cache.nixos.org"; \
+              snarf -s snarfd:9000 list-upstream-caches --token $SNARF_CLIENT_TOKEN | grep "https://cache.nixos.org"; \
               snarf -s snarfd:9000 add-closure --token $SNARF_CLIENT_TOKEN $(realpath $(which snarf))')
           '';
         };

proto/management.proto

@@ -14,7 +14,8 @@ service ManagementService {
   rpc CreateClientToken(NewClientTokenRequest) returns (ClientToken) {}
   rpc FilterHashes(stream NARHashRequest) returns (stream NARHashResponse);
   rpc AddUpstreamCache(AddUpstreamCacheRequest) returns (AddUpstreamCacheResponse);
-}
+  rpc ListUpstreamCaches(ListUpstreamCachesRequest) returns (ListUpstreamCachesResponse);
+ }
 
 message NARHashRequest {
   bytes digest = 1;
@@ -31,3 +32,9 @@ message AddUpstreamCacheRequest {
 message AddUpstreamCacheResponse {
   bool success = 1;
 }
+
+message ListUpstreamCachesRequest {}
+
+message ListUpstreamCachesResponse {
+  repeated string base_urls = 1;
+}

src/bin/client.rs

@@ -32,6 +32,12 @@ enum ClientCommand {
         /// The base_url of the upstream cache
         base_url: String,
     },
+    /// List all configured upstream caches
+    ListUpstreamCaches {
+        /// The authentication token
+        #[arg(short, long, env = "SNARF_CLIENT_TOKEN", required = true)]
+        token: String,
+    },
 }
 
 /// CLI arguments for the client
@@ -61,11 +67,14 @@ async fn main() -> anyhow::Result<(), Box<dyn std::error::Error + Send + Sync>>
 
     let client_cli = ClientCli::parse();
 
-    Ok(match &client_cli.command {
+    match &client_cli.command {
         ClientCommand::AddClosure { .. } => add_closure(&client_cli).await?,
         ClientCommand::CreateToken => create_token(&client_cli).await?,
         ClientCommand::AddUpstreamCache { .. } => add_upstream_cache(&client_cli).await?,
-    })
+        ClientCommand::ListUpstreamCaches { .. } => list_upstream_caches(&client_cli).await?,
+    };
+
+    Ok(())
 }
 
 /// Add the closure of a path in the store to the cache. The path can be of arbitrary depth, in any case
@@ -243,6 +252,23 @@ async fn add_upstream_cache(client_cli: &ClientCli) -> anyhow::Result<()> {
     Ok(())
 }
 
+async fn list_upstream_caches(client_cli: &ClientCli) -> anyhow::Result<()> {
+    let ClientCommand::ListUpstreamCaches { token } = &client_cli.command else {
+        bail!("Upstream cache called with the wrong command");
+    };
+
+    let response = get_client(&client_cli.server_address, Some(token))
+        .await?
+        .list_upstream_caches(tonic::Request::new(ListUpstreamCachesRequest {}))
+        .await?;
+
+    for base_url in response.into_inner().base_urls {
+        println!("{}", base_url);
+    }
+
+    Ok(())
+}
+
 /// Get a client to connect to the server at server_address.
 ///
 /// An optional token can be passed to authenticate the requests.

src/cache.rs

@@ -58,6 +58,10 @@ impl NARCache {
         debug!("Checking narinfo {} for existence", url);
         Ok(client.head(&url).send().await?.status().is_success())
     }
+
+    pub fn base_url(&self) -> &str {
+        &self.base_url
+    }
 }
 
 pub mod persistence {

src/server/services.rs

@@ -283,4 +283,18 @@ impl management_service_server::ManagementService for ManagementServiceWrapper {
             success: true,
         }))
     }
+
+    async fn list_upstream_caches(
+        &self,
+        _: tonic::Request<ListUpstreamCachesRequest>,
+    ) -> anyhow::Result<tonic::Response<ListUpstreamCachesResponse>, tonic::Status> {
+        Ok(tonic::Response::new(ListUpstreamCachesResponse {
+            base_urls: self
+                .upstream_caches
+                .caches()
+                .iter()
+                .map(|cache| cache.base_url().to_owned())
+                .collect(),
+        }))
+    }
 }