Add --content-base-url option to warg-server. (#136)

This PR adds a `--content-base-url` to `warg-server` that allows
specifying what the expected base URL to use when serving content.

By default, the base URL is derived from the listen address for the
server.

For servers that aren't bound to the address that will also be serving
the content, the `--content-base-url` option can be used to control the
URLs the server returns for fetching content.

crates/server/Cargo.toml

@@ -14,7 +14,7 @@ anyhow = { workspace = true }
 axum = { workspace = true }
 clap = { workspace = true }
 futures = { workspace = true }
-reqwest = { workspace = true }
+url = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true }
 tokio-util = { workspace = true }

crates/server/src/api/mod.rs

@@ -12,12 +12,13 @@ use tower_http::{
     LatencyUnit,
 };
 use tracing::{Level, Span};
+use url::Url;
 
 pub mod v1;
 
 /// Creates the router for the API.
 pub fn create_router(
-    base_url: String,
+    content_base_url: Url,
     core: Arc<CoreService>,
     temp_dir: PathBuf,
     files_dir: PathBuf,
@@ -28,7 +29,7 @@ pub fn create_router(
         .nest(
             "/v1",
             v1::create_router(
-                base_url,
+                content_base_url,
                 core,
                 temp_dir,
                 files_dir.clone(),

crates/server/src/api/v1/mod.rs

@@ -14,6 +14,7 @@ use axum::{
 };
 use serde::{Serialize, Serializer};
 use std::{path::PathBuf, sync::Arc};
+use url::Url;
 
 pub mod fetch;
 pub mod package;
@@ -89,7 +90,7 @@ pub async fn not_found() -> impl IntoResponse {
 }
 
 pub fn create_router(
-    base_url: String,
+    content_base_url: Url,
     core: Arc<CoreService>,
     temp_dir: PathBuf,
     files_dir: PathBuf,
@@ -99,7 +100,7 @@ pub fn create_router(
     let proof_config = proof::Config::new(core.log_data().clone(), core.map_data().clone());
     let package_config = package::Config::new(
         core.clone(),
-        base_url,
+        content_base_url,
         files_dir,
         temp_dir,
         content_policy,

crates/server/src/api/v1/package.rs

@@ -20,6 +20,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use tempfile::NamedTempFile;
 use tokio::io::AsyncWriteExt;
+use url::Url;
 use warg_api::v1::package::{
     ContentSource, PackageError, PackageRecord, PackageRecordState, PublishRecordRequest,
 };
@@ -33,7 +34,7 @@ use warg_protocol::{
 #[derive(Clone)]
 pub struct Config {
     core_service: Arc<CoreService>,
-    base_url: String,
+    content_base_url: Url,
     files_dir: PathBuf,
     temp_dir: PathBuf,
     content_policy: Option<Arc<dyn ContentPolicy>>,
@@ -43,15 +44,15 @@ pub struct Config {
 impl Config {
     pub fn new(
         core_service: Arc<CoreService>,
-        base_url: String,
+        content_base_url: Url,
         files_dir: PathBuf,
         temp_dir: PathBuf,
         content_policy: Option<Arc<dyn ContentPolicy>>,
         record_policy: Option<Arc<dyn RecordPolicy>>,
     ) -> Self {
         Self {
             core_service,
-            base_url,
+            content_base_url,
             files_dir,
             temp_dir,
             content_policy,
@@ -83,11 +84,12 @@ impl Config {
     }
 
     fn content_url(&self, digest: &AnyHash) -> String {
-        format!(
-            "{url}/content/{name}",
-            url = self.base_url,
-            name = self.content_file_name(digest)
-        )
+        self.content_base_url
+            .join("content/")
+            .unwrap()
+            .join(&self.content_file_name(digest))
+            .unwrap()
+            .to_string()
     }
 }
 

crates/server/src/bin/warg-server.rs

@@ -3,6 +3,7 @@ use clap::{Parser, ValueEnum};
 use std::{net::SocketAddr, path::PathBuf};
 use tokio::signal;
 use tracing_subscriber::filter::LevelFilter;
+use url::Url;
 use warg_crypto::signing::PrivateKey;
 use warg_server::{args::get_opt_content, Config, Server};
 
@@ -28,6 +29,10 @@ struct Args {
     #[arg(long, env = "WARG_CONTENT_DIR")]
     content_dir: PathBuf,
 
+    /// The base content URL to use; defaults to the server address.
+    #[arg(long, env = "WARG_CONTENT_BASE_URL")]
+    content_base_url: Option<Url>,
+
     /// The data store to use for the server.
     #[arg(long, env = "WARG_DATA_STORE", default_value = "memory")]
     data_store: DataStoreKind,
@@ -106,10 +111,14 @@ async fn main() -> Result<()> {
             .parse()
             .context("failed to parse operator key")?;
 
-    let config = Config::new(operator_key, args.content_dir)
+    let mut config = Config::new(operator_key, args.content_dir)
         .with_addr(args.listen)
         .with_shutdown(shutdown_signal());
 
+    if let Some(url) = args.content_base_url {
+        config = config.with_content_base_url(url);
+    }
+
     let config = match args.data_store {
         #[cfg(feature = "postgres")]
         DataStoreKind::Postgres => {

crates/server/src/lib.rs

@@ -12,6 +12,7 @@ use std::{
     sync::Arc,
     time::Duration,
 };
+use url::Url;
 use warg_crypto::signing::PrivateKey;
 
 pub mod api;
@@ -27,8 +28,9 @@ const DEFAULT_CHECKPOINT_INTERVAL: Duration = Duration::from_secs(5);
 pub struct Config {
     operator_key: PrivateKey,
     addr: Option<SocketAddr>,
-    data_store: Option<Box<dyn datastore::DataStore>>,
+    data_store: Option<Box<dyn DataStore>>,
     content_dir: PathBuf,
+    content_base_url: Option<Url>,
     shutdown: Option<Pin<Box<dyn Future<Output = ()> + Send + Sync>>>,
     checkpoint_interval: Option<Duration>,
     content_policy: Option<Arc<dyn ContentPolicy>>,
@@ -67,6 +69,7 @@ impl Config {
             addr: None,
             data_store: None,
             content_dir,
+            content_base_url: None,
             shutdown: None,
             checkpoint_interval: None,
             content_policy: None,
@@ -80,6 +83,14 @@ impl Config {
         self
     }
 
+    /// Specify the content base URL to use.
+    ///
+    /// If not set, the content base URL will be derived from the server address.
+    pub fn with_content_base_url(mut self, url: Url) -> Self {
+        self.content_base_url = Some(url);
+        self
+    }
+
     /// Specify the data store to use.
     ///
     /// If this is not specified, the server will use an in-memory data store.
@@ -206,10 +217,13 @@ impl Server {
             )
         })?;
 
-        let base_url = format!("http://{addr}", addr = self.config.addr.unwrap());
+        let content_base_url = self.config.content_base_url.unwrap_or_else(|| {
+            Url::parse(&format!("http://{addr}", addr = self.config.addr.unwrap())).unwrap()
+        });
+
         let server = axum::Server::from_tcp(listener)?.serve(
             create_router(
-                base_url,
+                content_base_url,
                 core,
                 temp_dir,
                 files_dir,

tests/client.rs

@@ -21,7 +21,7 @@ async fn client_incrementally_fetches() -> Result<()> {
     const RELEASE_COUNT: usize = 300;
     const PACKAGE_ID: &str = "test:package";
 
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
 
     let client = create_client(&config)?;
     let signing_key = support::test_signing_key().parse().unwrap();

tests/memory/mod.rs

@@ -7,13 +7,13 @@ use warg_crypto::signing::PrivateKey;
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn it_starts_with_initial_checkpoint() -> Result<()> {
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
     test_initial_checkpoint(&config).await
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn it_publishes_a_component() -> Result<()> {
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
     test_component_publishing(&config).await?;
 
     // There should be two log entries in the registry
@@ -30,7 +30,7 @@ async fn it_publishes_a_component() -> Result<()> {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn it_publishes_a_wit_package() -> Result<()> {
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
     test_wit_publishing(&config).await?;
 
     // There should be two log entries in the registry
@@ -47,7 +47,7 @@ async fn it_publishes_a_wit_package() -> Result<()> {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn it_rejects_non_wasm_content() -> Result<()> {
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
     test_wasm_content_policy(&config).await
 }
 
@@ -56,6 +56,7 @@ async fn it_rejects_unauthorized_signing_key() -> Result<()> {
     let (_server, config) = spawn_server(
         &root().await?,
         None,
+        None,
         Some(vec![(
             "test".to_string(),
             test_signing_key()
@@ -72,12 +73,24 @@ async fn it_rejects_unauthorized_signing_key() -> Result<()> {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn it_rejects_unknown_signing_key() -> Result<()> {
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
     test_unknown_signing_key(&config).await
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn it_rejects_invalid_signature() -> Result<()> {
-    let (_server, config) = spawn_server(&root().await?, None, None).await?;
+    let (_server, config) = spawn_server(&root().await?, None, None, None).await?;
     test_invalid_signature(&config).await
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn it_formats_custom_content_urls() -> Result<()> {
+    let (_server, config) = spawn_server(
+        &root().await?,
+        Some("https://example.com".parse().unwrap()),
+        None,
+        None,
+    )
+    .await?;
+    test_custom_content_url(&config).await
+}

tests/postgres/mod.rs

@@ -31,6 +31,7 @@ async fn it_works_with_postgres() -> TestResult {
     let root = root().await?;
     let (server, config) = spawn_server(
         &root,
+        None,
         Some(data_store()?),
         Some(vec![(
             "test".to_string(),
@@ -73,7 +74,7 @@ async fn it_works_with_postgres() -> TestResult {
     drop(server);
 
     // Restart the server and ensure the data is still there
-    let (_server, config) = spawn_server(&root, Some(data_store()?), None).await?;
+    let (server, config) = spawn_server(&root, None, Some(data_store()?), None).await?;
 
     test_unknown_signing_key(&config).await?;
 
@@ -103,5 +104,18 @@ async fn it_works_with_postgres() -> TestResult {
             .context("failed to resolve package")?;
     }
 
+    // Restart the server for the custom content URL test
+    drop(client);
+    drop(server);
+    let (_server, config) = spawn_server(
+        &root,
+        Some("https://example.com".parse().unwrap()),
+        Some(data_store()?),
+        None,
+    )
+    .await?;
+
+    test_custom_content_url(&config).await?;
+
     Ok(())
 }