12.4 KiB · text History e832e77
//! Shared iroh node bootstrap for devendra.
//!
//! Both the controller and the agent need the same substrate: an iroh
//! [`Endpoint`], a persistent blob store, a gossip instance, the docs protocol,
//! and a [`Router`] wiring those three ALPNs together. [`Node`] builds and owns
//! all of it, and exposes just the handful of document operations devendra
//! uses (read/write a keyed entry, share/import a doc). Keeping this in one
//! place means the two daemons can never drift on how they stand up iroh.
//!
//! Persistence matters: the controller must reuse the same document namespace
//! (and therefore the same join ticket) across restarts, and the agent avoids
//! re-downloading content it already has. Both point [`Node::spawn`] at a data
//! directory backed by [`iroh_blobs::store::fs::FsStore`] and
//! [`iroh_docs::protocol::Docs::persistent`].

use std::path::Path;

use anyhow::{Context, Result, bail};
use bytes::Bytes;
use iroh::{Endpoint, SecretKey, endpoint::presets, protocol::Router};
use iroh_blobs::{ALPN as BLOBS_ALPN, BlobsProtocol, store::fs::FsStore};
use iroh_docs::{
    ALPN as DOCS_ALPN, AuthorId,
    api::{Doc, DocsApi},
    protocol::Docs,
    store::Query,
};
use iroh_gossip::{ALPN as GOSSIP_ALPN, net::Gossip};

// Re-export the doc-facing types callers need so they don't each depend on the
// exact iroh-docs module paths.
pub use futures_lite::StreamExt;
pub use iroh_docs::{
    DocTicket, NamespaceId, api::Doc as DocHandle, api::protocol::AddrInfoOptions,
    api::protocol::ShareMode, engine::LiveEvent,
};

/// A running iroh node: endpoint, blob store, gossip, docs, and the router that
/// serves all three protocols. Dropping the `Node` drops the [`Router`], which
/// shuts the accept loop down.
pub struct Node {
    endpoint: Endpoint,
    blobs: FsStore,
    #[allow(dead_code)]
    gossip: Gossip,
    docs: Docs,
    #[allow(dead_code)]
    router: Router,
    author: AuthorId,
}

impl Node {
    /// Stand up a persistent node rooted at `data_dir`. Creates the directory
    /// (and `blobs/`, `docs/` subdirectories) if missing.
    ///
    /// `bind_port` fixes the UDP port the iroh endpoint listens on. Pass `Some`
    /// for a long-lived server (the controller) so its direct address is stable
    /// across restarts — otherwise every restart picks a new port, invalidating
    /// the addresses embedded in already-issued join tickets. Pass `None` for
    /// short-lived / client nodes (the agent), which only dial out.
    pub async fn spawn(data_dir: &Path, bind_port: Option<u16>) -> Result<Self> {
        // Create the store subdirectories up front: the blob store and the
        // docs redb file are created *inside* these paths, so the directories
        // must exist first.
        let blobs_dir = data_dir.join("blobs");
        let docs_dir = data_dir.join("docs");
        for dir in [data_dir, &blobs_dir, &docs_dir] {
            tokio::fs::create_dir_all(dir)
                .await
                .with_context(|| format!("creating {}", dir.display()))?;
        }

        // Load (or create + persist) the endpoint secret key so this node keeps
        // a stable node id across restarts. Without this, every run mints a new
        // identity and every previously-issued ticket points at a dead node.
        let secret_key = load_or_create_secret_key(data_dir).await?;
        let mut builder = Endpoint::builder(presets::N0).secret_key(secret_key);
        if let Some(port) = bind_port {
            let addr = std::net::SocketAddr::from((std::net::Ipv4Addr::UNSPECIFIED, port));
            builder = builder
                .bind_addr(addr)
                .map_err(|e| anyhow::anyhow!("invalid bind address 0.0.0.0:{port}: {e}"))?;
        }
        let endpoint = builder.bind().await.context("binding iroh endpoint")?;

        let blobs = FsStore::load(&blobs_dir)
            .await
            .context("opening blob store")?;

        let gossip = Gossip::builder().spawn(endpoint.clone());

        let docs = Docs::persistent(docs_dir)
            .spawn(endpoint.clone(), (*blobs).clone(), gossip.clone())
            .await
            .context("spawning docs protocol")?;

        let router = Router::builder(endpoint.clone())
            .accept(BLOBS_ALPN, BlobsProtocol::new(&blobs, None))
            .accept(GOSSIP_ALPN, gossip.clone())
            .accept(DOCS_ALPN, docs.clone())
            .spawn();

        // A stable default author, persisted in the docs store, so every entry
        // this node writes is attributed consistently across restarts.
        let author = docs
            .api()
            .author_default()
            .await
            .context("resolving default author")?;

        Ok(Self {
            endpoint,
            blobs,
            gossip,
            docs,
            router,
            author,
        })
    }

