atagen

vineyard

public
Star0
Watch0
25d14e49ab6861630869e2129d40652900bc5833
1
ssh://brewery@pls.lobotomise.me:33/atagen/vineyard.git
src/federation.rsRust650 lines21.8 KB
1
use std::collections::{BTreeMap, HashMap, HashSet};
2
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
3
use std::sync::Arc;
4
use std::time::{Duration, SystemTime, UNIX_EPOCH};
5
6
use anyhow::{Context as _, Result, bail};
7
use futures_util::{SinkExt as _, StreamExt as _};
8
use nostr::event::{Event, Kind};
9
use nostr::filter::{Filter, MatchEventOptions};
10
use nostr::key::PublicKey;
11
use nostr::message::MachineReadablePrefix;
12
use nostr_database::{NostrDatabase as _, RejectedReason, SaveEventStatus};
13
use serde_json::{Value, json};
14
use tokio::net::{TcpStream, lookup_host};
15
use tokio::time::{Instant, sleep, timeout, timeout_at};
16
use tokio_tungstenite::tungstenite::{Message, protocol::WebSocketConfig};
17
use tokio_tungstenite::{
18
    MaybeTlsStream, WebSocketStream, client_async_tls_with_config, connect_async_with_config,
19
};
20
use tokio_util::sync::CancellationToken;
21
use tokio_util::task::TaskTracker;
22
23
use crate::config::{Config, Peer, relay_url};
24
use crate::policy::{Policy, WritePolicyResult, coordinate};
25
use crate::relay::Relay;
26
use crate::store::Store;
27
28
const KINDS: [u16; 10] = [5, 1617, 1621, 1630, 1631, 1632, 10002, 30382, 30617, 30618];
29
const OVERLAP: u64 = 7 * 24 * 60 * 60;
30
const RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);
31
type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
32
33
#[derive(Clone)]
34
struct Link {
35
    peer: Peer,
36
    config: Arc<Config>,
37
    store: Arc<Store>,
38
    relay: Relay,
39
    policy: Policy,
40
    discovered: bool,
41
    stop: CancellationToken,
42
}
43
44
struct Connection {
45
    socket: Socket,
46
    stop: CancellationToken,
47
    write_timeout: Duration,
48
    request_sequence: u64,
49
}
50
51
pub fn spawn(
52
    config: Arc<Config>,
53
    store: Arc<Store>,
54
    relay: Relay,
55
    policy: Policy,
56
    task_tracker: &TaskTracker,
57
    stop: CancellationToken,
58
) {
59
    let tasks = task_tracker.clone();
60
    task_tracker
61
        .clone()
62
        .spawn(manage(config, store, relay, policy, tasks, stop));
63
}
64
65
async fn manage(
66
    config: Arc<Config>,
67
    store: Arc<Store>,
68
    relay: Relay,
69
    policy: Policy,
70
    tasks: TaskTracker,
71
    stop: CancellationToken,
72
) {
73
    let mut active = HashMap::<String, CancellationToken>::new();
74
    loop {
75
        let mut desired = BTreeMap::new();
76
        for configured_peer in &config.peers {
77
            let mut normalized_peer = configured_peer.clone();
78
            let Some(url) = relay_url(&normalized_peer.url) else {
79
                continue;
80
            };
81
            normalized_peer.url = url.to_string();
82
            desired.insert(normalized_peer.url.clone(), (normalized_peer, false));
83
        }
84
        match store
85
            .journal
86
            .discovered(config.limits.max_discovered_peers)
87
            .await
88
        {
89
            Ok(urls) => {
90
                for raw in urls {
91
                    if let Some(peer) = discovered_peer(&config, &raw) {
92
                        desired.entry(peer.url.clone()).or_insert((peer, true));
93
                    }
94
                }
95
            }
96
            Err(error) => tracing::error!(%error, "loading discovered peers failed"),
97
        }
98
99
        active.retain(|url, peer_stop| {
100
            if desired.contains_key(url) {
101
                true
102
            } else {
103
                peer_stop.cancel();
104
                false
105
            }
106
        });
107
        for (url, (peer, discovered)) in desired {
108
            if active.contains_key(&url) {
109
                continue;
110
            }
111
            let peer_stop = stop.child_token();
112
            let link = Link {
113
                peer,
114
                config: Arc::clone(&config),
115
                store: Arc::clone(&store),
116
                relay: relay.clone(),
117
                policy: policy.clone(),
118
                discovered,
119
                stop: peer_stop.clone(),
120
            };
121
            tasks.spawn(run(link.clone(), true));
122
            tasks.spawn(run(link, false));
123
            active.insert(url, peer_stop);
124
        }
125
126
        tokio::select! {
127
            () = stop.cancelled() => break,
128
            () = sleep(Duration::from_secs(10)) => {}
129
        }
130
    }
131
    for peer_stop in active.into_values() {
132
        peer_stop.cancel();
133
    }
134
}
135
136
fn discovered_peer(config: &Config, raw: &str) -> Option<Peer> {
137
    let url = relay_url(raw)?;
138
    if (!config.allow_private_discovery && url.scheme() != "wss")
139
        || url == relay_url(&config.public_url)?
140
    {
141
        return None;
142
    }
143
    Some(Peer {
144
        url: url.to_string(),
145
        query_limit: config.limits.max_query_results,
146
        rescan_from: None,
147
    })
148
}
149
150
async fn run(link: Link, outbound: bool) {
151
    let direction = if outbound { "outbound" } else { "inbound" };
152
    let mut backoff = 5;
153
    while !link.stop.is_cancelled() {
154
        let started = Instant::now();
155
        let result = if outbound {
156
            publish(&link).await
157
        } else {
158
            subscribe(&link).await
159
        };
160
        if link.stop.is_cancelled() {
161
            break;
162
        }
163
        if let Err(error) = result {
164
            tracing::warn!(peer = %link.peer.url, %direction, %error, "peer connection interrupted");
165
        }
166
        if started.elapsed() > Duration::from_secs(60) {
167
            backoff = 5;
168
        }
169
        tokio::select! {
170
            () = link.stop.cancelled() => break,
171
            () = sleep(Duration::from_secs(backoff)) => {},
172
        }
173
        backoff = (backoff * 2).min(60);
174
    }
175
}
176
177
impl Connection {
178
    async fn open(link: &Link) -> Result<Self> {
179
        let config = WebSocketConfig::default()
180
            .max_message_size(Some(link.config.limits.max_message_bytes))
181
            .max_frame_size(Some(link.config.limits.max_message_bytes));
182
        let connected = async {
183
            if link.discovered {
184
                connect_discovered(link, config).await
185
            } else {
186
                connect_async_with_config(&link.peer.url, Some(config), false)
187
                    .await
188
                    .map(|(socket, _)| socket)
189
                    .map_err(Into::into)
190
            }
191
        };
192
        let socket = tokio::select! {
193
            () = link.stop.cancelled() => bail!("stopping"),
194
            result = timeout(RESPONSE_TIMEOUT, connected) => {
195
                result.context("peer connection timed out")??
196
            }
197
        };
198
        Ok(Self {
199
            socket,
200
            stop: link.stop.clone(),
201
            write_timeout: Duration::from_secs(link.config.limits.write_timeout_seconds),
202
            request_sequence: 0,
203
        })
204
    }
205
206
    async fn send(&mut self, value: Value) -> Result<()> {
207
        self.send_message(Message::Text(value.to_string().into()))
208
            .await
209
    }
210
211
    async fn send_message(&mut self, message: Message) -> Result<()> {
212
        tokio::select! {
213
            () = self.stop.cancelled() => bail!("stopping"),
214
            result = timeout(self.write_timeout, self.socket.send(message)) => {
215
                result.context("peer write timed out")??;
216
                Ok(())
217
            }
218
        }
219
    }
220
221
    async fn next(&mut self, deadline: Instant) -> Result<Option<Value>> {
222
        loop {
223
            let message = tokio::select! {
224
                () = self.stop.cancelled() => bail!("stopping"),
225
                result = timeout_at(deadline, self.socket.next()) => match result {
226
                    Err(_) => return Ok(None),
227
                    Ok(Some(message)) => message?,
228
                    Ok(None) => bail!("peer disconnected"),
229
                }
230
            };
231
            match message {
232
                Message::Text(text) => return Ok(Some(serde_json::from_str(&text)?)),
233
                Message::Ping(payload) => self.send_message(Message::Pong(payload)).await?,
234
                Message::Close(_) => bail!("peer closed the connection"),
235
                Message::Binary(_) | Message::Pong(_) | Message::Frame(_) => {}
236
            }
237
        }
238
    }
239
}
240
241
async fn connect_discovered(link: &Link, config: WebSocketConfig) -> Result<Socket> {
242
    let url = url::Url::parse(&link.peer.url)?;
243
    let host = url.host_str().context("discovered peer has no host")?;
244
    let port = url
245
        .port_or_known_default()
246
        .context("discovered peer has no port")?;
247
    let mut last_error = None;
248
    for address in lookup_host((host, port)).await? {
249
        if !link.config.allow_private_discovery && !public_ip(address.ip()) {
250
            continue;
251
        }
252
        match TcpStream::connect(address).await {
253
            Ok(stream) => {
254
                return client_async_tls_with_config(&link.peer.url, stream, Some(config), None)
255
                    .await
256
                    .map(|(socket, _)| socket)
257
                    .map_err(Into::into);
258
            }
259
            Err(error) => last_error = Some(error),
260
        }
261
    }
262
    match last_error {
263
        Some(error) => Err(error.into()),
264
        None => bail!("discovered peer did not resolve to an allowed address"),
265
    }
266
}
267
268
fn public_ip(address: IpAddr) -> bool {
269
    match address {
270
        IpAddr::V4(ipv4) => public_ipv4(ipv4),
271
        IpAddr::V6(ipv6) => public_ipv6(ipv6),
272
    }
273
}
274
275
fn public_ipv4(ip: Ipv4Addr) -> bool {
276
    let [first, second, _, _] = ip.octets();
277
    !(ip.is_private()
278
        || ip.is_loopback()
279
        || ip.is_link_local()
280
        || ip.is_broadcast()
281
        || ip.is_documentation()
282
        || ip.is_unspecified()
283
        || ip.is_multicast()
284
        || first == 0
285
        || (first == 100 && (64..=127).contains(&second))
286
        || (first == 198 && matches!(second, 18 | 19)))
287
}
288
289
fn public_ipv6(ip: Ipv6Addr) -> bool {
290
    if let Some(ipv4) = ip.to_ipv4_mapped() {
291
        return public_ipv4(ipv4);
292
    }
293
    let segments = ip.segments();
294
    !(ip.is_loopback()
295
        || ip.is_unspecified()
296
        || ip.is_unique_local()
297
        || ip.is_unicast_link_local()
298
        || ip.is_multicast()
299
        || (segments[0] == 0x2001 && segments[1] == 0x0db8))
300
}
301
302
async fn publish(link: &Link) -> Result<()> {
303
    let mut connection = Connection::open(link).await?;
304
    let mut announced = HashSet::new();
305
    loop {
306
        let pending = link.store.journal.pending(&link.peer.url, 32).await?;
307
        if pending.is_empty() {
308
            connection
309
                .next(Instant::now() + Duration::from_secs(1))
310
                .await?;
311
            continue;
312
        }
313
        for (sequence, id) in pending {
314
            if announced.len() >= link.config.limits.max_repositories {
315
                announced.clear();
316
            }
317
            if let Some(event) = link.store.event_by_id(&id).await? {
318
                let repository_event = repository_coordinate(&event).is_some();
319
                if let Some(announcement) = announcement_for(link, &event).await? {
320
                    if announced.insert(announcement.id) {
321
                        if let Some(error) = publish_event(&mut connection, &announcement).await? {
322
                            announced.remove(&announcement.id);
323
                            defer(link, sequence, &error).await?;
324
                            continue;
325
                        }
326
                    }
327
                } else if repository_event {
328
                    link.store
329
                        .journal
330
                        .acknowledge(&link.peer.url, sequence)
331
                        .await?;
332
                    continue;
333
                }
334
                if let Some(error) = publish_event(&mut connection, &event).await? {
335
                    if link.store.event_by_id(&id).await?.is_some() {
336
                        defer(link, sequence, &error).await?;
337
                        continue;
338
                    }
339
                } else if event.kind == Kind::from(30617) {
340
                    announced.insert(event.id);
341
                }
342
            }
343
            link.store
344
                .journal
345
                .acknowledge(&link.peer.url, sequence)
346
                .await?;
347
        }
348
    }
349
}
350
351
async fn defer(link: &Link, sequence: u64, error: &str) -> Result<()> {
352
    let retry_at = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() + 300;
353
    link.store
354
        .journal
355
        .defer(&link.peer.url, sequence, error, retry_at)
356
        .await?;
357
    tracing::warn!(peer = %link.peer.url, sequence, %error, "peer delivery deferred");
358
    Ok(())
359
}
360
361
async fn publish_event(connection: &mut Connection, event: &Event) -> Result<Option<String>> {
362
    tokio::select! {
363
        () = connection.stop.cancelled() => bail!("stopping"),
364
        () = sleep(Duration::from_secs(1)) => {},
365
    }
366
    connection.send(json!(["EVENT", event])).await?;
367
    let id = event.id.to_hex();
368
    let deadline = Instant::now() + RESPONSE_TIMEOUT;
369
    loop {
370
        let frame = connection
371
            .next(deadline)
372
            .await?
373
            .context("peer acknowledgement timed out")?;
374
        if frame[0] == "OK" && frame[1].as_str() == Some(&id) {
375
            let reason = frame[3].as_str().unwrap_or("peer rejected event");
376
            return Ok(if frame[2] == true || reason.starts_with("duplicate:") {
377
                None
378
            } else {
379
                Some(reason.chars().take(512).collect())
380
            });
381
        }
382
    }
383
}
384
385
async fn announcement_for(link: &Link, event: &Event) -> Result<Option<Event>> {
386
    let Some((author, identifier)) = repository_coordinate(event) else {
387
        return Ok(None);
388
    };
389
    Ok(link
390
        .store
391
        .query(
392
            Filter::new()
393
                .kind(Kind::from(30617))
394
                .author(author)
395
                .identifier(identifier)
396
                .limit(1),
397
        )
398
        .await?
399
        .into_iter()
400
        .next())
401
}
402
403
fn repository_coordinate(event: &Event) -> Option<(PublicKey, String)> {
404
    if event.kind == Kind::from(30618) {
405
        event
406
            .tags
407
            .identifier()
408
            .map(|identifier| (event.pubkey, identifier))
409
    } else if matches!(event.kind.as_u16(), 1617 | 1621 | 1630..=1632) {
410
        event.tags.iter().find_map(|tag| {
411
            let values = tag.as_slice();
412
            if values.first().map(String::as_str) != Some("a") {
413
                return None;
414
            }
415
            let (kind, author, identifier) = coordinate(values.get(1)?)?;
416
            if kind != 30617 {
417
                return None;
418
            }
419
            Some((author, identifier.to_owned()))
420
        })
421
    } else {
422
        None
423
    }
424
}
425
426
async fn import(link: &Link, value: Value) -> Result<()> {
427
    let event: Event = serde_json::from_value(value)?;
428
    if event.verify().is_err() {
429
        tracing::debug!(peer = %link.peer.url, event = %event.id, "ignored invalid peer event");
430
        return Ok(());
431
    }
432
    match link.policy.admit_peer(&event).await {
433
        WritePolicyResult::Accept => {}
434
        WritePolicyResult::Respond { status: true, .. } => return Ok(()),
435
        WritePolicyResult::Respond {
436
            prefix: MachineReadablePrefix::Error | MachineReadablePrefix::RateLimited,
437
            message,
438
            ..
439
        } => bail!("peer import paused by {message}"),
440
        WritePolicyResult::Respond { message, .. } => {
441
            tracing::debug!(peer = %link.peer.url, event = %event.id, %message, "ignored out of scope peer event");
442
            return Ok(());
443
        }
444
    }
445
    match link.relay.add_event(event).await? {
446
        SaveEventStatus::Success
447
        | SaveEventStatus::Rejected(
448
            RejectedReason::Duplicate
449
            | RejectedReason::Deleted
450
            | RejectedReason::Expired
451
            | RejectedReason::Replaced
452
            | RejectedReason::InvalidDelete
453
            | RejectedReason::Vanished,
454
        ) => Ok(()),
455
        status @ SaveEventStatus::Rejected(_) => {
456
            bail!("peer event was not stored with status {status:?}")
457
        }
458
    }
459
}
460
461
async fn subscribe(link: &Link) -> Result<()> {
462
    let mut connection = Connection::open(link).await?;
463
    connection
464
        .send(json!(["REQ", "live", {"kinds": KINDS, "limit": 0_u8}]))
465
        .await?;
466
    let deadline = Instant::now() + RESPONSE_TIMEOUT;
467
    loop {
468
        let frame = connection
469
            .next(deadline)
470
            .await?
471
            .context("live subscription timed out")?;
472
        if frame[0] == "EOSE" && frame[1] == "live" {
473
            break;
474
        }
475
        handle_live(link, frame).await?;
476
    }
477
    let mut completed = HashSet::new();
478
    let mut next_replay = Instant::now();
479
    loop {
480
        let replay_all = Instant::now() >= next_replay;
481
        let until = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
482
        for (scope, filter) in [
483
            ("announcements".to_owned(), json!({"kinds": [30_617_u16]})),
484
            ("bindings".to_owned(), json!({"kinds": [30_382_u16]})),
485
            ("relays".to_owned(), json!({"kinds": [10_002_u16]})),
486
        ] {
487
            if replay_all || !completed.contains(&scope) {
488
                replay(link, &mut connection, &scope, filter, until).await?;
489
                completed.insert(scope);
490
            }
491
        }
492
        let mut active_scopes = HashSet::from([
493
            "announcements".to_owned(),
494
            "bindings".to_owned(),
495
            "relays".to_owned(),
496
        ]);
497
        for (author, identifier) in link
498
            .store
499
            .journal
500
            .repositories(link.config.limits.max_repositories)
501
            .await?
502
        {
503
            let address = format!("30617:{author}:{identifier}");
504
            for (scope, filter) in [
505
                (
506
                    format!("state:{address}"),
507
                    json!({"kinds": [30_618_u16], "authors": [author], "#d": [identifier]}),
508
                ),
509
                (
510
                    format!("contributions:{address}"),
511
                    json!({"kinds": [1_617_u16, 1_621_u16, 1_630_u16, 1_631_u16, 1_632_u16], "#a": [address]}),
512
                ),
513
            ] {
514
                active_scopes.insert(scope.clone());
515
                if replay_all || !completed.contains(&scope) {
516
                    replay(link, &mut connection, &scope, filter, until).await?;
517
                    completed.insert(scope);
518
                }
519
            }
520
        }
521
        completed.retain(|scope| active_scopes.contains(scope));
522
        if replay_all {
523
            replay(
524
                link,
525
                &mut connection,
526
                "deletions",
527
                json!({"kinds": [5_u16]}),
528
                until,
529
            )
530
            .await?;
531
            next_replay = Instant::now() + Duration::from_secs(300);
532
        }
533
        let refresh = Instant::now() + Duration::from_secs(30);
534
        while let Some(frame) = connection.next(refresh).await? {
535
            handle_live(link, frame).await?;
536
        }
537
    }
538
}
539
540
async fn handle_live(link: &Link, frame: Value) -> Result<()> {
541
    if frame[1] != "live" {
542
        return Ok(());
543
    }
544
    match frame[0].as_str() {
545
        Some("EVENT") => import(link, frame[2].clone()).await,
546
        Some("CLOSED") => bail!("peer closed live subscription with {}", frame[2]),
547
        _ => Ok(()),
548
    }
549
}
550
551
async fn replay(
552
    link: &Link,
553
    connection: &mut Connection,
554
    scope: &str,
555
    filter: Value,
556
    until: u64,
557
) -> Result<()> {
558
    let checkpoint = link
559
        .store
560
        .journal
561
        .checkpoint(&link.peer.url, scope)
562
        .await?
563
        .unwrap_or(0);
564
    let overlap_start = checkpoint.saturating_sub(OVERLAP);
565
    let replay_start = link
566
        .peer
567
        .rescan_from
568
        .map_or(overlap_start, |from| from.min(overlap_start));
569
    let mut ranges = vec![(replay_start.min(until), until)];
570
    while let Some((since, range_until)) = ranges.pop() {
571
        if replay_range(link, connection, filter.clone(), since, range_until).await? {
572
            link.store
573
                .journal
574
                .advance(&link.peer.url, scope, range_until)
575
                .await?;
576
        } else {
577
            if since == range_until {
578
                bail!(
579
                    "recovery required for {scope} at {since}, one second exceeds the peer history limit"
580
                );
581
            }
582
            let middle = since + (range_until - since) / 2;
583
            ranges.push((middle + 1, range_until));
584
            ranges.push((since, middle));
585
        }
586
    }
587
    Ok(())
588
}
589
590
async fn replay_range(
591
    link: &Link,
592
    connection: &mut Connection,
593
    mut filter: Value,
594
    since: u64,
595
    until: u64,
596
) -> Result<bool> {
597
    let pause = Instant::now() + Duration::from_secs(1);
598
    while let Some(frame) = connection.next(pause).await? {
599
        handle_live(link, frame).await?;
600
    }
601
    filter["since"] = json!(since);
602
    filter["until"] = json!(until);
603
    filter["limit"] = json!(link.peer.query_limit);
604
    let matching: Filter = serde_json::from_value(filter.clone())?;
605
    connection.request_sequence += 1;
606
    let subscription = format!("history-{}", connection.request_sequence);
607
    connection
608
        .send(json!(["REQ", subscription, filter]))
609
        .await?;
610
    let deadline = Instant::now() + RESPONSE_TIMEOUT;
611
    let mut received = HashSet::new();
612
    let complete = loop {
613
        let frame = connection
614
            .next(deadline)
615
            .await?
616
            .context("peer history timed out")?;
617
        if frame[1] != subscription {
618
            handle_live(link, frame).await?;
619
            continue;
620
        }
621
        match frame[0].as_str() {
622
            Some("EVENT") => {
623
                let event: Event = serde_json::from_value(frame[2].clone())?;
624
                if !matching.match_event(&event, MatchEventOptions::new()) {
625
                    bail!("peer returned an event outside the requested history");
626
                }
627
                received.insert(event.id);
628
                import(link, frame[2].clone()).await?;
629
                if received.len() >= link.peer.query_limit {
630
                    break false;
631
                }
632
            }
633
            Some("EOSE") => break true,
634
            Some("CLOSED") => {
635
                let message = frame[2].as_str().unwrap_or_default();
636
                if (message.starts_with("rate-limited:") || message.starts_with("error:"))
637
                    && (message.contains("byte limit")
638
                        || message.contains("result limit")
639
                        || message.contains("too many results"))
640
                {
641
                    break false;
642
                }
643
                bail!("peer closed history with {message}");
644
            }
645
            _ => {}
646
        }
647
    };
648
    connection.send(json!(["CLOSE", subscription])).await?;
649
    Ok(complete)
650
}