1use 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#[derive(Debug, Clone)]
26pub struct HttpServerConfig {
27 pub bind_addr: SocketAddr,
29 pub cors_enabled: bool,
31 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 pub fn with_addr(mut self, addr: SocketAddr) -> Self {
48 self.bind_addr = addr;
49 self
50 }
51
52 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
61pub struct HttpState<T> {
63 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#[derive(Debug, Serialize)]
77pub struct HealthResponse {
78 pub status: String,
80 pub version: String,
82 pub uptime_secs: u64,
84}
85
86#[derive(Debug, Serialize)]
88pub struct ErrorResponse {
89 pub error: String,
91 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
102pub 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, })
115}
116
117async fn ready_handler() -> StatusCode {
118 StatusCode::OK
119}
120
121pub struct HttpServer {
123 config: HttpServerConfig,
124 router: Router,
125}
126
127impl HttpServer {
128 pub fn new(config: HttpServerConfig) -> Self {
130 Self {
131 config,
132 router: Router::new(),
133 }
134 }
135
136 pub fn with_router(mut self, router: Router) -> Self {
138 self.router = router;
139 self
140 }
141
142 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#[cfg(feature = "quic")]
160#[derive(Debug, Clone)]
161pub struct QuicConfig {
162 pub bind_addr: SocketAddr,
164 pub server_name: String,
166 pub max_streams: u32,
168 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#[cfg(feature = "quic")]
186pub struct QuicTransport {
187 config: QuicConfig,
188 endpoint: Option<quinn::Endpoint>,
189}
190
191#[cfg(feature = "quic")]
192impl QuicTransport {
193 pub fn new(config: QuicConfig) -> Self {
195 Self {
196 config,
197 endpoint: None,
198 }
199 }
200
201 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 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 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 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 pub fn local_addr(&self) -> Option<SocketAddr> {
277 self.endpoint.as_ref().and_then(|e| e.local_addr().ok())
278 }
279
280 pub fn close(&self) {
282 if let Some(endpoint) = &self.endpoint {
283 endpoint.close(0u32.into(), b"shutdown");
284 }
285 }
286}
287
288#[derive(Debug, Clone, Serialize, Deserialize)]
290pub enum PeerMessage {
291 Ping {
293 node_id: String,
295 timestamp: u64,
297 },
298 Pong {
300 node_id: String,
302 timestamp: u64,
304 },
305 RouteRequest {
307 request_id: String,
309 input: String,
311 },
312 RouteResponse {
314 request_id: String,
316 expert_index: usize,
318 result: Vec<u8>,
320 },
321 ShardAssign {
323 shard_id: u64,
325 node_id: String,
327 },
328 ShardMigrate {
330 shard_id: u64,
332 from_node: String,
334 to_node: String,
336 },
337}
338
339impl PeerMessage {
340 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 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}