hatchery_c2_client/
drain.rs

1//! Core-side drain client for the A2 restreamer.
2//!
3//! Core opens the stream to C2; C2 never dials core.
4
5use hatchery_core::{
6    ChannelKind, DrainEnvelope, Frame, FrameError, OpenPayload,
7};
8use std::net::SocketAddr;
9use std::time::Duration;
10use tokio::io::{AsyncReadExt, AsyncWriteExt};
11use tokio::net::TcpStream;
12
13#[derive(Debug)]
14pub enum DrainError {
15    Io(std::io::Error),
16    Frame(FrameError),
17    Auth,
18    /// Peer closed the TCP stream (distinct from a read timeout).
19    Closed,
20}
21
22impl std::fmt::Display for DrainError {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        match self {
25            Self::Io(e) => write!(f, "{e}"),
26            Self::Frame(e) => write!(f, "{e}"),
27            Self::Auth => write!(f, "drain auth refused"),
28            Self::Closed => write!(f, "drain closed"),
29        }
30    }
31}
32
33impl std::error::Error for DrainError {}
34
35impl From<std::io::Error> for DrainError {
36    fn from(value: std::io::Error) -> Self {
37        Self::Io(value)
38    }
39}
40
41pub struct DrainClient {
42    addr: SocketAddr,
43    token: String,
44}
45
46impl DrainClient {
47    pub fn new(addr: SocketAddr, token: impl Into<String>) -> Self {
48        Self {
49            addr,
50            token: token.into(),
51        }
52    }
53
54    pub async fn connect(&self) -> Result<DrainSession, DrainError> {
55        let mut stream = TcpStream::connect(self.addr).await?;
56        let _ = stream.set_nodelay(true);
57        stream.write_all(self.token.as_bytes()).await?;
58        stream.write_all(b"\n").await?;
59        Ok(DrainSession { stream, buf: Vec::new() })
60    }
61}
62
63pub struct DrainSession {
64    stream: TcpStream,
65    /// Bytes read but not yet decoded (several envelopes may share one read).
66    buf: Vec<u8>,
67}
68
69impl DrainSession {
70    pub async fn open(
71        &mut self,
72        node_id: &str,
73        channel_id: u32,
74        kind: ChannelKind,
75    ) -> Result<(), DrainError> {
76        let frame = Frame::open(
77            channel_id,
78            &OpenPayload {
79                channel_kind: kind,
80                node_id: node_id.to_owned(),
81            },
82        )
83        .map_err(DrainError::Frame)?;
84        self.stream
85            .write_all(&frame.encode().map_err(DrainError::Frame)?)
86            .await?;
87        Ok(())
88    }
89
90    pub async fn resume(&mut self, after: u64) -> Result<(), DrainError> {
91        self.stream
92            .write_all(&Frame::resume(after).encode().map_err(DrainError::Frame)?)
93            .await?;
94        Ok(())
95    }
96
97    pub async fn ack(&mut self, seq: u64) -> Result<(), DrainError> {
98        self.stream
99            .write_all(&Frame::ack(seq).encode().map_err(DrainError::Frame)?)
100            .await?;
101        Ok(())
102    }
103
104    /// Downlink: send opaque bytes to the node selected by `open` (C2 relays as-is).
105    pub async fn send_data(&mut self, channel_id: u32, payload: Vec<u8>) -> Result<(), DrainError> {
106        let frame = Frame::data(channel_id, 0, payload).map_err(DrainError::Frame)?;
107        self.stream
108            .write_all(&frame.encode().map_err(DrainError::Frame)?)
109            .await?;
110        Ok(())
111    }
112
113    /// Next drained envelope, or `None` on timeout. EOF is `Err(Closed)`. Non-DATA
114    /// control frames from C2 (e.g. PONG) are skipped.
115    pub async fn recv_envelope(
116        &mut self,
117        deadline: Duration,
118    ) -> Result<Option<DrainEnvelope>, DrainError> {
119        let until = tokio::time::Instant::now() + deadline;
120        loop {
121            if !self.buf.is_empty() {
122                match DrainEnvelope::decode(&self.buf) {
123                    Ok((env, used)) => {
124                        self.buf.drain(..used);
125                        return Ok(Some(env));
126                    }
127                    Err(FrameError::Truncated) => {}
128                    Err(_) => {
129                        // A bare frame (GOAWAY on auth failure, PONG): surface auth refusal.
130                        if let Ok((frame, used)) = Frame::decode(&self.buf) {
131                            self.buf.drain(..used);
132                            if frame.frame_type == hatchery_core::FrameType::GoAway {
133                                return Err(DrainError::Auth);
134                            }
135                            continue;
136                        }
137                        return Err(DrainError::Frame(FrameError::Io("bad drain envelope".into())));
138                    }
139                }
140            }
141            let mut chunk = vec![0u8; 64 * 1024];
142            let left = until.saturating_duration_since(tokio::time::Instant::now());
143            match tokio::time::timeout(left, self.stream.read(&mut chunk)).await {
144                Ok(Ok(0)) => return Err(DrainError::Closed),
145                Ok(Ok(n)) => self.buf.extend_from_slice(&chunk[..n]),
146                Ok(Err(e)) => return Err(e.into()),
147                Err(_) => return Ok(None),
148            }
149        }
150    }
151}
hatchery_c2_client/drain.rs

hatchery_c2_client/
drain.rs
