opentracker

opentracker/crates/otserver/src/udp.rs
udp.rs
⎇
Raw
1//! The UDP receive loop.
2//!
3//! `SO_REUSEPORT` lets several tasks bind the *same* port and have the kernel
4//! spread datagrams across them, which is how one port scales past one core
5//! without a dispatcher task in the middle. Each worker owns its own socket, so
6//! there is no shared state on the hot path at all.
7//!
8//! The loop deliberately contains no `.await` between receiving and replying: the
9//! whole of [`Ingest::handle`] is synchronous, and a full writer channel is
10//! answered rather than waited on. Applying backpressure here would turn one slow
11//! disk into packet loss for every device at once.
12
13use std::net::SocketAddr;
14use std::sync::Arc;
15
16use anyhow::{Context, Result};
17use socket2::{Domain, Protocol, Socket, Type};
18use tokio::net::UdpSocket;
19use tracing::{error, info, warn};
20
21use crate::db::now;
22use crate::ingest::{Ingest, Peer, Transport};
23
24/// Receive buffer per socket. A queue flush from every device at once arrives as
25/// a burst, and the kernel's default (a few hundred kB) drops it on the floor
26/// before this process ever sees it.
27const RECV_BUFFER_BYTES: usize = 2 * 1024 * 1024;
28
29/// One more than the largest datagram, so an oversized packet is *seen* to be
30/// oversized instead of being silently truncated into something that might parse.
31const READ_BUFFER: usize = otproto::MAX_DATAGRAM + 1;
32
33fn bind_reuseport(addr: SocketAddr) -> Result<UdpSocket> {
34 let domain = if addr.is_ipv6() {
35 Domain::IPV6
36 } else {
37 Domain::IPV4
38 };
39 let socket =
40 Socket::new(domain, Type::DGRAM, Some(Protocol::UDP)).context("creating socket")?;
41 socket.set_reuse_address(true).context("SO_REUSEADDR")?;
42 socket.set_reuse_port(true).context("SO_REUSEPORT")?;
43
44 // Best effort: on Linux the kernel doubles the requested value and caps it at
45 // net.core.rmem_max, so a smaller buffer than asked for is normal and not
46 // worth failing startup over.
47 if let Err(e) = socket.set_recv_buffer_size(RECV_BUFFER_BYTES) {
48 warn!(error = %e, "could not enlarge the UDP receive buffer; bursts may be dropped");
49 }
50
51 socket.set_nonblocking(true).context("set_nonblocking")?;
52 socket
53 .bind(&addr.into())
54 .with_context(|| format!("binding {addr}/udp"))?;
55 UdpSocket::from_std(socket.into()).context("handing the socket to tokio")
56}
57
58/// Spawn `workers` receive tasks on `addr`.
59pub fn spawn(
60 addr: SocketAddr,
61 workers: usize,
62 ingest: Arc<Ingest>,
63) -> Result<Vec<tokio::task::JoinHandle<()>>> {
64 let mut tasks = Vec::with_capacity(workers);
65 for id in 0..workers {
66 let socket = bind_reuseport(addr)?;
67 tasks.push(tokio::spawn(run(id, socket, Arc::clone(&ingest))));
68 }
69 info!(%addr, workers, "OTP/1 UDP listener started");
70 // The trap worth stating in the log, because operators hit it on day one and
71 // the symptom (everything falls back to TLS) is far from the cause.
72 info!("reminder: HTTP reverse proxies do not forward UDP — {addr} needs its own firewall rule");
73 Ok(tasks)
74}
75
76async fn run(id: usize, socket: UdpSocket, ingest: Arc<Ingest>) {
77 let mut buf = vec![0u8; READ_BUFFER];
78 loop {
79 let (len, peer_addr) = match socket.recv_from(&mut buf).await {
80 Ok(v) => v,
81 Err(e) => {
82 // On UDP a send error can surface here as ICMP-driven
83 // ECONNREFUSED for a *previous* send. It says nothing about the
84 // socket's health, so log and keep going rather than exiting the
85 // worker and silently losing a quarter of the capacity.
86 warn!(worker = id, error = %e, "recv_from failed");
87 continue;
88 }
89 };
90
91 let peer = Peer {
92 addr: peer_addr,
93 transport: Transport::Udp,
94 };
95 if let Some(reply) = ingest.handle(&buf[..len], peer, now())
96 && let Err(e) = socket.send_to(&reply, peer_addr).await
97 {
98 // The phone will retry; there is nothing to recover here.
99 error!(worker = id, %peer_addr, error = %e, "sending reply failed");
100 }
101 }
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107 use otproto::msg::Direction;
108 use otproto::{Key, Message, Point, kdf};
109 use std::time::Duration;
110
111 const TOKEN_ID: u64 = 0xABCD_0123_4567_89EF;
112 const TOKEN_KEY: Key = [0x33; 32];
113
114 /// Round trip a real datagram over a real loopback socket. This is the only
115 /// test that exercises the socket options and the reply path together.
116 #[tokio::test]
117 async fn a_datagram_over_loopback_is_acked() {
118 let dir = tempfile::tempdir().expect("temp dir");
119 let db = crate::db::Db::open(&dir.path().join("t.db"))
120 .await
121 .expect("open");
122 sqlx::query(
123 "INSERT INTO users (id, username, pw_hash, display_name, created_at, pw_changed_at) \
124 VALUES (1, 'a', 'x', 'A', 0, 0)",
125 )
126 .execute(&db.write)
127 .await
128 .expect("user");
129
130 let (writer, _task) = crate::writer::spawn(db.write.clone());
131 let ingest = Arc::new(Ingest::new(writer, 30 * 86_400, None));
132 ingest.insert_token(crate::ingest::TokenSlot::new(TOKEN_ID, 1, &TOKEN_KEY, 1));
133
134 // Port 0 lets the OS choose; then read it back, because SO_REUSEPORT
135 // workers must all bind the *same* concrete port.
136 let probe = bind_reuseport("127.0.0.1:0".parse().expect("literal")).expect("bind");
137 let addr = probe.local_addr().expect("local addr");
138 drop(probe);
139
140 let _tasks = spawn(addr, 2, Arc::clone(&ingest)).expect("spawn");
141
142 let client = UdpSocket::bind("127.0.0.1:0").await.expect("client bind");
143 let k_up = kdf::derive(&TOKEN_KEY, Direction::Up);
144 let msg = Message::Loc(vec![Point {
145 acc_dm: Some(50),
146 ..Point::new(now() as u32, 525_200_080, 134_050_000)
147 }]);
148 let datagram = otproto::seal_message(&k_up, TOKEN_ID, [0x77; 12], &msg);
149 client.send_to(&datagram, addr).await.expect("send");
150
151 let mut buf = vec![0u8; READ_BUFFER];
152 let len = tokio::time::timeout(Duration::from_secs(2), client.recv(&mut buf))
153 .await
154 .expect("no reply within 2s")
155 .expect("recv");
156
157 let k_down = kdf::derive(&TOKEN_KEY, Direction::Down);
158 match otproto::open_message(&k_down, &buf[..len])
159 .expect("ack opens")
160 .1
161 {
162 Message::Ack(ack) => assert_eq!(ack.nonces, vec![[0x77u8; 12]]),
163 other => panic!("expected an ACK, got {other:?}"),
164 }
165 assert!(len <= datagram.len(), "the reply amplified the request");
166 }
167
168 #[tokio::test]
169 async fn several_workers_can_share_one_port() {
170 let probe = bind_reuseport("127.0.0.1:0".parse().expect("literal")).expect("first bind");
171 let addr = probe.local_addr().expect("local addr");
172 // The second bind on the same concrete port is the thing SO_REUSEPORT
173 // makes legal, and the thing the whole multi-worker design rests on.
174 let second = bind_reuseport(addr).expect("SO_REUSEPORT should allow a second bind");
175 assert_eq!(second.local_addr().expect("addr"), addr);
176 }
177}
178