1use crate::auth_lys::session::LysSession;
4use crate::auth_lys::{AccessDecisionError, HotAction, SubjectBinding};
5use crate::link_handshake::MachineIdentity;
6use liminal::protocol::{Frame, FrameType, codec, validate_stream};
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeMap;
9use std::sync::{Arc, Mutex, mpsc};
10
11type Device = [u8; 32];
12type Pair = (Device, Device);
13const STREAM_BOUND: usize = 256;
14
15#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
17#[serde(deny_unknown_fields)]
18pub struct RelayConfig {
19 pub enabled: bool,
21}
22
23#[derive(Debug, thiserror::Error)]
25pub enum RelayRefusal {
26 #[error(transparent)]
28 Decision(#[from] AccessDecisionError),
29 #[error("{0}")]
31 Named(&'static str),
32}
33
34pub trait RelaySink: Send + Sync + std::fmt::Debug {
36 fn send(&self, frame: Frame) -> Result<(), RelayRefusal>;
41}
42
43impl RelaySink for mpsc::SyncSender<Frame> {
44 fn send(&self, frame: Frame) -> Result<(), RelayRefusal> {
45 self.try_send(frame).map_err(|error| match error {
46 mpsc::TrySendError::Full(_) => RelayRefusal::Named("relay_outbound_full"),
47 mpsc::TrySendError::Disconnected(_) => RelayRefusal::Named("peer_not_connected"),
48 })
49 }
50}
51
52#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)]
54pub struct RelayCounters {
55 pub frames_relayed: u64,
57 pub bytes_relayed: u64,
59 pub refused: u64,
61}
62
63#[derive(Debug)]
64struct Endpoint {
65 generation: u64,
66 sink: Arc<dyn RelaySink>,
67 streams: BTreeMap<Device, u32>,
68}
69
70#[derive(Debug, Default)]
71struct State {
72 generation: u64,
73 endpoints: BTreeMap<Device, Endpoint>,
74 pending: BTreeMap<Pair, u32>,
75 counters: BTreeMap<Pair, RelayCounters>,
76}
77
78#[derive(Debug)]
80pub struct Relay {
81 network: String,
82 recorded: BTreeMap<String, MachineIdentity>,
83 state: Mutex<State>,
84}
85
86impl Relay {
87 #[must_use]
89 pub fn new(network: String, recorded: BTreeMap<String, MachineIdentity>) -> Arc<Self> {
90 Arc::new(Self {
91 network,
92 recorded,
93 state: Mutex::new(State::default()),
94 })
95 }
96
97 pub fn connect(
102 self: &Arc<Self>,
103 session: Arc<dyn LysSession>,
104 sink: Arc<dyn RelaySink>,
105 ) -> Result<RelayConnection, RelayRefusal> {
106 let device = *self
107 .recorded
108 .get(session.holder())
109 .ok_or(RelayRefusal::Named("machine_not_recorded"))?
110 .ed25519_public();
111 let mut state = self.lock()?;
112 if state.endpoints.contains_key(&device) {
113 return Err(RelayRefusal::Named("device_already_connected"));
114 }
115 state.generation = state
116 .generation
117 .checked_add(1)
118 .ok_or(RelayRefusal::Named("relay_ids_exhausted"))?;
119 let generation = state.generation;
120 state.endpoints.insert(
121 device,
122 Endpoint {
123 generation,
124 sink,
125 streams: BTreeMap::new(),
126 },
127 );
128 drop(state);
129 Ok(RelayConnection {
130 relay: Arc::clone(self),
131 device,
132 generation,
133 session,
134 })
135 }
136
137 pub fn counters(
142 &self,
143 from: MachineIdentity,
144 to: MachineIdentity,
145 ) -> Result<RelayCounters, RelayRefusal> {
146 Ok(self
147 .lock()?
148 .counters
149 .get(&(*from.ed25519_public(), *to.ed25519_public()))
150 .copied()
151 .unwrap_or_default())
152 }
153
154 fn lock(&self) -> Result<std::sync::MutexGuard<'_, State>, RelayRefusal> {
155 self.state
156 .lock()
157 .map_err(|_| RelayRefusal::Named("relay_state_poisoned"))
158 }
159}
160
161pub struct RelayConnection {
163 relay: Arc<Relay>,
164 device: Device,
165 generation: u64,
166 session: Arc<dyn LysSession>,
167}
168
169impl std::fmt::Debug for RelayConnection {
170 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171 formatter
172 .debug_struct("RelayConnection")
173 .field("generation", &self.generation)
174 .finish_non_exhaustive()
175 }
176}
177
178impl RelayConnection {
179 pub fn mark_stream(&self, stream: u32, peer: &str, now: u64) -> Result<(), RelayRefusal> {
184 validate_stream(FrameType::LinkMessage, stream)
185 .map_err(|_| RelayRefusal::Named("invalid_relay_stream"))?;
186 let mut device = [0; 32];
187 hex::decode_to_slice(peer, &mut device)
188 .map_err(|_| RelayRefusal::Named("invalid_peer_device"))?;
189 if hex::encode(device) != peer || device == self.device {
190 return Err(RelayRefusal::Named("invalid_peer_device"));
191 }
192 self.decide(device, now)?;
193 let mut state = self.relay.lock()?;
194 let endpoint = self.endpoint(&state)?;
195 if endpoint
196 .streams
197 .iter()
198 .any(|(to, id)| (*id == stream && *to != device) || (*to == device && *id != stream))
199 {
200 return Err(RelayRefusal::Named("relay_stream_in_use"));
201 }
202 if !endpoint.streams.contains_key(&device) && endpoint.streams.len() >= STREAM_BOUND {
203 return Err(RelayRefusal::Named("relay_stream_limit"));
204 }
205 state
206 .endpoints
207 .get_mut(&self.device)
208 .ok_or(RelayRefusal::Named("peer_not_connected"))?
209 .streams
210 .insert(device, stream);
211 drop(state);
212 Ok(())
213 }
214
215 pub fn receive(&self, mut frame: Frame, now: u64) -> Result<(), RelayRefusal> {
221 let stream = frame.stream_id();
222 let mut state = self.relay.lock()?;
223 let endpoint = self.endpoint(&state)?;
224 let peer = *endpoint
225 .streams
226 .iter()
227 .find(|(_, id)| **id == stream)
228 .map(|(peer, _)| peer)
229 .ok_or(RelayRefusal::Named("relay_stream_not_marked"))?;
230 let pair = (self.device, peer);
231 if matches!(frame, Frame::LinkMessage { .. }) {
232 if let Err(error) = self.decide(peer, now) {
233 return refuse(&mut state, pair, stream, &error.to_string());
234 }
235 let Some(target) = state.endpoints.get(&peer) else {
236 return refuse(&mut state, pair, stream, "peer_not_connected");
237 };
238 let Some(&target_stream) = target.streams.get(&self.device) else {
239 return refuse(&mut state, pair, stream, "peer_relay_stream_not_marked");
240 };
241 if state.pending.contains_key(&pair) {
242 return refuse(&mut state, pair, stream, "relay_in_flight");
243 }
244 remap(&mut frame, target_stream)?;
245 let bytes = codec::encoded_len(&frame)
246 .map_err(|_| RelayRefusal::Named("invalid_relay_frame"))?;
247 if bytes > crate::link_carrier::WIRE_FRAME_BOUND {
248 return refuse(&mut state, pair, stream, "relay_frame_too_large");
249 }
250 let previous = state.counters.get(&pair).copied().unwrap_or_default();
251 let frames_relayed = increment(previous.frames_relayed)?;
252 let bytes_relayed = previous
253 .bytes_relayed
254 .checked_add(
255 u64::try_from(bytes)
256 .map_err(|_| RelayRefusal::Named("relay_counter_overflow"))?,
257 )
258 .ok_or(RelayRefusal::Named("relay_counter_overflow"))?;
259 if let Err(error) = target.sink.send(frame) {
260 return refuse(&mut state, pair, stream, &error.to_string());
261 }
262 state.counters.insert(
263 pair,
264 RelayCounters {
265 frames_relayed,
266 bytes_relayed,
267 ..previous
268 },
269 );
270 state.pending.insert(pair, stream);
271 drop(state);
272 return Ok(());
273 }
274 let original = (peer, self.device);
275 let sender_stream = *state
276 .pending
277 .get(&original)
278 .ok_or(RelayRefusal::Named("unexpected_relay_reply"))?;
279 let refused = matches!(frame, Frame::PublishError { .. });
280 remap(&mut frame, sender_stream)?;
281 let count = if refused {
282 Some(increment(
283 state.counters.entry(original).or_default().refused,
284 )?)
285 } else {
286 None
287 };
288 state
289 .endpoints
290 .get(&peer)
291 .ok_or(RelayRefusal::Named("peer_not_connected"))?
292 .sink
293 .send(frame)?;
294 state.pending.remove(&original);
295 if let Some(count) = count {
296 state.counters.entry(original).or_default().refused = count;
297 }
298 drop(state);
299 Ok(())
300 }
301
302 pub fn disconnect(&self) -> Result<(), RelayRefusal> {
307 let mut state = self.relay.lock()?;
308 if state
309 .endpoints
310 .get(&self.device)
311 .is_none_or(|endpoint| endpoint.generation != self.generation)
312 {
313 return Ok(());
314 }
315 state.endpoints.remove(&self.device);
316 let affected: Vec<_> = state
317 .pending
318 .iter()
319 .filter(|((from, to), _)| *from == self.device || *to == self.device)
320 .map(|(pair, stream)| (*pair, *stream))
321 .collect();
322 let mut failed = None;
323 for (pair, stream) in affected {
324 state.pending.remove(&pair);
325 if pair.1 == self.device
326 && let Err(error) = refuse(&mut state, pair, stream, "peer_not_connected")
327 {
328 tracing::error!(%error, "relay disconnect notification refused");
329 failed = Some(error);
330 }
331 }
332 drop(state);
333 failed.map_or(Ok(()), Err)
334 }
335
336 fn endpoint<'a>(&self, state: &'a State) -> Result<&'a Endpoint, RelayRefusal> {
337 state
338 .endpoints
339 .get(&self.device)
340 .filter(|endpoint| endpoint.generation == self.generation)
341 .ok_or(RelayRefusal::Named("peer_not_connected"))
342 }
343
344 fn decide(&self, peer: Device, now: u64) -> Result<(), RelayRefusal> {
345 let subject = format!("net.relay.{}", hex::encode(peer));
346 let binding = SubjectBinding::new(self.relay.network.clone(), subject.clone(), subject)
347 .map_err(|_| RelayRefusal::Named("invalid_relay_subject"))?;
348 self.session.decide(&binding, HotAction::Call, now)?;
349 Ok(())
350 }
351}
352
353impl Drop for RelayConnection {
354 fn drop(&mut self) {
355 if let Err(error) = self.disconnect() {
356 tracing::error!(%error, "relay disconnect refused");
357 }
358 }
359}
360
361fn increment(value: u64) -> Result<u64, RelayRefusal> {
362 value
363 .checked_add(1)
364 .ok_or(RelayRefusal::Named("relay_counter_overflow"))
365}
366
367fn refuse(state: &mut State, pair: Pair, stream: u32, reason: &str) -> Result<(), RelayRefusal> {
368 let counters = state.counters.entry(pair).or_default();
369 counters.refused = increment(counters.refused)?;
370 state
371 .endpoints
372 .get(&pair.0)
373 .ok_or(RelayRefusal::Named("peer_not_connected"))?
374 .sink
375 .send(Frame::PublishError {
376 flags: 0,
377 stream_id: stream,
378 reason_code: 0x0302,
379 message: Some(reason.to_owned()),
380 })
381}
382
383const fn remap(frame: &mut Frame, stream: u32) -> Result<(), RelayRefusal> {
384 match frame {
385 Frame::LinkMessage { stream_id, .. }
386 | Frame::PublishAck { stream_id, .. }
387 | Frame::PublishError { stream_id, .. } => {
388 *stream_id = stream;
389 Ok(())
390 }
391 _ => Err(RelayRefusal::Named("invalid_relay_frame")),
392 }
393}