iroh_relay/server/
http_server.rs

1//! Low-level HTTP server components for embedding the relay service.
2//!
3//! This module provides [`RelayService`] which can be used to embed relay functionality
4//! into an existing HTTP server. It handles individual connections and provides
5//! the core relay protocol implementation.
6//!
7//! For a complete relay server implementation, see the parent [`server`](super) module.
8
9use std::{collections::HashMap, convert::Infallible, net::SocketAddr, sync::Arc, time::Duration};
10
11use bytes::Bytes;
12use derive_more::Debug;
13use http::{
14    header::{CONNECTION, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL, SEC_WEBSOCKET_VERSION},
15    response::Builder as ResponseBuilder,
16};
17use hyper::{
18    HeaderMap, Method, Request, Response, StatusCode,
19    body::Incoming,
20    header::{HeaderValue, SEC_WEBSOCKET_ACCEPT, UPGRADE},
21    service::Service,
22    upgrade::Upgraded,
23};
24use n0_error::{e, ensure, stack_error};
25use n0_future::MaybeFuture;
26use tokio::{
27    net::{TcpListener, TcpStream},
28    sync::{Notify, watch},
29    time::Instant,
30};
31use tokio_rustls_acme::AcmeAcceptor;
32use tokio_util::{sync::CancellationToken, task::AbortOnDropHandle};
33use tracing::{Instrument, debug, error, info, info_span, trace, warn};
34
35use super::{
36    AllowAll, ClientRequest, DynAccessControl, SpawnError, clients::Clients,
37    streams::InvalidBucketConfig,
38};
39use crate::{
40    KeyCache,
41    defaults::{DEFAULT_KEY_CACHE_CAPACITY, timeouts::SERVER_WRITE_TIMEOUT},
42    http::{
43        CLIENT_AUTH_HEADER, ProtocolVersion, RELAY_PATH, SUPPORTED_WEBSOCKET_VERSION,
44        WEBSOCKET_UPGRADE_PROTOCOL,
45    },
46    protos::{handshake, relay::MAX_FRAME_SIZE, streams::WsBytesFramed},
47    server::{
48        ClientRateLimit,
49        client::Config,
50        metrics::Metrics,
51        streams::{MaybeTlsStream, RateLimited, RelayedStream},
52    },
53};
54
55/// Boxed HTTP response body produced by [`RelayServiceWithNotify`].
56pub type BytesBody = Box<
57    dyn 'static + Send + Unpin + hyper::body::Body<Data = hyper::body::Bytes, Error = Infallible>,
58>;
59/// Boxed error type returned from [`RelayServiceWithNotify`]'s [`hyper::service::Service`] impl.
60pub type HyperError = Box<dyn std::error::Error + Send + Sync>;
61/// Result alias for HTTP responses produced by [`RelayServiceWithNotify`].
62pub type HyperResult<T> = std::result::Result<T, HyperError>;
63pub(super) type HyperHandler = Box<
64    dyn Fn(Request<Incoming>, ResponseBuilder) -> HyperResult<Response<BytesBody>>
65        + Send
66        + Sync
67        + 'static,
68>;
69
70/// WebSocket GUID needed for accepting websocket connections, see RFC 6455 (<https://www.rfc-editor.org/rfc/rfc6455>) section 1.3
71const SEC_WEBSOCKET_ACCEPT_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
72
73/// Timeout for a connection to finish the TLS and WebSocket upgrade handshakes.
74///
75/// The connection is aborted if the connection does not complete the TLS handshake
76/// and establishes relay protocol WebSocket stream within this timeout.
77const ESTABLISH_TIMEOUT: Duration = Duration::from_secs(30);
78
79/// Derives the accept key for WebSocket handshake according to RFC 6455.
80/// Takes the client's Sec-WebSocket-Key value and returns the calculated accept key.
81fn derive_accept_key(client_key: &HeaderValue) -> String {
82    use sha1::Digest;
83
84    let mut sha1 = sha1::Sha1::new();
85    sha1.update(client_key.as_bytes());
86    sha1.update(SEC_WEBSOCKET_ACCEPT_GUID);
87    data_encoding::BASE64.encode(&sha1.finalize())
88}
89
90/// Creates a new [`BytesBody`] with given content.
91fn body_full(content: impl Into<hyper::body::Bytes>) -> BytesBody {
92    Box::new(http_body_util::Full::new(content.into()))
93}
94
95#[allow(clippy::result_large_err)]
96fn downcast_upgrade(upgraded: Upgraded) -> Result<(MaybeTlsStream, Bytes), ConnectionHandlerError> {
97    match upgraded.downcast::<hyper_util::rt::TokioIo<MaybeTlsStream>>() {
98        Ok(parts) => Ok((parts.io.into_inner(), parts.read_buf)),
99        Err(_) => Err(e!(ConnectionHandlerError::DowncastUpgrade)),
100    }
101}
102
103/// The Relay HTTP server.
104///
105/// A running HTTP server serving the relay endpoint and optionally a number of additional
106/// HTTP services added with [`ServerBuilder::request_handler`].  If configured using
107/// [`ServerBuilder::tls_config`] the server will handle TLS as well.
108///
109/// Created using [`ServerBuilder::spawn`].
110#[derive(Debug)]
111pub(super) struct Server {
112    addr: SocketAddr,
113    http_server_task: AbortOnDropHandle<()>,
114    cancel_server_loop: CancellationToken,
115    service: RelayService,
116}
117
118impl Server {
119    /// Returns a handle for this server.
120    ///
121    /// The server runs in the background as several async tasks.  This allows controlling
122    /// the server, in particular it allows gracefully shutting down the server.
123    pub(super) fn handle(&self) -> ServerHandle {
124        ServerHandle {
125            cancel_token: self.cancel_server_loop.clone(),
126        }
127    }
128
129    /// Closes the underlying relay server and the HTTP(S) server tasks.
130    pub(super) fn shutdown(&self) {
131        self.cancel_server_loop.cancel();
132    }
133
134    /// Returns the [`AbortOnDropHandle`] for the supervisor task managing the server.
135    ///
136    /// This is the root of all the tasks for the server.  Aborting it will abort all the
137    /// other tasks for the server.  Awaiting it will complete when all the server tasks are
138    /// completed.
139    pub(super) fn task_handle(&mut self) -> &mut AbortOnDropHandle<()> {
140        &mut self.http_server_task
141    }
142
143    /// Returns the local address of this server.
144    pub(super) fn addr(&self) -> SocketAddr {
145        self.addr
146    }
147
148    /// Returns the [`RelayService`] driving this server.
149    pub(super) fn service(&self) -> &RelayService {
150        &self.service
151    }
152}
153
154/// A handle for the [`Server`].
155///
156/// This does not allow access to the task but can communicate with it.
157#[derive(Debug, Clone)]
158pub(super) struct ServerHandle {
159    cancel_token: CancellationToken,
160}
161
162impl ServerHandle {
163    /// Gracefully shut down the server.
164    pub(super) fn shutdown(&self) {
165        self.cancel_token.cancel()
166    }
167}
168
169/// Configuration to use for the TLS connection
170///
171/// This struct wraps a rustls server configuration and TLS acceptor for use with
172/// [`RelayService::handle_connection`].
173///
174/// # Example
175///
176/// ```
177/// use std::sync::Arc;
178///
179/// use iroh_relay::server::http_server::TlsConfig;
180/// use rustls::ServerConfig;
181/// use webpki_types::{CertificateDer, PrivateKeyDer};
182///
183/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
184/// // Set ring as the process-level default crypto provider
185/// rustls::crypto::ring::default_provider()
186///     .install_default()
187///     .ok();
188/// // Generate a self-signed certificate for testing
189/// let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()])?;
190/// let cert_der = cert.cert.der().to_vec();
191/// let private_key_der = cert.signing_key.serialize_der();
192///
193/// // Create rustls types
194/// let cert_chain = vec![CertificateDer::from(cert_der)];
195/// let private_key = PrivateKeyDer::try_from(private_key_der)?;
196///
197/// // Create a rustls ServerConfig
198/// let server_config = Arc::new(
199///     ServerConfig::builder()
200///         .with_no_client_auth()
201///         .with_single_cert(cert_chain, private_key)?,
202/// );
203///
204/// // Create TlsConfig for use with RelayService
205/// let tls_config = TlsConfig::new(server_config);
206/// # Ok(())
207/// # }
208/// ```
209#[derive(Debug, Clone)]
210pub struct TlsConfig {
211    /// The server config
212    pub(super) config: Arc<rustls::ServerConfig>,
213    /// The kind
214    pub(super) acceptor: TlsAcceptor,
215}
216
217impl TlsConfig {
218    /// Creates a new `TlsConfig` from a rustls `ServerConfig`.
219    ///
220    /// This creates a manual TLS acceptor using the provided server configuration.
221    /// The acceptor will handle TLS handshakes for incoming connections.
222    ///
223    /// # Example
224    ///
225    /// ```
226    /// use std::sync::Arc;
227    ///
228    /// use iroh_relay::server::http_server::TlsConfig;
229    /// use rustls::ServerConfig;
230    /// use webpki_types::{CertificateDer, PrivateKeyDer};
231    ///
232    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
233    /// // Set ring as the process-level default crypto provider
234    /// rustls::crypto::ring::default_provider()
235    ///     .install_default()
236    ///     .ok();
237    /// // Generate a self-signed certificate for testing
238    /// let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()])?;
239    /// let cert_der = cert.cert.der().to_vec();
240    /// let private_key_der = cert.signing_key.serialize_der();
241    ///
242    /// // Create rustls types
243    /// let cert_chain = vec![CertificateDer::from(cert_der)];
244    /// let private_key = PrivateKeyDer::try_from(private_key_der)?;
245    ///
246    /// let server_config = Arc::new(
247    ///     ServerConfig::builder()
248    ///         .with_no_client_auth()
249    ///         .with_single_cert(cert_chain, private_key)?,
250    /// );
251    ///
252    /// let tls_config = TlsConfig::new(server_config);
253    /// # Ok(())
254    /// # }
255    /// ```
256    pub fn new(config: Arc<rustls::ServerConfig>) -> Self {
257        let acceptor = tokio_rustls::TlsAcceptor::from(config.clone());
258        Self {
259            config,
260            acceptor: TlsAcceptor::Manual(acceptor),
261        }
262    }
263}
264
265/// Errors when attempting to upgrade and
266#[allow(missing_docs)]
267#[stack_error(derive, add_meta)]
268#[non_exhaustive]
269pub enum ServeConnectionError {
270    #[error("TLS[acme] handshake")]
271    TlsHandshake {
272        #[error(std_err)]
273        source: std::io::Error,
274    },
275    #[error("TLS[acme] serve connection")]
276    ServeConnection {
277        #[error(std_err)]
278        source: hyper::Error,
279    },
280    #[error("TLS[manual] accept")]
281    ManualAccept {
282        #[error(std_err)]
283        source: std::io::Error,
284    },
285    #[error("TLS[acme] accept")]
286    LetsEncryptAccept {
287        #[error(std_err)]
288        source: std::io::Error,
289    },
290    #[error("HTTPS connection")]
291    Https {
292        #[error(std_err)]
293        source: hyper::Error,
294    },
295    #[error("HTTP connection")]
296    Http {
297        #[error(std_err)]
298        source: hyper::Error,
299    },
300    #[error("Connection did not reach established state within timeout")]
301    EstablishTimeout,
302}
303
304/// Server accept errors.
305#[allow(missing_docs)]
306#[stack_error(derive, add_meta, from_sources)]
307#[non_exhaustive]
308pub enum AcceptError {
309    #[error(transparent)]
310    Handshake { source: handshake::Error },
311    #[error("rate limiting misconfigured")]
312    RateLimitingMisconfigured { source: InvalidBucketConfig },
313}
314
315/// Server connection errors, includes errors that can happen on `accept`.
316#[allow(missing_docs)]
317#[stack_error(derive, add_meta, from_sources)]
318#[non_exhaustive]
319pub enum ConnectionHandlerError {
320    #[error(transparent)]
321    Accept { source: AcceptError },
322    #[error("Could not downcast the upgraded connection to MaybeTlsStream")]
323    DowncastUpgrade {},
324    #[error("Cannot deal with buffered data yet: {buf:?}")]
325    BufferNotEmpty { buf: Bytes },
326}
327
328/// Builder for the Relay HTTP Server.
329///
330/// Defaults to handling relay requests on the "/relay" (and "/derp" for backwards compatibility) endpoint.
331/// Other HTTP endpoints can be added using [`ServerBuilder::request_handler`].
332#[derive(derive_more::Debug)]
333pub(super) struct ServerBuilder {
334    /// The ip + port combination for this server.
335    addr: SocketAddr,
336    /// Optional tls configuration/TlsAcceptor combination.
337    ///
338    /// When `None`, the server will serve HTTP, otherwise it will serve HTTPS.
339    tls_config: Option<TlsConfig>,
340    /// A map of request handlers to routes.
341    ///
342    /// Used when certain routes in your server should be made available at the same port as
343    /// the relay server, and so must be handled along side requests to the relay endpoint.
344    handlers: Handlers,
345    /// Headers to use for HTTP responses.
346    headers: HeaderMap,
347    /// Rate-limiting configuration for an individual client connection.
348    ///
349    /// Rate-limiting is enforced on received traffic from individual clients.  This
350    /// configuration applies to a single client connection.
351    client_rx_ratelimit: Option<ClientRateLimit>,
352    /// The capacity of the key cache.
353    key_cache_capacity: usize,
354    /// Access control for endpoints.
355    access: Arc<dyn DynAccessControl>,
356    metrics: Option<Arc<Metrics>>,
357    establish_timeout: Duration,
358}
359
360impl ServerBuilder {
361    /// Creates a new [ServerBuilder].
362    pub(super) fn new(addr: SocketAddr) -> Self {
363        Self {
364            addr,
365            tls_config: None,
366            handlers: Default::default(),
367            headers: HeaderMap::new(),
368            client_rx_ratelimit: None,
369            key_cache_capacity: DEFAULT_KEY_CACHE_CAPACITY,
370            access: Arc::new(AllowAll),
371            metrics: None,
372            establish_timeout: ESTABLISH_TIMEOUT,
373        }
374    }
375
376    /// Sets the metrics collector.
377    pub(super) fn metrics(mut self, metrics: Arc<Metrics>) -> Self {
378        self.metrics = Some(metrics);
379        self
380    }
381
382    /// Set the access control.
383    pub(super) fn access(mut self, access: Arc<dyn DynAccessControl>) -> Self {
384        self.access = access;
385        self
386    }
387
388    /// Serves all requests content using TLS.
389    pub(super) fn tls_config(mut self, config: Option<TlsConfig>) -> Self {
390        self.tls_config = config;
391        self
392    }
393
394    /// Sets the timeout after which connections are aborted if they don't become fully established.
395    ///
396    /// The timeout is started immediately after a TCP connection comes in, and cleared once
397    /// the connection has finished the TLS handshake and fully processed the WebSocket request
398    /// to initiate the relay protocol. If the timeout expires before being cleared, the
399    /// connection is aborted.
400    ///
401    /// Defaults to 30s.
402    #[cfg(test)]
403    pub(super) fn establish_timeout(mut self, timeout: Duration) -> Self {
404        self.establish_timeout = timeout;
405        self
406    }
407
408    /// Sets the per-client rate-limit configuration for incoming data.
409    ///
410    /// On each client connection the incoming data is rate-limited.  By default
411    /// no rate limit is enforced.
412    pub(super) fn client_rx_ratelimit(mut self, config: ClientRateLimit) -> Self {
413        self.client_rx_ratelimit = Some(config);
414        self
415    }
416
417    /// Adds a custom handler for a specific Method & URI.
418    pub(super) fn request_handler(
419        mut self,
420        method: Method,
421        uri_path: &'static str,
422        handler: HyperHandler,
423    ) -> Self {
424        self.handlers.insert((method, uri_path), handler);
425        self
426    }
427
428    /// Adds HTTP headers to responses.
429    pub(super) fn headers(mut self, headers: HeaderMap) -> Self {
430        for (k, v) in headers.iter() {
431            self.headers.insert(k.clone(), v.clone());
432        }
433        self
434    }
435
436    /// Set the capacity of the cache for public keys.
437    pub(super) fn key_cache_capacity(mut self, capacity: usize) -> Self {
438        self.key_cache_capacity = capacity;
439        self
440    }
441
442    /// Builds and spawns an HTTP(S) Relay Server.
443    pub(super) async fn spawn(self) -> Result<Server, SpawnError> {
444        let cancel_token = CancellationToken::new();
445
446        let service = RelayService::new(
447            self.handlers,
448            self.headers,
449            self.client_rx_ratelimit,
450            KeyCache::new(self.key_cache_capacity),
451            self.access,
452            self.metrics.unwrap_or_default(),
453        );
454
455        let addr = self.addr;
456        let tls_config = self.tls_config;
457
458        // Bind a TCP listener on `addr` and handles content using HTTPS.
459
460        let listener = TcpListener::bind(&addr)
461            .await
462            .map_err(|err| e!(super::SpawnError::BindTcpListener { addr }, err))?;
463
464        let addr = listener
465            .local_addr()
466            .map_err(|err| e!(super::SpawnError::NoLocalAddr, err))?;
467        let http_str = tls_config.as_ref().map_or("HTTP/WS", |_| "HTTPS/WSS");
468        info!("[{http_str}] relay: serving on {addr}");
469
470        let cancel = cancel_token.clone();
471        let loop_service = service.clone();
472        let task = tokio::task::spawn(
473            async move {
474                let service = loop_service;
475                // create a join set to track all our connection tasks
476                let mut set = tokio::task::JoinSet::new();
477                loop {
478                    tokio::select! {
479                        biased;
480                        _ = cancel.cancelled() => {
481                            break;
482                        }
483                        Some(res) = set.join_next() => {
484                            if let Err(err) = res
485                                && err.is_panic()
486                            {
487                                panic!("task panicked: {err:#?}");
488                            }
489                        }
490                        res = listener.accept() => match res {
491                            Ok((stream, peer_addr)) => {
492                                debug!("connection opened from {peer_addr}");
493                                let tls_config = tls_config.clone();
494                                let service = service.clone();
495                                // spawn a task to handle the connection
496                                set.spawn(async move {
497                                    service
498                                        .handle_connection(stream, tls_config, self.establish_timeout)
499                                        .await
500                                }.instrument(info_span!("conn", peer = %peer_addr)));
501                            }
502                            Err(err) => {
503                                error!("failed to accept connection: {err}");
504                            }
505                        }
506                    }
507                }
508                service.shutdown().await;
509                set.shutdown().await;
510                debug!("server has been shutdown.");
511            }
512            .instrument(info_span!("relay-http-serve")),
513        );
514
515        Ok(Server {
516            addr,
517            http_server_task: AbortOnDropHandle::new(task),
518            cancel_server_loop: cancel_token,
519            service,
520        })
521    }
522}
523
524/// The hyper Service that serves the actual relay endpoints.
525///
526/// This service can be used standalone or embedded into an existing HTTP server.
527#[derive(Clone, Debug)]
528pub struct RelayService(Arc<Inner>);
529
530#[derive(Debug)]
531struct Inner {
532    handlers: Handlers,
533    headers: HeaderMap,
534    clients: Clients,
535    write_timeout: Duration,
536    rate_limit: watch::Sender<Option<ClientRateLimit>>,
537    key_cache: KeyCache,
538    access: Arc<dyn DynAccessControl>,
539    metrics: Arc<Metrics>,
540}
541
542#[stack_error(derive, add_meta)]
543enum RelayUpgradeReqError {
544    #[error("missing header: {header}")]
545    MissingHeader { header: http::HeaderName },
546    #[error("invalid header value for {header}: {details}")]
547    InvalidHeader {
548        header: http::HeaderName,
549        details: String,
550    },
551    #[error(
552        "invalid header value for {SEC_WEBSOCKET_VERSION}: unsupported websocket version, only supporting {SUPPORTED_WEBSOCKET_VERSION}"
553    )]
554    UnsupportedWebsocketVersion,
555    #[error(
556        "invalid header value for {SEC_WEBSOCKET_PROTOCOL}: unsupported relay version: we support {we_support} but you only provide {you_support}"
557    )]
558    UnsupportedRelayVersion {
559        we_support: String,
560        you_support: String,
561    },
562}
563
564impl RelayServiceWithNotify {
565    fn build_response(&self) -> http::response::Builder {
566        let mut res = Response::builder();
567        for (key, value) in self.service.0.headers.iter() {
568            res = res.header(key, value);
569        }
570        res
571    }
572
573    /// Upgrades the HTTP connection to the relay protocol, runs relay client.
574    fn handle_relay_ws_upgrade(
575        &self,
576        mut req: Request<Incoming>,
577    ) -> Result<Response<BytesBody>, RelayUpgradeReqError> {
578        fn expect_header(
579            req: &Request<Incoming>,
580            header: http::HeaderName,
581        ) -> Result<&HeaderValue, RelayUpgradeReqError> {
582            req.headers()
583                .get(&header)
584                .ok_or_else(|| e!(RelayUpgradeReqError::MissingHeader { header }))
585        }
586
587        let upgrade_header = expect_header(&req, UPGRADE)?;
588        ensure!(
589            upgrade_header == HeaderValue::from_static(WEBSOCKET_UPGRADE_PROTOCOL),
590            RelayUpgradeReqError::InvalidHeader {
591                header: UPGRADE,
592                details: format!("value must be {WEBSOCKET_UPGRADE_PROTOCOL}")
593            }
594        );
595
596        let key = expect_header(&req, SEC_WEBSOCKET_KEY)?.clone();
597        let version = expect_header(&req, SEC_WEBSOCKET_VERSION)?.clone();
598
599        ensure!(
600            version.as_bytes() == SUPPORTED_WEBSOCKET_VERSION.as_bytes(),
601            RelayUpgradeReqError::UnsupportedWebsocketVersion
602        );
603
604        let subprotocols = expect_header(&req, SEC_WEBSOCKET_PROTOCOL)?
605            .to_str()
606            .ok()
607            .ok_or_else(|| {
608                e!(RelayUpgradeReqError::InvalidHeader {
609                    header: SEC_WEBSOCKET_PROTOCOL,
610                    details: "header value is not ascii".to_string()
611                })
612            })?;
613        let protocol_version = subprotocols
614            .split(",")
615            .map(|s| s.trim())
616            .filter_map(ProtocolVersion::match_from_str)
617            .max()
618            .ok_or_else(|| {
619                e!(RelayUpgradeReqError::UnsupportedRelayVersion {
620                    we_support: ProtocolVersion::all_joined(),
621                    you_support: subprotocols.to_string()
622                })
623            })?;
624
625        // Setup a future that will eventually receive the upgraded
626        // connection and talk a new protocol, and spawn the future
627        // into the runtime.
628        //
629        // Note: This can't possibly be fulfilled until the 101 response
630        // is returned below, so it's better to spawn this future instead
631        // waiting for it to complete to then return a response.
632        tokio::task::spawn({
633            let this = self.clone();
634            let deadline = Instant::now() + this.establish_timeout;
635            let establish = async move {
636                match hyper::upgrade::on(&mut req).await {
637                    Ok(upgraded) => {
638                        let (parts, _) = req.into_parts();
639                        if let Err(err) = this
640                            .service
641                            .0
642                            .relay_connection_handler(upgraded, parts, protocol_version)
643                            .await
644                        {
645                            warn!("error accepting upgraded connection: {err:#}",);
646                        } else {
647                            // We have passed the connection to the relay protocol handler,
648                            // thus we trigger the on_establish notification so that timeouts
649                            // on the upper layer will be cleared.
650                            this.on_establish.notify_waiters();
651                            debug!("upgraded connection completed");
652                        };
653                    }
654                    Err(err) => warn!("upgrade error: {err:#}"),
655                }
656            };
657            async move {
658                // After the WebSocket upgrade, the HTTP timeout no longer covers this socket.
659                // Give the client a deadline to finish authentication and the access check.
660                if tokio::time::timeout_at(deadline, establish).await.is_err() {
661                    debug!("relay connection establishment timed out");
662                }
663            }
664            .instrument(info_span!("handler"))
665        });
666
667        // Now return a 101 Response saying we agree to the upgrade to the
668        // websocket upgrade protocol
669        Ok(self
670            .build_response()
671            .status(StatusCode::SWITCHING_PROTOCOLS)
672            .header(
673                UPGRADE,
674                HeaderValue::from_static(WEBSOCKET_UPGRADE_PROTOCOL),
675            )
676            .header(SEC_WEBSOCKET_ACCEPT, derive_accept_key(&key))
677            .header(SEC_WEBSOCKET_PROTOCOL, protocol_version.to_header_value())
678            .header(CONNECTION, "upgrade")
679            .body(body_full("switching to websocket protocol"))
680            .expect("valid body"))
681    }
682}
683
684/// Combines [`RelayService`] with a notification token.
685///
686/// This struct implements [`Service`]. Note that the service has to be called with hyper's `io`
687/// argument set to [`MaybeTlsStream`] wrapped by [`hyper_util::rt::TokioIo`], otherwise handling
688/// WebSocket requests at `/relay` will fail at runtime with [`ConnectionHandlerError::DowncastUpgrade`].
689///
690/// The notification token is triggered once the relay connection is fully established. It can be used
691/// to cancel a timeout aborting the TCP connection if no upgrade request is received in some time.
692///
693/// ## Example
694///
695/// ```no_run
696/// # use std::{sync::Arc, time::Duration};
697/// # use http::HeaderMap;
698/// # use hyper::server::conn::http1;
699/// # use hyper_util::rt::TokioIo;
700/// # use tokio::{net::TcpListener, sync::Notify};
701/// # use iroh_relay::{
702/// #     KeyCache,
703/// #     server::{
704/// #         AllowAll, Metrics,
705/// #         http_server::{Handlers, RelayService, RelayServiceWithNotify},
706/// #         streams::MaybeTlsStream
707/// #     },
708/// # };
709/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
710/// let service = RelayService::new(
711///     Handlers::default(),
712///     HeaderMap::new(),
713///     None,
714///     KeyCache::new(1024),
715///     Arc::new(AllowAll),
716///     Arc::new(Metrics::default()),
717/// );
718/// let service = RelayServiceWithNotify::with_timeout(
719///     service,
720///     Arc::new(Notify::new()),
721///     Duration::from_secs(30),
722/// );
723///
724/// let listener = TcpListener::bind("127.0.0.1:0").await?;
725/// let (stream, _peer) = listener.accept().await?;
726/// // Wrap the TCP stream in `MaybeTlsStream`, otherwise the relay WebSocket handler will error at runtime
727/// // for all WebSocket requests to `/relay`.
728/// let stream = MaybeTlsStream::Plain(stream);
729/// http1::Builder::new()
730///     .serve_connection(TokioIo::new(stream), service)
731///     .with_upgrades()
732///     .await?;
733/// # Ok(())
734/// # }
735/// ```
736#[derive(Debug, Clone)]
737pub struct RelayServiceWithNotify {
738    service: RelayService,
739    on_establish: Arc<Notify>,
740    establish_timeout: Duration,
741}
742
743impl RelayServiceWithNotify {
744    /// Creates a new service wrapper for a connection.
745    ///
746    /// The `on_establish` notification is triggered once the connection is passed to the
747    /// relay protocol, i.e. after a WebSocket request on /relay is received and established.
748    /// Upgraded connections must finish authentication and authorization within 30 seconds
749    /// of receiving the upgrade request.
750    pub fn new(service: RelayService, on_establish: Arc<Notify>) -> Self {
751        Self::with_timeout(service, on_establish, ESTABLISH_TIMEOUT)
752    }
753
754    /// Creates a service wrapper with a relay establishment timeout.
755    ///
756    /// Each upgrade request starts a new timeout covering the WebSocket upgrade, authentication,
757    /// and authorization. The `on_establish` notification is triggered once the connection is
758    /// passed to the relay protocol. Established connections are not subject to this timeout.
759    ///
760    /// The caller must separately bound TLS handshakes and HTTP requests before an upgrade request
761    /// reaches this service.
762    pub fn with_timeout(
763        service: RelayService,
764        on_establish: Arc<Notify>,
765        timeout: Duration,
766    ) -> Self {
767        Self {
768            service,
769            on_establish,
770            establish_timeout: timeout,
771        }
772    }
773}
774
775impl Service<Request<Incoming>> for RelayServiceWithNotify {
776    type Response = Response<BytesBody>;
777    type Error = HyperError;
778    type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
779
780    fn call(&self, req: Request<Incoming>) -> Self::Future {
781        // Create a client if the request hits the relay endpoint.
782        if matches!(
783            (req.method(), req.uri().path()),
784            (&hyper::Method::GET, RELAY_PATH)
785        ) {
786            let response = match self.handle_relay_ws_upgrade(req) {
787                Ok(response) => Ok(response),
788                // It's convention to send back the version(s) we *do* support
789                Err(e @ RelayUpgradeReqError::UnsupportedWebsocketVersion { .. }) => self
790                    .build_response()
791                    .status(StatusCode::BAD_REQUEST)
792                    .header(SEC_WEBSOCKET_VERSION, SUPPORTED_WEBSOCKET_VERSION)
793                    .body(body_full(e.to_string())),
794                Err(e) => self
795                    .build_response()
796                    .status(StatusCode::BAD_REQUEST)
797                    .body(body_full(e.to_string())),
798            }
799            .map_err(Into::into);
800            return std::future::ready(response);
801        }
802        // Otherwise handle the relay connection as normal.
803
804        // Check all other possible endpoints.
805        let uri = req.uri().clone();
806        if let Some(handler) = self
807            .service
808            .0
809            .handlers
810            .get(&(req.method().clone(), uri.path()))
811        {
812            let response = handler(req, self.service.0.default_response());
813            return std::future::ready(response);
814        }
815
816        // Otherwise return 404
817        let response = self
818            .service
819            .0
820            .not_found_fn(req, self.service.0.default_response());
821        std::future::ready(response)
822    }
823}
824
825impl Inner {
826    fn default_response(&self) -> ResponseBuilder {
827        let mut response = Response::builder();
828        for (key, value) in self.headers.iter() {
829            response = response.header(key.clone(), value.clone());
830        }
831        response
832    }
833
834    fn not_found_fn(
835        &self,
836        _req: Request<Incoming>,
837        mut res: ResponseBuilder,
838    ) -> HyperResult<Response<BytesBody>> {
839        for (k, v) in self.headers.iter() {
840            res = res.header(k.clone(), v.clone());
841        }
842        let body = body_full("Not Found");
843        let r = res.status(StatusCode::NOT_FOUND).body(body)?;
844        HyperResult::Ok(r)
845    }
846
847    /// The server HTTP handler to do HTTP upgrades.
848    ///
849    /// This handler runs while doing the connection upgrade handshake.  Once the connection
850    /// is upgraded it sends the stream to the relay server which takes it over.  After
851    /// having sent off the connection this handler returns.
852    async fn relay_connection_handler(
853        &self,
854        upgraded: Upgraded,
855        request_parts: http::request::Parts,
856        protocol_version: ProtocolVersion,
857    ) -> Result<(), ConnectionHandlerError> {
858        debug!("relay_connection upgraded");
859        let (io, read_buf) = downcast_upgrade(upgraded)?;
860        if !read_buf.is_empty() {
861            return Err(e!(ConnectionHandlerError::BufferNotEmpty { buf: read_buf }));
862        }
863
864        self.accept(io, request_parts, protocol_version).await?;
865        Ok(())
866    }
867
868    /// Adds a new connection to the server and serves it.
869    ///
870    /// Will error if it takes too long (10 sec) to write or read to the connection, if there is
871    /// some read or write error to the connection,  if the server is meant to verify clients,
872    /// and is unable to verify this one, or if there is some issue communicating with the server.
873    ///
874    /// The provided [`AsyncRead`] and [`AsyncWrite`] must be already connected to the connection.
875    ///
876    /// [`AsyncRead`]: tokio::io::AsyncRead
877    /// [`AsyncWrite`]: tokio::io::AsyncWrite
878    async fn accept(
879        &self,
880        io: MaybeTlsStream,
881        request_parts: http::request::Parts,
882        protocol_version: ProtocolVersion,
883    ) -> Result<(), AcceptError> {
884        trace!("accept: start");
885
886        // Set the socket to NO_DELAY.
887        io.disable_nagle();
888
889        let io = RateLimited::from_watcher(io, self.rate_limit.subscribe(), self.metrics.clone())
890            .map_err(|err| e!(AcceptError::RateLimitingMisconfigured, err))?;
891        let limited_watcher = io.limited_watcher();
892
893        // Create a server builder with default config
894        let websocket = tokio_websockets::ServerBuilder::new()
895            .limits(tokio_websockets::Limits::default().max_payload_len(Some(MAX_FRAME_SIZE)))
896            // Serve will create a WebSocketStream on an already upgraded connection
897            .serve(io);
898
899        let mut io = WsBytesFramed { io: websocket };
900
901        let client_auth_header = request_parts.headers.get(CLIENT_AUTH_HEADER).cloned();
902        let authentication = handshake::serverside(&mut io, client_auth_header).await?;
903
904        trace!(?authentication.mechanism, "accept: verified authentication");
905
906        let request =
907            ClientRequest::new(authentication.client_key, protocol_version, request_parts);
908
909        // Authorize the request against the configured `AccessControl`.
910        let guard = authentication
911            .authorize_with(&request, &self.access, &mut io)
912            .await?;
913
914        trace!("accept: verified authorization");
915
916        let io = RelayedStream {
917            inner: io,
918            key_cache: self.key_cache.clone(),
919        };
920
921        trace!("accept: build client conn");
922        let mut client_conn_builder = Config::new(guard, io, protocol_version);
923        client_conn_builder.write_timeout = self.write_timeout;
924        client_conn_builder.rate_limited = Some(limited_watcher);
925        trace!(endpoint_id = %request.endpoint_id().fmt_short(), "create client");
926
927        // build and register client, starting up read & write loops for the client
928        // connection
929        self.clients
930            .register(client_conn_builder, self.metrics.clone());
931        Ok(())
932    }
933}
934
935/// TLS Certificate Authority acceptor.
936#[derive(Clone, derive_more::Debug)]
937pub(super) enum TlsAcceptor {
938    /// Uses Let's Encrypt as the Certificate Authority. This is used in production.
939    LetsEncrypt(#[debug("tokio_rustls_acme::AcmeAcceptor")] AcmeAcceptor),
940    /// Manually added tls acceptor. Generally used for tests or for when we've passed in
941    /// a certificate via a file.
942    Manual(#[debug("tokio_rustls::TlsAcceptor")] tokio_rustls::TlsAcceptor),
943}
944
945impl RelayService {
946    /// Creates a new RelayService.
947    ///
948    /// This allows embedding the relay service into an existing HTTP server.
949    pub fn new(
950        handlers: Handlers,
951        headers: HeaderMap,
952        rate_limit: Option<ClientRateLimit>,
953        key_cache: KeyCache,
954        access: Arc<dyn DynAccessControl>,
955        metrics: Arc<Metrics>,
956    ) -> Self {
957        Self(Arc::new(Inner {
958            handlers,
959            headers,
960            clients: Clients::default(),
961            write_timeout: SERVER_WRITE_TIMEOUT,
962            rate_limit: watch::Sender::new(rate_limit),
963            key_cache,
964            access,
965            metrics,
966        }))
967    }
968
969    /// Updates the per-client receive rate limit.
970    ///
971    /// The new rate limit will be applied to all current and future client connections.
972    /// Passing `None` will remove rate limiting from all connections.
973    pub fn set_client_rate_limit(&self, rate_limit: Option<ClientRateLimit>) {
974        self.0.rate_limit.send_replace(rate_limit);
975    }
976
977    /// Shuts down the relay service, disconnecting all clients.
978    pub async fn shutdown(&self) {
979        self.0.clients.shutdown().await;
980    }
981
982    /// Returns a reference to the registry of currently connected clients.
983    ///
984    /// The returned [`Clients`] handle can be used at runtime to disconnect a
985    /// connected endpoint via [`Clients::disconnect`].
986    pub fn clients(&self) -> &Clients {
987        &self.0.clients
988    }
989
990    /// Handle the incoming connection.
991    ///
992    /// If a `tls_config` is given, will serve the connection using HTTPS, otherwise HTTP.
993    ///
994    /// TLS and HTTP establishment are bounded by `establish_timeout`. An upgrade request starts
995    /// a separate timeout of the same duration for the WebSocket upgrade, authentication, and
996    /// authorization. The connection is aborted if either timeout expires.
997    ///
998    /// # Example
999    ///
1000    /// ```no_run
1001    /// # use std::{sync::Arc, time::Duration};
1002    /// # use tokio::net::TcpStream;
1003    /// # use http::HeaderMap;
1004    /// # use iroh_relay::server::http_server::{Handlers, RelayService, TlsConfig};
1005    /// # use iroh_relay::{KeyCache, server::{AllowAll, Metrics}};
1006    /// # use webpki_types::{CertificateDer, PrivateKeyDer};
1007    /// # async fn example(stream: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
1008    /// // Create a relay service
1009    /// let handlers = Handlers::default();
1010    /// let headers = HeaderMap::new();
1011    /// let key_cache = KeyCache::new(1024);
1012    /// let metrics = Arc::new(Metrics::default());
1013    /// let relay_service = RelayService::new(
1014    ///     handlers,
1015    ///     headers,
1016    ///     None, // No rate limiting
1017    ///     key_cache,
1018    ///     Arc::new(AllowAll),
1019    ///     metrics,
1020    /// );
1021    ///
1022    /// // Generate a self-signed certificate for HTTPS
1023    /// let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()])?;
1024    /// let cert_der = cert.cert.der().to_vec();
1025    /// let private_key_der = cert.signing_key.serialize_der();
1026    /// let cert_chain = vec![CertificateDer::from(cert_der)];
1027    /// let private_key = PrivateKeyDer::try_from(private_key_der)?;
1028    ///
1029    /// // Serve with HTTPS
1030    /// let server_config = Arc::new(
1031    ///     rustls::ServerConfig::builder()
1032    ///         .with_no_client_auth()
1033    ///         .with_single_cert(cert_chain, private_key)?,
1034    /// );
1035    /// let tls_config = TlsConfig::new(server_config);
1036    /// relay_service
1037    ///     .clone()
1038    ///     .handle_connection(stream, Some(tls_config), Duration::from_secs(30))
1039    ///     .await;
1040    ///
1041    /// // Or serve with plain HTTP
1042    /// # let stream = TcpStream::connect("127.0.0.1:0").await?;
1043    /// relay_service
1044    ///     .handle_connection(stream, None, Duration::from_secs(30))
1045    ///     .await;
1046    /// # Ok(())
1047    /// # }
1048    /// ```
1049    pub async fn handle_connection(
1050        self,
1051        stream: TcpStream,
1052        tls_config: Option<TlsConfig>,
1053        establish_timeout: Duration,
1054    ) {
1055        let metrics = self.0.metrics.clone();
1056        metrics.http_connections.inc();
1057        // We create a notification token to be triggered once the connection is fully established
1058        // and passed to the relay server.
1059        let on_establish = Arc::new(Notify::new());
1060        let service =
1061            RelayServiceWithNotify::with_timeout(self, on_establish.clone(), establish_timeout);
1062
1063        // This is the main connection future, driving the connection to completion.
1064        let serve_fut = async move {
1065            match tls_config {
1066                Some(tls_config) => {
1067                    debug!("HTTPS: serve connection");
1068                    service.tls_serve_connection(stream, tls_config).await
1069                }
1070                None => {
1071                    debug!("HTTP: serve connection");
1072                    let stream = MaybeTlsStream::Plain(stream);
1073                    service.serve_connection(stream).await
1074                }
1075            }
1076        };
1077
1078        // Bound TLS and HTTP establishment. Hyper completes this future when it hands off an
1079        // upgraded socket; the upgrade task separately bounds authentication and authorization.
1080        let res = clearable_timeout(establish_timeout, on_establish, serve_fut)
1081            .await
1082            .map_err(|_elapsed| e!(ServeConnectionError::EstablishTimeout))
1083            .flatten();
1084
1085        metrics.http_connections_closed.inc();
1086
1087        if let Err(error) = res {
1088            match error {
1089                ServeConnectionError::ManualAccept { source, .. }
1090                | ServeConnectionError::LetsEncryptAccept { source, .. }
1091                    if source.kind() == std::io::ErrorKind::UnexpectedEof =>
1092                {
1093                    debug!(reason=?source, "peer disconnected");
1094                }
1095                // From hyper: <https://github.com/hyperium/hyper/commit/271bba16672ff54a44e043c5cc1ae6b9345bb172>
1096                // `hyper::Error::IncompleteMessage` is hyper's equivalent of UnexpectedEof
1097                ServeConnectionError::Https { source, .. }
1098                | ServeConnectionError::Http { source, .. }
1099                    if source.is_incomplete_message() =>
1100                {
1101                    debug!(reason=?source, "peer disconnected");
1102                }
1103                _ => {
1104                    metrics.http_connections_errored.inc();
1105                    error!(?error, "failed to handle connection");
1106                }
1107            }
1108        }
1109    }
1110}
1111
1112impl RelayServiceWithNotify {
1113    /// Serves a TLS connection.
1114    async fn tls_serve_connection(
1115        self,
1116        stream: TcpStream,
1117        tls_config: TlsConfig,
1118    ) -> Result<(), ServeConnectionError> {
1119        let TlsConfig { acceptor, config } = tls_config;
1120        let stream = match acceptor {
1121            TlsAcceptor::LetsEncrypt(a) => {
1122                match a
1123                    .accept(stream)
1124                    .await
1125                    .map_err(|err| e!(ServeConnectionError::LetsEncryptAccept, err))?
1126                {
1127                    None => {
1128                        info!("TLS[acme]: received TLS-ALPN-01 validation request");
1129                        return Ok(());
1130                    }
1131                    Some(start_handshake) => {
1132                        debug!("TLS[acme]: start handshake");
1133                        let tls_stream = start_handshake
1134                            .into_stream(config)
1135                            .await
1136                            .map_err(|err| e!(ServeConnectionError::TlsHandshake, err))?;
1137                        MaybeTlsStream::Tls(tls_stream)
1138                    }
1139                }
1140            }
1141            TlsAcceptor::Manual(a) => {
1142                debug!("TLS[manual]: accept");
1143                let tls_stream = a
1144                    .accept(stream)
1145                    .await
1146                    .map_err(|err| e!(ServeConnectionError::ManualAccept, err))?;
1147                MaybeTlsStream::Tls(tls_stream)
1148            }
1149        };
1150        self.serve_connection(stream).await
1151    }
1152
1153    /// Wrapper for the actual http connection (with upgrades)
1154    async fn serve_connection(self, io: MaybeTlsStream) -> Result<(), ServeConnectionError> {
1155        hyper::server::conn::http1::Builder::new()
1156            .serve_connection(hyper_util::rt::TokioIo::new(io), self)
1157            .with_upgrades()
1158            .await
1159            .map_err(|err| e!(ServeConnectionError::ServeConnection, err))
1160    }
1161}
1162
1163/// A collection of HTTP request handlers for custom endpoints.
1164#[derive(Default)]
1165pub struct Handlers(HashMap<(Method, &'static str), HyperHandler>);
1166
1167impl std::fmt::Debug for Handlers {
1168    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1169        let s = self.0.keys().fold(String::new(), |curr, next| {
1170            let (method, uri) = next;
1171            format!("{curr}\n({method},{uri}): Box<Fn(ResponseBuilder) -> Result<Response<Body>> + Send + Sync + 'static>")
1172        });
1173        write!(f, "HashMap<{s}>")
1174    }
1175}
1176
1177impl std::ops::Deref for Handlers {
1178    type Target = HashMap<(Method, &'static str), HyperHandler>;
1179
1180    fn deref(&self) -> &Self::Target {
1181        &self.0
1182    }
1183}
1184
1185impl std::ops::DerefMut for Handlers {
1186    fn deref_mut(&mut self) -> &mut Self::Target {
1187        &mut self.0
1188    }
1189}
1190
1191/// Requires a future to complete before the specified duration elapses, unless the timeout is cleared.
1192///
1193/// If the future completes before the duration has elapsed, then the completed value is returned.
1194/// Otherwise, an error is returned and the future is canceled.
1195///
1196/// If `clear_timeout` is triggered, the timeout is cleared and the future is always run to completion.
1197async fn clearable_timeout<F: Future>(
1198    timeout: Duration,
1199    clear_timeout: Arc<Notify>,
1200    fut: F,
1201) -> Result<F::Output, Elapsed> {
1202    tokio::pin!(fut);
1203    let timeout = MaybeFuture::Some(tokio::time::sleep(timeout));
1204    tokio::pin!(timeout);
1205    loop {
1206        tokio::select! {
1207            biased;
1208            res = &mut fut => {
1209                return Ok(res);
1210            }
1211            _ = clear_timeout.notified() => {
1212                timeout.as_mut().set_none();
1213            },
1214            _ = &mut timeout => {
1215                return Err(Elapsed);
1216            }
1217        }
1218    }
1219}
1220
1221#[stack_error(derive)]
1222#[error("Timeout elapsed")]
1223struct Elapsed;
1224
1225#[cfg(test)]
1226mod tests {
1227    use std::sync::Arc;
1228
1229    use iroh_base::{PublicKey, SecretKey};
1230    use iroh_dns::dns::DnsResolver;
1231    use n0_error::{Result, StdResultExt, bail_any};
1232    use n0_future::{SinkExt, StreamExt};
1233    use n0_tracing_test::traced_test;
1234    use rand::{RngExt, SeedableRng};
1235    use reqwest::Url;
1236    use tokio::io::{AsyncReadExt, AsyncWriteExt};
1237    use tracing::info;
1238
1239    use super::*;
1240    use crate::{
1241        client::{Client, ClientBuilder, ConnectError, conn::Conn},
1242        protos::relay::{ClientToRelayMsg, Datagrams, RelayToClientMsg},
1243        tls::{CaTlsConfig, default_provider},
1244    };
1245
1246    pub(crate) fn make_tls_config() -> TlsConfig {
1247        let subject_alt_names = vec!["localhost".to_string()];
1248
1249        let cert = rcgen::generate_simple_self_signed(subject_alt_names).unwrap();
1250        let rustls_certificate = cert.cert.der().clone();
1251        let rustls_key =
1252            rustls::pki_types::PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der());
1253        let config = rustls::ServerConfig::builder_with_provider(Arc::new(
1254            rustls::crypto::ring::default_provider(),
1255        ))
1256        .with_safe_default_protocol_versions()
1257        .expect("protocols supported by ring")
1258        .with_no_client_auth()
1259        .with_single_cert(vec![(rustls_certificate)], rustls_key.into())
1260        .expect("cert is right");
1261
1262        TlsConfig::new(Arc::new(config))
1263    }
1264
1265    #[tokio::test]
1266    #[traced_test]
1267    async fn test_http_clients_and_server() -> Result {
1268        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0u64);
1269
1270        let a_key = SecretKey::from_bytes(&rng.random());
1271        let b_key = SecretKey::from_bytes(&rng.random());
1272
1273        // start server
1274        let server = ServerBuilder::new("127.0.0.1:0".parse().unwrap())
1275            .spawn()
1276            .await?;
1277
1278        let addr = server.addr();
1279
1280        // get dial info
1281        let port = addr.port();
1282        let addr = if let std::net::IpAddr::V4(ipv4_addr) = addr.ip() {
1283            ipv4_addr
1284        } else {
1285            bail_any!("cannot get ipv4 addr from socket addr {addr:?}");
1286        };
1287
1288        info!("addr: {addr}:{port}");
1289        let relay_addr: Url = format!("http://{addr}:{port}").parse().unwrap();
1290
1291        // create clients
1292        let (a_key, mut client_a) = create_test_client(a_key, relay_addr.clone()).await?;
1293        info!("created client {a_key:?}");
1294        let (b_key, mut client_b) = create_test_client(b_key, relay_addr).await?;
1295        info!("created client {b_key:?}");
1296
1297        info!("ping a");
1298        client_a.send(ClientToRelayMsg::Ping([1u8; 8])).await?;
1299        let pong = client_a.next().await.expect("eos")?;
1300        assert!(matches!(pong, RelayToClientMsg::Pong { .. }));
1301
1302        info!("ping b");
1303        client_b.send(ClientToRelayMsg::Ping([2u8; 8])).await?;
1304        let pong = client_b.next().await.expect("eos")?;
1305        assert!(matches!(pong, RelayToClientMsg::Pong { .. }));
1306
1307        info!("sending message from a to b");
1308        let msg = Datagrams::from(b"hi there, client b!");
1309        client_a
1310            .send(ClientToRelayMsg::Datagrams {
1311                dst_endpoint_id: b_key,
1312                datagrams: msg.clone(),
1313            })
1314            .await?;
1315        info!("waiting for message from a on b");
1316        let (got_key, got_msg) =
1317            process_msg(client_b.next().await).expect("expected message from client_a");
1318        assert_eq!(a_key, got_key);
1319        assert_eq!(msg, got_msg);
1320
1321        info!("sending message from b to a");
1322        let msg = Datagrams::from(b"right back at ya, client b!");
1323        client_b
1324            .send(ClientToRelayMsg::Datagrams {
1325                dst_endpoint_id: a_key,
1326                datagrams: msg.clone(),
1327            })
1328            .await?;
1329        info!("waiting for message b on a");
1330        let (got_key, got_msg) =
1331            process_msg(client_a.next().await).expect("expected message from client_b");
1332        assert_eq!(b_key, got_key);
1333        assert_eq!(msg, got_msg);
1334
1335        // Close before shutting down, otherwise we'll try to send close frames on broken pipes
1336        client_a.close().await?;
1337        client_b.close().await?;
1338        server.shutdown();
1339
1340        Ok(())
1341    }
1342
1343    async fn create_test_client(
1344        key: SecretKey,
1345        server_url: Url,
1346    ) -> Result<(PublicKey, Client), ConnectError> {
1347        let public_key = key.public();
1348        let client = ClientBuilder::new(server_url, key, DnsResolver::new()).tls_client_config(
1349            CaTlsConfig::insecure_skip_verify()
1350                .client_config(default_provider())
1351                .expect("infallible"),
1352        );
1353        let client = client.connect().await?;
1354
1355        Ok((public_key, client))
1356    }
1357
1358    fn process_msg(
1359        msg: Option<Result<RelayToClientMsg, crate::client::RecvError>>,
1360    ) -> Option<(PublicKey, Datagrams)> {
1361        match msg {
1362            Some(Err(e)) => {
1363                info!("client `recv` error {e}");
1364                None
1365            }
1366            Some(Ok(msg)) => {
1367                info!("got message on: {msg:?}");
1368                if let RelayToClientMsg::Datagrams {
1369                    remote_endpoint_id: source,
1370                    datagrams,
1371                } = msg
1372                {
1373                    Some((source, datagrams))
1374                } else {
1375                    None
1376                }
1377            }
1378            None => {
1379                info!("client end of stream");
1380                None
1381            }
1382        }
1383    }
1384
1385    #[tokio::test]
1386    #[traced_test]
1387    async fn test_subprotocol_negotiation_picks_latest() -> Result {
1388        let server = ServerBuilder::new("127.0.0.1:0".parse().unwrap())
1389            .spawn()
1390            .await?;
1391        let addr = server.addr();
1392
1393        for offered in [
1394            "iroh-relay-v2,iroh-relay-v1",
1395            "iroh-relay-v1,iroh-relay-v2",
1396            "baz, iroh-relay-v1, iroh-relay-v2, boo",
1397            "foo, iroh-relay-v2, bar",
1398        ] {
1399            let ws_uri = format!("ws://{addr}{RELAY_PATH}");
1400            let (_stream, response) = tokio_websockets::ClientBuilder::new()
1401                .uri(&ws_uri)
1402                .expect("valid websocket URI")
1403                .add_header(
1404                    SEC_WEBSOCKET_PROTOCOL,
1405                    HeaderValue::from_str(offered).expect("valid subprotocol header value"),
1406                )
1407                .expect("header accepted by websocket client")
1408                .connect()
1409                .await
1410                .expect("websocket upgrade succeeds");
1411            let negotiated = response
1412                .headers()
1413                .get(SEC_WEBSOCKET_PROTOCOL)
1414                .expect("Sec-WebSocket-Protocol response header is present");
1415            assert_eq!(negotiated, "iroh-relay-v2", "offered={offered}");
1416        }
1417
1418        server.shutdown();
1419        Ok(())
1420    }
1421
1422    #[tokio::test]
1423    #[traced_test]
1424    async fn test_https_clients_and_server() -> Result {
1425        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0u64);
1426
1427        let a_key = SecretKey::from_bytes(&rng.random());
1428        let b_key = SecretKey::from_bytes(&rng.random());
1429
1430        // create tls_config
1431        let tls_config = make_tls_config();
1432
1433        // start server
1434        let mut server = ServerBuilder::new("127.0.0.1:0".parse().unwrap())
1435            .tls_config(Some(tls_config))
1436            .spawn()
1437            .await?;
1438
1439        let addr = server.addr();
1440
1441        // get dial info
1442        let port = addr.port();
1443        let addr = if let std::net::IpAddr::V4(ipv4_addr) = addr.ip() {
1444            ipv4_addr
1445        } else {
1446            bail_any!("cannot get ipv4 addr from socket addr {addr:?}");
1447        };
1448
1449        info!("Relay listening on: {addr}:{port}");
1450
1451        let url: Url = format!("https://localhost:{port}").parse().unwrap();
1452
1453        // create clients
1454        let (a_key, mut client_a) = create_test_client(a_key, url.clone()).await?;
1455        info!("created client {a_key:?}");
1456        let (b_key, mut client_b) = create_test_client(b_key, url).await?;
1457        info!("created client {b_key:?}");
1458
1459        info!("ping a");
1460        client_a.send(ClientToRelayMsg::Ping([1u8; 8])).await?;
1461        let pong = client_a.next().await.expect("eos")?;
1462        assert!(matches!(pong, RelayToClientMsg::Pong { .. }));
1463
1464        info!("ping b");
1465        client_b.send(ClientToRelayMsg::Ping([2u8; 8])).await?;
1466        let pong = client_b.next().await.expect("eos")?;
1467        assert!(matches!(pong, RelayToClientMsg::Pong { .. }));
1468
1469        info!("sending message from a to b");
1470        let msg = Datagrams::from(b"hi there, client b!");
1471        client_a
1472            .send(ClientToRelayMsg::Datagrams {
1473                dst_endpoint_id: b_key,
1474                datagrams: msg.clone(),
1475            })
1476            .await?;
1477        info!("waiting for message from a on b");
1478        let (got_key, got_msg) =
1479            process_msg(client_b.next().await).expect("expected message from client_a");
1480        assert_eq!(a_key, got_key);
1481        assert_eq!(msg, got_msg);
1482
1483        info!("sending message from b to a");
1484        let msg = Datagrams::from(b"right back at ya, client b!");
1485        client_b
1486            .send(ClientToRelayMsg::Datagrams {
1487                dst_endpoint_id: a_key,
1488                datagrams: msg.clone(),
1489            })
1490            .await?;
1491        info!("waiting for message b on a");
1492        let (got_key, got_msg) =
1493            process_msg(client_a.next().await).expect("expected message from client_b");
1494        assert_eq!(b_key, got_key);
1495        assert_eq!(msg, got_msg);
1496
1497        // Close before shutting down, otherwise we'll try to send close frames on broken pipes
1498        client_a.close().await?;
1499        client_b.close().await?;
1500        server.shutdown();
1501        server.task_handle().await.std_context("join")?;
1502
1503        Ok(())
1504    }
1505
1506    async fn make_test_client(client: tokio::io::DuplexStream, key: &SecretKey) -> Result<Conn> {
1507        let client = crate::client::streams::MaybeTlsStream::Test(client);
1508        let client = tokio_websockets::ClientBuilder::new().take_over(client);
1509        let client = Conn::new(client, KeyCache::test(), key, Default::default()).await?;
1510        Ok(client)
1511    }
1512
1513    #[tokio::test]
1514    #[traced_test]
1515    async fn test_server_basic() -> Result {
1516        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0u64);
1517
1518        info!("Create the server.");
1519        let metrics = Arc::new(Metrics::default());
1520        let service = RelayService::new(
1521            Default::default(),
1522            Default::default(),
1523            None,
1524            KeyCache::test(),
1525            Arc::new(crate::server::AllowAll),
1526            metrics.clone(),
1527        );
1528
1529        info!("Create client A and connect it to the server.");
1530        let key_a = SecretKey::from_bytes(&rng.random());
1531        let public_key_a = key_a.public();
1532        let (client_a, rw_a) = tokio::io::duplex(10);
1533        let s = service.clone();
1534        let handler_task = tokio::spawn(async move {
1535            s.0.accept(
1536                MaybeTlsStream::Test(rw_a),
1537                Request::new(()).into_parts().0,
1538                Default::default(),
1539            )
1540            .await
1541        });
1542        let mut client_a = make_test_client(client_a, &key_a).await?;
1543        handler_task.await.std_context("join")??;
1544
1545        info!("Create client B and connect it to the server.");
1546        let key_b = SecretKey::from_bytes(&rng.random());
1547        let public_key_b = key_b.public();
1548        let (client_b, rw_b) = tokio::io::duplex(10);
1549        let s = service.clone();
1550        let handler_task = tokio::spawn(async move {
1551            s.0.accept(
1552                MaybeTlsStream::Test(rw_b),
1553                Request::new(()).into_parts().0,
1554                Default::default(),
1555            )
1556            .await
1557        });
1558        let mut client_b = make_test_client(client_b, &key_b).await?;
1559        handler_task.await.std_context("join")??;
1560
1561        info!("Send message from A to B.");
1562        let msg = Datagrams::from(b"hello client b!!");
1563        client_a
1564            .send(ClientToRelayMsg::Datagrams {
1565                dst_endpoint_id: public_key_b,
1566                datagrams: msg.clone(),
1567            })
1568            .await?;
1569        match client_b.next().await.unwrap()? {
1570            RelayToClientMsg::Datagrams {
1571                remote_endpoint_id,
1572                datagrams,
1573            } => {
1574                assert_eq!(public_key_a, remote_endpoint_id);
1575                assert_eq!(msg, datagrams);
1576            }
1577            msg => {
1578                bail_any!("expected ReceivedDatagrams msg, got {msg:?}");
1579            }
1580        }
1581
1582        info!("Send message from B to A.");
1583        let msg = Datagrams::from(b"nice to meet you client a!!");
1584        client_b
1585            .send(ClientToRelayMsg::Datagrams {
1586                dst_endpoint_id: public_key_a,
1587                datagrams: msg.clone(),
1588            })
1589            .await?;
1590        match client_a.next().await.unwrap()? {
1591            RelayToClientMsg::Datagrams {
1592                remote_endpoint_id,
1593                datagrams,
1594            } => {
1595                assert_eq!(public_key_b, remote_endpoint_id);
1596                assert_eq!(msg, datagrams);
1597            }
1598            msg => {
1599                bail_any!("expected ReceivedDatagrams msg, got {msg:?}");
1600            }
1601        }
1602
1603        info!("Close the server and clients");
1604        service.shutdown().await;
1605        tokio::time::sleep(Duration::from_secs(1)).await;
1606
1607        info!("Fail to send message from A to B.");
1608        let res = client_a
1609            .send(ClientToRelayMsg::Datagrams {
1610                dst_endpoint_id: public_key_b,
1611                datagrams: Datagrams::from(b"try to send"),
1612            })
1613            .await;
1614        assert!(res.is_err());
1615        assert!(client_b.next().await.is_none());
1616
1617        drop(client_a);
1618        drop(client_b);
1619
1620        service.shutdown().await;
1621
1622        assert_eq!(metrics.accepts.get(), metrics.disconnects.get());
1623
1624        Ok(())
1625    }
1626
1627    #[tokio::test]
1628    async fn test_server_replace_client() -> Result {
1629        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0u64);
1630
1631        info!("Create the server.");
1632        let service = RelayService::new(
1633            Default::default(),
1634            Default::default(),
1635            None,
1636            KeyCache::test(),
1637            Arc::new(crate::server::AllowAll),
1638            Default::default(),
1639        );
1640
1641        info!("Create client A and connect it to the server.");
1642        let key_a = SecretKey::from_bytes(&rng.random());
1643        let public_key_a = key_a.public();
1644        let (client_a, rw_a) = tokio::io::duplex(10);
1645        let s = service.clone();
1646        let handler_task = tokio::spawn(async move {
1647            s.0.accept(
1648                MaybeTlsStream::Test(rw_a),
1649                Request::new(()).into_parts().0,
1650                Default::default(),
1651            )
1652            .await
1653        });
1654        let mut client_a = make_test_client(client_a, &key_a).await?;
1655        handler_task.await.std_context("join")??;
1656
1657        info!("Create client B and connect it to the server.");
1658        let key_b = SecretKey::from_bytes(&rng.random());
1659        let public_key_b = key_b.public();
1660        let (client_b, rw_b) = tokio::io::duplex(10);
1661        let s = service.clone();
1662        let handler_task = tokio::spawn(async move {
1663            s.0.accept(
1664                MaybeTlsStream::Test(rw_b),
1665                Request::new(()).into_parts().0,
1666                Default::default(),
1667            )
1668            .await
1669        });
1670        let mut client_b = make_test_client(client_b, &key_b).await?;
1671        handler_task.await.std_context("join")??;
1672
1673        info!("Send message from A to B.");
1674        let msg = Datagrams::from(b"hello client b!!");
1675        client_a
1676            .send(ClientToRelayMsg::Datagrams {
1677                dst_endpoint_id: public_key_b,
1678                datagrams: msg.clone(),
1679            })
1680            .await?;
1681        match client_b.next().await.expect("eos")? {
1682            RelayToClientMsg::Datagrams {
1683                remote_endpoint_id,
1684                datagrams,
1685            } => {
1686                assert_eq!(public_key_a, remote_endpoint_id);
1687                assert_eq!(msg, datagrams);
1688            }
1689            msg => {
1690                bail_any!("expected ReceivedDatagrams msg, got {msg:?}");
1691            }
1692        }
1693
1694        info!("Send message from B to A.");
1695        let msg = Datagrams::from(b"nice to meet you client a!!");
1696        client_b
1697            .send(ClientToRelayMsg::Datagrams {
1698                dst_endpoint_id: public_key_a,
1699                datagrams: msg.clone(),
1700            })
1701            .await?;
1702        match client_a.next().await.expect("eos")? {
1703            RelayToClientMsg::Datagrams {
1704                remote_endpoint_id,
1705                datagrams,
1706            } => {
1707                assert_eq!(public_key_b, remote_endpoint_id);
1708                assert_eq!(msg, datagrams);
1709            }
1710            msg => {
1711                bail_any!("expected ReceivedDatagrams msg, got {msg:?}");
1712            }
1713        }
1714
1715        info!("Create client B and connect it to the server");
1716        let (new_client_b, new_rw_b) = tokio::io::duplex(10);
1717        let s = service.clone();
1718        let handler_task = tokio::spawn(async move {
1719            s.0.accept(
1720                MaybeTlsStream::Test(new_rw_b),
1721                Request::new(()).into_parts().0,
1722                Default::default(),
1723            )
1724            .await
1725        });
1726        let mut new_client_b = make_test_client(new_client_b, &key_b).await?;
1727        handler_task.await.std_context("join")??;
1728
1729        // assert!(client_b.recv().await.is_err());
1730
1731        info!("Send message from A to B.");
1732        let msg = Datagrams::from(b"are you still there, b?!");
1733        client_a
1734            .send(ClientToRelayMsg::Datagrams {
1735                dst_endpoint_id: public_key_b,
1736                datagrams: msg.clone(),
1737            })
1738            .await?;
1739        match new_client_b.next().await.expect("eos")? {
1740            RelayToClientMsg::Datagrams {
1741                remote_endpoint_id,
1742                datagrams,
1743            } => {
1744                assert_eq!(public_key_a, remote_endpoint_id);
1745                assert_eq!(msg, datagrams);
1746            }
1747            msg => {
1748                bail_any!("expected ReceivedDatagrams msg, got {msg:?}");
1749            }
1750        }
1751
1752        info!("Send message from B to A.");
1753        let msg = Datagrams::from(b"just had a spot of trouble but I'm back now,a!!");
1754        new_client_b
1755            .send(ClientToRelayMsg::Datagrams {
1756                dst_endpoint_id: public_key_a,
1757                datagrams: msg.clone(),
1758            })
1759            .await?;
1760        match client_a.next().await.expect("eos")? {
1761            RelayToClientMsg::Datagrams {
1762                remote_endpoint_id,
1763                datagrams,
1764            } => {
1765                assert_eq!(public_key_b, remote_endpoint_id);
1766                assert_eq!(msg, datagrams);
1767            }
1768            msg => {
1769                bail_any!("expected ReceivedDatagrams msg, got {msg:?}");
1770            }
1771        }
1772
1773        info!("Close the server and clients");
1774        service.shutdown().await;
1775
1776        info!("Sending message from A to B fails");
1777        let res = client_a
1778            .send(ClientToRelayMsg::Datagrams {
1779                dst_endpoint_id: public_key_b,
1780                datagrams: Datagrams::from(b"try to send"),
1781            })
1782            .await;
1783        assert!(res.is_err());
1784        assert!(new_client_b.next().await.is_none());
1785        Ok(())
1786    }
1787
1788    #[tokio::test]
1789    #[traced_test]
1790    async fn test_establish_timeout() -> Result {
1791        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42u64);
1792
1793        // Start server with a very short establish timeout.
1794        let server = ServerBuilder::new("127.0.0.1:0".parse().unwrap())
1795            .establish_timeout(Duration::from_millis(500))
1796            .spawn()
1797            .await?;
1798
1799        let addr = server.addr();
1800        let port = addr.port();
1801        let addr = if let std::net::IpAddr::V4(ipv4_addr) = addr.ip() {
1802            ipv4_addr
1803        } else {
1804            bail_any!("cannot get ipv4 addr from socket addr {addr:?}");
1805        };
1806        let relay_url: Url = format!("http://{addr}:{port}").parse().unwrap();
1807
1808        // 1. A lingering connection that never upgrades should be aborted by the timeout.
1809        info!("opening lingering TCP connection (no upgrade)");
1810        let mut lingering = TcpStream::connect(format!("{addr}:{port}")).await?;
1811        // Write a partial HTTP request but never complete the upgrade.
1812        lingering
1813            .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n")
1814            .await?;
1815        // Wait for the server to abort this connection.
1816        let mut buf = [0u8; 1];
1817        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
1818        let read = tokio::time::timeout_at(deadline, lingering.read(&mut buf)).await;
1819        // The server should close the connection, resulting in a read of 0 bytes or an error.
1820        match read {
1821            Ok(Ok(0)) => info!("lingering connection closed by server (EOF)"),
1822            Ok(Err(e)) => info!("lingering connection closed by server (error: {e})"),
1823            other => bail_any!("expected lingering connection to be closed, got {other:?}"),
1824        }
1825
1826        // 2. A properly established client should NOT be aborted by the timeout.
1827        info!("connecting a proper relay client");
1828        let key = SecretKey::from_bytes(&rng.random());
1829        let (_, mut client) = create_test_client(key, relay_url).await?;
1830
1831        // Wait longer than the establish timeout to prove the connection survives.
1832        tokio::time::sleep(Duration::from_millis(1000)).await;
1833
1834        // Ping should still work.
1835        client.send(ClientToRelayMsg::Ping([7u8; 8])).await?;
1836        let pong = client.next().await.expect("expected pong")?;
1837        assert!(matches!(pong, RelayToClientMsg::Pong { .. }));
1838        info!("established connection survived past the timeout");
1839
1840        client.close().await?;
1841        server.shutdown();
1842        Ok(())
1843    }
1844}
iroh_relay/server/http_server.rs

iroh_relay/server/
http_server.rs
