refactor(iroh): Allow to register custom protocols (#2358)

## Description

* feat(iroh-net): Allow to set the list of accepted ALPN protocols at
runtime
* refactor(iroh): Spawning a node can now be performed in two stages:
First `Builder::build()` is called, which returns a Future that resolves
to a new type `ProtocolBuilder`. The `ProtocolBuilder` is then spawned
into the actual, running `Node`. If the intermediate step is not needed,
`Builder::spawn` performs both in one call, therefore this change is not
breaking.
* feat(iroh): Allow to accept custom ALPN protocols in an Iroh node.
Introduce a `ProtocolHandler` trait for accepting incoming connections
and adds `ProtocolBuilder::accept` to register these handlers per ALPN.
* refactor(iroh): Move towards more structured concurrency by spawning
tasks into a `JoinSet`
* refactor(iroh): Improve shutdown flow and perform more things
concurently.

originally based on #2357 but now directly to main

## Breaking Changes

* `iroh_net::endpoint::make_server_config` takes
`Arc<quinn::TransportConfig>` instead of
`Option<quinn::TransportConfig>`. If you used the `None` case, replace
with `quinn::TransportConfig::default()`.

## Notes & open questions

<!-- Any notes, remarks or open questions you have to make about the PR.
-->

## Change checklist

- [x] Self-review.
- [x] Documentation updates if relevant.
- [ ] Tests if relevant.
- [x] All breaking changes documented.

Author: Frando

iroh-blobs/src/store/fs.rs

@@ -1486,6 +1486,8 @@ impl Actor {
         let mut msgs = PeekableFlumeReceiver::new(self.state.msgs.clone());
         while let Some(msg) = msgs.recv() {
             if let ActorMessage::Shutdown { tx } = msg {
+                // Make sure the database is dropped before we send the reply.
+                drop(self);
                 if let Some(tx) = tx {
                     tx.send(()).ok();
                 }

iroh-blobs/src/store/traits.rs

@@ -295,7 +295,7 @@ pub trait ReadableStore: Map {
 }
 
 /// The mutable part of a Bao store.
-pub trait Store: ReadableStore + MapMut {
+pub trait Store: ReadableStore + MapMut + std::fmt::Debug {
     /// This trait method imports a file from a local path.
     ///
     /// `data` is the path to the file.

iroh-net/src/endpoint.rs

@@ -125,15 +125,12 @@ impl Builder {
             }
         };
         let secret_key = self.secret_key.unwrap_or_else(SecretKey::generate);
-        let mut server_config = make_server_config(
-            &secret_key,
-            self.alpn_protocols,
-            self.transport_config,
-            self.keylog,
-        )?;
-        if let Some(c) = self.concurrent_connections {
-            server_config.concurrent_connections(c);
-        }
+        let static_config = StaticConfig {
+            transport_config: Arc::new(self.transport_config.unwrap_or_default()),
+            keylog: self.keylog,
+            concurrent_connections: self.concurrent_connections,
+            secret_key: secret_key.clone(),
+        };
         let dns_resolver = self
             .dns_resolver
             .unwrap_or_else(|| default_resolver().clone());
@@ -149,7 +146,7 @@ impl Builder {
             #[cfg(any(test, feature = "test-utils"))]
             insecure_skip_relay_cert_verify: self.insecure_skip_relay_cert_verify,
         };
-        Endpoint::bind(Some(server_config), msock_opts, self.keylog).await
+        Endpoint::bind(static_config, msock_opts, self.alpn_protocols).await
     }
 
     // # The very common methods everyone basically needs.
@@ -296,17 +293,41 @@ impl Builder {
     }
 }
 
+/// Configuration for a [`quinn::Endpoint`] that cannot be changed at runtime.
+#[derive(Debug)]
+struct StaticConfig {
+    secret_key: SecretKey,
+    transport_config: Arc<quinn::TransportConfig>,
+    keylog: bool,
+    concurrent_connections: Option<u32>,
+}
+
+impl StaticConfig {
+    /// Create a [`quinn::ServerConfig`] with the specified ALPN protocols.
+    fn create_server_config(&self, alpn_protocols: Vec<Vec<u8>>) -> Result<quinn::ServerConfig> {
+        let mut server_config = make_server_config(
+            &self.secret_key,
+            alpn_protocols,
+            self.transport_config.clone(),
+            self.keylog,
+        )?;
+        if let Some(c) = self.concurrent_connections {
+            server_config.concurrent_connections(c);
+        }
+        Ok(server_config)
+    }
+}
+
 /// Creates a [`quinn::ServerConfig`] with the given secret key and limits.
 pub fn make_server_config(
     secret_key: &SecretKey,
     alpn_protocols: Vec<Vec<u8>>,
-    transport_config: Option<quinn::TransportConfig>,
+    transport_config: Arc<quinn::TransportConfig>,
     keylog: bool,
 ) -> Result<quinn::ServerConfig> {
     let tls_server_config = tls::make_server_config(secret_key, alpn_protocols, keylog)?;
     let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(tls_server_config));
-    server_config.transport_config(Arc::new(transport_config.unwrap_or_default()));
-
+    server_config.transport_config(transport_config);
     Ok(server_config)
 }
 
@@ -334,12 +355,11 @@ pub fn make_server_config(
 /// [QUIC]: https://quicwg.org
 #[derive(Clone, Debug)]
 pub struct Endpoint {
-    secret_key: Arc<SecretKey>,
     msock: Handle,
     endpoint: quinn::Endpoint,
     rtt_actor: Arc<rtt_actor::RttHandle>,
-    keylog: bool,
     cancel_token: CancellationToken,
+    static_config: Arc<StaticConfig>,
 }
 
 impl Endpoint {
@@ -359,16 +379,17 @@ impl Endpoint {
     /// This is for internal use, the public interface is the [`Builder`] obtained from
     /// [Self::builder]. See the methods on the builder for documentation of the parameters.
     async fn bind(
-        server_config: Option<quinn::ServerConfig>,
+        static_config: StaticConfig,
         msock_opts: magicsock::Options,
-        keylog: bool,
+        initial_alpns: Vec<Vec<u8>>,
     ) -> Result<Self> {
-        let secret_key = msock_opts.secret_key.clone();
-        let span = info_span!("magic_ep", me = %secret_key.public().fmt_short());
+        let span = info_span!("magic_ep", me = %static_config.secret_key.public().fmt_short());
         let _guard = span.enter();
         let msock = magicsock::MagicSock::spawn(msock_opts).await?;
         trace!("created magicsock");
 
+        let server_config = static_config.create_server_config(initial_alpns)?;
+
         let mut endpoint_config = quinn::EndpointConfig::default();
         // Setting this to false means that quinn will ignore packets that have the QUIC fixed bit
         // set to 0. The fixed bit is the 3rd bit of the first byte of a packet.
@@ -379,22 +400,31 @@ impl Endpoint {
 
         let endpoint = quinn::Endpoint::new_with_abstract_socket(
             endpoint_config,
-            server_config,
+            Some(server_config),
             msock.clone(),
             Arc::new(quinn::TokioRuntime),
         )?;
         trace!("created quinn endpoint");
 
         Ok(Self {
-            secret_key: Arc::new(secret_key),
             msock,
             endpoint,
             rtt_actor: Arc::new(rtt_actor::RttHandle::new()),
-            keylog,
             cancel_token: CancellationToken::new(),
+            static_config: Arc::new(static_config),
         })
     }
 
+    /// Set the list of accepted ALPN protocols.
+    ///
+    /// This will only affect new incoming connections.
+    /// Note that this *overrides* the current list of ALPNs.
+    pub fn set_alpns(&self, alpns: Vec<Vec<u8>>) -> Result<()> {
+        let server_config = self.static_config.create_server_config(alpns)?;
+        self.endpoint.set_server_config(Some(server_config));
+        Ok(())
+    }
+
     // # Methods for establishing connectivity.
 
     /// Connects to a remote [`Endpoint`].
@@ -480,10 +510,10 @@ impl Endpoint {
         let client_config = {
             let alpn_protocols = vec![alpn.to_vec()];
             let tls_client_config = tls::make_client_config(
-                &self.secret_key,
+                &self.static_config.secret_key,
                 Some(*node_id),
                 alpn_protocols,
-                self.keylog,
+                self.static_config.keylog,
             )?;
             let mut client_config = quinn::ClientConfig::new(Arc::new(tls_client_config));
             let mut transport_config = quinn::TransportConfig::default();
@@ -579,15 +609,15 @@ impl Endpoint {
 
     /// Returns the secret_key of this endpoint.
     pub fn secret_key(&self) -> &SecretKey {
-        &self.secret_key
+        &self.static_config.secret_key
     }
 
     /// Returns the node id of this endpoint.
     ///
     /// This ID is the unique addressing information of this node and other peers must know
     /// it to be able to connect to this node.
     pub fn node_id(&self) -> NodeId {
-        self.secret_key.public()
+        self.static_config.secret_key.public()
     }
 
     /// Returns the current [`NodeAddr`] for this endpoint.

iroh/Cargo.toml

@@ -102,3 +102,7 @@ required-features = ["examples"]
 [[example]]
 name = "client"
 required-features = ["examples"]
+
+[[example]]
+name = "custom-protocol"
+required-features = ["examples"]

iroh/examples/custom-protocol.rs

@@ -0,0 +1,127 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use clap::Parser;
+use futures_lite::future::Boxed as BoxedFuture;
+use iroh::{
+    client::MemIroh,
+    net::{
+        endpoint::{get_remote_node_id, Connecting},
+        Endpoint, NodeId,
+    },
+    node::ProtocolHandler,
+};
+use tracing_subscriber::{prelude::*, EnvFilter};
+
+#[derive(Debug, Parser)]
+pub struct Cli {
+    #[clap(subcommand)]
+    command: Command,
+}
+
+#[derive(Debug, Parser)]
+pub enum Command {
+    Accept,
+    Connect { node: NodeId },
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    setup_logging();
+    let args = Cli::parse();
+    // create a new node
+    let builder = iroh::node::Node::memory().build().await?;
+    let proto = ExampleProto::new(builder.client().clone(), builder.endpoint().clone());
+    let node = builder
+        .accept(EXAMPLE_ALPN, Arc::new(proto.clone()))
+        .spawn()
+        .await?;
+
+    // print the ticket if this is the accepting side
+    match args.command {
+        Command::Accept => {
+            let node_id = node.node_id();
+            println!("node id: {node_id}");
+            // wait until ctrl-c
+            tokio::signal::ctrl_c().await?;
+        }
+        Command::Connect { node: node_id } => {
+            proto.connect(node_id).await?;
+        }
+    }
+
+    node.shutdown().await?;
+
+    Ok(())
+}
+
+const EXAMPLE_ALPN: &[u8] = b"example-proto/0";
+
+#[derive(Debug, Clone)]
+struct ExampleProto {
+    client: MemIroh,
+    endpoint: Endpoint,
+}
+
+impl ProtocolHandler for ExampleProto {
+    fn accept(self: Arc<Self>, connecting: Connecting) -> BoxedFuture<Result<()>> {
+        Box::pin(async move {
+            let connection = connecting.await?;
+            let peer = get_remote_node_id(&connection)?;
+            println!("accepted connection from {peer}");
+            let mut send_stream = connection.open_uni().await?;
+            // Let's create a new blob for each incoming connection.
+            // This functions as an example of using existing iroh functionality within a protocol
+            // (you likely don't want to create a new blob for each connection for real)
+            let content = format!("this blob is created for my beloved peer {peer} ♥");
+            let hash = self
+                .client
+                .blobs()
+                .add_bytes(content.as_bytes().to_vec())
+                .await?;
+            // Send the hash over our custom protocol.
+            send_stream.write_all(hash.hash.as_bytes()).await?;
+            send_stream.finish().await?;
+            println!("closing connection from {peer}");
+            Ok(())
+        })
+    }
+}
+
+impl ExampleProto {
+    pub fn new(client: MemIroh, endpoint: Endpoint) -> Self {
+        Self { client, endpoint }
+    }
+
+    pub async fn connect(&self, remote_node_id: NodeId) -> Result<()> {
+        println!("our node id: {}", self.endpoint.node_id());
+        println!("connecting to {remote_node_id}");
+        let conn = self
+            .endpoint
+            .connect_by_node_id(&remote_node_id, EXAMPLE_ALPN)
+            .await?;
+        let mut recv_stream = conn.accept_uni().await?;
+        let hash_bytes = recv_stream.read_to_end(32).await?;
+        let hash = iroh::blobs::Hash::from_bytes(hash_bytes.try_into().unwrap());
+        println!("received hash: {hash}");
+        self.client
+            .blobs()
+            .download(hash, remote_node_id.into())
+            .await?
+            .await?;
+        println!("blob downloaded");
+        let content = self.client.blobs().read_to_bytes(hash).await?;
+        let message = String::from_utf8(content.to_vec())?;
+        println!("blob content: {message}");
+        Ok(())
+    }
+}
+
+/// Set the RUST_LOG env var to one of {debug,info,warn} to see logging.
+fn setup_logging() {
+    tracing_subscriber::registry()
+        .with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr))
+        .with(EnvFilter::from_default_env())
+        .try_init()
+        .ok();
+}