1use crate::session::{ControlReader, ControlWriter, PathStats, SendError, Session};
22use bytes::{Bytes, BytesMut};
23use futures_util::{SinkExt, StreamExt};
24use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
25use std::sync::{Arc, Mutex};
26use std::time::{Duration, Instant};
27use tokio::io::{AsyncReadExt, AsyncWriteExt};
28use tokio::sync::{mpsc, oneshot, watch};
29use tokio_tungstenite::tungstenite::Message;
30
31pub const KIND_DATAGRAM: u8 = 1;
32pub const KIND_CONTROL: u8 = 2;
33
34const ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
36const MAX_DATAGRAM: usize = 60_000;
39
40enum Out {
41 Frame(Bytes, usize),
44 Close(String),
45}
46
47#[derive(Default)]
48struct RttSample {
49 rtt_ms: f64,
50 min_rtt_ms: f64,
51}
52
53struct Rate {
54 at: Instant,
55 bytes: u64,
56}
57
58pub struct Relay {
59 out: mpsc::UnboundedSender<Out>,
60 queued: Arc<AtomicUsize>,
62 capacity: usize,
63 bytes_tx: Arc<AtomicU64>,
64 rtt: Arc<Mutex<RttSample>>,
65 rate: Mutex<Rate>,
66 closed: watch::Receiver<bool>,
67}
68
69impl Relay {
70 pub fn max_datagram_size(&self) -> usize {
71 MAX_DATAGRAM
72 }
73
74 pub fn send_buffer_space(&self) -> usize {
75 self.capacity
76 .saturating_sub(self.queued.load(Ordering::Relaxed))
77 }
78
79 pub fn send_buffer_queued(&self) -> usize {
80 self.queued.load(Ordering::Relaxed)
81 }
82
83 pub fn send_datagram(&self, payload: Bytes) -> Result<(), SendError> {
84 if *self.closed.borrow() {
85 return Err(SendError::Closed);
86 }
87 if payload.len() > MAX_DATAGRAM {
88 return Err(SendError::TooLarge);
89 }
90 let len = payload.len();
93 let queued = self.queued.fetch_add(len, Ordering::Relaxed);
94 if queued + len > self.capacity {
95 self.queued.fetch_sub(len, Ordering::Relaxed);
96 return Err(SendError::Blocked);
97 }
98 let mut frame = BytesMut::with_capacity(len + 1);
99 frame.extend_from_slice(&[KIND_DATAGRAM]);
100 frame.extend_from_slice(&payload);
101 if self.out.send(Out::Frame(frame.freeze(), len)).is_err() {
102 self.queued.fetch_sub(len, Ordering::Relaxed);
103 return Err(SendError::Closed);
104 }
105 Ok(())
106 }
107
108 pub fn stats(&self) -> PathStats {
109 let rtt = self.rtt.lock().unwrap();
110 let now = Instant::now();
111 let mut rate = self.rate.lock().unwrap();
112 let bytes = self.bytes_tx.load(Ordering::Relaxed);
113 let dt = now.duration_since(rate.at).as_secs_f64();
114 let delivery_kbps = if dt > 0.05 {
115 ((bytes.saturating_sub(rate.bytes)) as f64 * 8.0 / dt / 1000.0) as u64
116 } else {
117 0
118 };
119 if dt > 0.05 {
120 rate.at = now;
121 rate.bytes = bytes;
122 }
123 PathStats {
124 rtt_ms: rtt.rtt_ms,
125 min_rtt_ms: rtt.min_rtt_ms,
126 cwnd: 0,
127 delivery_kbps,
128 lost_packets: 0,
129 current_mtu: 0,
130 }
131 }
132
133 pub fn close(&self, reason: &str) {
134 let _ = self.out.send(Out::Close(reason.to_string()));
135 }
136
137 pub async fn closed(&self) {
138 let mut rx = self.closed.clone();
139 while !*rx.borrow() {
140 if rx.changed().await.is_err() {
141 return;
142 }
143 }
144 }
145}
146
147pub(crate) async fn connect(
150 id: u64,
151 url: &str,
152 session_id: &str,
153 token: &str,
154 capacity: usize,
155) -> anyhow::Result<Session> {
156 let (ws, _) = tokio_tungstenite::connect_async(url).await?;
157 let (mut sink, mut stream) = ws.split();
158 let attach = serde_json::json!({
159 "type": "relay-attach", "session_id": session_id, "token": token, "party": "robot",
160 });
161 sink.send(Message::Text(attach.to_string().into())).await?;
162 let reply = tokio::time::timeout(ATTACH_TIMEOUT, async {
163 loop {
164 match stream.next().await {
165 Some(Ok(Message::Text(t))) => {
166 let v: serde_json::Value = serde_json::from_str(&t)?;
167 return Ok::<_, anyhow::Error>(v);
168 }
169 Some(Ok(Message::Close(_))) | None => {
170 anyhow::bail!("relay closed before attaching")
171 }
172 Some(Ok(_)) => continue,
173 Some(Err(e)) => return Err(e.into()),
174 }
175 }
176 })
177 .await
178 .map_err(|_| anyhow::anyhow!("relay attach timed out"))??;
179 match reply.get("type").and_then(|t| t.as_str()) {
180 Some("relay-attached") => {}
181 Some("denied") => anyhow::bail!(
182 "relay denied: {}",
183 reply
184 .get("reason")
185 .and_then(|r| r.as_str())
186 .unwrap_or("unknown")
187 ),
188 other => anyhow::bail!("unexpected relay reply {other:?}"),
189 }
190
191 let (out_tx, mut out_rx) = mpsc::unbounded_channel::<Out>();
192 let (dg_tx, dg_rx) = mpsc::channel::<Bytes>(1024);
193 let (ctl_tx, ctl_rx) = oneshot::channel::<(ControlReader, ControlWriter)>();
194 let (closed_tx, closed_rx) = watch::channel(false);
195 let queued = Arc::new(AtomicUsize::new(0));
196 let bytes_tx = Arc::new(AtomicU64::new(0));
197 let rtt = Arc::new(Mutex::new(RttSample::default()));
198
199 let (ctl_in_engine, mut ctl_in_task) = tokio::io::duplex(64 * 1024);
203 let (ctl_out_engine, mut ctl_out_task) = tokio::io::duplex(64 * 1024);
204 let reader: ControlReader = tokio::io::BufReader::new(Box::new(ctl_in_engine));
205 let writer: ControlWriter = Box::new(ctl_out_engine);
206 let _ = ctl_tx.send((reader, writer));
207
208 let relay = Relay {
209 out: out_tx,
210 queued: queued.clone(),
211 capacity,
212 bytes_tx: bytes_tx.clone(),
213 rtt: rtt.clone(),
214 rate: Mutex::new(Rate {
215 at: Instant::now(),
216 bytes: 0,
217 }),
218 closed: closed_rx,
219 };
220
221 tokio::spawn(async move {
222 let mut ping = tokio::time::interval(Duration::from_secs(1));
223 ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
224 let epoch = Instant::now();
225 let mut ctl_buf = vec![0u8; 16 * 1024];
226 let mut dropped = 0u64;
227 let reason: String = loop {
228 tokio::select! {
229 out = out_rx.recv() => {
230 match out {
231 Some(Out::Frame(frame, release)) => {
232 let len = frame.len() as u64;
233 let res = sink.send(Message::Binary(frame)).await;
234 if release > 0 {
235 queued.fetch_sub(release, Ordering::Relaxed);
236 }
237 if res.is_err() {
238 break "write failed".to_string();
239 }
240 bytes_tx.fetch_add(len, Ordering::Relaxed);
241 }
242 Some(Out::Close(reason)) => {
243 let _ = sink.send(Message::Close(None)).await;
244 break reason;
245 }
246 None => break "session dropped".to_string(),
247 }
248 }
249 n = ctl_out_task.read(&mut ctl_buf) => {
250 match n {
251 Ok(0) | Err(_) => break "control writer closed".to_string(),
252 Ok(n) => {
253 let mut frame = BytesMut::with_capacity(n + 1);
254 frame.extend_from_slice(&[KIND_CONTROL]);
255 frame.extend_from_slice(&ctl_buf[..n]);
256 if sink.send(Message::Binary(frame.freeze())).await.is_err() {
257 break "write failed".to_string();
258 }
259 }
260 }
261 }
262 msg = stream.next() => {
263 match msg {
264 Some(Ok(Message::Binary(b))) => {
265 match b.first().copied() {
266 Some(KIND_DATAGRAM) => {
267 if dg_tx.try_send(b.slice(1..)).is_err() {
268 dropped += 1;
269 }
270 }
271 Some(KIND_CONTROL) => {
272 let ok = ctl_in_task.write_all(&b[1..]).await.is_ok();
273 if !ok {
274 break "control reader closed".to_string();
275 }
276 }
277 _ => {}
278 }
279 }
280 Some(Ok(Message::Pong(p))) => {
281 if p.len() == 8 {
282 let sent = u64::from_be_bytes(p[..8].try_into().unwrap());
283 let now = epoch.elapsed().as_micros() as u64;
284 let ms = now.saturating_sub(sent) as f64 / 1000.0;
285 let mut r = rtt.lock().unwrap();
286 r.rtt_ms = ms;
287 if r.min_rtt_ms == 0.0 || ms < r.min_rtt_ms {
288 r.min_rtt_ms = ms;
289 }
290 }
291 }
292 Some(Ok(Message::Text(t))) => {
293 tracing::debug!(session = id, msg = %t, "relay text");
295 if t.contains("relay-closed") || t.contains("denied") {
296 break "relay closed".to_string();
297 }
298 }
299 Some(Ok(Message::Close(_))) | None => break "relay socket closed".to_string(),
300 Some(Ok(_)) => {}
301 Some(Err(e)) => break format!("relay error: {e}"),
302 }
303 }
304 _ = ping.tick() => {
305 let ts = (epoch.elapsed().as_micros() as u64).to_be_bytes();
306 if sink.send(Message::Ping(Bytes::copy_from_slice(&ts))).await.is_err() {
307 break "ping failed".to_string();
308 }
309 }
310 }
311 };
312 if dropped > 0 {
313 tracing::debug!(
314 session = id,
315 dropped,
316 "relayed datagrams dropped (receiver slow)"
317 );
318 }
319 tracing::info!(session = id, %reason, "relay session ended");
320 let _ = closed_tx.send(true);
321 });
322
323 Ok(Session::from_relay(id, relay, dg_rx, ctl_rx))
324}