hatchery_server/
drain_consumer.rs

1//! Core drain consumer (A2 / A3): core dials C2, drains one node stream, and answers on
2//! the same stream. C2 never dials core and never sees plaintext.
3//!
4//! Per stream: Noise_NK msg1 arrives as a DATA envelope, core answers msg2 as downlink
5//! DATA, then every envelope is Noise transport. Each plaintext goes to a
6//! [`CoreHandler`] together with the node id taken from the stream (never from the
7//! payload); the reply is encrypted back down. Every processed seq is ACKed, and a
8//! reconnect RESUMEs after the last ACK.
9
10use hatchery_c2_client::{DrainClient, DrainSession};
11use hatchery_core::inner::{decode_inner, encode_inner, InnerKind};
12use hatchery_core::nested::{HqRecipientSecret, NoiseResponder, NoiseTransport};
13use hatchery_core::ChannelKind;
14use std::sync::Arc;
15use std::time::Duration;
16use tokio::sync::watch;
17
18/// Core logic for one decrypted inner message. Blocking work is fine: it runs on a
19/// blocking thread. Return `None` for messages that need no reply.
20pub type CoreHandler = Arc<dyn Fn(&str, ChannelKind, &[u8]) -> Option<Vec<u8>> + Send + Sync>;
21
22pub struct NodeStream {
23    pub node_id: String,
24    pub channel_id: u32,
25    pub kind: ChannelKind,
26}
27
28pub struct DrainConsumer {
29    client: DrainClient,
30    recipient: HqRecipientSecret,
31    stream: NodeStream,
32    handler: CoreHandler,
33    last_acked: u64,
34    transport: Option<NoiseTransport>,
35    /// Core-initiated messages (operator control). Sent once the node's handshake landed.
36    outbound: Option<tokio::sync::mpsc::Receiver<Vec<u8>>>,
37}
38
39impl DrainConsumer {
40    pub fn new(client: DrainClient, recipient: HqRecipientSecret, stream: NodeStream, handler: CoreHandler) -> Self {
41        Self { client, recipient, stream, handler, last_acked: 0, transport: None, outbound: None }
42    }
43
44    /// Attach a queue of core-initiated plaintexts (see `node_controller::OperatorLink`).
45    pub fn with_outbound(mut self, rx: tokio::sync::mpsc::Receiver<Vec<u8>>) -> Self {
46        self.outbound = Some(rx);
47        self
48    }
49
50    pub fn last_acked(&self) -> u64 {
51        self.last_acked
52    }
53
54    /// Drain until shutdown, reconnecting with RESUME after the last ACK.
55    pub async fn run(mut self, mut shutdown: watch::Receiver<bool>) {
56        let mut backoff = Duration::from_millis(50);
57        while !*shutdown.borrow() {
58            match self.session(&mut shutdown).await {
59                Ok(()) => backoff = Duration::from_millis(50),
60                Err(err) => tracing::debug!(error = %err, "drain session ended"),
61            }
62            if *shutdown.borrow() {
63                break;
64            }
65            tokio::select! {
66                _ = tokio::time::sleep(backoff) => {}
67                _ = shutdown.changed() => {}
68            }
69            backoff = (backoff * 2).min(Duration::from_secs(5));
70        }
71    }
72
73    async fn session(&mut self, shutdown: &mut watch::Receiver<bool>) -> Result<(), String> {
74        // Prior Noise keys are invalid after a TCP reconnect; wait for a fresh msg1.
75        self.transport = None;
76        let mut s = self.client.connect().await.map_err(|e| e.to_string())?;
77        s.open(&self.stream.node_id, self.stream.channel_id, self.stream.kind).await.map_err(|e| e.to_string())?;
78        s.resume(self.last_acked).await.map_err(|e| e.to_string())?;
79        loop {
80            if *shutdown.borrow() {
81                return Ok(());
82            }
83            let ready = self.transport.is_some();
84            let env = tokio::select! {
85                env = s.recv_envelope(Duration::from_millis(250)) => env.map_err(|e| e.to_string())?,
86                Some(plain) = async {
87                    match self.outbound.as_mut() {
88                        Some(rx) if ready => rx.recv().await,
89                        _ => std::future::pending().await,
90                    }
91                } => {
92                    let ct = self.transport.as_mut().expect("transport").encrypt(&plain).map_err(|e| e.to_string())?;
93                    s.send_data(self.stream.channel_id, encode_inner(InnerKind::Transport, &ct))
94                        .await
95                        .map_err(|e| e.to_string())?;
96                    continue;
97                }
98                _ = shutdown.changed() => return Ok(()),
99            };
100            let Some(env) = env else { continue };
101            if env.node_id != self.stream.node_id {
102                continue;
103            }
104            let seq = env.frame.seq;
105            if seq != 0 && seq <= self.last_acked {
106                continue;
107            }
108            self.handle(&mut s, &env.frame.payload).await?;
109            if seq != 0 {
110                s.ack(seq).await.map_err(|e| e.to_string())?;
111                self.last_acked = seq;
112            }
113        }
114    }
115
116    async fn handle(&mut self, s: &mut DrainSession, inner: &[u8]) -> Result<(), String> {
117        match decode_inner(inner) {
118            Some((InnerKind::Handshake1, msg1)) => {
119                let responder = NoiseResponder::new(&self.recipient).map_err(|e| e.to_string())?;
120                match responder.complete(msg1) {
121                    Ok((msg2, transport)) => {
122                        self.transport = Some(transport);
123                        s.send_data(self.stream.channel_id, encode_inner(InnerKind::Handshake2, &msg2))
124                            .await
125                            .map_err(|e| e.to_string())?;
126                    }
127                    // Wrong recipient key or garbage: drop it, keep draining.
128                    Err(err) => tracing::debug!(error = %err, "handshake refused"),
129                }
130            }
131            Some((InnerKind::Transport, ct)) => {
132                let Some(transport) = self.transport.as_mut() else { return Ok(()) };
133                let Ok(plain) = transport.decrypt(ct) else {
134                    // Desync (often after C2 dropped a frame under credit pressure).
135                    // Forget the session so the next Handshake1 can recover.
136                    tracing::warn!(node = %self.stream.node_id, channel = self.stream.channel_id, "inner decrypt failed; resetting Noise session");
137                    self.transport = None;
138                    return Ok(());
139                };
140                let handler = Arc::clone(&self.handler);
141                let node = self.stream.node_id.clone();
142                let kind = self.stream.kind;
143                let reply = tokio::task::spawn_blocking(move || handler(&node, kind, &plain))
144                    .await
145                    .map_err(|e| e.to_string())?;
146                if let Some(reply) = reply {
147                    let ct = self.transport.as_mut().expect("transport").encrypt(&reply).map_err(|e| e.to_string())?;
148                    s.send_data(self.stream.channel_id, encode_inner(InnerKind::Transport, &ct))
149                        .await
150                        .map_err(|e| e.to_string())?;
151                }
152            }
153            _ => {}
154        }
155        Ok(())
156    }
157}
hatchery_server/drain_consumer.rs

hatchery_server/
drain_consumer.rs