    /// The node's public key, printable for operators.
    pub fn node_id(&self) -> iroh::EndpointId {
        self.endpoint.id()
    }

    /// Flush stores and tear down the network cleanly. Call before a process
    /// exits: the blob store buffers content, and an abrupt exit can leave a
    /// document entry pointing at a blob whose bytes never hit disk. Essential
    /// for the one-shot CLI commands, which otherwise exit mid-write.
    pub async fn shutdown(&self) -> Result<()> {
        // Flush the blob store's database to disk first — this is the operation
        // that guarantees written content survives the exit.
        self.blobs.sync_db().await.context("syncing blob store")?;
        // Stop accepting new connections.
        if let Err(e) = self.router.shutdown().await {
            tracing::warn!("router shutdown: {e}");
        }
        // Ask the store actor to stop. It may drop its ack as it exits, so this
        // is best-effort — the sync above already persisted our writes.
        if let Err(e) = self.blobs.shutdown().await {
            tracing::debug!("blob store shutdown ack: {e}");
        }
        self.endpoint.close().await;
        Ok(())
    }

    /// The docs API handle.
    pub fn docs(&self) -> &DocsApi {
        self.docs.api()
    }

    /// The author id this node writes as.
    pub fn author(&self) -> AuthorId {
        self.author
    }

    /// Create a brand-new document (namespace). The caller owns write
    /// capability. Used by the controller on first boot.
    pub async fn create_doc(&self) -> Result<Doc> {
        self.docs().create().await.context("creating doc")
    }

    /// Open a previously-created document by namespace id, if this node's store
    /// knows it. Returns `None` if not present.
    pub async fn open_doc(&self, id: NamespaceId) -> Result<Option<Doc>> {
        self.docs().open(id).await.context("opening doc")
    }

    /// Import a document from a join ticket and begin syncing with its peers,
    /// returning the doc plus a live event stream. Used by the agent.
    pub async fn import_doc(
        &self,
        ticket: DocTicket,
    ) -> Result<(Doc, impl futures_lite::Stream<Item = Result<LiveEvent>>)> {
        self.docs()
            .import_and_subscribe(ticket)
            .await
            .context("importing doc from ticket")
    }

    /// Re-initiate sync with the ticket's peers. Cheap and idempotent — used as
    /// a periodic nudge to recover from a dropped connection to the controller.
    pub async fn resync(&self, doc: &Doc, ticket: &DocTicket) -> Result<()> {
        doc.start_sync(ticket.nodes.clone())
            .await
            .context("re-syncing with controller")?;
        Ok(())
    }

    /// Produce a shareable join ticket for `doc` at the given capability level.
    /// Agents get [`ShareMode::Write`] so they can post status back; a
    /// read-only observer would get [`ShareMode::Read`].
    pub async fn share(&self, doc: &Doc, mode: ShareMode) -> Result<DocTicket> {
        doc.share(mode, AddrInfoOptions::RelayAndAddresses)
            .await
            .context("sharing doc")
    }

