hatchery_c2_client/
lib.rs

1//! Bounded typed HTTP/1.1 client for the local Gate4Agent C2 observer API.
2
3use hatchery_core::{HealthResponse, ReadyResponse, StatusResponse, MAX_C2_RESPONSE_BYTES};
4use serde::de::DeserializeOwned;
5use std::io;
6use std::net::SocketAddr;
7use std::time::Duration;
8use thiserror::Error;
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tokio::net::TcpStream;
11use tokio::time::timeout;
12
13#[cfg(any(windows, unix))]
14mod runtime;
15#[cfg(any(windows, unix))]
16mod reconnect;
17#[cfg(unix)]
18mod hq_wg;
19mod drain;
20pub use drain::{DrainClient, DrainError, DrainSession};
21
22#[cfg(any(windows, unix))]
23/// Dial the **core** operator control socket (A2). Same wire as the former C2
24/// control endpoint; the socket now lives on hatchery-server, not on C2.
25pub use runtime::connect_local as connect_core_operator;
26#[cfg(any(windows, unix))]
27/// Dial a reconnecting core operator control socket (A2).
28pub use reconnect::connect_local_reconnecting as connect_core_operator_reconnecting;
29#[cfg(any(windows, unix))]
30pub use runtime::{
31    connect_local, connect_tunnel_tcp, C2ConnectionLossReason, C2ControlError, C2ControlHandle,
32    C2EventReceiver, C2PendingRequest,
33};
34#[cfg(any(windows, unix))]
35pub use reconnect::{
36    connect_local_reconnecting, connect_tunnel_reconnecting, C2LinkState,
37    C2ReconnectingEventReceiver, C2ReconnectingHandle,
38};
39#[cfg(unix)]
40pub use hq_wg::{
41    bring_up_hq_wireguard, HqWgAllowedIp, HqWgClientConfig, HqWgClientError, HqWgPeer,
42};
43
44const MAX_RESPONSE_HEADERS: usize = 16 * 1024;
45
46#[derive(Clone)]
47pub struct C2Client {
48    address: SocketAddr,
49    token: String,
50    deadline: Duration,
51}
52
53impl C2Client {
54    pub fn new(address: SocketAddr, token: impl Into<String>) -> Result<Self, C2ClientError> {
55        let token = token.into();
56        validate_token(&token)?;
57        Ok(Self { address, token, deadline: Duration::from_secs(3) })
58    }
59
60    pub fn with_deadline(mut self, deadline: Duration) -> Self {
61        self.deadline = deadline;
62        self
63    }
64
65    pub async fn health(&self) -> Result<HealthResponse, C2ClientError> {
66        self.get("/health", false, false).await
67    }
68
69    pub async fn ready(&self) -> Result<ReadyResponse, C2ClientError> {
70        self.get("/ready", false, true).await
71    }
72
73    pub async fn status(&self) -> Result<StatusResponse, C2ClientError> {
74        self.get("/status", true, false).await
75    }
76
77    async fn get<T: DeserializeOwned>(&self, path: &str, authenticated: bool, allow_unready: bool) -> Result<T, C2ClientError> {
78        timeout(self.deadline, self.get_inner(path, authenticated, allow_unready))
79            .await
80            .map_err(|_| C2ClientError::Deadline)?
81    }
82
83    async fn get_inner<T: DeserializeOwned>(&self, path: &str, authenticated: bool, allow_unready: bool) -> Result<T, C2ClientError> {
84        let mut stream = TcpStream::connect(self.address).await?;
85        let auth = if authenticated {
86            format!("Authorization: Bearer {}\r\n", self.token)
87        } else {
88            String::new()
89        };
90        let request = format!("GET {path} HTTP/1.1\r\nHost: {}\r\n{auth}Connection: close\r\n\r\n", self.address);
91        stream.write_all(request.as_bytes()).await?;
92        let mut bytes = Vec::with_capacity(4096);
93        let mut chunk = [0_u8; 4096];
94        loop {
95            let count = stream.read(&mut chunk).await?;
96            if count == 0 { break; }
97            if bytes.len().saturating_add(count) > MAX_C2_RESPONSE_BYTES + MAX_RESPONSE_HEADERS {
98                return Err(C2ClientError::Oversized);
99            }
100            bytes.extend_from_slice(&chunk[..count]);
101        }
102        let header_end = bytes.windows(4).position(|window| window == b"\r\n\r\n")
103            .ok_or(C2ClientError::Malformed)?;
104        if header_end > MAX_RESPONSE_HEADERS { return Err(C2ClientError::Oversized); }
105        let headers = std::str::from_utf8(&bytes[..header_end]).map_err(|_| C2ClientError::Malformed)?;
106        let mut lines = headers.split("\r\n");
107        let status_line = lines.next().ok_or(C2ClientError::Malformed)?;
108        let mut status_parts = status_line.split_whitespace();
109        if status_parts.next() != Some("HTTP/1.1") { return Err(C2ClientError::Malformed); }
110        let status = status_parts.next().ok_or(C2ClientError::Malformed)?
111            .parse::<u16>().map_err(|_| C2ClientError::Malformed)?;
112        let mut content_length = None;
113        for line in lines {
114            let (name, value) = line.split_once(':').ok_or(C2ClientError::Malformed)?;
115            if name.eq_ignore_ascii_case("content-length") {
116                if content_length.is_some() { return Err(C2ClientError::Malformed); }
117                content_length = Some(value.trim().parse::<usize>().map_err(|_| C2ClientError::Malformed)?);
118            }
119        }
120        let content_length = content_length.ok_or(C2ClientError::Malformed)?;
121        if content_length > MAX_C2_RESPONSE_BYTES { return Err(C2ClientError::Oversized); }
122        let body = &bytes[header_end + 4..];
123        if body.len() != content_length { return Err(C2ClientError::Malformed); }
124        if status == 401 { return Err(C2ClientError::Unauthorized); }
125        if !(200..300).contains(&status) && !(allow_unready && status == 503) { return Err(C2ClientError::HttpStatus(status)); }
126        serde_json::from_slice(body).map_err(C2ClientError::Json)
127    }
128}
129
130impl std::fmt::Debug for C2Client {
131    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132        formatter.debug_struct("C2Client")
133            .field("address", &self.address)
134            .field("token", &"[REDACTED]")
135            .field("deadline", &self.deadline)
136            .finish()
137    }
138}
139
140fn validate_token(token: &str) -> Result<(), C2ClientError> {
141    if token.is_empty() || token.len() > 4096 || !token.bytes().all(|byte| matches!(byte, 0x21..=0x7e)) {
142        return Err(C2ClientError::InvalidToken);
143    }
144    Ok(())
145}
146
147#[derive(Debug, Error)]
148pub enum C2ClientError {
149    #[error("C2 bearer token must contain 1..=4096 visible ASCII bytes without whitespace")]
150    InvalidToken,
151    #[error("C2 HTTP I/O failed: {0}")]
152    Io(#[from] io::Error),
153    #[error("C2 request exceeded its deadline")]
154    Deadline,
155    #[error("C2 returned HTTP 401")]
156    Unauthorized,
157    #[error("C2 returned HTTP status {0}")]
158    HttpStatus(u16),
159    #[error("C2 returned malformed HTTP")]
160    Malformed,
161    #[error("C2 response exceeded the 8 MiB body bound")]
162    Oversized,
163    #[error("C2 returned malformed JSON: {0}")]
164    Json(serde_json::Error),
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170    use tokio::net::TcpListener;
171
172    async fn one_response(response: Vec<u8>) -> SocketAddr {
173        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
174        let address = listener.local_addr().unwrap();
175        tokio::spawn(async move {
176            let (mut stream, _) = listener.accept().await.unwrap();
177            let mut request = [0_u8; 1024];
178            let _ = stream.read(&mut request).await.unwrap();
179            stream.write_all(&response).await.unwrap();
180            stream.shutdown().await.unwrap();
181        });
182        address
183    }
184
185    #[test]
186    fn bearer_rejects_header_injection_and_debug_redacts_secret() {
187        let address = "127.0.0.1:1".parse().unwrap();
188        assert!(matches!(C2Client::new(address, "safe\r\nInjected: yes"), Err(C2ClientError::InvalidToken)));
189        let client = C2Client::new(address, "unique-secret").unwrap();
190        let debug = format!("{client:?}");
191        assert!(!debug.contains("unique-secret"));
192        assert!(debug.contains("[REDACTED]"));
193    }
194
195    #[tokio::test]
196    async fn http_client_rejects_unauthorized_malformed_and_oversized_responses() {
197        let unauthorized = one_response(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_vec()).await;
198        assert!(matches!(C2Client::new(unauthorized, "token").unwrap().status().await, Err(C2ClientError::Unauthorized)));
199
200        let malformed = one_response(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n{}".to_vec()).await;
201        assert!(matches!(C2Client::new(malformed, "token").unwrap().health().await, Err(C2ClientError::Malformed)));
202
203        let oversized = one_response(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", MAX_C2_RESPONSE_BYTES + 1).into_bytes()).await;
204        assert!(matches!(C2Client::new(oversized, "token").unwrap().health().await, Err(C2ClientError::Oversized)));
205    }
206
207    #[tokio::test]
208    async fn typed_ready_decodes_bounded_503_initializing_response() {
209        let body = br#"{"ready":false,"api_version":1,"configured_nodes":2,"attempted_nodes":0,"online_nodes":0,"offline_nodes":2,"parked_nodes":0}"#;
210        let mut response = format!("HTTP/1.1 503 Service Unavailable\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).into_bytes();
211        response.extend_from_slice(body);
212        let address = one_response(response).await;
213        let ready = C2Client::new(address, "token").unwrap().ready().await.unwrap();
214        assert!(!ready.ready);
215        assert_eq!(ready.configured_nodes, 2);
216        assert_eq!(ready.attempted_nodes, 0);
217    }
218}
hatchery_c2_client/lib.rs

hatchery_c2_client/
lib.rs
