ri/gateway/
server.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![cfg(feature = "gateway")]
19#![allow(non_snake_case)]
20
21use crate::core::{RiResult, RiError};
22use crate::gateway::{RiGateway, RiGatewayConfig, RiGatewayRequest};
23use hyper::{Body, Request as HyperRequest, Response as HyperResponse, Server, StatusCode};
24use hyper::service::{make_service_fn, service_fn};
25use std::collections::HashMap as FxHashMap;
26use std::convert::Infallible;
27use std::net::SocketAddr;
28use std::sync::Arc;
29use tokio::sync::RwLock;
30use tokio_rustls::rustls::ServerConfig;
31
32/// Default maximum request body size (10 MB)
33const DEFAULT_MAX_BODY_SIZE: usize = 10 * 1024 * 1024;
34
35/// Maximum URL path length
36const MAX_PATH_LENGTH: usize = 2048;
37
38/// Maximum number of headers
39const MAX_HEADERS: usize = 100;
40
41pub struct RiGatewayServer {
42    gateway: Arc<RiGateway>,
43    config: Arc<RwLock<RiGatewayConfig>>,
44    addr: SocketAddr,
45    tls_config: Option<ServerConfig>,
46    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
47}
48
49impl RiGatewayServer {
50    pub fn new(gateway: Arc<RiGateway>, config: Arc<RwLock<RiGatewayConfig>>, addr: SocketAddr) -> Self {
51        Self {
52            gateway,
53            config,
54            addr,
55            tls_config: None,
56            shutdown_tx: None,
57        }
58    }
59
60    pub fn with_tls(mut self, tls_config: ServerConfig) -> Self {
61        self.tls_config = Some(tls_config);
62        self
63    }
64
65    pub async fn serve(&mut self) -> RiResult<()> {
66        let addr = self.addr;
67        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
68        self.shutdown_tx = Some(shutdown_tx);
69
70        let gateway = self.gateway.clone();
71        let config = self.config.clone();
72
73        let service = make_service_fn(move |_conn| {
74            let gateway = gateway.clone();
75            let config = config.clone();
76            async move {
77                Ok::<_, Infallible>(service_fn(move |req: HyperRequest<Body>| {
78                    Self::handle_request(req, gateway.clone(), config.clone())
79                }))
80            }
81        });
82
83        let server = Server::bind(&addr)
84            .http1_pipeline_flush(true)
85            .serve(service);
86
87        let graceful = server.with_graceful_shutdown(async {
88            shutdown_rx.await.ok();
89        });
90
91        graceful.await.map_err(|e| RiError::Other(format!("Server error: {}", e)))
92    }
93
94    async fn handle_request(
95        req: HyperRequest<Body>,
96        gateway: Arc<RiGateway>,
97        config: Arc<RwLock<RiGatewayConfig>>,
98    ) -> Result<HyperResponse<Body>, Infallible> {
99        let request_id = uuid::Uuid::new_v4().to_string();
100        let start = std::time::Instant::now();
101
102        let method = req.method().to_string();
103        let path = req.uri().path().to_string();
104        
105        // Security: Validate path length to prevent buffer overflow attacks
106        if path.len() > MAX_PATH_LENGTH {
107            log::warn!(
108                "[Ri.Gateway] Path too long: {} chars (max {})",
109                path.len(), MAX_PATH_LENGTH
110            );
111            return Ok(HyperResponse::builder()
112                .status(StatusCode::URI_TOO_LONG)
113                .body(Body::from("URI Too Long"))
114                .unwrap_or_else(|_| HyperResponse::default()));
115        }
116        
117        // Security: Validate X-Forwarded-For header to prevent IP spoofing
118        // Only trust X-Forwarded-For if the request comes from a trusted proxy
119        // For now, we use the direct connection address as the source of truth
120        // and log the X-Forwarded-For for reference only
121        let direct_remote_addr = req.extensions()
122            .get::<SocketAddr>()
123            .map(|a| a.to_string())
124            .unwrap_or_else(|| "unknown".to_string());
125        
126        let remote_addr = if let Some(xff) = req.headers().get("X-Forwarded-For") {
127            if let Ok(xff_str) = xff.to_str() {
128                // Parse X-Forwarded-For: client, proxy1, proxy2
129                // The rightmost (last) IP is the most recent proxy
130                // The leftmost (first) IP is the original client (if chain is trusted)
131                let ips: Vec<&str> = xff_str.split(',').map(|s| s.trim()).collect();
132                
133                // Validate each IP address format
134                let mut valid_ips = Vec::new();
135                for ip in &ips {
136                    // Basic IP validation: must contain only valid characters
137                    if ip.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == ':' || c == '-') {
138                        // Additional validation: IP must not be empty and have reasonable length
139                        if !ip.is_empty() && ip.len() <= 45 { // Max IPv6 length is 45 chars
140                            valid_ips.push(*ip);
141                        }
142                    }
143                }
144                
145                // Use the first valid IP as the client IP (if any valid IPs exist)
146                // This assumes the proxy chain is trusted
147                // In production, you should validate against a list of trusted proxy IPs
148                if let Some(&client_ip) = valid_ips.first() {
149                    // Log both the direct connection and the forwarded IP for audit
150                    log::debug!(
151                        "[Ri.Gateway] Request from {} (X-Forwarded-For: {}, direct: {})",
152                        client_ip, xff_str, direct_remote_addr
153                    );
154                    client_ip.to_string()
155                } else {
156                    log::warn!("[Ri.Gateway] Invalid X-Forwarded-For format: {}", xff_str);
157                    direct_remote_addr
158                }
159            } else {
160                direct_remote_addr
161            }
162        } else {
163            direct_remote_addr
164        };
165
166        let mut headers = FxHashMap::default();
167        for (key, value) in req.headers() {
168            // Security: Limit number of headers to prevent memory exhaustion
169            if headers.len() >= MAX_HEADERS {
170                log::warn!("[Ri.Gateway] Too many headers, ignoring remaining");
171                break;
172            }
173            if let Ok(v) = value.to_str() {
174                headers.insert(key.as_str().to_string(), v.to_string());
175            }
176        }
177
178        // Security: URL-decode query parameters to prevent encoding-based attacks
179        let query_params = {
180            let uri = req.uri();
181            let query = uri.query().unwrap_or("");
182            let mut params = FxHashMap::default();
183            for pair in query.split('&') {
184                if let Some((key, value)) = pair.split_once('=') {
185                    let decoded_key = urlencoding::decode(key).map(|cow| cow.into_owned()).unwrap_or_else(|_| key.to_string());
186                    let decoded_value = urlencoding::decode(value).map(|cow| cow.into_owned()).unwrap_or_else(|_| value.to_string());
187                    params.insert(decoded_key, decoded_value);
188                }
189            }
190            params
191        };
192
193        // Security: Limit request body size to prevent memory exhaustion attacks
194        let body = match hyper::body::to_bytes(req.into_body()).await {
195            Ok(bytes) => {
196                if bytes.len() > DEFAULT_MAX_BODY_SIZE {
197                    log::warn!(
198                        "[Ri.Gateway] Request body too large: {} bytes (max {} bytes)",
199                        bytes.len(), DEFAULT_MAX_BODY_SIZE
200                    );
201                    return Ok(HyperResponse::builder()
202                        .status(StatusCode::PAYLOAD_TOO_LARGE)
203                        .body(Body::from("Payload Too Large"))
204                        .unwrap_or_else(|_| HyperResponse::default()));
205                }
206                if bytes.is_empty() {
207                    None
208                } else {
209                    Some(bytes.to_vec())
210                }
211            }
212            Err(e) => {
213                log::warn!("[Ri.Gateway] Failed to read request body: {}", e);
214                None
215            }
216        };
217
218        let ri_request = RiGatewayRequest::new(
219            method.clone(),
220            path.clone(),
221            headers,
222            query_params,
223            body,
224            remote_addr.clone(),
225        );
226
227        let response = gateway.handle_request(ri_request).await;
228
229        let duration_ms = start.elapsed().as_secs_f64() * 1000.0;
230
231        if config.read().await.enable_logging {
232            let log_level = &config.read().await.log_level;
233            match log_level.as_str() {
234                "debug" => {
235                    log::debug!(
236                        target: "Ri.Gateway",
237                        "{} {} {} {} {}ms",
238                        method,
239                        path,
240                        response.status_code,
241                        request_id,
242                        duration_ms
243                    );
244                }
245                "info" => {
246                    log::info!(
247                        target: "Ri.Gateway",
248                        "{} {} {} {}ms",
249                        method,
250                        path,
251                        response.status_code,
252                        duration_ms
253                    );
254                }
255                "warn" => {
256                    log::warn!(
257                        target: "Ri.Gateway",
258                        "{} {} {} {}ms",
259                        method,
260                        path,
261                        response.status_code,
262                        duration_ms
263                    );
264                }
265                "error" => {
266                    log::error!(
267                        target: "Ri.Gateway",
268                        "{} {} {} {}ms",
269                        method,
270                        path,
271                        response.status_code,
272                        duration_ms
273                    );
274                }
275                _ => {}
276            }
277        }
278
279        let mut hyper_response = HyperResponse::builder()
280            .status(StatusCode::from_u16(response.status_code).unwrap_or(StatusCode::OK));
281
282        for (key, value) in response.headers {
283            if let (Ok(k), Ok(v)) = (key.parse::<hyper::header::HeaderName>(), value.parse::<hyper::header::HeaderValue>()) {
284                hyper_response = hyper_response.header(k, v);
285            }
286        }
287
288        let body = Body::from(response.body);
289        Ok(hyper_response.body(body).unwrap_or_else(|_| HyperResponse::default()))
290    }
291
292    pub async fn shutdown(&mut self) {
293        if let Some(tx) = self.shutdown_tx.take() {
294            let _ = tx.send(());
295        }
296    }
297}
298
299impl Drop for RiGatewayServer {
300    fn drop(&mut self) {
301        if let Some(tx) = self.shutdown_tx.take() {
302            let _ = tx.send(());
303        }
304    }
305}
306
307pub fn load_tls_config(
308    cert_path: &str,
309    key_path: &str,
310) -> RiResult<ServerConfig> {
311    let cert = std::fs::read(cert_path)
312        .map_err(|e| RiError::Config(format!("Failed to read TLS certificate: {}", e)))?;
313    let key = std::fs::read(key_path)
314        .map_err(|e| RiError::Config(format!("Failed to read TLS key: {}", e)))?;
315
316    let cert_chain = tokio_rustls::rustls::Certificate(cert);
317    let private_key = tokio_rustls::rustls::PrivateKey(key);
318
319    let mut server_config = tokio_rustls::rustls::ServerConfig::builder()
320        .with_safe_defaults()
321        .with_no_client_auth()
322        .with_single_cert(vec![cert_chain], private_key)
323        .map_err(|e| RiError::Config(format!("Failed to build TLS config: {}", e)))?;
324
325    server_config.alpn_protocols = vec!["h2".into(), "http/1.1".into()];
326
327    Ok(server_config)
328}
ri/gateway/server.rs

ri/gateway/
server.rs
