iota_http/
connection_info.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2025 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::HashMap,
7    sync::{Arc, Mutex, RwLock},
8};
9
10use tokio_rustls::rustls::pki_types::CertificateDer;
11
12use crate::config::{OnPeerConnectionEvent, PeerConnectionEvent};
13
14pub(crate) type ActiveConnections<A = std::net::SocketAddr> =
15    Arc<RwLock<HashMap<ConnectionId, ConnectionInfo<A>>>>;
16
17pub type ConnectionId = usize;
18
19#[derive(Debug)]
20pub struct ConnectionInfo<A>(Arc<Inner<A>>);
21
22#[derive(Clone, Debug)]
23pub struct PeerCertificates(Arc<Vec<tokio_rustls::rustls::pki_types::CertificateDer<'static>>>);
24
25impl PeerCertificates {
26    pub fn peer_certs(&self) -> &[tokio_rustls::rustls::pki_types::CertificateDer<'static>] {
27        self.0.as_ref()
28    }
29}
30
31impl<A> ConnectionInfo<A> {
32    pub(crate) fn new(
33        address: A,
34        peer_certificates: Option<Arc<Vec<CertificateDer<'static>>>>,
35        graceful_shutdown_token: tokio_util::sync::CancellationToken,
36    ) -> Self {
37        Self(Arc::new(Inner {
38            address,
39            time_established: std::time::Instant::now(),
40            peer_certificates: peer_certificates.map(PeerCertificates),
41            graceful_shutdown_token,
42        }))
43    }
44
45    /// The peer's remote address
46    pub fn remote_address(&self) -> &A {
47        &self.0.address
48    }
49
50    /// Time the Connection was established
51    pub fn time_established(&self) -> std::time::Instant {
52        self.0.time_established
53    }
54
55    pub fn peer_certificates(&self) -> Option<&PeerCertificates> {
56        self.0.peer_certificates.as_ref()
57    }
58
59    /// A stable identifier for this connection
60    pub fn id(&self) -> ConnectionId {
61        &*self.0 as *const _ as usize
62    }
63
64    /// Trigger a graceful shutdown of this connection
65    pub fn close(&self) {
66        self.0.graceful_shutdown_token.cancel()
67    }
68}
69
70#[derive(Debug)]
71struct Inner<A = std::net::SocketAddr> {
72    address: A,
73
74    // Time that the connection was established
75    time_established: std::time::Instant,
76
77    peer_certificates: Option<PeerCertificates>,
78    graceful_shutdown_token: tokio_util::sync::CancellationToken,
79}
80
81#[derive(Debug, Clone)]
82pub struct ConnectInfo<A = std::net::SocketAddr> {
83    /// Returns the local address of this connection.
84    pub local_addr: A,
85    /// Returns the remote (peer) address of this connection.
86    pub remote_addr: A,
87}
88
89impl<A> ConnectInfo<A> {
90    /// Return the local address the IO resource is connected.
91    pub fn local_addr(&self) -> &A {
92        &self.local_addr
93    }
94
95    /// Return the remote address the IO resource is connected too.
96    pub fn remote_addr(&self) -> &A {
97        &self.remote_addr
98    }
99}
100
101/// Number of established connections held by each authenticated peer, keyed by
102/// the peer's public key.
103#[derive(Clone, Debug)]
104pub(crate) struct PeerConnectionCounts {
105    counts: Arc<Mutex<HashMap<Vec<u8>, usize>>>,
106    on_event: Option<OnPeerConnectionEvent>,
107}
108
109impl PeerConnectionCounts {
110    pub(crate) fn new(on_event: Option<OnPeerConnectionEvent>) -> Self {
111        Self {
112            counts: Arc::default(),
113            on_event,
114        }
115    }
116
117    /// Counts one more connection for `peer`, or returns `None` if the peer
118    /// already holds `max` of them.
119    pub(crate) fn register(&self, peer: &[u8], max: usize) -> Option<PeerConnectionGuard> {
120        let held = {
121            let mut counts = self.counts.lock().unwrap();
122            // A zero `max` is rejected by `Config::validate`, so a peer with no
123            // entry yet is always below the limit.
124            if counts.get(peer).is_some_and(|count| *count >= max) {
125                None
126            } else {
127                let count = counts.entry(peer.to_vec()).or_insert(0);
128                *count += 1;
129                Some(*count)
130            }
131        };
132
133        let Some(held) = held else {
134            self.notify(peer, PeerConnectionEvent::RefusedAtLimit { held: max });
135            return None;
136        };
137        self.notify(peer, PeerConnectionEvent::Established { held });
138        Some(PeerConnectionGuard {
139            counts: self.clone(),
140            peer: peer.to_vec(),
141        })
142    }
143
144    fn notify(&self, peer: &[u8], event: PeerConnectionEvent) {
145        if let Some(on_event) = &self.on_event {
146            on_event.call(peer, event);
147        }
148    }
149}
150
151/// Gives the peer its connection back when dropped.
152pub(crate) struct PeerConnectionGuard {
153    counts: PeerConnectionCounts,
154    peer: Vec<u8>,
155}
156
157impl Drop for PeerConnectionGuard {
158    fn drop(&mut self) {
159        let held = {
160            let mut counts = self.counts.counts.lock().unwrap();
161            let Some(count) = counts.get_mut(&self.peer) else {
162                return;
163            };
164            *count -= 1;
165            let held = *count;
166            if held == 0 {
167                counts.remove(&self.peer);
168            }
169            held
170        };
171        self.counts
172            .notify(&self.peer, PeerConnectionEvent::Closed { held });
173    }
174}
iota_http/connection_info.rs

iota_http/
connection_info.rs