    /// Read the latest value for `key`, returning its bytes, or `None` if the
    /// key has no live entry. Content is pulled from the local blob store,
    /// which the sync layer keeps populated.
    pub async fn read(&self, doc: &Doc, key: &str) -> Result<Option<Vec<u8>>> {
        let entry = doc
            .get_one(Query::single_latest_per_key().key_exact(key).build())
            .await
            .with_context(|| format!("querying key {key}"))?;
        let Some(entry) = entry else {
            return Ok(None);
        };
        // A tombstone (deleted entry) has an empty hash / zero length.
        if entry.content_len() == 0 {
            return Ok(Some(Vec::new()));
        }
        // The entry can sync before its content blob finishes downloading.
        // Treat "not local yet" the same as "absent" so callers defer cleanly
        // rather than erroring; a later ContentReady event retriggers the read.
        let hash = entry.content_hash();
        if !self
            .blobs
            .blobs()
            .has(hash)
            .await
            .context("checking blob presence")?
        {
            return Ok(None);
        }
        let bytes = self
            .blobs
            .blobs()
            .get_bytes(hash)
            .await
            .with_context(|| format!("reading content for key {key}"))?;
        Ok(Some(bytes.to_vec()))
    }

    /// Write `value` under `key`, storing the bytes in the blob store and
    /// recording the entry. Overwrites any previous value for the key.
    pub async fn write(&self, doc: &Doc, key: &str, value: impl Into<Bytes>) -> Result<()> {
        doc.set_bytes(self.author, key.to_string().into_bytes(), value.into())
            .await
            .with_context(|| format!("writing key {key}"))?;
        Ok(())
    }

    /// List the latest live keys under `prefix` (UTF-8), skipping tombstones.
    /// Used to enumerate hosts with pending adoption proposals.
    pub async fn list_keys(&self, doc: &Doc, prefix: &str) -> Result<Vec<String>> {
        let stream = doc
            .get_many(Query::single_latest_per_key().key_prefix(prefix).build())
            .await
            .with_context(|| format!("listing keys under {prefix}"))?;
        let mut stream = std::pin::pin!(stream);
        let mut keys = Vec::new();
        while let Some(entry) = stream.next().await {
            let entry = entry?;
            // Skip tombstones (deleted entries have zero-length content).
            if entry.content_len() == 0 {
                continue;
            }
            if let Ok(k) = std::str::from_utf8(entry.key()) {
                keys.push(k.to_string());
            }
        }
        Ok(keys)
    }

    /// Delete every entry under `prefix` (writes tombstones). Used to clear an
    /// adoption proposal once it's been accepted or dismissed.
    pub async fn delete_prefix(&self, doc: &Doc, prefix: &str) -> Result<()> {
        doc.del(self.author, prefix.to_string().into_bytes())
            .await
            .with_context(|| format!("deleting prefix {prefix}"))?;
        Ok(())
    }
}

/// Read the persisted endpoint secret key from `<data_dir>/secret.key`, or
/// generate a new one and persist it (mode 0600). The 32 raw key bytes are
/// stored verbatim.
async fn load_or_create_secret_key(data_dir: &Path) -> Result<SecretKey> {
    let path = data_dir.join("secret.key");
    match tokio::fs::read(&path).await {
        Ok(bytes) => {
            let arr: [u8; 32] = bytes.as_slice().try_into().map_err(|_| {
                anyhow::anyhow!(
                    "secret key file {} is {} bytes, expected 32",
                    path.display(),
                    bytes.len()
                )
            })?;
            Ok(SecretKey::from_bytes(&arr))
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
            let secret_key = SecretKey::generate();
            tokio::fs::write(&path, secret_key.to_bytes())
                .await
                .with_context(|| format!("writing {}", path.display()))?;
            // Tighten permissions: this is a private key.
            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                tokio::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))
                    .await
                    .with_context(|| format!("chmod 600 {}", path.display()))?;
            }
            tracing::info!(path = %path.display(), "generated new node secret key");
            Ok(secret_key)
        }
        Err(e) => bail!("reading {}: {e}", path.display()),
    }
}