1use crate::uplink_journal::UplinkJournal;
17use hatchery_core::inner::{decode_inner, encode_inner, InnerKind};
18use hatchery_core::nested::{node_seal_outer, NoiseInitiator, NoiseTransport, PublicKeyBytes};
19use hatchery_core::{grow_read_buf, ChannelKind, Frame, FrameError, FrameType, OpenPayload};
20use serde_json::{json, Value};
21use std::collections::{BTreeMap, VecDeque};
22use std::net::SocketAddr;
23use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
24use std::sync::{Arc, Mutex};
25use std::time::Duration;
26use tokio::io::{AsyncReadExt, AsyncWriteExt};
27use tokio::net::TcpStream;
28use tokio::sync::{mpsc, Notify};
29use tokio::task::JoinHandle;
30
31const UPLINK_PING_EVERY: Duration = Duration::from_secs(10);
33const UPLINK_DEAD_AFTER: Duration = Duration::from_secs(35);
35
36#[derive(Clone, Copy, Debug)]
38pub struct UpstreamKeys {
39 pub c2_public: PublicKeyBytes,
40 pub recipient_public: PublicKeyBytes,
41}
42
43#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
45pub enum SendPrio {
46 Control = 0,
48 Pty = 1,
50 Mail = 2,
52 Browser = 3,
54}
55
56const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
57const CHAN_CAP_CONTROL: usize = 64;
58const CHAN_CAP_PTY: usize = 128;
59const CHAN_CAP_MAIL: usize = 16;
60const CHAN_CAP_BROWSER: usize = 8;
61
62struct Queued {
63 prio: SendPrio,
64 frame: Frame,
65}
66
67struct ChanQ {
68 q: VecDeque<Queued>,
69 credit: Option<u32>,
71}
72
73impl ChanQ {
74 fn new() -> Self {
75 Self { q: VecDeque::new(), credit: None }
76 }
77
78 fn cap_for(prio: SendPrio) -> usize {
79 match prio {
80 SendPrio::Control => CHAN_CAP_CONTROL,
81 SendPrio::Pty => CHAN_CAP_PTY,
82 SendPrio::Mail => CHAN_CAP_MAIL,
83 SendPrio::Browser => CHAN_CAP_BROWSER,
84 }
85 }
86}
87
88struct OutState {
89 chans: BTreeMap<u32, ChanQ>,
90}
91
92pub struct NodeUplink {
94 enqueue: mpsc::UnboundedSender<(u32, SendPrio, Frame)>,
95 routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>>,
96 credits: Arc<Mutex<BTreeMap<u32, u32>>>,
98 credit_pulse: Arc<Notify>,
99 journal: Option<Arc<UplinkJournal>>,
100 alive: Arc<AtomicBool>,
101 dead: Arc<Notify>,
102 tasks: Vec<JoinHandle<()>>,
103}
104
105impl Drop for NodeUplink {
106 fn drop(&mut self) {
107 self.alive.store(false, Ordering::SeqCst);
108 for t in &self.tasks {
109 t.abort();
110 }
111 }
112}
113
114impl NodeUplink {
115 pub async fn connect(uplink: SocketAddr, journal: Option<Arc<UplinkJournal>>) -> Result<Arc<Self>, String> {
116 let stream = TcpStream::connect(uplink).await.map_err(|e| e.to_string())?;
117 let _ = stream.set_nodelay(true);
118 let (mut rd, mut wr) = stream.into_split();
119 let (enq_tx, mut enq_rx) = mpsc::unbounded_channel::<(u32, SendPrio, Frame)>();
120 let routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>> = Arc::default();
121 let credits: Arc<Mutex<BTreeMap<u32, u32>>> = Arc::default();
122 let credit_pulse = Arc::new(Notify::new());
123 let alive = Arc::new(AtomicBool::new(true));
124 let dead = Arc::new(Notify::new());
125 let wake = Arc::new(Notify::new());
126
127 let out_state = Arc::new(Mutex::new(OutState { chans: BTreeMap::new() }));
128 let born = std::time::Instant::now();
133 let last_rx = Arc::new(AtomicU64::new(0));
134
135 let state_in = Arc::clone(&out_state);
137 let wake_in = Arc::clone(&wake);
138 let alive_in = Arc::clone(&alive);
139 let ingress = tokio::spawn(async move {
140 while let Some((ch, prio, frame)) = enq_rx.recv().await {
141 if !alive_in.load(Ordering::SeqCst) {
142 break;
143 }
144 let mut st = state_in.lock().expect("out");
145 let q = st.chans.entry(ch).or_insert_with(ChanQ::new);
146 let cap = ChanQ::cap_for(prio);
147 if prio == SendPrio::Browser {
148 q.q.retain(|x| x.prio != SendPrio::Browser);
150 if q.q.len() >= cap {
151 continue;
153 }
154 } else {
155 while q.q.len() >= cap {
156 if let Some(i) = q
158 .q
159 .iter()
160 .enumerate()
161 .filter(|(_, x)| x.prio >= SendPrio::Mail)
162 .max_by_key(|(_, x)| x.prio)
163 .map(|(i, _)| i)
164 {
165 q.q.remove(i);
166 } else if prio <= SendPrio::Pty {
167 q.q.pop_front();
169 } else {
170 break;
171 }
172 }
173 if q.q.len() >= cap {
174 continue;
175 }
176 }
177 q.q.push_back(Queued { prio, frame });
178 drop(st);
179 wake_in.notify_one();
180 }
181 });
182
183 let state_wr = Arc::clone(&out_state);
186 let wake_wr = Arc::clone(&wake);
187 let alive_wr = Arc::clone(&alive);
188 let dead_wr = Arc::clone(&dead);
189 let routes_wr = Arc::clone(&routes);
190 let writer = tokio::spawn(async move {
191 loop {
192 if !alive_wr.load(Ordering::SeqCst) {
193 break;
194 }
195 let next = {
196 let mut st = state_wr.lock().expect("out");
197 pick_next(&mut st)
198 };
199 let Some((ch, is_data, frame)) = next else {
200 wake_wr.notified().await;
201 continue;
202 };
203 let Ok(bytes) = frame.encode() else {
204 kill(&alive_wr, &dead_wr, &routes_wr);
205 break;
206 };
207 match tokio::time::timeout(WRITE_TIMEOUT, wr.write_all(&bytes)).await {
208 Ok(Ok(())) => {
209 if is_data {
210 let mut st = state_wr.lock().expect("out");
211 if let Some(q) = st.chans.get_mut(&ch) {
212 if let Some(c) = q.credit.as_mut() {
213 *c = c.saturating_sub(1);
214 }
215 }
216 }
217 }
218 Ok(Err(_)) | Err(_) => {
219 tracing::warn!(%ch, "uplink write timed out or failed; reconnecting");
220 kill(&alive_wr, &dead_wr, &routes_wr);
221 break;
222 }
223 }
224 }
225 });
226
227 let enq_reader = enq_tx.clone();
228 let reader_journal = journal.clone();
229 let reader_routes = Arc::clone(&routes);
230 let alive_rd = Arc::clone(&alive);
231 let dead_rd = Arc::clone(&dead);
232 let state_rd = Arc::clone(&out_state);
233 let wake_rd = Arc::clone(&wake);
234 let credits_rd = Arc::clone(&credits);
235 let credit_pulse_rd = Arc::clone(&credit_pulse);
236 let last_rx_rd = Arc::clone(&last_rx);
237 let reader = tokio::spawn(async move {
238 let mut buf = vec![0u8; 128 * 1024];
239 let mut filled = 0usize;
240 loop {
241 let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
242 if n == 0 {
243 break;
244 }
245 filled += n;
246 last_rx_rd.store(born.elapsed().as_millis() as u64, Ordering::SeqCst);
247 loop {
248 match Frame::decode(&buf[..filled]) {
249 Ok((frame, used)) => {
250 buf.copy_within(used..filled, 0);
251 filled -= used;
252 match frame.frame_type {
253 FrameType::Data => {
254 let tx = reader_routes.lock().expect("routes").get(&frame.channel_id).cloned();
255 if let Some(tx) = tx {
256 let _ = tx.send(frame.payload);
257 }
258 }
259 FrameType::Window => {
260 let credit = frame.seq.min(u64::from(u32::MAX)) as u32;
261 {
262 let mut st = state_rd.lock().expect("out");
263 st.chans.entry(frame.channel_id).or_insert_with(ChanQ::new).credit = Some(credit);
264 }
265 credits_rd.lock().expect("credits").insert(frame.channel_id, credit);
266 credit_pulse_rd.notify_waiters();
267 wake_rd.notify_one();
268 }
269 FrameType::JournalReq => {
270 let Some(j) = reader_journal.as_ref() else { continue };
271 let limit = frame.journal_req_limit().unwrap_or(64) as usize;
272 if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
273 if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
274 let _ = enq_reader.send((frame.channel_id, SendPrio::Control, reply));
275 }
276 }
277 }
278 FrameType::GoAway | FrameType::Close => {
279 kill(&alive_rd, &dead_rd, &reader_routes);
280 return;
281 }
282 _ => {}
283 }
284 }
285 Err(FrameError::Truncated) => break,
286 Err(_) => {
287 kill(&alive_rd, &dead_rd, &reader_routes);
288 return;
289 }
290 }
291 }
292 if filled == buf.len() {
293 if grow_read_buf(&mut buf, filled).is_err() {
296 break;
297 }
298 }
299 }
300 kill(&alive_rd, &dead_rd, &reader_routes);
301 });
302
303 let ka_tx = enq_tx.clone();
304 let alive_ka = Arc::clone(&alive);
305 let dead_ka = Arc::clone(&dead);
306 let routes_ka = Arc::clone(&routes);
307 let keepalive = tokio::spawn(async move {
308 let mut nonce = 0u64;
309 let mut tick = tokio::time::interval(UPLINK_PING_EVERY);
310 tick.tick().await;
311 loop {
312 tick.tick().await;
313 if !alive_ka.load(Ordering::SeqCst) {
314 break;
315 }
316 let now = born.elapsed().as_millis() as u64;
317 let idle = now.saturating_sub(last_rx.load(Ordering::SeqCst));
318 if idle > UPLINK_DEAD_AFTER.as_millis() as u64 {
319 tracing::warn!(idle_ms = idle, "uplink silent past keepalive deadline; reconnecting");
320 kill(&alive_ka, &dead_ka, &routes_ka);
321 break;
322 }
323 nonce = nonce.wrapping_add(1);
324 if ka_tx.send((0, SendPrio::Control, Frame::ping(nonce))).is_err() {
325 break;
326 }
327 }
328 });
329
330 Ok(Arc::new(Self {
331 enqueue: enq_tx,
332 routes,
333 credits,
334 credit_pulse,
335 journal,
336 alive,
337 dead,
338 tasks: vec![ingress, writer, reader, keepalive],
339 }))
340 }
341
342 pub fn is_alive(&self) -> bool {
343 self.alive.load(Ordering::SeqCst)
344 }
345
346 pub async fn closed(&self) {
348 if !self.is_alive() {
349 return;
350 }
351 self.dead.notified().await;
352 }
353
354 fn push(&self, channel_id: u32, prio: SendPrio, frame: Frame) -> Result<(), String> {
355 if !self.is_alive() {
356 return Err("uplink closed".into());
357 }
358 self.enqueue
359 .send((channel_id, prio, frame))
360 .map_err(|_| "uplink closed".to_owned())
361 }
362
363 pub async fn open(
365 self: &Arc<Self>,
366 node_id: &str,
367 channel_id: u32,
368 kind: ChannelKind,
369 keys: UpstreamKeys,
370 handshake_timeout: Duration,
371 ) -> Result<NodeSession, String> {
372 let (down_tx, mut down_rx) = mpsc::unbounded_channel::<Vec<u8>>();
373 self.routes.lock().expect("routes").insert(channel_id, down_tx);
374 self.credits.lock().expect("credits").remove(&channel_id);
375 let open = Frame::open(channel_id, &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() })
376 .map_err(|e| e.to_string())?;
377 self.push(channel_id, SendPrio::Control, open)?;
378
379 tokio::time::timeout(handshake_timeout, async {
382 loop {
383 if !self.is_alive() {
384 return Err("uplink closed".to_owned());
385 }
386 if self.credits.lock().expect("credits").get(&channel_id).copied().unwrap_or(0) > 0 {
387 return Ok(());
388 }
389 tokio::select! {
390 _ = self.credit_pulse.notified() => {}
391 _ = self.closed() => return Err("uplink closed".to_owned()),
392 }
393 }
394 })
395 .await
396 .map_err(|_| "window credit timeout after OPEN".to_owned())??;
397
398 let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
399 let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
400 self.send_inner(channel_id, keys, SendPrio::Control, &encode_inner(InnerKind::Handshake1, &msg1))?;
401
402 let msg2 = tokio::time::timeout(handshake_timeout, async {
403 loop {
404 let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
405 if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
406 return Ok(body.to_vec());
407 }
408 }
409 })
410 .await
411 .map_err(|_| "handshake timeout".to_owned())??;
412 let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
413 Ok(NodeSession {
414 channel_id,
415 keys,
416 default_prio: match kind {
417 ChannelKind::M4a => SendPrio::Mail,
418 ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents => SendPrio::Pty,
419 ChannelKind::Vault | ChannelKind::Observation => SendPrio::Mail,
420 },
421 down: down_rx,
422 transport,
423 journal: self.journal.clone(),
424 uplink: Arc::clone(self),
425 })
426 }
427
428 fn send_inner(&self, channel_id: u32, keys: UpstreamKeys, prio: SendPrio, inner: &[u8]) -> Result<(), String> {
429 if let Some(j) = &self.journal {
430 j.append(channel_id, inner)?;
431 }
432 let sealed = node_seal_outer(keys.c2_public, inner).map_err(|e| e.to_string())?;
433 let frame = Frame::data(channel_id, 0, sealed).map_err(|e| e.to_string())?;
434 self.push(channel_id, prio, frame)
435 }
436}
437
438fn kill(alive: &AtomicBool, dead: &Notify, routes: &Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>) {
439 if alive.swap(false, Ordering::SeqCst) {
440 routes.lock().expect("routes").clear();
441 dead.notify_waiters();
442 }
443}
444
445fn pick_next(st: &mut OutState) -> Option<(u32, bool, Frame)> {
447 let mut best: Option<(SendPrio, u32, usize)> = None;
449 for (&ch, cq) in &st.chans {
450 for (i, item) in cq.q.iter().enumerate() {
451 let needs_credit = matches!(item.frame.frame_type, FrameType::Data);
452 if needs_credit {
453 match cq.credit {
454 Some(0) => continue,
455 None => {
456 if item.prio > SendPrio::Control {
459 continue;
460 }
461 }
462 Some(_) => {}
463 }
464 }
465 let cand = (item.prio, ch, i);
466 if best.is_none_or(|b| cand.0 < b.0 || (cand.0 == b.0 && cand.1 < b.1)) {
467 best = Some(cand);
468 }
469 }
470 }
471 let (prio, ch, idx) = best?;
472 let cq = st.chans.get_mut(&ch)?;
473 let item = cq.q.remove(idx)?;
474 debug_assert_eq!(item.prio, prio);
475 let is_data = matches!(item.frame.frame_type, FrameType::Data);
476 Some((ch, is_data, item.frame))
477}
478
479pub struct NodeSession {
480 channel_id: u32,
481 keys: UpstreamKeys,
482 default_prio: SendPrio,
483 down: mpsc::UnboundedReceiver<Vec<u8>>,
484 transport: NoiseTransport,
485 journal: Option<Arc<UplinkJournal>>,
486 uplink: Arc<NodeUplink>,
487}
488
489impl NodeSession {
490 pub async fn connect(
491 uplink: SocketAddr,
492 node_id: &str,
493 channel_id: u32,
494 kind: ChannelKind,
495 keys: UpstreamKeys,
496 journal: Option<Arc<UplinkJournal>>,
497 handshake_timeout: Duration,
498 ) -> Result<Self, String> {
499 let link = NodeUplink::connect(uplink, journal).await?;
500 link.open(node_id, channel_id, kind, keys, handshake_timeout).await
501 }
502
503 pub fn uplink(&self) -> Arc<NodeUplink> {
504 Arc::clone(&self.uplink)
505 }
506
507 pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
508 self.send_prio(plaintext, self.default_prio).await
509 }
510
511 pub async fn send_prio(&mut self, plaintext: &[u8], prio: SendPrio) -> Result<(), String> {
512 if !self.uplink.is_alive() {
513 return Err("uplink closed".into());
514 }
515 let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
516 let inner = encode_inner(InnerKind::Transport, &ct);
517 self.uplink.send_inner(self.channel_id, self.keys, prio, &inner)
518 }
519
520 pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
521 tokio::time::timeout(timeout, async {
522 loop {
523 tokio::select! {
524 bytes = self.down.recv() => {
525 let Some(bytes) = bytes else { return Err("uplink closed".to_owned()) };
526 if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
527 return self.transport.decrypt(ct).map_err(|e| e.to_string());
528 }
529 }
530 _ = self.uplink.closed() => return Err("uplink closed".to_owned()),
531 }
532 }
533 })
534 .await
535 .map_err(|_| "reply timeout".to_owned())?
536 }
537
538 pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
539 let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
540 self.send_prio(&req, SendPrio::Mail).await?;
541 let reply = self.recv(timeout).await?;
542 let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
543 let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
544 let body = reply.get("body").cloned().unwrap_or(Value::Null);
545 if status == 200 {
546 Ok(body)
547 } else {
548 let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
549 Err(format!("{status}: {msg}"))
550 }
551 }
552}
553
554#[cfg(test)]
555mod tests {
556 use super::*;
557
558 #[test]
559 fn pick_prefers_control_over_browser() {
560 let mut st = OutState { chans: BTreeMap::new() };
561 let q = st.chans.entry(1).or_insert_with(ChanQ::new);
562 q.credit = Some(10);
563 q.q.push_back(Queued {
564 prio: SendPrio::Browser,
565 frame: Frame::data(1, 0, vec![1]).unwrap(),
566 });
567 q.q.push_back(Queued {
568 prio: SendPrio::Control,
569 frame: Frame::data(1, 0, vec![2]).unwrap(),
570 });
571 q.q.push_back(Queued {
572 prio: SendPrio::Pty,
573 frame: Frame::data(1, 0, vec![3]).unwrap(),
574 });
575 let (_, _, f) = pick_next(&mut st).unwrap();
576 assert_eq!(f.payload, vec![2]);
577 let (_, _, f) = pick_next(&mut st).unwrap();
578 assert_eq!(f.payload, vec![3]);
579 let (_, _, f) = pick_next(&mut st).unwrap();
580 assert_eq!(f.payload, vec![1]);
581 }
582
583 #[test]
584 fn pick_holds_bulk_until_window() {
585 let mut st = OutState { chans: BTreeMap::new() };
586 let q = st.chans.entry(2).or_insert_with(ChanQ::new);
587 q.credit = None;
588 q.q.push_back(Queued {
589 prio: SendPrio::Mail,
590 frame: Frame::data(2, 0, vec![9]).unwrap(),
591 });
592 q.q.push_back(Queued {
593 prio: SendPrio::Control,
594 frame: Frame::data(2, 0, vec![1]).unwrap(),
595 });
596 let (_, _, f) = pick_next(&mut st).unwrap();
597 assert_eq!(f.payload, vec![1]);
598 assert!(pick_next(&mut st).is_none());
599 st.chans.get_mut(&2).unwrap().credit = Some(4);
600 let (_, _, f) = pick_next(&mut st).unwrap();
601 assert_eq!(f.payload, vec![9]);
602 }
603}