Skip to main content

forge_orchestration/
networking.rs

1//! Networking module for Forge
2//!
3//! ## Table of Contents
4//! - **QuicTransport**: QUIC-based peer communication (requires `quic` feature)
5//! - **HttpServer**: Axum-based HTTP/REST API server
6//! - **GrpcServer**: Tonic-based gRPC server (placeholder)
7
8use crate::error::{ForgeError, Result};
9use axum::{
10    http::StatusCode,
11    response::IntoResponse,
12    routing::get,
13    Json, Router,
14};
15use serde::{Deserialize, Serialize};
16use std::net::SocketAddr;
17use std::sync::Arc;
18use tokio::sync::RwLock;
19use tracing::info;
20
21#[cfg(feature = "quic")]
22use std::sync::Arc as QuicArc;
23
24/// HTTP server configuration
25#[derive(Debug, Clone)]
26pub struct HttpServerConfig {
27    /// Bind address
28    pub bind_addr: SocketAddr,
29    /// Enable CORS
30    pub cors_enabled: bool,
31    /// Request timeout in seconds
32    pub timeout_secs: u64,
33}
34
35impl Default for HttpServerConfig {
36    fn default() -> Self {
37        Self {
38            bind_addr: ([0, 0, 0, 0], 8080).into(),
39            cors_enabled: true,
40            timeout_secs: 30,
41        }
42    }
43}
44
45impl HttpServerConfig {
46    /// Create with custom bind address
47    pub fn with_addr(mut self, addr: SocketAddr) -> Self {
48        self.bind_addr = addr;
49        self
50    }
51
52    /// Parse from string address
53    pub fn with_addr_str(mut self, addr: &str) -> Result<Self> {
54        self.bind_addr = addr
55            .parse()
56            .map_err(|e| ForgeError::config(format!("Invalid address: {}", e)))?;
57        Ok(self)
58    }
59}
60
61/// Shared state for HTTP handlers
62pub struct HttpState<T> {
63    /// Application state
64    pub app: Arc<RwLock<T>>,
65}
66
67impl<T> Clone for HttpState<T> {
68    fn clone(&self) -> Self {
69        Self {
70            app: Arc::clone(&self.app),
71        }
72    }
73}
74
75/// Health check response
76#[derive(Debug, Serialize)]
77pub struct HealthResponse {
78    /// Health status string (e.g. `"healthy"`).
79    pub status: String,
80    /// Crate/build version.
81    pub version: String,
82    /// Process uptime in seconds.
83    pub uptime_secs: u64,
84}
85
86/// Error response
87#[derive(Debug, Serialize)]
88pub struct ErrorResponse {
89    /// Human-readable error message.
90    pub error: String,
91    /// HTTP status code.
92    pub code: u16,
93}
94
95impl IntoResponse for ErrorResponse {
96    fn into_response(self) -> axum::response::Response {
97        let status = StatusCode::from_u16(self.code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
98        (status, Json(self)).into_response()
99    }
100}
101
102/// Create the base router with health endpoints
103pub fn base_router<T: Send + Sync + 'static>() -> Router<HttpState<T>> {
104    Router::new()
105        .route("/health", get(health_handler))
106        .route("/ready", get(ready_handler))
107}
108
109async fn health_handler() -> Json<HealthResponse> {
110    Json(HealthResponse {
111        status: "healthy".to_string(),
112        version: env!("CARGO_PKG_VERSION").to_string(),
113        uptime_secs: 0, // TODO: Track actual uptime
114    })
115}
116
117async fn ready_handler() -> StatusCode {
118    StatusCode::OK
119}
120
121/// HTTP server wrapper
122pub struct HttpServer {
123    config: HttpServerConfig,
124    router: Router,
125}
126
127impl HttpServer {
128    /// Create a new HTTP server
129    pub fn new(config: HttpServerConfig) -> Self {
130        Self {
131            config,
132            router: Router::new(),
133        }
134    }
135
136    /// Set the router
137    pub fn with_router(mut self, router: Router) -> Self {
138        self.router = router;
139        self
140    }
141
142    /// Start the server
143    pub async fn serve(self) -> Result<()> {
144        let listener = tokio::net::TcpListener::bind(self.config.bind_addr)
145            .await
146            .map_err(|e| ForgeError::network(format!("Failed to bind: {}", e)))?;
147
148        info!(addr = %self.config.bind_addr, "HTTP server starting");
149
150        axum::serve(listener, self.router)
151            .await
152            .map_err(|e| ForgeError::network(format!("Server error: {}", e)))?;
153
154        Ok(())
155    }
156}
157
158/// QUIC transport configuration
159#[cfg(feature = "quic")]
160#[derive(Debug, Clone)]
161pub struct QuicConfig {
162    /// Bind address
163    pub bind_addr: SocketAddr,
164    /// Server name for TLS
165    pub server_name: String,
166    /// Max concurrent streams
167    pub max_streams: u32,
168    /// Idle timeout in seconds
169    pub idle_timeout_secs: u64,
170}
171
172#[cfg(feature = "quic")]
173impl Default for QuicConfig {
174    fn default() -> Self {
175        Self {
176            bind_addr: ([0, 0, 0, 0], 4433).into(),
177            server_name: "forge".to_string(),
178            max_streams: 100,
179            idle_timeout_secs: 30,
180        }
181    }
182}
183
184/// QUIC transport for peer communication
185#[cfg(feature = "quic")]
186pub struct QuicTransport {
187    config: QuicConfig,
188    endpoint: Option<quinn::Endpoint>,
189}
190
191#[cfg(feature = "quic")]
192impl QuicTransport {
193    /// Create a new QUIC transport
194    pub fn new(config: QuicConfig) -> Self {
195        Self {
196            config,
197            endpoint: None,
198        }
199    }
200
201    /// Generate self-signed certificate for development
202    fn generate_self_signed_cert() -> Result<(rustls::pki_types::CertificateDer<'static>, rustls::pki_types::PrivateKeyDer<'static>)> {
203        let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()])
204            .map_err(|e| ForgeError::network(format!("Failed to generate cert: {}", e)))?;
205
206        let cert_der = rustls::pki_types::CertificateDer::from(cert.cert.der().to_vec());
207        let key_der = rustls::pki_types::PrivateKeyDer::try_from(cert.key_pair.serialize_der())
208            .map_err(|e| ForgeError::network(format!("Failed to serialize key: {}", e)))?;
209
210        Ok((cert_der, key_der))
211    }
212
213    /// Start as server
214    pub async fn start_server(&mut self) -> Result<()> {
215        let (cert, key) = Self::generate_self_signed_cert()?;
216
217        let mut server_crypto = rustls::ServerConfig::builder()
218            .with_no_client_auth()
219            .with_single_cert(vec![cert], key)
220            .map_err(|e| ForgeError::network(format!("TLS config error: {}", e)))?;
221
222        server_crypto.alpn_protocols = vec![b"forge".to_vec()];
223
224        let server_config = quinn::ServerConfig::with_crypto(Arc::new(
225            quinn::crypto::rustls::QuicServerConfig::try_from(server_crypto)
226                .map_err(|e| ForgeError::network(format!("QUIC config error: {}", e)))?,
227        ));
228
229        let endpoint = quinn::Endpoint::server(server_config, self.config.bind_addr)
230            .map_err(|e| ForgeError::network(format!("Failed to create endpoint: {}", e)))?;
231
232        info!(addr = %self.config.bind_addr, "QUIC server started");
233        self.endpoint = Some(endpoint);
234
235        Ok(())
236    }
237
238    /// Accept incoming connections
239    pub async fn accept(&self) -> Result<quinn::Connection> {
240        let endpoint = self
241            .endpoint
242            .as_ref()
243            .ok_or_else(|| ForgeError::network("Endpoint not started"))?;
244
245        let incoming = endpoint
246            .accept()
247            .await
248            .ok_or_else(|| ForgeError::network("Endpoint closed"))?;
249
250        let conn = incoming
251            .await
252            .map_err(|e| ForgeError::network(format!("Connection error: {}", e)))?;
253
254        info!(remote = %conn.remote_address(), "QUIC connection accepted");
255        Ok(conn)
256    }
257
258    /// Connect to a peer
259    pub async fn connect(&self, addr: SocketAddr) -> Result<quinn::Connection> {
260        let endpoint = self
261            .endpoint
262            .as_ref()
263            .ok_or_else(|| ForgeError::network("Endpoint not started"))?;
264
265        let conn = endpoint
266            .connect(addr, &self.config.server_name)
267            .map_err(|e| ForgeError::network(format!("Connect error: {}", e)))?
268            .await
269            .map_err(|e| ForgeError::network(format!("Connection error: {}", e)))?;
270
271        info!(remote = %addr, "QUIC connection established");
272        Ok(conn)
273    }
274
275    /// Get local address
276    pub fn local_addr(&self) -> Option<SocketAddr> {
277        self.endpoint.as_ref().and_then(|e| e.local_addr().ok())
278    }
279
280    /// Close the transport
281    pub fn close(&self) {
282        if let Some(endpoint) = &self.endpoint {
283            endpoint.close(0u32.into(), b"shutdown");
284        }
285    }
286}
287
288/// Message types for peer communication
289#[derive(Debug, Clone, Serialize, Deserialize)]
290pub enum PeerMessage {
291    /// Heartbeat ping
292    Ping {
293        /// Sender node id.
294        node_id: String,
295        /// Send timestamp (epoch ms).
296        timestamp: u64,
297    },
298    /// Heartbeat pong
299    Pong {
300        /// Responder node id.
301        node_id: String,
302        /// Send timestamp (epoch ms).
303        timestamp: u64,
304    },
305    /// Route request to expert
306    RouteRequest {
307        /// Correlation id for the request.
308        request_id: String,
309        /// Input payload to route.
310        input: String,
311    },
312    /// Route response from expert
313    RouteResponse {
314        /// Correlation id matching the request.
315        request_id: String,
316        /// Expert index that handled the request.
317        expert_index: usize,
318        /// Serialized result bytes.
319        result: Vec<u8>,
320    },
321    /// Shard assignment notification
322    ShardAssign {
323        /// Shard being assigned.
324        shard_id: u64,
325        /// Node receiving the shard.
326        node_id: String,
327    },
328    /// Shard migration request
329    ShardMigrate {
330        /// Shard being migrated.
331        shard_id: u64,
332        /// Source node id.
333        from_node: String,
334        /// Destination node id.
335        to_node: String,
336    },
337}
338
339impl PeerMessage {
340    /// Serialize to bytes
341    pub fn to_bytes(&self) -> Result<Vec<u8>> {
342        serde_json::to_vec(self).map_err(|e| ForgeError::network(format!("Serialize error: {}", e)))
343    }
344
345    /// Deserialize from bytes
346    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
347        serde_json::from_slice(bytes)
348            .map_err(|e| ForgeError::network(format!("Deserialize error: {}", e)))
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355
356    #[test]
357    fn test_http_config_default() {
358        let config = HttpServerConfig::default();
359        assert_eq!(config.bind_addr.port(), 8080);
360        assert!(config.cors_enabled);
361    }
362
363    #[test]
364    #[cfg(feature = "quic")]
365    fn test_quic_config_default() {
366        let config = QuicConfig::default();
367        assert_eq!(config.bind_addr.port(), 4433);
368        assert_eq!(config.server_name, "forge");
369    }
370
371    #[test]
372    fn test_peer_message_serialization() {
373        let msg = PeerMessage::Ping {
374            node_id: "node-1".to_string(),
375            timestamp: 12345,
376        };
377
378        let bytes = msg.to_bytes().unwrap();
379        let decoded = PeerMessage::from_bytes(&bytes).unwrap();
380
381        match decoded {
382            PeerMessage::Ping { node_id, timestamp } => {
383                assert_eq!(node_id, "node-1");
384                assert_eq!(timestamp, 12345);
385            }
386            _ => panic!("Wrong message type"),
387        }
388    }
389}