feat(iroh-relay)!: implement authentication (#3086)

## Description

Design RFC:
https://www.notion.so/number-zero/Relay-Authentication-16f5df1306fb80ac9e31c1ccb04e026b?pvs=4

## Breaking Changes

- added: field `access` to `iroh_relay::server::RelayConfig`

## Notes & open questions

- This reuses the `health` frame in the relay protocol to indicate the
authentcation issue. The frame was previously never sent in our code,
but is now more explicitly interpreted to disconnect from the server.
- ~~Should the callback be `async`? If this does more complex things
like access a DB we probably want this to be `async`.~~ it is now async

## Change checklist

- [ ] Self-review.
- [ ] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [ ] Tests if relevant.
- [ ] All breaking changes documented.

Author: dignifiedquire

iroh-relay/src/main.rs

@@ -11,6 +11,8 @@ use std::{
 
 use anyhow::{bail, Context as _, Result};
 use clap::Parser;
+use futures_lite::FutureExt;
+use iroh_base::NodeId;
 use iroh_relay::{
     defaults::{
         DEFAULT_HTTPS_PORT, DEFAULT_HTTP_PORT, DEFAULT_METRICS_PORT, DEFAULT_RELAY_QUIC_PORT,
@@ -170,6 +172,59 @@ struct Config {
     metrics_bind_addr: Option<SocketAddr>,
     /// The capacity of the key cache.
     key_cache_capacity: Option<usize>,
+    /// Access control for relaying connections.
+    ///
+    /// This controls which nodes are allowed to relay connections, other endpoints, like STUN are not controlled by this.
+    #[serde(default)]
+    access: AccessConfig,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
+#[serde(rename_all = "lowercase")]
+enum AccessConfig {
+    /// Allows everyone
+    #[default]
+    Everyone,
+    /// Allows only these nodes.
+    Allowlist(Vec<NodeId>),
+    /// Allows everyone, except these nodes.
+    Denylist(Vec<NodeId>),
+}
+
+impl From<AccessConfig> for iroh_relay::server::AccessConfig {
+    fn from(cfg: AccessConfig) -> Self {
+        match cfg {
+            AccessConfig::Everyone => iroh_relay::server::AccessConfig::Everyone,
+            AccessConfig::Allowlist(allow_list) => {
+                let allow_list = Arc::new(allow_list);
+                iroh_relay::server::AccessConfig::Restricted(Box::new(move |node_id| {
+                    let allow_list = allow_list.clone();
+                    async move {
+                        if allow_list.contains(&node_id) {
+                            iroh_relay::server::Access::Allow
+                        } else {
+                            iroh_relay::server::Access::Deny
+                        }
+                    }
+                    .boxed()
+                }))
+            }
+            AccessConfig::Denylist(deny_list) => {
+                let deny_list = Arc::new(deny_list);
+                iroh_relay::server::AccessConfig::Restricted(Box::new(move |node_id| {
+                    let deny_list = deny_list.clone();
+                    async move {
+                        if deny_list.contains(&node_id) {
+                            iroh_relay::server::Access::Deny
+                        } else {
+                            iroh_relay::server::Access::Allow
+                        }
+                    }
+                    .boxed()
+                }))
+            }
+        }
+    }
 }
 
 impl Config {
@@ -202,6 +257,7 @@ impl Default for Config {
             enable_metrics: cfg_defaults::enable_metrics(),
             metrics_bind_addr: None,
             key_cache_capacity: Default::default(),
+            access: AccessConfig::Everyone,
         }
     }
 }
@@ -574,7 +630,9 @@ async fn build_relay_config(cfg: Config) -> Result<relay::ServerConfig<std::io::
         tls: relay_tls.and_then(|tls| if dangerous_http_only { None } else { Some(tls) }),
         limits,
         key_cache_capacity: cfg.key_cache_capacity,
+        access: cfg.access.clone().into(),
     };
+
     let stun_config = relay::StunConfig {
         bind_addr: cfg.stun_bind_addr(),
     };
@@ -639,6 +697,9 @@ mod metrics {
 mod tests {
     use std::num::NonZeroU32;
 
+    use iroh_base::SecretKey;
+    use rand::SeedableRng;
+    use rand_chacha::ChaCha8Rng;
     use testresult::TestResult;
 
     use super::*;
@@ -676,4 +737,28 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_access_config() -> TestResult {
+        let config = "
+            access = \"everyone\"
+        ";
+        let config = Config::from_str(config)?;
+        assert_eq!(config.access, AccessConfig::Everyone);
+
+        let mut rng = ChaCha8Rng::seed_from_u64(0);
+        let node_id = SecretKey::generate(&mut rng).public();
+
+        let config = format!(
+            "
+            access.allowlist = [
+              \"{node_id}\",
+            ]
+        "
+        );
+        let config = Config::from_str(dbg!(&config))?;
+        assert_eq!(config.access, AccessConfig::Allowlist(vec![node_id]));
+
+        Ok(())
+    }
 }

iroh-relay/src/protos/relay.rs

@@ -96,10 +96,9 @@ pub(crate) enum FrameType {
     /// 8 byte payload, the contents of ping being replied to
     Pong = 13,
     /// Sent from server to client to tell the client if their connection is
-    /// unhealthy somehow. Currently the only unhealthy state is whether the
-    /// connection is detected as a duplicate.
-    /// The entire frame body is the text of the error message. An empty message
-    /// clears the error state.
+    /// unhealthy somehow.
+    ///
+    /// Currently this is used to indicate that the connection was closed because of authentication issues.
     Health = 14,
 
     /// Sent from server to client for the server to declare that it's restarting.

iroh-relay/src/server.rs

@@ -20,11 +20,12 @@ use std::{fmt, future::Future, net::SocketAddr, num::NonZeroU32, pin::Pin, sync:
 
 use anyhow::{anyhow, bail, Context, Result};
 use derive_more::Debug;
-use futures_lite::StreamExt;
+use futures_lite::{future::Boxed, StreamExt};
 use http::{
     response::Builder as ResponseBuilder, HeaderMap, Method, Request, Response, StatusCode,
 };
 use hyper::body::Incoming;
+use iroh_base::NodeId;
 #[cfg(feature = "test-utils")]
 use iroh_base::RelayUrl;
 use iroh_metrics::inc;
@@ -120,6 +121,40 @@ pub struct RelayConfig<EC: fmt::Debug, EA: fmt::Debug = EC> {
     pub limits: Limits,
     /// Key cache capacity.
     pub key_cache_capacity: Option<usize>,
+    /// Access configuration.
+    pub access: AccessConfig,
+}
+
+/// Controls which nodes are allowed to use the relay.
+#[derive(derive_more::Debug)]
+pub enum AccessConfig {
+    /// Everyone
+    Everyone,
+    /// Only nodes for which the function returns `Access::Allow`.
+    #[debug("restricted")]
+    Restricted(Box<dyn Fn(NodeId) -> Boxed<Access> + Send + Sync + 'static>),
+}
+
+impl AccessConfig {
+    /// Is this node allowed?
+    pub async fn is_allowed(&self, node: NodeId) -> bool {
+        match self {
+            Self::Everyone => true,
+            Self::Restricted(check) => {
+                let res = check(node).await;
+                matches!(res, Access::Allow)
+            }
+        }
+    }
+}
+
+/// Access restriction for a node.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum Access {
+    /// Access is allowed.
+    Allow,
+    /// Access is denied.
+    Deny,
 }
 
 /// Configuration for the STUN server.
@@ -318,6 +353,7 @@ impl Server {
                 let mut builder = http_server::ServerBuilder::new(relay_bind_addr)
                     .headers(headers)
                     .key_cache_capacity(key_cache_capacity)
+                    .access(relay_config.access)
                     .request_handler(Method::GET, "/", Box::new(root_handler))
                     .request_handler(Method::GET, "/index.html", Box::new(root_handler))
                     .request_handler(Method::GET, RELAY_PROBE_PATH, Box::new(probe_handler))
@@ -772,6 +808,7 @@ mod tests {
     use std::{net::Ipv4Addr, time::Duration};
 
     use bytes::Bytes;
+    use futures_lite::FutureExt;
     use futures_util::SinkExt;
     use http::header::UPGRADE;
     use iroh_base::{NodeId, SecretKey};
@@ -790,6 +827,7 @@ mod tests {
                 tls: None,
                 limits: Default::default(),
                 key_cache_capacity: Some(1024),
+                access: AccessConfig::Everyone,
             }),
             quic: None,
             stun: None,
@@ -840,6 +878,7 @@ mod tests {
                 tls: None,
                 limits: Default::default(),
                 key_cache_capacity: Some(1024),
+                access: AccessConfig::Everyone,
             }),
             stun: None,
             quic: None,
@@ -1106,4 +1145,96 @@ mod tests {
         assert_eq!(txid, txid_back);
         assert_eq!(response_addr, socket.local_addr().unwrap());
     }
+
+    #[tokio::test]
+    async fn test_relay_access_control() -> Result<()> {
+        let _guard = iroh_test::logging::setup();
+
+        let a_secret_key = SecretKey::generate(rand::thread_rng());
+        let a_key = a_secret_key.public();
+
+        let server = Server::spawn(ServerConfig::<(), ()> {
+            relay: Some(RelayConfig::<(), ()> {
+                http_bind_addr: (Ipv4Addr::LOCALHOST, 0).into(),
+                tls: None,
+                limits: Default::default(),
+                key_cache_capacity: Some(1024),
+                access: AccessConfig::Restricted(Box::new(move |node_id| {
+                    async move {
+                        info!("checking {}", node_id);
+                        // reject node a
+                        if node_id == a_key {
+                            Access::Deny
+                        } else {
+                            Access::Allow
+                        }
+                    }
+                    .boxed()
+                })),
+            }),
+            quic: None,
+            stun: None,
+            metrics_addr: None,
+        })
+        .await
+        .unwrap();
+        let relay_url = format!("http://{}", server.http_addr().unwrap());
+        let relay_url: RelayUrl = relay_url.parse()?;
+
+        // set up client a
+        let resolver = crate::dns::default_resolver().clone();
+        let mut client_a = ClientBuilder::new(relay_url.clone(), a_secret_key, resolver)
+            .connect()
+            .await?;
+
+        // the next message should be the rejection of the connection
+        tokio::time::timeout(Duration::from_millis(500), async move {
+            match client_a.next().await.unwrap().unwrap() {
+                ReceivedMessage::Health { problem } => {
+                    assert_eq!(problem, Some("not authenticated".to_string()));
+                }
+                msg => {
+                    panic!("other msg: {:?}", msg);
+                }
+            }
+        })
+        .await?;
+
+        // test that another client has access
+
+        // set up client b
+        let b_secret_key = SecretKey::generate(rand::thread_rng());
+        let b_key = b_secret_key.public();
+
+        let resolver = crate::dns::default_resolver().clone();
+        let mut client_b = ClientBuilder::new(relay_url.clone(), b_secret_key, resolver)
+            .connect()
+            .await?;
+
+        // set up client c
+        let c_secret_key = SecretKey::generate(rand::thread_rng());
+        let c_key = c_secret_key.public();
+
+        let resolver = crate::dns::default_resolver().clone();
+        let mut client_c = ClientBuilder::new(relay_url.clone(), c_secret_key, resolver)
+            .connect()
+            .await?;
+
+        // send message from b to c
+        let msg = Bytes::from("hello, c");
+        let res = try_send_recv(&mut client_b, &mut client_c, c_key, msg.clone()).await?;
+
+        if let ReceivedMessage::ReceivedPacket {
+            remote_node_id,
+            data,
+        } = res
+        {
+            assert_eq!(b_key, remote_node_id);
+            assert_eq!(msg, data);
+        } else {
+            panic!("client_c received unexpected message {res:?}");
+        }
+
+        Ok(())
+    }
 }

iroh-relay/src/server/client.rs

@@ -2,7 +2,7 @@
 
 use std::{future::Future, num::NonZeroU32, pin::Pin, sync::Arc, task::Poll, time::Duration};
 
-use anyhow::{Context, Result};
+use anyhow::{bail, Context, Result};
 use bytes::Bytes;
 use futures_lite::FutureExt;
 use futures_sink::Sink;
@@ -333,8 +333,8 @@ impl Actor {
                 self.write_frame(Frame::Pong { data }).await?;
                 inc!(Metrics, sent_pong);
             }
-            Frame::Health { .. } => {
-                inc!(Metrics, other_packets_recv);
+            Frame::Health { problem } => {
+                bail!("server issue: {:?}", problem);
             }
             _ => {
                 inc!(Metrics, unknown_frames);