karyon_p2p/discovery/kademlia/
refresh.rs

1use std::{
2    collections::HashMap,
3    net::{IpAddr, SocketAddr},
4    sync::Arc,
5    time::{Duration, Instant},
6};
7
8use log::{error, info, trace};
9use rand::{rngs::OsRng, TryRngCore};
10
11use karyon_core::{
12    async_runtime::Executor,
13    async_util::{sleep, timeout, Backoff, TaskGroup, TaskResult},
14};
15
16use karyon_net::{udp, Endpoint};
17
18use crate::{
19    access_control::{Action, Subject},
20    discovery::kademlia::{
21        messages::RefreshMsg,
22        routing_table::{BucketEntry, Entry, RoutingTable, PENDING_ENTRY, UNREACHABLE_ENTRY},
23    },
24    message::{pick_endpoint, Protocol},
25    monitor::{ConnectionKind, DiscoveryKind, Monitor},
26    peer::ConnDirection,
27    util::{decode, encode},
28    Config, Error, Result,
29};
30
31/// Maximum failures for an entry before removing it from the routing table.
32pub const MAX_FAILURES: u32 = 3;
33
34// Max UDP datagram payload size for refresh messages.
35const MAX_UDP_BUF: usize = 1024;
36
37// Max entries pulled from each bucket per refresh round.
38const REFRESH_PER_BUCKET: usize = 8;
39
40// Token-bucket parameters for the per-IP rate limit on incoming pings.
41// Capacity allows small bursts; refill rate caps sustained traffic.
42const RL_CAPACITY: u32 = 5;
43const RL_REFILL_PER_SEC: f64 = 0.5;
44
45/// Per-IP token bucket used by `listen_to_ping_msg` to drop floods.
46struct RateBucket {
47    tokens: f64,
48    last_refill: Instant,
49}
50
51impl RateBucket {
52    fn new() -> Self {
53        Self {
54            tokens: RL_CAPACITY as f64,
55            last_refill: Instant::now(),
56        }
57    }
58
59    /// Refill and try to take one token. Returns true if allowed.
60    fn allow(&mut self) -> bool {
61        let now = Instant::now();
62        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
63        self.tokens = (self.tokens + elapsed * RL_REFILL_PER_SEC).min(RL_CAPACITY as f64);
64        self.last_refill = now;
65        if self.tokens >= 1.0 {
66            self.tokens -= 1.0;
67            true
68        } else {
69            false
70        }
71    }
72}
73
74pub struct RefreshService {
75    /// Routing table
76    table: Arc<RoutingTable>,
77
78    /// UDP listen endpoint (None when no udp discovery endpoint is configured).
79    listen_endpoint: Option<Endpoint>,
80
81    /// Managing spawned tasks.
82    task_group: TaskGroup,
83
84    /// Holds the configuration for the P2P network.
85    config: Arc<Config>,
86
87    /// Responsible for network and system monitoring.
88    monitor: Arc<Monitor>,
89}
90
91impl RefreshService {
92    /// Creates a new refresh service
93    pub fn new(
94        config: Arc<Config>,
95        table: Arc<RoutingTable>,
96        monitor: Arc<Monitor>,
97        listen_endpoint: Option<Endpoint>,
98        executor: Executor,
99    ) -> Self {
100        Self {
101            table,
102            listen_endpoint,
103            task_group: TaskGroup::with_executor(executor.clone()),
104            config,
105            monitor,
106        }
107    }
108
109    /// Start the refresh service
110    pub async fn start(self: &Arc<Self>) -> Result<()> {
111        if let Some(endpoint) = self.listen_endpoint.clone() {
112            self.task_group.spawn_then(
113                {
114                    let this = self.clone();
115                    async move { this.listen_loop(endpoint).await }
116                },
117                |res| async move {
118                    if let TaskResult::Completed(Err(err)) = res {
119                        error!("Listen loop stopped: {err}");
120                    }
121                },
122            );
123        }
124
125        self.task_group.spawn_then(
126            {
127                let this = self.clone();
128                async move { this.refresh_loop().await }
129            },
130            |res| async move {
131                if let TaskResult::Completed(Err(err)) = res {
132                    error!("Refresh loop stopped: {err}");
133                }
134            },
135        );
136
137        Ok(())
138    }
139
140    /// Shuts down the refresh service
141    pub async fn shutdown(&self) {
142        self.task_group.cancel().await;
143    }
144
145    /// Periodically refresh the routing table by pinging the oldest entries.
146    async fn refresh_loop(self: Arc<Self>) -> Result<()> {
147        loop {
148            sleep(Duration::from_secs(self.config.refresh_interval)).await;
149            trace!("Start refreshing the routing table...");
150
151            self.monitor.notify(DiscoveryKind::RefreshStarted).await;
152
153            let entries = self.table.refresh_candidates(REFRESH_PER_BUCKET);
154            let succeeded = self.clone().do_refresh(&entries).await;
155
156            // Empty sweep counts as success with 0 — fail only when at
157            // least one entry was tried and all failed.
158            if !entries.is_empty() && succeeded == 0 {
159                self.monitor.notify(DiscoveryKind::RefreshFailed).await;
160            } else {
161                self.monitor
162                    .notify(DiscoveryKind::RefreshSucceeded(succeeded))
163                    .await;
164            }
165        }
166    }
167
168    /// Iterates over the entries and initiates a connection. Returns
169    /// the number of entries refreshed successfully.
170    async fn do_refresh(self: Arc<Self>, entries: &[BucketEntry]) -> usize {
171        use futures_util::stream::{FuturesUnordered, StreamExt};
172        let mut succeeded = 0;
173        for chunk in entries.chunks(16) {
174            let mut tasks = FuturesUnordered::new();
175            for bucket_entry in chunk {
176                if bucket_entry.failures >= MAX_FAILURES {
177                    let pid = bucket_entry.entry.key.into();
178                    self.table.remove_entry(&bucket_entry.entry.key);
179                    self.monitor.notify(DiscoveryKind::EntryEvicted(pid)).await;
180                    continue;
181                }
182                tasks.push(self.clone().refresh_entry(bucket_entry.clone()))
183            }
184            while let Some(ok) = tasks.next().await {
185                if ok {
186                    succeeded += 1;
187                }
188            }
189        }
190        succeeded
191    }
192
193    /// Refresh a specific entry by pinging it over UDP. Returns true
194    /// if the ping succeeded.
195    async fn refresh_entry(self: Arc<Self>, bucket_entry: BucketEntry) -> bool {
196        let key = &bucket_entry.entry.key;
197        match self.connect(&bucket_entry.entry).await {
198            Ok(_) => {
199                self.table.update_entry(key, PENDING_ENTRY);
200                true
201            }
202            Err(err) => {
203                trace!("Failed to refresh entry {key:?}: {err}");
204                if bucket_entry.failures >= MAX_FAILURES {
205                    let pid = (*key).into();
206                    self.table.remove_entry(key);
207                    self.monitor.notify(DiscoveryKind::EntryEvicted(pid)).await;
208                    return false;
209                }
210                self.table.update_entry(key, UNREACHABLE_ENTRY);
211                false
212            }
213        }
214    }
215
216    /// Send a ping over UDP with retries.
217    async fn connect(&self, entry: &Entry) -> Result<()> {
218        let mut retry = 0;
219        let supported = [Protocol::Udp];
220        let endpoint = pick_endpoint(&entry.discovery_addrs, &supported)
221            .ok_or(Error::Lookup("No UDP discovery address available".into()))?;
222
223        if !self.config.access_control.allow(
224            &Subject::Endpoint(&endpoint),
225            Action::Probe(ConnDirection::Outbound),
226        ) {
227            return Err(Error::AccessDenied);
228        }
229
230        let conn = udp::dial(&endpoint, Default::default()).await?;
231        let peer_addr = SocketAddr::try_from(endpoint.clone())?;
232        let backoff = Backoff::new(100, 5000);
233        while retry < self.config.refresh_connect_retries {
234            match self.send_ping_msg(&conn, peer_addr).await {
235                Ok(()) => return Ok(()),
236                Err(Error::Timeout) => {
237                    retry += 1;
238                    backoff.sleep().await;
239                }
240                Err(err) => {
241                    return Err(err);
242                }
243            }
244        }
245        Err(Error::Timeout)
246    }
247
248    /// Listen on UDP for Ping messages.
249    async fn listen_loop(self: Arc<Self>, endpoint: Endpoint) -> Result<()> {
250        let conn = match udp::listen(&endpoint, Default::default()).await {
251            Ok(c) => {
252                self.monitor
253                    .notify(ConnectionKind::Listening(endpoint.clone()))
254                    .await;
255                c
256            }
257            Err(err) => {
258                self.monitor
259                    .notify(ConnectionKind::ListenFailed(endpoint.clone()))
260                    .await;
261                return Err(err.into());
262            }
263        };
264        info!("Start listening on {endpoint}");
265
266        // Per-IP rate limit. Lives on the listen task, no lock needed:
267        // listen_to_ping_msg is only ever called from this loop.
268        let mut rate_limiter: HashMap<IpAddr, RateBucket> = HashMap::new();
269
270        loop {
271            let res = self.listen_to_ping_msg(&conn, &mut rate_limiter).await;
272            if let Err(err) = res {
273                trace!("Failed to handle ping msg {err}");
274                self.monitor.notify(ConnectionKind::AcceptFailed).await;
275            }
276        }
277    }
278
279    /// Listen for a Ping message and respond with a Pong message.
280    /// Drops pings from sources not in the routing table or that
281    /// exceed the per-IP rate limit.
282    async fn listen_to_ping_msg(
283        &self,
284        conn: &udp::UdpConn,
285        rate_limiter: &mut HashMap<IpAddr, RateBucket>,
286    ) -> Result<()> {
287        let mut buf = vec![0u8; MAX_UDP_BUF];
288        let (n, sender) = conn.recv_from(&mut buf).await?;
289
290        let sender_ip = sender.ip();
291        if !self.table.has_discovery_ip(&sender_ip) {
292            trace!("Drop refresh ping from unknown source {sender_ip}");
293            return Ok(());
294        }
295
296        let allowed = rate_limiter
297            .entry(sender_ip)
298            .or_insert_with(RateBucket::new)
299            .allow();
300        if !allowed {
301            trace!("Drop rate-limited refresh ping from {sender_ip}");
302            return Ok(());
303        }
304
305        let sender_ep = Endpoint::new_udp_addr(sender);
306
307        if !self.config.access_control.allow(
308            &Subject::Endpoint(&sender_ep),
309            Action::Probe(ConnDirection::Inbound),
310        ) {
311            trace!("Drop refresh ping from denied source {sender_ip}");
312            return Ok(());
313        }
314
315        self.monitor
316            .notify(ConnectionKind::Accepted(sender_ep.clone()))
317            .await;
318
319        let (msg, _) = decode::<RefreshMsg>(&buf[..n])?;
320        match msg {
321            RefreshMsg::Ping(m) => {
322                let pong_msg = RefreshMsg::Pong(m);
323                let encoded = encode(&pong_msg)?;
324                conn.send_to(&encoded, sender).await?;
325            }
326            RefreshMsg::Pong(_) => return Err(Error::InvalidMsg("Unexpected pong msg".into())),
327        }
328
329        self.monitor
330            .notify(ConnectionKind::Disconnected(sender_ep))
331            .await;
332        Ok(())
333    }
334
335    /// Sends a Ping msg and waits for the Pong response.
336    async fn send_ping_msg(&self, conn: &udp::UdpConn, peer_addr: SocketAddr) -> Result<()> {
337        let mut nonce: [u8; 32] = [0; 32];
338        OsRng.try_fill_bytes(&mut nonce)?;
339
340        let ping = RefreshMsg::Ping(nonce);
341        let encoded = encode(&ping)?;
342        conn.send_to(&encoded, peer_addr).await?;
343
344        let t = Duration::from_secs(self.config.refresh_response_timeout);
345        let mut buf = vec![0u8; MAX_UDP_BUF];
346        let (n, _) = timeout(t, conn.recv_from(&mut buf)).await??;
347        let (msg, _) = decode::<RefreshMsg>(&buf[..n])?;
348
349        match msg {
350            RefreshMsg::Pong(n) => {
351                if n != nonce {
352                    return Err(Error::InvalidPongMsg);
353                }
354                Ok(())
355            }
356            _ => Err(Error::InvalidMsg("Unexpected ping msg".into())),
357        }
358    }
359}
karyon_p2p/discovery/kademlia/refresh.rs

karyon_p2p/discovery/kademlia/
refresh.rs
