1 files changed, 83 insertions, 26 deletions

diff --git a/src/main.rs b/src/main.rs
index 44f9dd4..96758ec 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,27 +15,22 @@ use anyhow::{Context, Result};
 use automerge::AutoCommit;
 use clap::Parser;
 use tokio::sync::{broadcast, mpsc, watch, RwLock};
-use tracing::{debug, error, info, trace};
+use tracing::{debug, error, info, trace, warn};
 use url::Url;
 
 mod api;
-mod auth;
-mod bootstrap;
-mod compaction;
-mod firewall;
-mod mycelium;
-mod peer;
-mod replica;
+mod document;
+mod remake;
 mod state;
-mod store;
-mod sync;
+mod swarm;
 mod transport;
 
-use auth::AuthState;
-use firewall::{Firewall, FirewallTransport};
-use mycelium::{MyceliumSetup, MyceliumSpawnConfig};
-use replica::{Epoch, Replica};
+use api::auth::AuthState;
+use document::replica::{Epoch, Replica};
 use state::{AppState, NodeLifecycle, Timing};
+use swarm::{bootstrap, peer, sync};
+use transport::firewall::{Firewall, FirewallTransport};
+use transport::mycelium::{self, MyceliumSetup, MyceliumSpawnConfig};
 use transport::Transport;
 
 #[derive(Parser, Debug)]
@@ -80,6 +75,11 @@ struct Args {
     #[arg(long, value_name = "PATH")]
     htpasswd: String,
 
+    /// Pre-shared swarm key: only peers that know it may join (raw bytes, at least 16;
+    /// e.g. `head -c 32 /dev/urandom`). Without it, anyone who knows the swarm ID can.
+    #[arg(long, value_name = "PATH")]
+    swarm_key_file: Option<PathBuf>,
+
     /// Keepalive interval; dormancy, resync and probe timings derive from it.
     #[arg(long, value_name = "SECONDS", default_value_t = 60)]
     keepalive_secs: u64,
@@ -96,6 +96,19 @@ async fn main() -> Result<()> {
 
     let args = Args::parse();
 
+    // Read before anything starts, so a bad key file fails right away.
+    let swarm_key = match &args.swarm_key_file {
+        Some(path) => Some(Arc::new(swarm::auth::SwarmKey::from_file(path)?)),
+        None => {
+            warn!("No swarm key: any peer that learns the swarm ID can join");
+            None
+        }
+    };
+
+    // First thing: until a handler is installed, SIGINT and SIGTERM kill the process on
+    // the spot, and startup (spawning Mycelium, bootstrap) can take minutes.
+    let stop = watch_for_shutdown().context("Installing signal handlers")?;
+
     let mycelium_setup = {
         let mut spawn_config = MyceliumSpawnConfig {
             rest_url_opt: args.mycelium_rest_url,
@@ -107,7 +120,10 @@ async fn main() -> Result<()> {
         if let Some(net_name) = args.mycelium_network_name {
             spawn_config.network_name = net_name;
         }
-        MyceliumSetup::bootstrap(spawn_config).await?
+        tokio::select! {
+            setup = MyceliumSetup::bootstrap(spawn_config) => setup?,
+            _ = stopped(stop.clone()) => return Ok(()),
+        }
     };
 
     let mycelium_rest_url = mycelium_setup.rest_url.clone();
@@ -147,7 +163,7 @@ async fn main() -> Result<()> {
         doc: Arc::new(RwLock::new(Replica::new(initial_epoch, AutoCommit::new()))),
         epoch_tx,
         lifecycle: Arc::new(RwLock::new(NodeLifecycle::Joining {
-            discovered_cluster_id: None,
+            discovered_swarm_id: None,
         })),
         peer_states: Arc::new(RwLock::new(HashMap::new())),
         transport: Arc::new(FirewallTransport::new(
@@ -160,7 +176,9 @@ async fn main() -> Result<()> {
         initial_sync_done: Arc::new(AtomicBool::new(false)),
         change_events,
         timing: Timing::from_keepalive(Duration::from_secs(args.keepalive_secs)),
-        compaction: Arc::new(compaction::Compaction::default()),
+        remake: Arc::new(remake::RemakeState::default()),
+        swarm_key,
+        stale: Default::default(),
     };
 
     trace!("Application state initialized");
@@ -191,14 +209,18 @@ async fn main() -> Result<()> {
     trace!("Peer worker task spawned");
 
     // Switches to committed epochs and catches up with newer ones. Starts before
-    // bootstrapping: a node joining a compacted cluster can only finish its initial sync
+    // bootstrapping: a node joining a compacted swarm can only finish its initial sync
     // after fetching the current epoch from a peer.
-    tokio::spawn(compaction::run_worker(state.clone()));
-    trace!("Compaction worker task spawned");
+    tokio::spawn(remake::run_worker(state.clone()));
+    trace!("Remake worker task spawned");
 
-    bootstrap::bootstrap(&state, &seeds)
-        .await
-        .context("Joining cluster")?;
+    tokio::select! {
+        joined = bootstrap::bootstrap(&state, &seeds) => joined.context("Joining swarm")?,
+        _ = stopped(stop.clone()) => {
+            info!("Stopped while joining the swarm");
+            return Ok(());
+        }
+    }
 
     // Start the change publisher after bootstrapping: streams only exist from then on,
     // and it avoids diffing the initial sync.
@@ -206,7 +228,7 @@ async fn main() -> Result<()> {
     tokio::spawn(async move { sync::run_change_publisher(publisher_state).await });
 
     info!(
-        cluster_id = ?state.cluster_id().await,
+        swarm_id = ?state.swarm_id().await,
         ready = state.is_ready().await,
         "Bootstrap completed"
     );
@@ -246,8 +268,7 @@ async fn main() -> Result<()> {
             std::mem::drop(state);
         }
 
-        _ = tokio::signal::ctrl_c() => {
-            info!("Received manual interrupt signal. Shutting down...");
+        _ = stopped(stop.clone()) => {
             let _ = shutdown_tx.send(());
             let _ = server_fut.as_mut().await;
         }
@@ -257,5 +278,41 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+/// Turns SIGINT and SIGTERM into a shutdown request. A second signal exits at once (the
+/// supervised Mycelium still gets its parent-death signal then).
+fn watch_for_shutdown() -> Result<watch::Receiver<bool>> {
+    use tokio::signal::unix::{signal, SignalKind};
+
+    let mut interrupt = signal(SignalKind::interrupt())?;
+    let mut terminate = signal(SignalKind::terminate())?;
+    let (stop_tx, stop_rx) = watch::channel(false);
+
+    tokio::spawn(async move {
+        let name = tokio::select! {
+            _ = interrupt.recv() => "SIGINT",
+            _ = terminate.recv() => "SIGTERM",
+        };
+        info!(
+            signal = name,
+            "Shutting down (send it again to exit immediately)"
+        );
+        let _ = stop_tx.send(true);
+
+        tokio::select! {
+            _ = interrupt.recv() => {}
+            _ = terminate.recv() => {}
+        }
+        error!("Second signal; exiting immediately");
+        std::process::exit(130);
+    });
+
+    Ok(stop_rx)
+}
+
+/// Resolves once a shutdown was requested.
+async fn stopped(mut stop: watch::Receiver<bool>) {
+    let _ = stop.wait_for(|stop| *stop).await;
+}
+
 // === Footnotes ===============================================================
 // [concurrent]: these comments are in place so we remember what to do if we go from async to concurrent.