0dbd480
commit 0dbd480
8 files changed
+3046 -0
Showing 8 changed files+3046-0
▼crates/justapi-core/Cargo.toml
+79-0
| @@ -0,0 +1,79 @@ | ||
1 | +[package] | |
2 | +name = "justapi-core" | |
3 | +version.workspace = true | |
4 | +edition.workspace = true | |
5 | +license.workspace = true | |
6 | + | |
7 | +[features] | |
8 | +default = ["opentelemetry"] | |
9 | +tls = ["rustls", "tokio-rustls", "rustls-pemfile"] | |
10 | +ws = ["tokio-tungstenite", "sha1", "base64"] | |
11 | +compression = [] | |
12 | +brotli-compression = ["compression", "brotli"] | |
13 | +zstd-compression = ["compression", "zstd"] | |
14 | +opentelemetry = ["dep:opentelemetry", "dep:opentelemetry_sdk", "dep:opentelemetry-stdout", "dep:opentelemetry-otlp", "dep:tracing-opentelemetry"] | |
15 | +db = ["sqlx", "chrono"] | |
16 | +inference = ["justapi-inference"] | |
17 | + | |
18 | +[dependencies] | |
19 | +tokio.workspace = true | |
20 | +hyper.workspace = true | |
21 | +hyper-util.workspace = true | |
22 | +http-body-util.workspace = true | |
23 | +serde = { workspace = true, features = ["derive"] } | |
24 | +serde_json.workspace = true | |
25 | +serde_urlencoded.workspace = true | |
26 | +tracing.workspace = true | |
27 | +tracing-subscriber = { workspace = true, features = ["env-filter", "json"] } | |
28 | +tracing-appender.workspace = true | |
29 | +anyhow.workspace = true | |
30 | +matchit.workspace = true | |
31 | +async-trait.workspace = true | |
32 | +jsonwebtoken.workspace = true | |
33 | +governor.workspace = true | |
34 | +futures = "0.3" | |
35 | +tokio-util.workspace = true | |
36 | +flate2.workspace = true | |
37 | +brotli = { workspace = true, optional = true } | |
38 | +zstd = { workspace = true, optional = true } | |
39 | +opentelemetry = { workspace = true, optional = true } | |
40 | +opentelemetry_sdk = { workspace = true, optional = true } | |
41 | +opentelemetry-stdout = { workspace = true, optional = true } | |
42 | +opentelemetry-otlp = { workspace = true, optional = true } | |
43 | +tracing-opentelemetry = { workspace = true, optional = true } | |
44 | +simd-json = { workspace = true, optional = true } | |
45 | +tokio-tungstenite = { workspace = true, optional = true } | |
46 | +tokio-stream.workspace = true | |
47 | +sha1 = { workspace = true, optional = true } | |
48 | +base64 = { workspace = true, optional = true } | |
49 | +jsonschema.workspace = true | |
50 | +rustls = { workspace = true, optional = true } | |
51 | +tokio-rustls = { workspace = true, optional = true } | |
52 | +rustls-pemfile = { workspace = true, optional = true } | |
53 | +sqlx = { workspace = true, optional = true } | |
54 | +chrono = { version = "0.4", optional = true } | |
55 | +rand = "0.8" | |
56 | +inventory = "0.3.24" | |
57 | +tonic = "0.14.6" | |
58 | +prost = "0.14.4" | |
59 | +tower = "0.5.3" | |
60 | +http = "1.4.2" | |
61 | +http-body = "1.0" | |
62 | +multer = "3.1" | |
63 | +bytes = "1.12.0" | |
64 | +wasmtime = "46.0.1" | |
65 | +redis = { version = "1.3.0", features = ["tokio-comp", "connection-manager"] } | |
66 | +notify = "8.2.0" | |
67 | +arc-swap = "1.9.2" | |
68 | +async-graphql = "7.2.1" | |
69 | +justapi-inference = { path = "../justapi-inference", optional = true } | |
70 | +tempfile = "3.27.0" | |
71 | + | |
72 | +[dev-dependencies] | |
73 | +proptest = "1" | |
74 | +futures = "0.3" | |
75 | +serde_json.workspace = true | |
76 | +tokio-tungstenite = { workspace = true, features = [] } | |
77 | + | |
78 | +[build-dependencies] | |
79 | +tonic-build = "0.14.6" | |
0 | 80 | |
▼crates/justapi-core/src/dummy_extract.rs
+14-0
| @@ -0,0 +1,14 @@ | ||
1 | +use hyper::HeaderMap; | |
2 | +use opentelemetry::propagation::Extractor; | |
3 | + | |
4 | +pub struct HeaderExtractor<'a>(pub &'a HeaderMap); | |
5 | + | |
6 | +impl<'a> Extractor for HeaderExtractor<'a> { | |
7 | + fn get(&self, key: &str) -> Option<&str> { | |
8 | + self.0.get(key).and_then(|v| v.to_str().ok()) | |
9 | + } | |
10 | + | |
11 | + fn keys(&self) -> Vec<&str> { | |
12 | + self.0.keys().map(|k| k.as_str()).collect() | |
13 | + } | |
14 | +} | |
0 | 15 | |
▼crates/justapi-core/src/error_catalog.rs
+182-0
| @@ -0,0 +1,182 @@ | ||
1 | +//! Error catalog: well-known JustAPI error codes with helpful messages. | |
2 | +//! | |
3 | +//! Each [`CatalogEntry`] bundles a machine-readable code, a short title, a | |
4 | +//! template message, and default suggestions so that CLI/server code can | |
5 | +//! construct rich [`super::dx::Diagnostic`] values with a single call. | |
6 | + | |
7 | +use super::dx::{DiagLevel, Diagnostic}; | |
8 | + | |
9 | +/// A catalogue entry for a well-known error. | |
10 | +#[derive(Debug, Clone)] | |
11 | +pub struct CatalogEntry { | |
12 | + /// Machine-readable code, e.g. `"E001"`. | |
13 | + pub code: &'static str, | |
14 | + /// Short title, e.g. `"Invalid address format"`. | |
15 | + pub title: &'static str, | |
16 | + /// Longer template message. | |
17 | + pub template: &'static str, | |
18 | + /// Default suggestions shown to the user. | |
19 | + pub suggestions: &'static [&'static str], | |
20 | +} | |
21 | + | |
22 | +impl CatalogEntry { | |
23 | + /// Build a [`Diagnostic`] from this catalogue entry. | |
24 | + /// | |
25 | + /// The diagnostic is always at [`DiagLevel::Error`] and carries the | |
26 | + /// catalogue's code, template message, and default suggestions. | |
27 | + pub fn into_diagnostic(&self) -> Diagnostic { | |
28 | + let mut d = Diagnostic::new(DiagLevel::Error, self.template).code(self.code); | |
29 | + for s in self.suggestions { | |
30 | + d = d.suggestion(*s); | |
31 | + } | |
32 | + d | |
33 | + } | |
34 | + | |
35 | + /// Build a diagnostic and attach context (e.g. a file path or value). | |
36 | + pub fn with_context(&self, ctx: impl Into<String>) -> Diagnostic { | |
37 | + self.into_diagnostic().context(ctx) | |
38 | + } | |
39 | +} | |
40 | + | |
41 | +// ── Catalog constants ──────────────────────────────────────────────── | |
42 | + | |
43 | +/// E001 — Invalid address format. | |
44 | +pub const E001_INVALID_ADDRESS: CatalogEntry = CatalogEntry { | |
45 | + code: "E001", | |
46 | + title: "Invalid address format", | |
47 | + template: "the address could not be parsed as HOST:PORT", | |
48 | + suggestions: &[ | |
49 | + "use the format HOST:PORT, e.g. 127.0.0.1:8080", | |
50 | + "ensure the port number is between 1 and 65535", | |
51 | + ], | |
52 | +}; | |
53 | + | |
54 | +/// E002 — Port already in use. | |
55 | +pub const E002_PORT_IN_USE: CatalogEntry = CatalogEntry { | |
56 | + code: "E002", | |
57 | + title: "Port already in use", | |
58 | + template: "the requested port is already bound by another process", | |
59 | + suggestions: &[ | |
60 | + "choose a different port with --addr 127.0.0.1:<PORT>", | |
61 | + "find the blocking process with `lsof -i :<PORT>` or `ss -tlnp`", | |
62 | + ], | |
63 | +}; | |
64 | + | |
65 | +/// E003 — TLS certificate not found. | |
66 | +pub const E003_TLS_CERT_NOT_FOUND: CatalogEntry = CatalogEntry { | |
67 | + code: "E003", | |
68 | + title: "TLS certificate not found", | |
69 | + template: "could not read the TLS certificate file", | |
70 | + suggestions: &[ | |
71 | + "verify the path passed to --tls-cert exists", | |
72 | + "generate a self-signed cert: openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes", | |
73 | + ], | |
74 | +}; | |
75 | + | |
76 | +/// E004 — Database URL invalid. | |
77 | +pub const E004_DB_URL_INVALID: CatalogEntry = CatalogEntry { | |
78 | + code: "E004", | |
79 | + title: "Database URL invalid", | |
80 | + template: "the database connection URL could not be parsed", | |
81 | + suggestions: &[ | |
82 | + "use the format: postgres://user:password@host:5432/dbname", | |
83 | + "or for SQLite: sqlite://path/to/db.sqlite", | |
84 | + ], | |
85 | +}; | |
86 | + | |
87 | +/// E005 — Migration file parse error. | |
88 | +pub const E005_MIGRATION_PARSE: CatalogEntry = CatalogEntry { | |
89 | + code: "E005", | |
90 | + title: "Migration file parse error", | |
91 | + template: "a migration file could not be parsed", | |
92 | + suggestions: &[ | |
93 | + "migration files must be named V<number>__<name>.sql", | |
94 | + "ensure the SQL syntax is valid for your database", | |
95 | + ], | |
96 | +}; | |
97 | + | |
98 | +/// E006 — Python not found. | |
99 | +pub const E006_PYTHON_NOT_FOUND: CatalogEntry = CatalogEntry { | |
100 | + code: "E006", | |
101 | + title: "Python not found", | |
102 | + template: "python3 is not available on your PATH", | |
103 | + suggestions: &[ | |
104 | + "install Python 3.9+ from https://www.python.org/downloads/", | |
105 | + "or use pyenv: `pyenv install 3.12`", | |
106 | + "ensure python3 is on your PATH", | |
107 | + ], | |
108 | +}; | |
109 | + | |
110 | +/// E007 — Route conflict. | |
111 | +pub const E007_ROUTE_CONFLICT: CatalogEntry = CatalogEntry { | |
112 | + code: "E007", | |
113 | + title: "Route conflict", | |
114 | + template: "two or more routes match the same path pattern", | |
115 | + suggestions: &[ | |
116 | + "check your route definitions for duplicate paths", | |
117 | + "use distinct path segments or HTTP methods to disambiguate", | |
118 | + ], | |
119 | +}; | |
120 | + | |
121 | +/// E008 — Invalid configuration. | |
122 | +pub const E008_INVALID_CONFIG: CatalogEntry = CatalogEntry { | |
123 | + code: "E008", | |
124 | + title: "Invalid configuration", | |
125 | + template: "the configuration file contains invalid or missing values", | |
126 | + suggestions: &[ | |
127 | + "check the config file for typos or missing required keys", | |
128 | + "refer to the JustAPI configuration reference for valid options", | |
129 | + ], | |
130 | +}; | |
131 | + | |
132 | +/// All catalogue entries, for enumeration in tests or help text. | |
133 | +pub const ALL_ENTRIES: &[&CatalogEntry] = &[ | |
134 | + &E001_INVALID_ADDRESS, | |
135 | + &E002_PORT_IN_USE, | |
136 | + &E003_TLS_CERT_NOT_FOUND, | |
137 | + &E004_DB_URL_INVALID, | |
138 | + &E005_MIGRATION_PARSE, | |
139 | + &E006_PYTHON_NOT_FOUND, | |
140 | + &E007_ROUTE_CONFLICT, | |
141 | + &E008_INVALID_CONFIG, | |
142 | +]; | |
143 | + | |
144 | +// ── Tests ──────────────────────────────────────────────────────────── | |
145 | + | |
146 | +#[cfg(test)] | |
147 | +mod tests { | |
148 | + use super::*; | |
149 | + | |
150 | + #[test] | |
151 | + fn all_entries_have_unique_codes() { | |
152 | + let mut codes: Vec<&str> = ALL_ENTRIES.iter().map(|e| e.code).collect(); | |
153 | + codes.sort(); | |
154 | + codes.dedup(); | |
155 | + assert_eq!(codes.len(), ALL_ENTRIES.len()); | |
156 | + } | |
157 | + | |
158 | + #[test] | |
159 | + fn catalog_entry_produces_diagnostic() { | |
160 | + let d = E001_INVALID_ADDRESS.into_diagnostic(); | |
161 | + let text = d.render_plain(); | |
162 | + assert!(text.contains("error[E001]")); | |
163 | + assert!(text.contains("= help:")); | |
164 | + } | |
165 | + | |
166 | + #[test] | |
167 | + fn catalog_with_context() { | |
168 | + let d = E003_TLS_CERT_NOT_FOUND.with_context("/etc/ssl/cert.pem"); | |
169 | + let text = d.render_plain(); | |
170 | + assert!(text.contains(" --> /etc/ssl/cert.pem")); | |
171 | + assert!(text.contains("E003")); | |
172 | + } | |
173 | + | |
174 | + #[test] | |
175 | + fn all_entries_render_without_panic() { | |
176 | + for entry in ALL_ENTRIES { | |
177 | + let d = entry.into_diagnostic(); | |
178 | + let _plain = d.render_plain(); | |
179 | + let _colored = d.render_colored(); | |
180 | + } | |
181 | + } | |
182 | +} | |
0 | 183 | |
▼crates/justapi-core/src/extract.rs
+365-0
| @@ -0,0 +1,365 @@ | ||
1 | +use std::collections::HashMap; | |
2 | +use std::str::FromStr; | |
3 | + | |
4 | +use anyhow::Result; | |
5 | +use bytes::Bytes; | |
6 | +use http_body_util::BodyExt; | |
7 | +use hyper::Request; | |
8 | +use serde::de::DeserializeOwned; | |
9 | + | |
10 | +// --------------------------------------------------------------------------- | |
11 | +// Typed form data (application/x-www-form-urlencoded) | |
12 | +// --------------------------------------------------------------------------- | |
13 | + | |
14 | +/// Typed form data extracted from an `application/x-www-form-urlencoded` body. | |
15 | +/// | |
16 | +/// # Example | |
17 | +/// | |
18 | +/// ```ignore | |
19 | +/// #[derive(Deserialize)] | |
20 | +/// struct Login { username: String, password: String } | |
21 | +/// | |
22 | +/// async fn handler(req: Request<Incoming>) -> Result<Response<ResponseBody>> { | |
23 | +/// let form = Form::<Login>::from_request(req).await?; | |
24 | +/// Ok(json_response(StatusCode::OK, &form.username)) | |
25 | +/// } | |
26 | +/// ``` | |
27 | +pub struct Form<T>(pub T); | |
28 | + | |
29 | +impl<T: DeserializeOwned> Form<T> { | |
30 | + /// Parse from raw `application/x-www-form-urlencoded` bytes. | |
31 | + pub fn from_bytes(bytes: &[u8]) -> Result<Self> { | |
32 | + let value: T = serde_urlencoded::from_bytes(bytes) | |
33 | + .map_err(|e| anyhow::anyhow!("form error: {}", e))?; | |
34 | + Ok(Form(value)) | |
35 | + } | |
36 | + | |
37 | + /// Parse from a request. Consumes the request body. Content-Type is NOT | |
38 | + /// checked here — callers should verify it themselves or use the | |
39 | + /// middleware-level guard. | |
40 | + pub async fn from_request<B>(req: Request<B>) -> Result<Self> | |
41 | + where | |
42 | + B: http_body::Body<Data = Bytes> + Send + Unpin + 'static, | |
43 | + B::Error: std::error::Error + Send + Sync + 'static, | |
44 | + { | |
45 | + let (_parts, body) = req.into_parts(); | |
46 | + let bytes = body | |
47 | + .collect() | |
48 | + .await | |
49 | + .map_err(|e| anyhow::anyhow!(e))? | |
50 | + .to_bytes(); | |
51 | + Self::from_bytes(&bytes) | |
52 | + } | |
53 | +} | |
54 | + | |
55 | +// --------------------------------------------------------------------------- | |
56 | +// Typed query parameters | |
57 | +// --------------------------------------------------------------------------- | |
58 | + | |
59 | +/// Typed query parameters extracted from the request URL. | |
60 | +/// | |
61 | +/// # Example | |
62 | +/// | |
63 | +/// ```ignore | |
64 | +/// #[derive(Deserialize)] | |
65 | +/// struct Page { page: Option<u32>, limit: Option<u32> } | |
66 | +/// | |
67 | +/// fn handler(req: Request<Incoming>) -> ... { | |
68 | +/// let q = Query::<Page>::from_request(&req)?; | |
69 | +/// } | |
70 | +/// ``` | |
71 | +pub struct Query<T>(pub T); | |
72 | + | |
73 | +impl<T: DeserializeOwned> Query<T> { | |
74 | + /// Parse from a raw query string (without the leading `?`). | |
75 | + pub fn parse(query: &str) -> Result<Self> { | |
76 | + let cleaned = query.trim_start_matches('?'); | |
77 | + let value: T = serde_urlencoded::from_str(cleaned) | |
78 | + .map_err(|e| anyhow::anyhow!("query error: {}", e))?; | |
79 | + Ok(Query(value)) | |
80 | + } | |
81 | + | |
82 | + /// Parse from a request's URI query string. | |
83 | + pub fn from_request<B>(req: &Request<B>) -> Result<Self> { | |
84 | + Self::parse(req.uri().query().unwrap_or("")) | |
85 | + } | |
86 | +} | |
87 | + | |
88 | +// --------------------------------------------------------------------------- | |
89 | +// Cookie jar | |
90 | +// --------------------------------------------------------------------------- | |
91 | + | |
92 | +/// Parsed cookies from the `Cookie` request header. | |
93 | +/// | |
94 | +/// # Example | |
95 | +/// | |
96 | +/// ```ignore | |
97 | +/// fn handler(req: Request<Incoming>) -> ... { | |
98 | +/// let jar = CookieJar::from_request(&req); | |
99 | +/// if let Some(session) = jar.get("session_id") { ... } | |
100 | +/// } | |
101 | +/// ``` | |
102 | +pub struct CookieJar { | |
103 | + cookies: HashMap<String, String>, | |
104 | +} | |
105 | + | |
106 | +impl CookieJar { | |
107 | + /// Parse the `Cookie` header into a key-value map. | |
108 | + pub fn from_request<B>(req: &Request<B>) -> Self { | |
109 | + let mut cookies = HashMap::new(); | |
110 | + if let Some(header) = req.headers().get("cookie").and_then(|v| v.to_str().ok()) { | |
111 | + for pair in header.split(';') { | |
112 | + let pair = pair.trim(); | |
113 | + if let Some((name, value)) = pair.split_once('=') { | |
114 | + cookies.insert(url_decode(name.trim()), url_decode(value.trim())); | |
115 | + } | |
116 | + } | |
117 | + } | |
118 | + CookieJar { cookies } | |
119 | + } | |
120 | + | |
121 | + pub fn get(&self, name: &str) -> Option<&str> { | |
122 | + self.cookies.get(name).map(|s| s.as_str()) | |
123 | + } | |
124 | + | |
125 | + pub fn is_empty(&self) -> bool { | |
126 | + self.cookies.is_empty() | |
127 | + } | |
128 | + | |
129 | + pub fn len(&self) -> usize { | |
130 | + self.cookies.len() | |
131 | + } | |
132 | + | |
133 | + pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> { | |
134 | + self.cookies.iter().map(|(k, v)| (k.as_str(), v.as_str())) | |
135 | + } | |
136 | +} | |
137 | + | |
138 | +// --------------------------------------------------------------------------- | |
139 | +// Header extraction helpers | |
140 | +// --------------------------------------------------------------------------- | |
141 | + | |
142 | +/// Access a header value as a string. | |
143 | +pub fn get_header<'a, B>(req: &'a Request<B>, name: &str) -> Option<&'a str> { | |
144 | + let header_name = hyper::header::HeaderName::from_bytes(name.as_bytes()).ok()?; | |
145 | + req.headers().get(&header_name)?.to_str().ok() | |
146 | +} | |
147 | + | |
148 | +/// Access a header value parsed into a specific type. | |
149 | +pub fn get_header_typed<T: FromStr, B>(req: &Request<B>, name: &str) -> Result<Option<T>> | |
150 | +where | |
151 | + T::Err: std::fmt::Display, | |
152 | +{ | |
153 | + match get_header(req, name) { | |
154 | + Some(val) => val | |
155 | + .parse::<T>() | |
156 | + .map(Some) | |
157 | + .map_err(|e| anyhow::anyhow!("invalid header '{}': {}", name, e)), | |
158 | + None => Ok(None), | |
159 | + } | |
160 | +} | |
161 | + | |
162 | +// --------------------------------------------------------------------------- | |
163 | +// Helpers | |
164 | +// --------------------------------------------------------------------------- | |
165 | + | |
166 | +/// Minimal URL-decoding (percent-encoding and '+') for cookie values. | |
167 | +fn url_decode(s: &str) -> String { | |
168 | + let mut result = String::with_capacity(s.len()); | |
169 | + let mut chars = s.chars(); | |
170 | + while let Some(c) = chars.next() { | |
171 | + if c == '%' { | |
172 | + let hi = chars.next().and_then(|c| c.to_digit(16)).unwrap_or(0); | |
173 | + let lo = chars.next().and_then(|c| c.to_digit(16)).unwrap_or(0); | |
174 | + result.push((hi as u8 * 16 + lo as u8) as char); | |
175 | + } else if c == '+' { | |
176 | + result.push(' '); | |
177 | + } else { | |
178 | + result.push(c); | |
179 | + } | |
180 | + } | |
181 | + result | |
182 | +} | |
183 | + | |
184 | +// --------------------------------------------------------------------------- | |
185 | +// Tests | |
186 | +// --------------------------------------------------------------------------- | |
187 | + | |
188 | +#[cfg(test)] | |
189 | +mod tests { | |
190 | + use super::*; | |
191 | + use hyper::Method; | |
192 | + use serde::Deserialize; | |
193 | + | |
194 | + type TestBody = http_body_util::Full<Bytes>; | |
195 | + | |
196 | + fn test_req(method: Method, uri: &str) -> Request<TestBody> { | |
197 | + Request::builder() | |
198 | + .method(method) | |
199 | + .uri(uri) | |
200 | + .body(TestBody::new(Bytes::new())) | |
201 | + .unwrap() | |
202 | + } | |
203 | + | |
204 | + fn form_req(body: &str) -> Request<TestBody> { | |
205 | + Request::builder() | |
206 | + .method(Method::POST) | |
207 | + .uri("/submit") | |
208 | + .header("content-type", "application/x-www-form-urlencoded") | |
209 | + .body(TestBody::new(Bytes::from(body.to_string()))) | |
210 | + .unwrap() | |
211 | + } | |
212 | + | |
213 | + // --- Form tests --- | |
214 | + | |
215 | + #[derive(Deserialize, Debug, PartialEq)] | |
216 | + struct LoginForm { | |
217 | + username: String, | |
218 | + password: String, | |
219 | + } | |
220 | + | |
221 | + #[test] | |
222 | + fn form_from_bytes_parses() { | |
223 | + let form = Form::<LoginForm>::from_bytes(b"username=alice&password=secret").unwrap(); | |
224 | + assert_eq!(form.0.username, "alice"); | |
225 | + assert_eq!(form.0.password, "secret"); | |
226 | + } | |
227 | + | |
228 | + #[test] | |
229 | + fn form_from_bytes_missing_field_errors() { | |
230 | + let result = Form::<LoginForm>::from_bytes(b"username=alice"); | |
231 | + assert!(result.is_err()); | |
232 | + } | |
233 | + | |
234 | + #[tokio::test] | |
235 | + async fn form_from_request_parses() { | |
236 | + let req = form_req("username=bob&password=pass"); | |
237 | + let form = Form::<LoginForm>::from_request(req).await.unwrap(); | |
238 | + assert_eq!(form.0.username, "bob"); | |
239 | + } | |
240 | + | |
241 | + #[tokio::test] | |
242 | + async fn form_from_request_empty_errors() { | |
243 | + let req = form_req(""); | |
244 | + let result = Form::<LoginForm>::from_request(req).await; | |
245 | + assert!(result.is_err()); | |
246 | + } | |
247 | + | |
248 | + // --- Query tests --- | |
249 | + | |
250 | + #[derive(Deserialize, Debug, PartialEq)] | |
251 | + struct Pagination { | |
252 | + page: Option<u32>, | |
253 | + limit: Option<u32>, | |
254 | + } | |
255 | + | |
256 | + #[test] | |
257 | + fn query_parse_parses() { | |
258 | + let q = Query::<Pagination>::parse("page=1&limit=20").unwrap(); | |
259 | + assert_eq!(q.0.page, Some(1)); | |
260 | + assert_eq!(q.0.limit, Some(20)); | |
261 | + } | |
262 | + | |
263 | + #[test] | |
264 | + fn query_from_request_parses() { | |
265 | + let req = test_req(Method::GET, "/items?page=3&limit=10"); | |
266 | + let q = Query::<Pagination>::from_request(&req).unwrap(); | |
267 | + assert_eq!(q.0.page, Some(3)); | |
268 | + } | |
269 | + | |
270 | + #[test] | |
271 | + fn query_empty_returns_defaults() { | |
272 | + let req = test_req(Method::GET, "/items"); | |
273 | + let q = Query::<Pagination>::from_request(&req).unwrap(); | |
274 | + assert_eq!(q.0.page, None); | |
275 | + assert_eq!(q.0.limit, None); | |
276 | + } | |
277 | + | |
278 | + #[test] | |
279 | + fn query_parse_with_leading_qmark() { | |
280 | + let q = Query::<Pagination>::parse("?page=5").unwrap(); | |
281 | + assert_eq!(q.0.page, Some(5)); | |
282 | + } | |
283 | + | |
284 | + // --- CookieJar tests --- | |
285 | + | |
286 | + #[test] | |
287 | + fn cookies_single() { | |
288 | + let mut req = test_req(Method::GET, "/"); | |
289 | + req.headers_mut() | |
290 | + .insert("cookie", "session_id=abc123".parse().unwrap()); | |
291 | + let jar = CookieJar::from_request(&req); | |
292 | + assert_eq!(jar.get("session_id"), Some("abc123")); | |
293 | + } | |
294 | + | |
295 | + #[test] | |
296 | + fn cookies_multiple() { | |
297 | + let mut req = test_req(Method::GET, "/"); | |
298 | + req.headers_mut().insert( | |
299 | + "cookie", | |
300 | + "session=xyz; theme=dark; lang=en-US".parse().unwrap(), | |
301 | + ); | |
302 | + let jar = CookieJar::from_request(&req); | |
303 | + assert_eq!(jar.get("session"), Some("xyz")); | |
304 | + assert_eq!(jar.get("theme"), Some("dark")); | |
305 | + assert_eq!(jar.get("lang"), Some("en-US")); | |
306 | + assert_eq!(jar.len(), 3); | |
307 | + } | |
308 | + | |
309 | + #[test] | |
310 | + fn cookies_no_header_empty() { | |
311 | + let req = test_req(Method::GET, "/"); | |
312 | + let jar = CookieJar::from_request(&req); | |
313 | + assert!(jar.is_empty()); | |
314 | + } | |
315 | + | |
316 | + #[test] | |
317 | + fn cookies_url_decoded() { | |
318 | + let mut req = test_req(Method::GET, "/"); | |
319 | + req.headers_mut() | |
320 | + .insert("cookie", "name=Hello%20World".parse().unwrap()); | |
321 | + let jar = CookieJar::from_request(&req); | |
322 | + assert_eq!(jar.get("name"), Some("Hello World")); | |
323 | + } | |
324 | + | |
325 | + #[test] | |
326 | + fn cookies_plus_decoded() { | |
327 | + let mut req = test_req(Method::GET, "/"); | |
328 | + req.headers_mut() | |
329 | + .insert("cookie", "q=rust+lang".parse().unwrap()); | |
330 | + let jar = CookieJar::from_request(&req); | |
331 | + assert_eq!(jar.get("q"), Some("rust lang")); | |
332 | + } | |
333 | + | |
334 | + // --- Header helpers --- | |
335 | + | |
336 | + #[test] | |
337 | + fn header_value_present() { | |
338 | + let mut req = test_req(Method::GET, "/"); | |
339 | + req.headers_mut() | |
340 | + .insert("x-api-key", "secret".parse().unwrap()); | |
341 | + assert_eq!(get_header(&req, "x-api-key"), Some("secret")); | |
342 | + } | |
343 | + | |
344 | + #[test] | |
345 | + fn header_value_missing() { | |
346 | + let req = test_req(Method::GET, "/"); | |
347 | + assert_eq!(get_header(&req, "x-api-key"), None); | |
348 | + } | |
349 | + | |
350 | + #[test] | |
351 | + fn header_typed_parses() { | |
352 | + let mut req = test_req(Method::GET, "/"); | |
353 | + req.headers_mut().insert("x-count", "42".parse().unwrap()); | |
354 | + let val: Option<u32> = get_header_typed(&req, "x-count").unwrap(); | |
355 | + assert_eq!(val, Some(42)); | |
356 | + } | |
357 | + | |
358 | + #[test] | |
359 | + fn header_typed_invalid_errors() { | |
360 | + let mut req = test_req(Method::GET, "/"); | |
361 | + req.headers_mut().insert("x-count", "abc".parse().unwrap()); | |
362 | + let result: Result<Option<u32>> = get_header_typed(&req, "x-count"); | |
363 | + assert!(result.is_err()); | |
364 | + } | |
365 | +} | |
0 | 366 | |
▼crates/justapi-core/src/lib.rs
+92-0
| @@ -0,0 +1,92 @@ | ||
1 | +pub mod alerting; | |
2 | +pub mod audit; | |
3 | +pub mod batching; | |
4 | +pub mod coalesce; | |
5 | +pub mod compress; | |
6 | +#[cfg(feature = "db")] | |
7 | +pub mod db; | |
8 | +pub mod dx; | |
9 | +pub mod error_catalog; | |
10 | +pub mod extract; | |
11 | +pub mod gateway; | |
12 | +pub mod graphql; | |
13 | +pub mod health; | |
14 | +pub mod memory; | |
15 | +pub mod metrics; | |
16 | +pub mod middleware; | |
17 | +pub mod multipart; | |
18 | +#[cfg(feature = "inference")] | |
19 | +pub mod openai; | |
20 | +pub mod openapi; | |
21 | +pub mod panic; | |
22 | +pub mod plugin; | |
23 | +pub mod rate_limit; | |
24 | +pub mod resilience; | |
25 | +pub mod router; | |
26 | +pub mod secrets; | |
27 | +pub mod serialize; | |
28 | +pub mod server; | |
29 | +pub mod static_files; | |
30 | +pub mod testing; | |
31 | +pub mod trace_context; | |
32 | +pub mod tracing_setup; | |
33 | +pub mod validate; | |
34 | +pub mod wasm; | |
35 | + | |
36 | +pub use server::{serve, Server}; | |
37 | + | |
38 | +use std::str::FromStr; | |
39 | + | |
40 | +use hyper::Method; | |
41 | + | |
42 | +/// The HTTP QUERY method (RFC 10008): a safe, idempotent request that | |
43 | +/// carries a representation (like POST) but must not change state (like GET). | |
44 | +/// | |
45 | +/// `hyper` does not ship a `Method::QUERY` constant, so JustAPI provides a | |
46 | +/// helper that returns one. | |
47 | +pub fn query_method() -> Method { | |
48 | + Method::from_str("QUERY").expect("QUERY is a valid HTTP method name") | |
49 | +} | |
50 | + | |
51 | +use anyhow::Result; | |
52 | +use futures::StreamExt; | |
53 | +use http_body_util::combinators::UnsyncBoxBody; | |
54 | +use http_body_util::BodyExt; | |
55 | +use http_body_util::Full; | |
56 | +use hyper::body::{Bytes, Frame}; | |
57 | +use hyper::Response; | |
58 | + | |
59 | +/// Unified response body type supporting both static and streaming payloads. | |
60 | +pub type ResponseBody = UnsyncBoxBody<Bytes, anyhow::Error>; | |
61 | + | |
62 | +/// Lift a static string into a JSON response. | |
63 | +pub fn json_response(status: hyper::StatusCode, body: &str) -> Response<ResponseBody> { | |
64 | + Response::builder() | |
65 | + .status(status) | |
66 | + .header("content-type", "application/json") | |
67 | + .header("content-length", body.len().to_string()) | |
68 | + .body(UnsyncBoxBody::new( | |
69 | + Full::new(Bytes::from(body.to_string())) | |
70 | + .map_err(|e: std::convert::Infallible| -> anyhow::Error { match e {} }), | |
71 | + )) | |
72 | + .unwrap() | |
73 | +} | |
74 | + | |
75 | +/// Lift an `anyhow::Result<Bytes>` stream into a streaming response. | |
76 | +pub fn streaming_response( | |
77 | + status: hyper::StatusCode, | |
78 | + content_type: &str, | |
79 | + stream: impl futures::Stream<Item = Result<Bytes>> + Send + 'static, | |
80 | +) -> Response<ResponseBody> { | |
81 | + let frame_stream = stream.map(|r| r.map(Frame::data)); | |
82 | + Response::builder() | |
83 | + .status(status) | |
84 | + .header("content-type", content_type) | |
85 | + .body(UnsyncBoxBody::new(http_body_util::StreamBody::new( | |
86 | + frame_stream, | |
87 | + ))) | |
88 | + .unwrap() | |
89 | +} | |
90 | +pub mod dummy_extract; | |
91 | +pub mod grpc; | |
92 | +pub mod test_codec; | |
0 | 93 | |
▼crates/justapi-core/src/router.rs
+271-0
| @@ -0,0 +1,271 @@ | ||
1 | +use std::collections::HashMap; | |
2 | + | |
3 | +use hyper::Method; | |
4 | +use matchit::Router as MatchitRouter; | |
5 | + | |
6 | +pub use matchit::{InsertError, MatchError}; | |
7 | + | |
8 | +#[derive(Debug, PartialEq, Eq)] | |
9 | +pub enum RouterError { | |
10 | + NotFound, | |
11 | + MethodNotAllowed, | |
12 | +} | |
13 | + | |
14 | +#[derive(Debug, Clone)] | |
15 | +pub struct Match<'a, 'p, T> { | |
16 | + pub handler: &'a T, | |
17 | + pub params: matchit::Params<'a, 'p>, | |
18 | +} | |
19 | + | |
20 | +#[derive(Debug, Clone)] | |
21 | +pub struct Router<T> { | |
22 | + routes: HashMap<Method, MatchitRouter<T>>, | |
23 | + fallback: Option<T>, | |
24 | + route_list: Vec<(Method, String)>, | |
25 | +} | |
26 | + | |
27 | +impl<T: Clone> Router<T> { | |
28 | + pub fn new() -> Self { | |
29 | + Self { | |
30 | + routes: HashMap::new(), | |
31 | + fallback: None, | |
32 | + route_list: Vec::new(), | |
33 | + } | |
34 | + } | |
35 | + | |
36 | + pub fn insert(&mut self, method: Method, path: &str, handler: T) -> Result<(), InsertError> { | |
37 | + let is_get = method == Method::GET; | |
38 | + let router = self.routes.entry(method.clone()).or_default(); | |
39 | + router.insert(path, handler.clone())?; | |
40 | + self.route_list.push((method, path.to_string())); | |
41 | + | |
42 | + if is_get { | |
43 | + let head_router = self.routes.entry(Method::HEAD).or_default(); | |
44 | + let _ = head_router.insert(path, handler); | |
45 | + self.route_list.push((Method::HEAD, path.to_string())); | |
46 | + } | |
47 | + Ok(()) | |
48 | + } | |
49 | + | |
50 | + /// List all registered routes as `(method, path)` pairs. | |
51 | + pub fn list_routes(&self) -> &[(Method, String)] { | |
52 | + &self.route_list | |
53 | + } | |
54 | + | |
55 | + pub fn at<'a, 'p>( | |
56 | + &'a self, | |
57 | + method: &Method, | |
58 | + path: &'p str, | |
59 | + ) -> Result<Match<'a, 'p, T>, RouterError> { | |
60 | + if let Some(router) = self.routes.get(method) { | |
61 | + if let Ok(matched) = router.at(path) { | |
62 | + return Ok(Match { | |
63 | + handler: matched.value, | |
64 | + params: matched.params, | |
65 | + }); | |
66 | + } | |
67 | + } | |
68 | + | |
69 | + for (other_method, router) in &self.routes { | |
70 | + if other_method != method && router.at(path).is_ok() { | |
71 | + return Err(RouterError::MethodNotAllowed); | |
72 | + } | |
73 | + } | |
74 | + | |
75 | + Err(RouterError::NotFound) | |
76 | + } | |
77 | + | |
78 | + pub fn set_fallback(&mut self, handler: T) { | |
79 | + self.fallback = Some(handler); | |
80 | + } | |
81 | + | |
82 | + pub fn fallback(&self) -> Option<&T> { | |
83 | + self.fallback.as_ref() | |
84 | + } | |
85 | +} | |
86 | + | |
87 | +impl<T: Clone> Default for Router<T> { | |
88 | + fn default() -> Self { | |
89 | + Self::new() | |
90 | + } | |
91 | +} | |
92 | + | |
93 | +#[cfg(test)] | |
94 | +mod tests { | |
95 | + use super::*; | |
96 | + | |
97 | + #[test] | |
98 | + fn test_basic_route() { | |
99 | + let mut router = Router::new(); | |
100 | + router | |
101 | + .insert(Method::GET, "/hello", "hello_handler") | |
102 | + .unwrap(); | |
103 | + let m = router.at(&Method::GET, "/hello").unwrap(); | |
104 | + assert_eq!(*m.handler, "hello_handler"); | |
105 | + } | |
106 | + | |
107 | + #[test] | |
108 | + fn test_route_not_found() { | |
109 | + let mut router = Router::new(); | |
110 | + router | |
111 | + .insert(Method::GET, "/hello", "hello_handler") | |
112 | + .unwrap(); | |
113 | + assert!(router.at(&Method::GET, "/unknown").is_err()); | |
114 | + } | |
115 | + | |
116 | + #[test] | |
117 | + fn test_wrong_method() { | |
118 | + let mut router = Router::new(); | |
119 | + router | |
120 | + .insert(Method::GET, "/hello", "hello_handler") | |
121 | + .unwrap(); | |
122 | + assert_eq!( | |
123 | + router.at(&Method::POST, "/hello").err().unwrap(), | |
124 | + RouterError::MethodNotAllowed | |
125 | + ); | |
126 | + } | |
127 | + | |
128 | + #[test] | |
129 | + fn test_path_param() { | |
130 | + let mut router = Router::new(); | |
131 | + router | |
132 | + .insert(Method::GET, "/users/{id}", "user_handler") | |
133 | + .unwrap(); | |
134 | + let m = router.at(&Method::GET, "/users/42").unwrap(); | |
135 | + assert_eq!(*m.handler, "user_handler"); | |
136 | + assert_eq!(m.params.get("id"), Some("42")); | |
137 | + } | |
138 | + | |
139 | + #[test] | |
140 | + fn test_catch_all() { | |
141 | + let mut router = Router::new(); | |
142 | + router | |
143 | + .insert(Method::GET, "/files/{*path}", "file_handler") | |
144 | + .unwrap(); | |
145 | + let m = router.at(&Method::GET, "/files/a/b/c").unwrap(); | |
146 | + assert_eq!(*m.handler, "file_handler"); | |
147 | + assert_eq!(m.params.get("path"), Some("a/b/c")); | |
148 | + } | |
149 | + | |
150 | + #[test] | |
151 | + fn test_multiple_routes() { | |
152 | + let mut router = Router::new(); | |
153 | + router | |
154 | + .insert(Method::GET, "/users/{id}", "user_handler") | |
155 | + .unwrap(); | |
156 | + router.insert(Method::GET, "/users", "users_list").unwrap(); | |
157 | + let m = router.at(&Method::GET, "/users/42").unwrap(); | |
158 | + assert_eq!(*m.handler, "user_handler"); | |
159 | + let m = router.at(&Method::GET, "/users").unwrap(); | |
160 | + assert_eq!(*m.handler, "users_list"); | |
161 | + } | |
162 | + | |
163 | + #[test] | |
164 | + fn test_post_route() { | |
165 | + let mut router = Router::new(); | |
166 | + router | |
167 | + .insert(Method::POST, "/data", "data_handler") | |
168 | + .unwrap(); | |
169 | + let m = router.at(&Method::POST, "/data").unwrap(); | |
170 | + assert_eq!(*m.handler, "data_handler"); | |
171 | + } | |
172 | + | |
173 | + #[test] | |
174 | + fn test_query_method_route() { | |
175 | + let mut router = Router::new(); | |
176 | + router | |
177 | + .insert(crate::query_method(), "/search", "search_handler") | |
178 | + .unwrap(); | |
179 | + let m = router.at(&crate::query_method(), "/search").unwrap(); | |
180 | + assert_eq!(*m.handler, "search_handler"); | |
181 | + // QUERY is distinct from POST on the same path. | |
182 | + assert_eq!( | |
183 | + router.at(&Method::POST, "/search").err().unwrap(), | |
184 | + RouterError::MethodNotAllowed | |
185 | + ); | |
186 | + } | |
187 | + | |
188 | + #[test] | |
189 | + fn test_fallback() { | |
190 | + let mut router = Router::new(); | |
191 | + router.set_fallback("not_found"); | |
192 | + assert_eq!(*router.fallback().unwrap(), "not_found"); | |
193 | + } | |
194 | + | |
195 | + #[test] | |
196 | + fn test_route_conflict() { | |
197 | + let mut router = Router::new(); | |
198 | + router.insert(Method::GET, "/users/{id}", "a").unwrap(); | |
199 | + let err = router.insert(Method::GET, "/users/{name}", "b"); | |
200 | + assert!(err.is_err()); | |
201 | + } | |
202 | + | |
203 | + #[test] | |
204 | + fn bench_500_route_lookup() { | |
205 | + let mut router = Router::new(); | |
206 | + // Insert 500 routes with various patterns | |
207 | + for i in 0..500 { | |
208 | + router | |
209 | + .insert( | |
210 | + Method::GET, | |
211 | + &format!("/api/v{}/users/{}/posts/{}", i % 5, i, i * 2), | |
212 | + format!("handler_{}", i), | |
213 | + ) | |
214 | + .unwrap(); | |
215 | + } | |
216 | + // Add some parameterized routes | |
217 | + for i in 0..100 { | |
218 | + router | |
219 | + .insert( | |
220 | + Method::GET, | |
221 | + &format!("/api/users/{{id}}/posts/{{post_id}}/comments/{}", i), | |
222 | + format!("param_handler_{}", i), | |
223 | + ) | |
224 | + .unwrap(); | |
225 | + } | |
226 | + // Catch-all | |
227 | + router | |
228 | + .insert(Method::GET, "/static/{*path}", "static_handler".to_string()) | |
229 | + .unwrap(); | |
230 | + | |
231 | + // Warmup | |
232 | + for _ in 0..1000 { | |
233 | + let _ = router.at(&Method::GET, "/api/v3/users/42/posts/84"); | |
234 | + } | |
235 | + | |
236 | + let iterations = 100_000u64; | |
237 | + let mut total_ns = 0u64; | |
238 | + | |
239 | + let test_paths = [ | |
240 | + "/api/v3/users/42/posts/84", | |
241 | + "/api/v0/users/100/posts/200", | |
242 | + "/api/users/999/posts/888/comments/50", | |
243 | + "/static/css/main.css", | |
244 | + "/unknown/path", | |
245 | + ]; | |
246 | + | |
247 | + for _ in 0..iterations { | |
248 | + for path in &test_paths { | |
249 | + let start = std::time::Instant::now(); | |
250 | + let _ = router.at(&Method::GET, path); | |
251 | + total_ns += start.elapsed().as_nanos() as u64; | |
252 | + } | |
253 | + } | |
254 | + | |
255 | + let total_lookups = iterations * test_paths.len() as u64; | |
256 | + let avg_ns = total_ns / total_lookups; | |
257 | + println!( | |
258 | + "500-route table: {} lookups, avg {}ns", | |
259 | + total_lookups, avg_ns, | |
260 | + ); | |
261 | + | |
262 | + // Benchmark gate: average lookup should be under 100ns (release) / 1us (debug) | |
263 | + let target = if cfg!(debug_assertions) { 1000 } else { 100 }; | |
264 | + assert!( | |
265 | + avg_ns < target, | |
266 | + "Route lookup too slow: {}ns (target < {}ns)", | |
267 | + avg_ns, | |
268 | + target | |
269 | + ); | |
270 | + } | |
271 | +} | |
0 | 272 | |
▼crates/justapi-core/src/serialize.rs
+247-0
| @@ -0,0 +1,247 @@ | ||
1 | +use anyhow::Result; | |
2 | +use bytes::Bytes; | |
3 | +use http_body_util::combinators::UnsyncBoxBody; | |
4 | +use http_body_util::BodyExt; | |
5 | +use hyper::header::HeaderName; | |
6 | +use hyper::{Response, StatusCode}; | |
7 | + | |
8 | +use crate::ResponseBody; | |
9 | + | |
10 | +// --------------------------------------------------------------------------- | |
11 | +// Low-level serialization helpers | |
12 | +// --------------------------------------------------------------------------- | |
13 | + | |
14 | +/// Serialize a value to a JSON string. | |
15 | +/// | |
16 | +/// Uses `simd-json` when the `simd-json` feature is enabled, | |
17 | +/// falling back to `serde_json` otherwise. | |
18 | +#[cfg(feature = "simd-json")] | |
19 | +pub fn to_json_string<T: serde::Serialize>(value: &T) -> Result<String> { | |
20 | + Ok(simd_json::serde::to_string(value)?) | |
21 | +} | |
22 | + | |
23 | +#[cfg(not(feature = "simd-json"))] | |
24 | +pub fn to_json_string<T: serde::Serialize>(value: &T) -> Result<String> { | |
25 | + Ok(serde_json::to_string(value)?) | |
26 | +} | |
27 | + | |
28 | +/// Serialize a value to a JSON byte vector. | |
29 | +#[cfg(feature = "simd-json")] | |
30 | +pub fn to_json_vec<T: serde::Serialize>(value: &T) -> Result<Vec<u8>> { | |
31 | + Ok(simd_json::serde::to_vec(value)?) | |
32 | +} | |
33 | + | |
34 | +#[cfg(not(feature = "simd-json"))] | |
35 | +pub fn to_json_vec<T: serde::Serialize>(value: &T) -> Result<Vec<u8>> { | |
36 | + Ok(serde_json::to_vec(value)?) | |
37 | +} | |
38 | + | |
39 | +// --------------------------------------------------------------------------- | |
40 | +// JsonResponse — typed response-model shaping | |
41 | +// --------------------------------------------------------------------------- | |
42 | + | |
43 | +/// A JSON response shaped from any `Serialize` type (FastAPI's `response_model` | |
44 | +/// equivalent, but zero-copy Rust + serde — no Python GIL overhead). | |
45 | +/// | |
46 | +/// Use the builder methods to set status code and headers, then call | |
47 | +/// [`into_response`](Self::into_response) to produce a `Response<ResponseBody>`. | |
48 | +/// | |
49 | +/// # Example | |
50 | +/// | |
51 | +/// ```ignore | |
52 | +/// #[derive(Serialize)] | |
53 | +/// struct User { id: u64, name: String } | |
54 | +/// | |
55 | +/// let resp = JsonResponse::new(User { id: 1, name: "Alice".into() }) | |
56 | +/// .with_status(StatusCode::CREATED) | |
57 | +/// .with_header("x-request-id", "abc-123") | |
58 | +/// .into_response(); | |
59 | +/// ``` | |
60 | +pub struct JsonResponse<T: serde::Serialize> { | |
61 | + data: T, | |
62 | + status: StatusCode, | |
63 | + headers: Vec<(HeaderName, String)>, | |
64 | +} | |
65 | + | |
66 | +impl<T: serde::Serialize> JsonResponse<T> { | |
67 | + /// Create a new `JsonResponse` with status 200 OK and no extra headers. | |
68 | + pub fn new(data: T) -> Self { | |
69 | + Self { | |
70 | + data, | |
71 | + status: StatusCode::OK, | |
72 | + headers: Vec::new(), | |
73 | + } | |
74 | + } | |
75 | + | |
76 | + /// Set the HTTP status code. | |
77 | + pub fn with_status(mut self, status: StatusCode) -> Self { | |
78 | + self.status = status; | |
79 | + self | |
80 | + } | |
81 | + | |
82 | + /// Add a response header. | |
83 | + pub fn with_header(mut self, name: &str, value: &str) -> Self { | |
84 | + if let Ok(h) = HeaderName::from_bytes(name.as_bytes()) { | |
85 | + self.headers.push((h, value.to_string())); | |
86 | + } | |
87 | + self | |
88 | + } | |
89 | + | |
90 | + /// Serialize the wrapped data and produce a `Response<ResponseBody>`. | |
91 | + pub fn into_response(self) -> Result<Response<ResponseBody>> { | |
92 | + let body_bytes = to_json_vec(&self.data)?; | |
93 | + let mut builder = Response::builder().status(self.status); | |
94 | + builder = builder.header("content-type", "application/json"); | |
95 | + builder = builder.header("content-length", body_bytes.len().to_string()); | |
96 | + for (name, value) in &self.headers { | |
97 | + builder = builder.header(name, value); | |
98 | + } | |
99 | + let resp = builder | |
100 | + .body(UnsyncBoxBody::new( | |
101 | + http_body_util::Full::new(Bytes::from(body_bytes)) | |
102 | + .map_err(|e: std::convert::Infallible| -> anyhow::Error { match e {} }), | |
103 | + )) | |
104 | + .unwrap(); | |
105 | + Ok(resp) | |
106 | + } | |
107 | +} | |
108 | + | |
109 | +/// Convenience: create a `JsonResponse` with status 200 from any Serialize value. | |
110 | +pub fn json_response_from<T: serde::Serialize>(data: T) -> JsonResponse<T> { | |
111 | + JsonResponse::new(data) | |
112 | +} | |
113 | + | |
114 | +#[cfg(test)] | |
115 | +mod tests { | |
116 | + use super::*; | |
117 | + use http_body_util::BodyExt; | |
118 | + use serde::Serialize; | |
119 | + | |
120 | + #[derive(Serialize)] | |
121 | + struct TestPayload { | |
122 | + message: String, | |
123 | + count: u32, | |
124 | + } | |
125 | + | |
126 | + #[test] | |
127 | + fn test_to_json_string() { | |
128 | + let payload = TestPayload { | |
129 | + message: "hello".into(), | |
130 | + count: 42, | |
131 | + }; | |
132 | + let json = to_json_string(&payload).unwrap(); | |
133 | + assert!(json.contains(r#""message":"hello""#)); | |
134 | + assert!(json.contains(r#""count":42"#)); | |
135 | + } | |
136 | + | |
137 | + #[test] | |
138 | + fn test_to_json_vec() { | |
139 | + let payload = TestPayload { | |
140 | + message: "hello".into(), | |
141 | + count: 42, | |
142 | + }; | |
143 | + let vec = to_json_vec(&payload).unwrap(); | |
144 | + let json = String::from_utf8(vec).unwrap(); | |
145 | + assert!(json.contains(r#""message":"hello""#)); | |
146 | + } | |
147 | + | |
148 | + // --- JsonResponse tests --- | |
149 | + | |
150 | + #[derive(Serialize)] | |
151 | + struct UserResp { | |
152 | + id: u64, | |
153 | + name: String, | |
154 | + #[serde(skip_serializing_if = "Option::is_none")] | |
155 | + email: Option<String>, | |
156 | + } | |
157 | + | |
158 | + #[tokio::test] | |
159 | + async fn json_response_default_status() { | |
160 | + let user = UserResp { | |
161 | + id: 1, | |
162 | + name: "Alice".into(), | |
163 | + email: None, | |
164 | + }; | |
165 | + let resp = JsonResponse::new(user).into_response().unwrap(); | |
166 | + assert_eq!(resp.status(), StatusCode::OK); | |
167 | + let body = resp.into_body().collect().await.unwrap().to_bytes(); | |
168 | + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); | |
169 | + assert_eq!(json["id"], 1); | |
170 | + assert_eq!(json["name"], "Alice"); | |
171 | + // email field skipped because None | |
172 | + assert!(json.get("email").is_none()); | |
173 | + } | |
174 | + | |
175 | + #[tokio::test] | |
176 | + async fn json_response_custom_status_and_headers() { | |
177 | + let user = UserResp { | |
178 | + id: 2, | |
179 | + name: "Bob".into(), | |
180 | + email: Some("[email protected]".into()), | |
181 | + }; | |
182 | + let resp = JsonResponse::new(user) | |
183 | + .with_status(StatusCode::CREATED) | |
184 | + .with_header("x-request-id", "req-999") | |
185 | + .with_header("x-version", "2.0") | |
186 | + .into_response() | |
187 | + .unwrap(); | |
188 | + assert_eq!(resp.status(), StatusCode::CREATED); | |
189 | + assert_eq!(resp.headers()["x-request-id"], "req-999"); | |
190 | + assert_eq!(resp.headers()["x-version"], "2.0"); | |
191 | + assert_eq!(resp.headers()["content-type"], "application/json"); | |
192 | + let body = resp.into_body().collect().await.unwrap().to_bytes(); | |
193 | + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); | |
194 | + assert_eq!(json["email"], "[email protected]"); | |
195 | + } | |
196 | + | |
197 | + #[tokio::test] | |
198 | + async fn json_response_from_convenience() { | |
199 | + let user = UserResp { | |
200 | + id: 3, | |
201 | + name: "Carol".into(), | |
202 | + email: None, | |
203 | + }; | |
204 | + let resp = json_response_from(user).into_response().unwrap(); | |
205 | + assert_eq!(resp.status(), StatusCode::OK); | |
206 | + } | |
207 | + | |
208 | + #[tokio::test] | |
209 | + async fn json_response_serde_rename() { | |
210 | + #[derive(Serialize)] | |
211 | + struct Renamed { | |
212 | + #[serde(rename = "userId")] | |
213 | + user_id: u64, | |
214 | + } | |
215 | + let resp = JsonResponse::new(Renamed { user_id: 42 }) | |
216 | + .into_response() | |
217 | + .unwrap(); | |
218 | + let body = resp.into_body().collect().await.unwrap().to_bytes(); | |
219 | + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); | |
220 | + assert_eq!(json["userId"], 42); | |
221 | + assert!(json.get("user_id").is_none()); | |
222 | + } | |
223 | + | |
224 | + #[tokio::test] | |
225 | + async fn json_response_flatten() { | |
226 | + use std::collections::HashMap; | |
227 | + | |
228 | + #[derive(Serialize)] | |
229 | + struct WithMeta { | |
230 | + data: String, | |
231 | + #[serde(flatten)] | |
232 | + meta: HashMap<String, serde_json::Value>, | |
233 | + } | |
234 | + | |
235 | + let mut meta = HashMap::new(); | |
236 | + meta.insert("page".into(), serde_json::json!(1)); | |
237 | + let val = WithMeta { | |
238 | + data: "content".into(), | |
239 | + meta, | |
240 | + }; | |
241 | + let resp = JsonResponse::new(val).into_response().unwrap(); | |
242 | + let body = resp.into_body().collect().await.unwrap().to_bytes(); | |
243 | + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); | |
244 | + assert_eq!(json["data"], "content"); | |
245 | + assert_eq!(json["page"], 1); | |
246 | + } | |
247 | +} | |
0 | 248 | |
▼crates/justapi-core/src/server/mod.rs
+1796-0
| @@ -0,0 +1,1796 @@ | ||
1 | +use std::net::SocketAddr; | |
2 | +use std::path::PathBuf; | |
3 | +use std::sync::Arc; | |
4 | +use std::sync::LazyLock; | |
5 | + | |
6 | +use anyhow::Result; | |
7 | +use http_body_util::BodyExt; | |
8 | +use hyper::body::{Bytes, Incoming}; | |
9 | +use hyper::service::service_fn; | |
10 | +use hyper::{Method, Request, Response, StatusCode}; | |
11 | +use hyper_util::rt::TokioIo; | |
12 | +use tokio::net::TcpListener; | |
13 | + | |
14 | +#[cfg(feature = "compression")] | |
15 | +use crate::compress::CompressionMiddleware; | |
16 | +use crate::health::HealthRegistry; | |
17 | +use crate::memory::{BufferPool, SharedArena}; | |
18 | +use crate::metrics::{self, Metrics}; | |
19 | +use crate::middleware::{ | |
20 | + ApiKeyAuth, Cors, JwtAuth, MiddlewareChain, OAuth2Password, RateLimiter, SecurityHeaders, | |
21 | +}; | |
22 | +use crate::openapi; | |
23 | +use crate::plugin::PluginRegistry; | |
24 | +use crate::router::{Match, Router}; | |
25 | +use crate::static_files::StaticDir; | |
26 | +use crate::{json_response, streaming_response, ResponseBody}; | |
27 | +#[cfg(feature = "ws")] | |
28 | +use futures::StreamExt; | |
29 | +use serde_json; | |
30 | +use tokio_util::sync::CancellationToken; | |
31 | + | |
32 | +// --------------------------------------------------------------------------- | |
33 | +// WebSocket handler (feature-gated on `ws`) | |
34 | +// --------------------------------------------------------------------------- | |
35 | + | |
36 | +#[cfg(feature = "ws")] | |
37 | +pub type WsRead = Box< | |
38 | + dyn futures::Stream< | |
39 | + Item = Result< | |
40 | + tokio_tungstenite::tungstenite::Message, | |
41 | + tokio_tungstenite::tungstenite::Error, | |
42 | + >, | |
43 | + > + Unpin | |
44 | + + Send, | |
45 | +>; | |
46 | + | |
47 | +#[cfg(feature = "ws")] | |
48 | +pub type WsWrite = Box< | |
49 | + dyn futures::Sink< | |
50 | + tokio_tungstenite::tungstenite::Message, | |
51 | + Error = tokio_tungstenite::tungstenite::Error, | |
52 | + > + Unpin | |
53 | + + Send, | |
54 | +>; | |
55 | + | |
56 | +/// A handler invoked when a WebSocket upgrade is accepted on a registered path. | |
57 | +/// Receives the request path plus the split read/write halves of the accepted | |
58 | +/// WebSocket stream. The handler owns the connection for its lifetime. | |
59 | +#[cfg(feature = "ws")] | |
60 | +pub type WsHandler = std::sync::Arc< | |
61 | + dyn Fn( | |
62 | + String, | |
63 | + WsRead, | |
64 | + WsWrite, | |
65 | + ) | |
66 | + -> std::pin::Pin<Box<dyn std::future::Future<Output = anyhow::Result<()>> + Send>> | |
67 | + + Send | |
68 | + + Sync, | |
69 | +>; | |
70 | + | |
71 | +#[cfg(feature = "ws")] | |
72 | +use tokio_tungstenite::tungstenite::Message; | |
73 | + | |
74 | +/// Default WebSocket handler used when no application handler is registered. | |
75 | +/// Echoes text/binary frames (mirroring the legacy raw-TCP behavior) so that | |
76 | +/// standalone servers remain WebSocket-compatible out of the box. `with_ws` | |
77 | +/// replaces this with an application-provided handler. | |
78 | +#[cfg(feature = "ws")] | |
79 | +fn default_ws_echo() -> WsHandler { | |
80 | + std::sync::Arc::new(|_path, mut read, mut write| { | |
81 | + Box::pin(async move { | |
82 | + use futures::{SinkExt, StreamExt}; | |
83 | + while let Some(msg) = read.next().await { | |
84 | + match msg { | |
85 | + Ok(m @ Message::Text(_)) | Ok(m @ Message::Binary(_)) => { | |
86 | + if write.send(m).await.is_err() { | |
87 | + break; | |
88 | + } | |
89 | + } | |
90 | + Ok(Message::Ping(p)) => { | |
91 | + let _ = write.send(Message::Pong(p)).await; | |
92 | + } | |
93 | + Ok(Message::Close(_)) | Err(_) => { | |
94 | + let _ = write.send(Message::Close(None)).await; | |
95 | + break; | |
96 | + } | |
97 | + _ => {} | |
98 | + } | |
99 | + } | |
100 | + let _ = write.close().await; | |
101 | + Ok(()) | |
102 | + }) | |
103 | + }) | |
104 | +} | |
105 | + | |
106 | +// --------------------------------------------------------------------------- | |
107 | +// Handler enum | |
108 | +// --------------------------------------------------------------------------- | |
109 | + | |
110 | +#[derive(Clone)] | |
111 | +#[allow(dead_code)] | |
112 | +pub(crate) enum Handler { | |
113 | + Static { | |
114 | + status: StatusCode, | |
115 | + body: &'static str, | |
116 | + }, | |
117 | + Echo, | |
118 | + ParamsEcho, | |
119 | + Sse, | |
120 | + Health, | |
121 | + Ready, | |
122 | + Live, | |
123 | + Prometheus, | |
124 | + OpenApiJson, | |
125 | + SwaggerUi, | |
126 | + Redoc, | |
127 | + GraphQL, | |
128 | + Custom(crate::middleware::HandlerFn), | |
129 | +} | |
130 | + | |
131 | +static BUILTIN_SPEC: LazyLock<String> = LazyLock::new(|| { | |
132 | + let spec = serde_json::json!({ | |
133 | + "openapi": "3.1.0", | |
134 | + "info": { | |
135 | + "title": "JustAPI", | |
136 | + "version": "0.1.0", | |
137 | + "description": "JustAPI Runtime — Python application server" | |
138 | + }, | |
139 | + "servers": [{"url": "http://localhost:8080", "description": "Local development"}], | |
140 | + "tags": [{"name": "builtin", "description": "Built-in routes"}], | |
141 | + "paths": { | |
142 | + "/hello": { | |
143 | + "get": { | |
144 | + "summary": "Hello world", | |
145 | + "description": "Returns a greeting message", | |
146 | + "operationId": "get_hello", | |
147 | + "tags": ["builtin"], | |
148 | + "responses": {"200": {"description": "Greeting"}} | |
149 | + } | |
150 | + }, | |
151 | + "/echo": { | |
152 | + "post": { | |
153 | + "summary": "Echo request body", | |
154 | + "description": "Returns the request body as-is", | |
155 | + "operationId": "post_echo", | |
156 | + "tags": ["builtin"], | |
157 | + "requestBody": { | |
158 | + "description": "Any JSON payload", | |
159 | + "content": {"application/json": {"schema": {}}}, | |
160 | + "required": true | |
161 | + }, | |
162 | + "responses": {"200": {"description": "Echoed body"}} | |
163 | + } | |
164 | + }, | |
165 | + "/users/{id}": { | |
166 | + "get": { | |
167 | + "summary": "Get user by ID", | |
168 | + "operationId": "get_users_id", | |
169 | + "tags": ["builtin"], | |
170 | + "parameters": [{ | |
171 | + "name": "id", | |
172 | + "in": "path", | |
173 | + "description": "User ID", | |
174 | + "required": true, | |
175 | + "schema": {"type": "string"} | |
176 | + }], | |
177 | + "responses": { | |
178 | + "200": {"description": "User found"}, | |
179 | + "404": {"description": "Not found"} | |
180 | + } | |
181 | + } | |
182 | + }, | |
183 | + "/events": { | |
184 | + "get": { | |
185 | + "summary": "Server-Sent Events stream", | |
186 | + "description": "Streams count events via SSE", | |
187 | + "operationId": "get_events", | |
188 | + "tags": ["builtin"], | |
189 | + "responses": {"200": {"description": "SSE event stream"}} | |
190 | + } | |
191 | + }, | |
192 | + "/health": { | |
193 | + "get": { | |
194 | + "summary": "Health check", | |
195 | + "description": "Returns health status of all registered components", | |
196 | + "operationId": "get_health", | |
197 | + "tags": ["builtin"], | |
198 | + "responses": { | |
199 | + "200": {"description": "Healthy"}, | |
200 | + "503": {"description": "Unhealthy"} | |
201 | + } | |
202 | + } | |
203 | + }, | |
204 | + "/ready": { | |
205 | + "get": { | |
206 | + "summary": "Readiness probe", | |
207 | + "description": "Returns readiness status for Kubernetes probes", | |
208 | + "operationId": "get_ready", | |
209 | + "tags": ["builtin"], | |
210 | + "responses": {"200": {"description": "Ready"}} | |
211 | + } | |
212 | + }, | |
213 | + "/live": { | |
214 | + "get": { | |
215 | + "summary": "Liveness probe", | |
216 | + "description": "Returns liveness status for Kubernetes probes", | |
217 | + "operationId": "get_live", | |
218 | + "tags": ["builtin"], | |
219 | + "responses": {"200": {"description": "Alive"}} | |
220 | + } | |
221 | + }, | |
222 | + "/metrics": { | |
223 | + "get": { | |
224 | + "summary": "Prometheus metrics", | |
225 | + "description": "Returns Prometheus-formatted metrics", | |
226 | + "operationId": "get_metrics", | |
227 | + "tags": ["builtin"], | |
228 | + "responses": {"200": {"description": "Metrics in Prometheus text format"}} | |
229 | + } | |
230 | + }, | |
231 | + "/openapi.json": { | |
232 | + "get": { | |
233 | + "summary": "OpenAPI 3.1 spec", | |
234 | + "description": "Returns the OpenAPI 3.1 specification for this server", | |
235 | + "operationId": "get_openapi_json", | |
236 | + "tags": ["builtin"], | |
237 | + "responses": {"200": {"description": "OpenAPI spec"}} | |
238 | + } | |
239 | + }, | |
240 | + "/docs": { | |
241 | + "get": { | |
242 | + "summary": "Swagger UI", | |
243 | + "description": "Interactive API documentation via Swagger UI", | |
244 | + "operationId": "get_docs", | |
245 | + "tags": ["builtin"], | |
246 | + "responses": {"200": {"description": "Swagger UI HTML"}} | |
247 | + } | |
248 | + }, | |
249 | + "/redoc": { | |
250 | + "get": { | |
251 | + "summary": "ReDoc", | |
252 | + "description": "Interactive API documentation via ReDoc", | |
253 | + "operationId": "get_redoc", | |
254 | + "tags": ["builtin"], | |
255 | + "responses": {"200": {"description": "ReDoc HTML"}} | |
256 | + } | |
257 | + } | |
258 | + } | |
259 | + }); | |
260 | + serde_json::to_string_pretty(&spec).unwrap() | |
261 | +}); | |
262 | + | |
263 | +// --------------------------------------------------------------------------- | |
264 | +// TLS config | |
265 | +// --------------------------------------------------------------------------- | |
266 | + | |
267 | +#[cfg(feature = "tls")] | |
268 | +pub struct TlsConfig { | |
269 | + pub cert_path: String, | |
270 | + pub key_path: String, | |
271 | +} | |
272 | + | |
273 | +// --------------------------------------------------------------------------- | |
274 | +// Server builder | |
275 | +// --------------------------------------------------------------------------- | |
276 | + | |
277 | +pub struct Server { | |
278 | + addr: SocketAddr, | |
279 | + chain: MiddlewareChain, | |
280 | + static_dir: Option<StaticDir>, | |
281 | + metrics: Metrics, | |
282 | + health_registry: Arc<HealthRegistry>, | |
283 | + #[cfg(feature = "tls")] | |
284 | + tls: Option<TlsConfig>, | |
285 | + shutdown: Option<CancellationToken>, | |
286 | + plugin_registry: PluginRegistry, | |
287 | + wasm_middleware: Option<Arc<crate::wasm::WasmEngine>>, | |
288 | + grpc_addr: Option<SocketAddr>, | |
289 | + grpc_handler: Option<crate::grpc::GrpcHandler>, | |
290 | + #[cfg(feature = "ws")] | |
291 | + ws_handler: Option<WsHandler>, | |
292 | + router: Option<Router<Handler>>, | |
293 | + openapi_spec: Option<Arc<String>>, | |
294 | +} | |
295 | + | |
296 | +impl Server { | |
297 | + pub fn new(addr: SocketAddr) -> Self { | |
298 | + let router = Router::new(); | |
299 | + let metrics = Metrics::new(); | |
300 | + let health_registry = Arc::new(HealthRegistry::new()); | |
301 | + // We initialize chain with a dummy handler. The real one is set in run(). | |
302 | + let chain = MiddlewareChain::new(Arc::new(|_| { | |
303 | + Box::pin(async { | |
304 | + Ok(Response::builder() | |
305 | + .status(StatusCode::NOT_FOUND) | |
306 | + .body(crate::UnsyncBoxBody::new( | |
307 | + http_body_util::Full::new(Bytes::from("Not Found")) | |
308 | + .map_err(|e: std::convert::Infallible| -> anyhow::Error { match e {} }), | |
309 | + )) | |
310 | + .unwrap()) | |
311 | + }) | |
312 | + })); | |
313 | + Self { | |
314 | + addr, | |
315 | + chain, | |
316 | + static_dir: None, | |
317 | + metrics, | |
318 | + health_registry, | |
319 | + shutdown: None, | |
320 | + plugin_registry: PluginRegistry::new(), | |
321 | + wasm_middleware: None, | |
322 | + #[cfg(feature = "tls")] | |
323 | + tls: None, | |
324 | + grpc_addr: None, | |
325 | + grpc_handler: None, | |
326 | + #[cfg(feature = "ws")] | |
327 | + ws_handler: Some(default_ws_echo()), | |
328 | + router: Some(router), | |
329 | + openapi_spec: None, | |
330 | + } | |
331 | + } | |
332 | + | |
333 | + /// Replace the built-in OpenAPI spec with a dynamically generated one | |
334 | + /// (e.g., from a Python JustAPIApp's registered routes). | |
335 | + pub fn with_openapi_spec(mut self, spec: String) -> Self { | |
336 | + self.openapi_spec = Some(Arc::new(spec)); | |
337 | + self | |
338 | + } | |
339 | + | |
340 | + /// Register a WebSocket handler. When a WebSocket upgrade is accepted on any | |
341 | + /// path, the handler is invoked with the path and the split stream halves. | |
342 | + #[cfg(feature = "ws")] | |
343 | + pub fn with_ws(mut self, handler: WsHandler) -> Self { | |
344 | + self.ws_handler = Some(handler); | |
345 | + self | |
346 | + } | |
347 | + | |
348 | + /// Set a custom fallback handler (used by the ASGI shim). | |
349 | + pub fn with_handler(mut self, handler: crate::middleware::HandlerFn) -> Self { | |
350 | + // If the user specifies a custom handler, it replaces the entire router behavior. | |
351 | + self.router = None; | |
352 | + self.chain.set_handler(handler); | |
353 | + self | |
354 | + } | |
355 | + | |
356 | + pub fn with_default_routes(mut self) -> Self { | |
357 | + let mut router = self.router.unwrap_or_default(); | |
358 | + router | |
359 | + .insert( | |
360 | + Method::GET, | |
361 | + "/hello", | |
362 | + Handler::Static { | |
363 | + status: StatusCode::OK, | |
364 | + body: r#"{"message":"hello"}"#, | |
365 | + }, | |
366 | + ) | |
367 | + .expect("valid"); | |
368 | + router | |
369 | + .insert(Method::POST, "/echo", Handler::Echo) | |
370 | + .expect("valid"); | |
371 | + router | |
372 | + .insert(Method::GET, "/echo/{*rest}", Handler::ParamsEcho) | |
373 | + .expect("valid"); | |
374 | + router | |
375 | + .insert( | |
376 | + Method::GET, | |
377 | + "/users/{id}", | |
378 | + Handler::Static { | |
379 | + status: StatusCode::OK, | |
380 | + body: r#"{"message":"user lookup"}"#, | |
381 | + }, | |
382 | + ) | |
383 | + .expect("valid"); | |
384 | + router | |
385 | + .insert(Method::GET, "/events", Handler::Sse) | |
386 | + .expect("valid"); | |
387 | + router | |
388 | + .insert(Method::GET, "/graphql", Handler::GraphQL) | |
389 | + .expect("valid"); | |
390 | + router | |
391 | + .insert(Method::POST, "/graphql", Handler::GraphQL) | |
392 | + .expect("valid"); | |
393 | + router.set_fallback(Handler::Static { | |
394 | + status: StatusCode::NOT_FOUND, | |
395 | + body: r#"{"error":"not found"}"#, | |
396 | + }); | |
397 | + self.router = Some(router); | |
398 | + self | |
399 | + } | |
400 | + | |
401 | + /// Register the OpenAI-compatible inference endpoints, bound to `engine`. | |
402 | + /// | |
403 | + /// Routes added: | |
404 | + /// - `POST /v1/chat/completions` | |
405 | + /// - `POST /v1/completions` | |
406 | + /// - `POST /v1/embeddings` | |
407 | + /// - `GET /v1/models` | |
408 | + #[cfg(feature = "inference")] | |
409 | + pub fn with_openai(mut self, engine: std::sync::Arc<justapi_inference::Engine>) -> Self { | |
410 | + if let Some(ref mut r) = self.router { | |
411 | + r.insert( | |
412 | + hyper::Method::POST, | |
413 | + "/v1/chat/completions", | |
414 | + Handler::Custom(crate::openai::chat_completions_handler(engine.clone())), | |
415 | + ) | |
416 | + .expect("valid route: /v1/chat/completions"); | |
417 | + r.insert( | |
418 | + hyper::Method::POST, | |
419 | + "/v1/completions", | |
420 | + Handler::Custom(crate::openai::completions_handler(engine.clone())), | |
421 | + ) | |
422 | + .expect("valid route: /v1/completions"); | |
423 | + r.insert( | |
424 | + hyper::Method::POST, | |
425 | + "/v1/embeddings", | |
426 | + Handler::Custom(crate::openai::embeddings_handler(engine.clone())), | |
427 | + ) | |
428 | + .expect("valid route: /v1/embeddings"); | |
429 | + r.insert( | |
430 | + hyper::Method::GET, | |
431 | + "/v1/models", | |
432 | + Handler::Custom(crate::openai::models_handler(engine.clone())), | |
433 | + ) | |
434 | + .expect("valid route: /v1/models"); | |
435 | + } | |
436 | + self | |
437 | + } | |
438 | + | |
439 | + /// Register the OpenAI-compatible inference endpoints backed by the | |
440 | + /// scheduler (continuous batching + RadixAttention prefix cache). | |
441 | + /// | |
442 | + /// Same routes as [`Server::with_openai`], but each request is admitted | |
443 | + /// through the scheduler's wait queue, scheduled for prefill/decode steps, | |
444 | + /// and prefix-cached when finished. Scheduler prefix-cache statistics are | |
445 | + /// exposed via the `/metrics` endpoint. | |
446 | + /// | |
447 | + /// The [`SchedulerEngine`](justapi_inference::SchedulerEngine) wraps the | |
448 | + /// engine and scheduler into a single generation interface. | |
449 | + #[cfg(feature = "inference")] | |
450 | + pub fn with_openai_scheduled( | |
451 | + mut self, | |
452 | + scheduler_engine: std::sync::Arc<justapi_inference::SchedulerEngine>, | |
453 | + ) -> Self { | |
454 | + // Register a metric provider that renders scheduler stats. | |
455 | + let stats_provider = { | |
456 | + let se = scheduler_engine.clone(); | |
457 | + move || { | |
458 | + let s = se.stats(); | |
459 | + format!( | |
460 | + concat!( | |
461 | + "# HELP justapi_scheduler_waiting Current number of waiting requests.\n", | |
462 | + "# TYPE justapi_scheduler_waiting gauge\n", | |
463 | + "justapi_scheduler_waiting {}\n", | |
464 | + "# HELP justapi_scheduler_running Current number of running sequences.\n", | |
465 | + "# TYPE justapi_scheduler_running gauge\n", | |
466 | + "justapi_scheduler_running {}\n", | |
467 | + "# HELP justapi_scheduler_prefilling Current number of prefilling sequences.\n", | |
468 | + "# TYPE justapi_scheduler_prefilling gauge\n", | |
469 | + "justapi_scheduler_prefilling {}\n", | |
470 | + "# HELP justapi_scheduler_back_pressure Whether the scheduler is under back-pressure.\n", | |
471 | + "# TYPE justapi_scheduler_back_pressure gauge\n", | |
472 | + "justapi_scheduler_back_pressure {}\n", | |
473 | + "# HELP justapi_scheduler_prefix_hits Total prefix-cache hits.\n", | |
474 | + "# TYPE justapi_scheduler_prefix_hits counter\n", | |
475 | + "justapi_scheduler_prefix_hits {}\n", | |
476 | + "# HELP justapi_scheduler_prefix_misses Total prefix-cache misses.\n", | |
477 | + "# TYPE justapi_scheduler_prefix_misses counter\n", | |
478 | + "justapi_scheduler_prefix_misses {}\n", | |
479 | + "# HELP justapi_scheduler_prefix_tokens_saved Tokens of prefill saved by prefix cache.\n", | |
480 | + "# TYPE justapi_scheduler_prefix_tokens_saved counter\n", | |
481 | + "justapi_scheduler_prefix_tokens_saved {}\n", | |
482 | + "# HELP justapi_scheduler_cached_kv_blocks KV blocks resident in prefix cache.\n", | |
483 | + "# TYPE justapi_scheduler_cached_kv_blocks gauge\n", | |
484 | + "justapi_scheduler_cached_kv_blocks {}\n", | |
485 | + ), | |
486 | + s.num_waiting, | |
487 | + s.num_running, | |
488 | + s.num_prefilling, | |
489 | + s.back_pressure as u8, | |
490 | + s.prefix.hits, | |
491 | + s.prefix.misses, | |
492 | + s.prefix.tokens_saved, | |
493 | + s.cached_kv_blocks, | |
494 | + ) | |
495 | + } | |
496 | + }; | |
497 | + self.metrics | |
498 | + .register_extra_provider(Box::new(stats_provider)); | |
499 | + | |
500 | + if let Some(ref mut r) = self.router { | |
501 | + r.insert( | |
502 | + hyper::Method::POST, | |
503 | + "/v1/chat/completions", | |
504 | + Handler::Custom(crate::openai::scheduled_chat_completions_handler( | |
505 | + scheduler_engine.clone(), | |
506 | + )), | |
507 | + ) | |
508 | + .expect("valid route: /v1/chat/completions"); | |
509 | + r.insert( | |
510 | + hyper::Method::POST, | |
511 | + "/v1/completions", | |
512 | + Handler::Custom(crate::openai::scheduled_completions_handler( | |
513 | + scheduler_engine.clone(), | |
514 | + )), | |
515 | + ) | |
516 | + .expect("valid route: /v1/completions"); | |
517 | + r.insert( | |
518 | + hyper::Method::POST, | |
519 | + "/v1/embeddings", | |
520 | + Handler::Custom(crate::openai::embeddings_handler( | |
521 | + scheduler_engine.engine().clone(), | |
522 | + )), | |
523 | + ) | |
524 | + .expect("valid route: /v1/embeddings"); | |
525 | + r.insert( | |
526 | + hyper::Method::GET, | |
527 | + "/v1/models", | |
528 | + Handler::Custom(crate::openai::models_handler( | |
529 | + scheduler_engine.engine().clone(), | |
530 | + )), | |
531 | + ) | |
532 | + .expect("valid route: /v1/models"); | |
533 | + } | |
534 | + self | |
535 | + } | |
536 | + | |
537 | + /// Register the OpenAI-compatible inference endpoints with LoRA-aware / | |
538 | + /// KV-aware routing via a [`justapi_inference::ControlPlane`] and | |
539 | + /// [`justapi_inference::Router`]. | |
540 | + /// | |
541 | + /// Each request is admitted through the control plane + router before | |
542 | + /// generation: if no replica can serve it (all at capacity / no route), the | |
543 | + /// endpoint returns `503`. The resolved model name is then used for the | |
544 | + /// actual generation. Same routes as [`Server::with_openai`]. | |
545 | + #[cfg(feature = "inference")] | |
546 | + pub fn with_openai_routed( | |
547 | + mut self, | |
548 | + engine: std::sync::Arc<justapi_inference::Engine>, | |
549 | + control_plane: std::sync::Arc<justapi_inference::ControlPlane>, | |
550 | + router: std::sync::Arc<justapi_inference::Router>, | |
551 | + ) -> Self { | |
552 | + if let Some(ref mut r) = self.router { | |
553 | + r.insert( | |
554 | + hyper::Method::POST, | |
555 | + "/v1/chat/completions", | |
556 | + Handler::Custom(crate::openai::routed_chat_completions_handler( | |
557 | + engine.clone(), | |
558 | + control_plane.clone(), | |
559 | + router.clone(), | |
560 | + )), | |
561 | + ) | |
562 | + .expect("valid route: /v1/chat/completions"); | |
563 | + r.insert( | |
564 | + hyper::Method::POST, | |
565 | + "/v1/completions", | |
566 | + Handler::Custom(crate::openai::routed_completions_handler( | |
567 | + engine.clone(), | |
568 | + control_plane.clone(), | |
569 | + router.clone(), | |
570 | + )), | |
571 | + ) | |
572 | + .expect("valid route: /v1/completions"); | |
573 | + r.insert( | |
574 | + hyper::Method::POST, | |
575 | + "/v1/embeddings", | |
576 | + Handler::Custom(crate::openai::routed_embeddings_handler( | |
577 | + engine.clone(), | |
578 | + control_plane.clone(), | |
579 | + router.clone(), | |
580 | + )), | |
581 | + ) | |
582 | + .expect("valid route: /v1/embeddings"); | |
583 | + r.insert( | |
584 | + hyper::Method::GET, | |
585 | + "/v1/models", | |
586 | + Handler::Custom(crate::openai::models_handler(engine.clone())), | |
587 | + ) | |
588 | + .expect("valid route: /v1/models"); | |
589 | + } | |
590 | + self | |
591 | + } | |
592 | + /// new connections and the `run()` future completes. | |
593 | + pub fn with_shutdown(mut self, token: CancellationToken) -> Self { | |
594 | + self.shutdown = Some(token); | |
595 | + self | |
596 | + } | |
597 | + | |
598 | + /// Attach a Gateway configuration with hot reloading | |
599 | + pub fn add_gateway(mut self, state: std::sync::Arc<crate::gateway::GatewayState>) -> Self { | |
600 | + self.chain | |
601 | + .add(crate::gateway::GatewayMiddleware::new(state)); | |
602 | + self | |
603 | + } | |
604 | + | |
605 | + pub fn with_static_dir(mut self, dir: impl Into<PathBuf>) -> Self { | |
606 | + self.static_dir = Some(StaticDir::new(dir)); | |
607 | + self | |
608 | + } | |
609 | + | |
610 | + /// Register a plugin with the server. | |
611 | + /// This calls the `build` hook immediately to allow the plugin to modify the Server configuration. | |
612 | + pub fn register_plugin(mut self, plugin: Box<dyn crate::plugin::Plugin>) -> Result<Self> { | |
613 | + plugin.build(&mut self)?; | |
614 | + self.plugin_registry.register(plugin); | |
615 | + Ok(self) | |
616 | + } | |
617 | + | |
618 | + pub fn add_cors(mut self, cors: Cors) -> Self { | |
619 | + self.chain.add(cors); | |
620 | + self | |
621 | + } | |
622 | + | |
623 | + pub fn add_security_headers(mut self, sh: SecurityHeaders) -> Self { | |
624 | + self.chain.add(sh); | |
625 | + self | |
626 | + } | |
627 | + | |
628 | + pub fn add_rate_limiter(mut self, rl: RateLimiter) -> Self { | |
629 | + self.chain.add(rl); | |
630 | + self | |
631 | + } | |
632 | + | |
633 | + /// Add JWT authentication middleware. | |
634 | + pub fn add_jwt(mut self, jwt: JwtAuth) -> Self { | |
635 | + self.chain.add(jwt); | |
636 | + self | |
637 | + } | |
638 | + | |
639 | + /// Add API-key authentication middleware (the "API-key scheme"). | |
640 | + pub fn add_api_key_auth(mut self, auth: ApiKeyAuth) -> Self { | |
641 | + self.chain.add(auth); | |
642 | + self | |
643 | + } | |
644 | + | |
645 | + /// Register the OAuth2 password flow: token endpoint + JwtAuth middleware. | |
646 | + /// | |
647 | + /// This registers `{token_path}` (default `/token`) as a POST handler that | |
648 | + /// accepts `application/x-www-form-urlencoded` with `grant_type=password`, | |
649 | + /// `username`, and `password`. On success it returns an `access_token` | |
650 | + /// signed by this provider. It also adds a [`JwtAuth`] middleware that | |
651 | + /// verifies every request (except the token path) against the same key. | |
652 | + /// | |
653 | + /// Use [`OAuth2Password::jwt_auth`] directly if you need to customise the | |
654 | + /// JWT verification rules (e.g. per-path exemptions). | |
655 | + pub fn with_oauth2_password(mut self, oauth2: OAuth2Password) -> Self { | |
656 | + let path = oauth2.token_path.clone(); | |
657 | + let handler = oauth2.token_handler::<Incoming>(); | |
658 | + if let Some(ref mut r) = self.router { | |
659 | + r.insert(hyper::Method::POST, &path, Handler::Custom(handler)) | |
660 | + .expect("valid oauth2 token path"); | |
661 | + } | |
662 | + let jwt = oauth2 | |
663 | + .jwt_auth() | |
664 | + .require_for(&path, crate::middleware::JwtRequirement::None); | |
665 | + self.chain.add(jwt); | |
666 | + self | |
667 | + } | |
668 | + | |
669 | + pub fn add_middleware(mut self, mw: impl crate::middleware::Middleware + 'static) -> Self { | |
670 | + self.chain.add(mw); | |
671 | + self | |
672 | + } | |
673 | + | |
674 | + /// Add response compression middleware (Gzip, feature-gated Brotli/Zstd). | |
675 | + #[cfg(feature = "compression")] | |
676 | + pub fn add_compression(mut self) -> Self { | |
677 | + self.chain.add(CompressionMiddleware::new()); | |
678 | + self | |
679 | + } | |
680 | + | |
681 | + pub fn add_circuit_breaker(mut self, config: crate::resilience::CircuitBreakerConfig) -> Self { | |
682 | + self.chain | |
683 | + .add(crate::resilience::PerRouteCircuitBreakerMiddleware::new( | |
684 | + config, | |
685 | + )); | |
686 | + self | |
687 | + } | |
688 | + | |
689 | + #[cfg(feature = "tls")] | |
690 | + pub fn with_tls(mut self, config: TlsConfig) -> Self { | |
691 | + self.tls = Some(config); | |
692 | + self | |
693 | + } | |
694 | + | |
695 | + pub fn metrics(&self) -> &Metrics { | |
696 | + &self.metrics | |
697 | + } | |
698 | + | |
699 | + pub fn health_registry(&self) -> &HealthRegistry { | |
700 | + &self.health_registry | |
701 | + } | |
702 | + | |
703 | + pub fn with_health_registry(mut self, registry: HealthRegistry) -> Self { | |
704 | + self.health_registry = Arc::new(registry); | |
705 | + self | |
706 | + } | |
707 | + | |
708 | + pub fn with_plugin(mut self, plugin: Box<dyn crate::plugin::Plugin>) -> Self { | |
709 | + self.plugin_registry.register(plugin); | |
710 | + self | |
711 | + } | |
712 | + | |
713 | + pub fn with_grpc(mut self, addr: SocketAddr, handler: crate::grpc::GrpcHandler) -> Self { | |
714 | + self.grpc_addr = Some(addr); | |
715 | + self.grpc_handler = Some(handler); | |
716 | + self | |
717 | + } | |
718 | + pub fn with_wasm_middleware(mut self, wasm_bytes: &[u8]) -> Result<Self> { | |
719 | + let engine = crate::wasm::WasmEngine::new(wasm_bytes)?; | |
720 | + self.wasm_middleware = Some(Arc::new(engine)); | |
721 | + Ok(self) | |
722 | + } | |
723 | + | |
724 | + pub async fn run(mut self) -> Result<()> { | |
725 | + if let Some(router) = self.router.take() { | |
726 | + let pool = Arc::new(BufferPool::new()); | |
727 | + let handler = make_handler( | |
728 | + Arc::new(router), | |
729 | + pool, | |
730 | + self.metrics.clone(), | |
731 | + self.health_registry.clone(), | |
732 | + self.openapi_spec.clone(), | |
733 | + ); | |
734 | + self.chain.set_handler(handler); | |
735 | + } | |
736 | + | |
737 | + let listener = tokio::net::TcpListener::bind(self.addr).await?; | |
738 | + let local_addr = listener.local_addr()?; | |
739 | + | |
740 | + if let (Some(grpc_addr), Some(grpc_handler)) = (self.grpc_addr, self.grpc_handler.take()) { | |
741 | + tracing::info!("Starting gRPC server on {}", grpc_addr); | |
742 | + let grpc_service = crate::grpc::DynamicGrpcService::new(grpc_handler); | |
743 | + let shutdown = self.shutdown.clone(); | |
744 | + tokio::spawn(async move { | |
745 | + if let Ok(grpc_listener) = tokio::net::TcpListener::bind(grpc_addr).await { | |
746 | + loop { | |
747 | + if let Some(ref token) = shutdown { | |
748 | + tokio::select! { | |
749 | + result = grpc_listener.accept() => { | |
750 | + if let Ok((stream, _)) = result { | |
751 | + let io = hyper_util::rt::TokioIo::new(stream); | |
752 | + let svc = hyper_util::service::TowerToHyperService::new(grpc_service.clone()); | |
753 | + tokio::spawn(async move { | |
754 | + let _ = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new()) | |
755 | + .serve_connection_with_upgrades(io, svc) | |
756 | + .await; | |
757 | + }); | |
758 | + } | |
759 | + } | |
760 | + _ = token.cancelled() => { | |
761 | + break; | |
762 | + } | |
763 | + } | |
764 | + } else { | |
765 | + if let Ok((stream, _)) = grpc_listener.accept().await { | |
766 | + let io = hyper_util::rt::TokioIo::new(stream); | |
767 | + let svc = hyper_util::service::TowerToHyperService::new( | |
768 | + grpc_service.clone(), | |
769 | + ); | |
770 | + tokio::spawn(async move { | |
771 | + let _ = hyper_util::server::conn::auto::Builder::new( | |
772 | + hyper_util::rt::TokioExecutor::new(), | |
773 | + ) | |
774 | + .serve_connection_with_upgrades(io, svc) | |
775 | + .await; | |
776 | + }); | |
777 | + } | |
778 | + } | |
779 | + } | |
780 | + } | |
781 | + }); | |
782 | + } | |
783 | + | |
784 | + self.plugin_registry.on_startup_all().await?; | |
785 | + | |
786 | + let chain = Arc::new(self.chain); | |
787 | + let static_dir = self.static_dir; | |
788 | + let metrics = self.metrics; | |
789 | + let plugin_registry = Arc::new(self.plugin_registry); | |
790 | + let wasm_middleware = self.wasm_middleware.clone(); | |
791 | + #[cfg(feature = "ws")] | |
792 | + let ws_handler = self.ws_handler.clone(); | |
793 | + | |
794 | + #[cfg(feature = "tls")] | |
795 | + if let Some(tls_config) = self.tls { | |
796 | + #[cfg(feature = "ws")] | |
797 | + let res = serve_with_tls( | |
798 | + listener, | |
799 | + chain, | |
800 | + tls_config, | |
801 | + static_dir, | |
802 | + metrics, | |
803 | + self.shutdown, | |
804 | + wasm_middleware, | |
805 | + ws_handler, | |
806 | + ) | |
807 | + .await; | |
808 | + #[cfg(not(feature = "ws"))] | |
809 | + let res = serve_with_tls( | |
810 | + listener, | |
811 | + chain, | |
812 | + tls_config, | |
813 | + static_dir, | |
814 | + metrics, | |
815 | + self.shutdown, | |
816 | + wasm_middleware, | |
817 | + ) | |
818 | + .await; | |
819 | + plugin_registry.on_shutdown_all().await?; | |
820 | + return res; | |
821 | + } | |
822 | + | |
823 | + tracing::info!("Listening on {} (plain HTTP/1.1)", local_addr); | |
824 | + #[cfg(feature = "ws")] | |
825 | + let res = serve_http( | |
826 | + listener, | |
827 | + chain, | |
828 | + static_dir, | |
829 | + metrics, | |
830 | + self.shutdown, | |
831 | + wasm_middleware, | |
832 | + ws_handler, | |
833 | + ) | |
834 | + .await; | |
835 | + #[cfg(not(feature = "ws"))] | |
836 | + let res = serve_http( | |
837 | + listener, | |
838 | + chain, | |
839 | + static_dir, | |
840 | + metrics, | |
841 | + self.shutdown, | |
842 | + wasm_middleware, | |
843 | + ) | |
844 | + .await; | |
845 | + plugin_registry.on_shutdown_all().await?; | |
846 | + res | |
847 | + } | |
848 | +} | |
849 | + | |
850 | +// --------------------------------------------------------------------------- | |
851 | +// Standalone helper | |
852 | +// --------------------------------------------------------------------------- | |
853 | + | |
854 | +pub async fn serve(listener: TcpListener) -> Result<()> { | |
855 | + let mut router: Router<Handler> = Router::new(); | |
856 | + router.set_fallback(Handler::Static { | |
857 | + status: StatusCode::NOT_FOUND, | |
858 | + body: r#"{"error":"not found"}"#, | |
859 | + }); | |
860 | + router.insert(hyper::Method::GET, "/hello", Handler::Static { status: StatusCode::OK, body: r#"{"message":"hello"}"# }).unwrap(); | |
861 | + router.insert(hyper::Method::POST, "/echo", Handler::Echo).unwrap(); | |
862 | + router.insert(hyper::Method::GET, "/health", Handler::Health).unwrap(); | |
863 | + router.insert(hyper::Method::GET, "/ready", Handler::Ready).unwrap(); | |
864 | + router.insert(hyper::Method::GET, "/live", Handler::Live).unwrap(); | |
865 | + router.insert(hyper::Method::GET, "/metrics", Handler::Prometheus).unwrap(); | |
866 | + router.insert(hyper::Method::GET, "/events", Handler::Sse).unwrap(); | |
867 | + | |
868 | + let router = Arc::new(router); | |
869 | + let pool = Arc::new(BufferPool::new()); | |
870 | + let metrics = Metrics::new(); | |
871 | + let health_registry = Arc::new(HealthRegistry::new()); | |
872 | + let handler = make_handler(router, pool, metrics.clone(), health_registry, None); | |
873 | + let chain = Arc::new(MiddlewareChain::new(handler)); | |
874 | + | |
875 | + #[cfg(feature = "ws")] | |
876 | + let res = serve_http( | |
877 | + listener, | |
878 | + chain, | |
879 | + None, | |
880 | + metrics, | |
881 | + None, | |
882 | + None, | |
883 | + Some(default_ws_echo()), | |
884 | + ) | |
885 | + .await; | |
886 | + #[cfg(not(feature = "ws"))] | |
887 | + let res = serve_http(listener, chain, None, metrics, None, None).await; | |
888 | + res | |
889 | +} | |
890 | + | |
891 | +// --------------------------------------------------------------------------- | |
892 | +// Handler factory | |
893 | +// --------------------------------------------------------------------------- | |
894 | + | |
895 | +fn make_handler( | |
896 | + router: Arc<Router<Handler>>, | |
897 | + pool: Arc<BufferPool>, | |
898 | + metrics: Metrics, | |
899 | + health_registry: Arc<HealthRegistry>, | |
900 | + openapi_spec: Option<Arc<String>>, | |
901 | +) -> crate::middleware::HandlerFn { | |
902 | + // Build the global GraphQL schema | |
903 | + let graphql_schema = Arc::new(crate::graphql::create_schema()); | |
904 | + | |
905 | + Arc::new(move |req: Request<Incoming>| { | |
906 | + let router = router.clone(); | |
907 | + let pool = pool.clone(); | |
908 | + let metrics = metrics.clone(); | |
909 | + let health_registry = health_registry.clone(); | |
910 | + let graphql_schema = graphql_schema.clone(); | |
911 | + let openapi_spec = openapi_spec.clone(); | |
912 | + Box::pin(async move { | |
913 | + let method = req.method().clone(); | |
914 | + let path = req.uri().path().to_string(); | |
915 | + match router.at(&method, &path) { | |
916 | + Ok(m) => { | |
917 | + execute_handler( | |
918 | + m, | |
919 | + req, | |
920 | + &pool, | |
921 | + &metrics, | |
922 | + Some(&health_registry), | |
923 | + Some(&graphql_schema), | |
924 | + openapi_spec.as_deref().map(|s| s.as_str()), | |
925 | + ) | |
926 | + .await | |
927 | + } | |
928 | + Err(crate::router::RouterError::MethodNotAllowed) => Ok(json_response( | |
929 | + StatusCode::METHOD_NOT_ALLOWED, | |
930 | + r#"{"error":"method not allowed"}"#, | |
931 | + )), | |
932 | + Err(crate::router::RouterError::NotFound) => { | |
933 | + let fb = router.fallback().unwrap(); | |
934 | + let m = Match { | |
935 | + handler: fb, | |
936 | + params: matchit::Params::new(), | |
937 | + }; | |
938 | + execute_handler( | |
939 | + m, | |
940 | + req, | |
941 | + &pool, | |
942 | + &metrics, | |
943 | + Some(&health_registry), | |
944 | + Some(&graphql_schema), | |
945 | + openapi_spec.as_deref().map(|s| s.as_str()), | |
946 | + ) | |
947 | + .await | |
948 | + } | |
949 | + } | |
950 | + }) | |
951 | + }) | |
952 | +} | |
953 | + | |
954 | +// --------------------------------------------------------------------------- | |
955 | +// HTTP server loop | |
956 | +// --------------------------------------------------------------------------- | |
957 | + | |
958 | +async fn serve_http( | |
959 | + listener: TcpListener, | |
960 | + chain: Arc<MiddlewareChain>, | |
961 | + static_dir: Option<StaticDir>, | |
962 | + metrics: Metrics, | |
963 | + shutdown: Option<CancellationToken>, | |
964 | + wasm_middleware: Option<Arc<crate::wasm::WasmEngine>>, | |
965 | + #[cfg(feature = "ws")] ws_handler: Option<WsHandler>, | |
966 | +) -> Result<()> { | |
967 | + let mut connections = tokio::task::JoinSet::new(); | |
968 | + | |
969 | + loop { | |
970 | + let token = shutdown.as_ref().cloned(); | |
971 | + | |
972 | + let (stream, peer) = if let Some(token) = &token { | |
973 | + tokio::select! { | |
974 | + result = listener.accept() => result?, | |
975 | + _ = token.cancelled() => { | |
976 | + tracing::info!("Shutdown signal received, stopping accept loop"); | |
977 | + break; | |
978 | + } | |
979 | + } | |
980 | + } else { | |
981 | + listener.accept().await? | |
982 | + }; | |
983 | + tracing::debug!("Accepted connection from {}", peer); | |
984 | + | |
985 | + // --- Normal HTTP via hyper ----------------------------------------- | |
986 | + let chain = chain.clone(); | |
987 | + let static_dir = static_dir.clone(); | |
988 | + let conn_metrics = metrics.clone(); | |
989 | + conn_metrics.connection_opened(); | |
990 | + let spawn_metrics = conn_metrics.clone(); | |
991 | + let wasm_middleware = wasm_middleware.clone(); | |
992 | + let peer_addr = peer; | |
993 | + | |
994 | + #[cfg(feature = "ws")] | |
995 | + let ws_handler = ws_handler.clone(); | |
996 | + | |
997 | + let token_clone = shutdown.clone(); | |
998 | + connections.spawn(async move { | |
999 | + let io = TokioIo::new(stream); | |
1000 | + let arena = Arc::new(SharedArena::new()); | |
1001 | + let svc = service_fn(move |mut req| { | |
1002 | + arena.reset(); | |
1003 | + let chain = chain.clone(); | |
1004 | + let arena = arena.clone(); | |
1005 | + let static_dir = static_dir.clone(); | |
1006 | + let metrics = spawn_metrics.clone(); | |
1007 | + let wasm_middleware = wasm_middleware.clone(); | |
1008 | + #[cfg(feature = "ws")] | |
1009 | + let ws_handler = ws_handler.clone(); | |
1010 | + async move { | |
1011 | + let path = req.uri().path().to_string(); | |
1012 | + let method = req.method().clone(); | |
1013 | + let _arena_path = arena.alloc_str(&path); | |
1014 | + let start = std::time::Instant::now(); | |
1015 | + req.extensions_mut().insert(peer_addr); | |
1016 | + | |
1017 | + let span = tracing::info_span!( | |
1018 | + "http.request", | |
1019 | + http.method = method.as_str(), | |
1020 | + http.path = %path, | |
1021 | + http.status_code = tracing::field::Empty, | |
1022 | + ); | |
1023 | + | |
1024 | + metrics.record_request(); | |
1025 | + | |
1026 | + #[cfg(feature = "ws")] | |
1027 | + if let Some(ref handler) = ws_handler { | |
1028 | + let is_ws = req.headers().get(hyper::header::UPGRADE) | |
1029 | + .and_then(|v| v.to_str().ok()) | |
1030 | + .map(|s| s.eq_ignore_ascii_case("websocket")) | |
1031 | + .unwrap_or(false); | |
1032 | + | |
1033 | + if is_ws { | |
1034 | + if let Some(key) = req.headers().get("sec-websocket-key") { | |
1035 | + let key_str = key.to_str().unwrap_or(""); | |
1036 | + let mut sha1 = sha1::Sha1::new(); | |
1037 | + use sha1::Digest; | |
1038 | + sha1.update(key_str.as_bytes()); | |
1039 | + sha1.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); | |
1040 | + use base64::Engine; | |
1041 | + let accept_key = base64::engine::general_purpose::STANDARD.encode(sha1.finalize()); | |
1042 | + | |
1043 | + let mut res = crate::json_response(StatusCode::SWITCHING_PROTOCOLS, ""); | |
1044 | + res.headers_mut().insert(hyper::header::UPGRADE, hyper::header::HeaderValue::from_static("websocket")); | |
1045 | + res.headers_mut().insert(hyper::header::CONNECTION, hyper::header::HeaderValue::from_static("upgrade")); | |
1046 | + res.headers_mut().insert("sec-websocket-accept", hyper::header::HeaderValue::from_str(&accept_key).unwrap()); | |
1047 | + | |
1048 | + let handler = handler.clone(); | |
1049 | + let path_clone = path.clone(); | |
1050 | + | |
1051 | + tokio::task::spawn(async move { | |
1052 | + match hyper::upgrade::on(&mut req).await { | |
1053 | + Ok(upgraded) => { | |
1054 | + let upgraded_io = hyper_util::rt::TokioIo::new(upgraded); | |
1055 | + let ws_stream = tokio_tungstenite::WebSocketStream::from_raw_socket( | |
1056 | + upgraded_io, | |
1057 | + tokio_tungstenite::tungstenite::protocol::Role::Server, | |
1058 | + None | |
1059 | + ).await; | |
1060 | + let (write, read) = ws_stream.split(); | |
1061 | + dispatch_ws(Box::new(read), Box::new(write), path_clone, &handler).await; | |
1062 | + } | |
1063 | + Err(e) => tracing::error!("WebSocket upgrade error: {}", e), | |
1064 | + } | |
1065 | + }); | |
1066 | + return Ok(res); | |
1067 | + } | |
1068 | + } | |
1069 | + } | |
1070 | + | |
1071 | + // WASM middleware (synchronous setup + one await) | |
1072 | + if let Some(ref wasm) = wasm_middleware { | |
1073 | + let wasm_span = tracing::debug_span!("wasm.middleware"); | |
1074 | + let mut hdrs = serde_json::Map::new(); | |
1075 | + for (k, v) in req.headers() { | |
1076 | + if let Ok(v_str) = v.to_str() { | |
1077 | + hdrs.insert( | |
1078 | + k.as_str().to_string(), | |
1079 | + serde_json::Value::String(v_str.to_string()), | |
1080 | + ); | |
1081 | + } | |
1082 | + } | |
1083 | + let req_json = serde_json::json!({ | |
1084 | + "path": path, | |
1085 | + "method": req.method().as_str(), | |
1086 | + "headers": hdrs, | |
1087 | + }) | |
1088 | + .to_string(); | |
1089 | + | |
1090 | + let wasm_result = { | |
1091 | + let _ws = wasm_span.enter(); | |
1092 | + wasm.execute_middleware(&req_json).await | |
1093 | + }; | |
1094 | + | |
1095 | + match wasm_result { | |
1096 | + Ok(res_json) => { | |
1097 | + if let Ok(parsed) = | |
1098 | + serde_json::from_str::<serde_json::Value>(&res_json) | |
1099 | + { | |
1100 | + if let Some(status) = parsed.get("status") { | |
1101 | + if let Some(status_code) = status.as_u64() { | |
1102 | + if status_code != 200 { | |
1103 | + let body = parsed | |
1104 | + .get("body") | |
1105 | + .and_then(|v| v.as_str()) | |
1106 | + .unwrap_or("") | |
1107 | + .to_string(); | |
1108 | + span.record("http.status_code", status_code); | |
1109 | + metrics.record_status( | |
1110 | + StatusCode::from_u16(status_code as u16) | |
1111 | + .unwrap_or(StatusCode::FORBIDDEN), | |
1112 | + ); | |
1113 | + metrics.record_latency( | |
1114 | + start.elapsed().as_secs_f64() * 1000.0, | |
1115 | + ); | |
1116 | + return Ok(json_response( | |
1117 | + StatusCode::from_u16(status_code as u16) | |
1118 | + .unwrap_or(StatusCode::FORBIDDEN), | |
1119 | + &body, | |
1120 | + )); | |
1121 | + } | |
1122 | + } | |
1123 | + } | |
1124 | + } | |
1125 | + } | |
1126 | + Err(e) => { | |
1127 | + tracing::error!("WASM execution error: {:?}", e); | |
1128 | + } | |
1129 | + } | |
1130 | + } | |
1131 | + | |
1132 | + // Run middleware chain (async, no entered span across await) | |
1133 | + let resp = chain.run(req).await; | |
1134 | + | |
1135 | + match resp { | |
1136 | + Ok(r) => { | |
1137 | + let status = r.status(); | |
1138 | + span.record("http.status_code", status.as_u16()); | |
1139 | + metrics.record_status(status); | |
1140 | + metrics.record_latency(start.elapsed().as_secs_f64() * 1000.0); | |
1141 | + | |
1142 | + // If the response is 404 and we have a static dir, try serving files | |
1143 | + if status == StatusCode::NOT_FOUND { | |
1144 | + if let Some(ref sd) = static_dir { | |
1145 | + if let Some(file_path) = sd.resolve(&path) { | |
1146 | + if tokio::fs::metadata(&file_path) | |
1147 | + .await | |
1148 | + .map(|m| m.is_file()) | |
1149 | + .unwrap_or(false) | |
1150 | + { | |
1151 | + return sd.serve_file(&file_path).await; | |
1152 | + } | |
1153 | + } | |
1154 | + } | |
1155 | + } | |
1156 | + Ok(r) | |
1157 | + } | |
1158 | + Err(_) => { | |
1159 | + span.record("http.status_code", 404u16); | |
1160 | + metrics.record_status(StatusCode::NOT_FOUND); | |
1161 | + metrics.record_latency(start.elapsed().as_secs_f64() * 1000.0); | |
1162 | + | |
1163 | + // Middleware error — try static files | |
1164 | + if let Some(ref sd) = static_dir { | |
1165 | + if let Some(file_path) = sd.resolve(&path) { | |
1166 | + if tokio::fs::metadata(&file_path) | |
1167 | + .await | |
1168 | + .map(|m| m.is_file()) | |
1169 | + .unwrap_or(false) | |
1170 | + { | |
1171 | + return sd.serve_file(&file_path).await; | |
1172 | + } | |
1173 | + } | |
1174 | + } | |
1175 | + Ok(json_response( | |
1176 | + StatusCode::NOT_FOUND, | |
1177 | + r#"{"error":"not found"}"#, | |
1178 | + )) | |
1179 | + } | |
1180 | + } | |
1181 | + } | |
1182 | + }); | |
1183 | + let builder = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new()); | |
1184 | + let mut conn = std::pin::pin!(builder.serve_connection_with_upgrades(io, svc)); | |
1185 | + | |
1186 | + if let Some(token) = token_clone { | |
1187 | + tokio::select! { | |
1188 | + res = &mut conn => { | |
1189 | + if let Err(e) = res { | |
1190 | + tracing::warn!("Connection error: {}", e); | |
1191 | + } | |
1192 | + } | |
1193 | + _ = token.cancelled() => { | |
1194 | + conn.as_mut().graceful_shutdown(); | |
1195 | + if let Err(e) = conn.await { | |
1196 | + tracing::warn!("Connection error during shutdown: {}", e); | |
1197 | + } | |
1198 | + } | |
1199 | + } | |
1200 | + } else { | |
1201 | + if let Err(e) = conn.await { | |
1202 | + tracing::warn!("Connection error: {}", e); | |
1203 | + } | |
1204 | + } | |
1205 | + conn_metrics.connection_closed(); | |
1206 | + }); | |
1207 | + } | |
1208 | + | |
1209 | + if shutdown.is_some() { | |
1210 | + tracing::info!( | |
1211 | + "Waiting for {} active connections to drain...", | |
1212 | + connections.len() | |
1213 | + ); | |
1214 | + tokio::select! { | |
1215 | + _ = async { while connections.join_next().await.is_some() {} } => { | |
1216 | + tracing::debug!("all connections closed gracefully"); | |
1217 | + } | |
1218 | + _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => { | |
1219 | + tracing::warn!("Graceful shutdown timeout exceeded. Dropping remaining connections."); | |
1220 | + } | |
1221 | + } | |
1222 | + } else { | |
1223 | + while connections.join_next().await.is_some() {} | |
1224 | + } | |
1225 | + | |
1226 | + Ok(()) | |
1227 | +} | |
1228 | + | |
1229 | +// --------------------------------------------------------------------------- | |
1230 | +// Route handler | |
1231 | +// --------------------------------------------------------------------------- | |
1232 | + | |
1233 | +async fn execute_handler( | |
1234 | + m: Match<'_, '_, Handler>, | |
1235 | + req: Request<Incoming>, | |
1236 | + pool: &BufferPool, | |
1237 | + metrics: &Metrics, | |
1238 | + health_registry: Option<&HealthRegistry>, | |
1239 | + graphql_schema: Option<&crate::graphql::AppSchema>, | |
1240 | + openapi_spec: Option<&str>, | |
1241 | +) -> Result<Response<ResponseBody>> { | |
1242 | + let handler_name = match m.handler { | |
1243 | + Handler::Static { .. } => "static", | |
1244 | + Handler::Echo => "echo", | |
1245 | + Handler::ParamsEcho => "params_echo", | |
1246 | + Handler::Sse => "sse", | |
1247 | + Handler::Health => "health", | |
1248 | + Handler::Ready => "ready", | |
1249 | + Handler::Live => "live", | |
1250 | + Handler::Prometheus => "prometheus", | |
1251 | + Handler::OpenApiJson => "openapi_json", | |
1252 | + Handler::SwaggerUi => "swagger_ui", | |
1253 | + Handler::Redoc => "redoc", | |
1254 | + Handler::GraphQL => "graphql", | |
1255 | + Handler::Custom(_) => "custom", | |
1256 | + }; | |
1257 | + let _handler_span = tracing::debug_span!("handler.execute", name = handler_name); | |
1258 | + | |
1259 | + match m.handler { | |
1260 | + Handler::Static { status, body } => Ok(json_response(*status, body)), | |
1261 | + Handler::Echo => { | |
1262 | + let body_bytes = match http_body_util::Limited::new(req.into_body(), 50 * 1024 * 1024) | |
1263 | + .collect() | |
1264 | + .await | |
1265 | + { | |
1266 | + Ok(collected) => { | |
1267 | + let bytes = collected.to_bytes(); | |
1268 | + metrics.add_bytes_in(bytes.len() as u64); | |
1269 | + bytes | |
1270 | + } | |
1271 | + Err(_) => { | |
1272 | + return Ok(json_response( | |
1273 | + StatusCode::BAD_REQUEST, | |
1274 | + r#"{"error":"bad request"}"#, | |
1275 | + )) | |
1276 | + } | |
1277 | + }; | |
1278 | + let mut buf = pool.acquire(body_bytes.len()); | |
1279 | + buf.extend_from_slice(&body_bytes); | |
1280 | + let body_str = String::from_utf8_lossy(&buf).to_string(); | |
1281 | + pool.release(buf); | |
1282 | + Ok(json_response(StatusCode::OK, &body_str)) | |
1283 | + } | |
1284 | + Handler::ParamsEcho => { | |
1285 | + let params_str: Vec<String> = m | |
1286 | + .params | |
1287 | + .iter() | |
1288 | + .map(|(k, v)| format!(r#""{}":"{}""#, k, v)) | |
1289 | + .collect(); | |
1290 | + let body = format!("{{{}}}", params_str.join(",")); | |
1291 | + Ok(json_response(StatusCode::OK, &body)) | |
1292 | + } | |
1293 | + Handler::Sse => Ok(sse_response()), | |
1294 | + Handler::Health => Ok(metrics::health_response()), | |
1295 | + Handler::Ready => { | |
1296 | + if let Some(reg) = health_registry { | |
1297 | + if reg.is_empty() { | |
1298 | + Ok(metrics::ready_response()) | |
1299 | + } else { | |
1300 | + Ok(reg.health_response().await) | |
1301 | + } | |
1302 | + } else { | |
1303 | + Ok(metrics::ready_response()) | |
1304 | + } | |
1305 | + } | |
1306 | + Handler::Live => Ok(metrics::live_response()), | |
1307 | + Handler::Prometheus => Ok(metrics::metrics_response(metrics)), | |
1308 | + Handler::OpenApiJson => { | |
1309 | + let body: String = if let Some(spec) = openapi_spec { | |
1310 | + spec.to_string() | |
1311 | + } else { | |
1312 | + BUILTIN_SPEC.clone() | |
1313 | + }; | |
1314 | + Ok(json_response(StatusCode::OK, &body)) | |
1315 | + } | |
1316 | + Handler::SwaggerUi => { | |
1317 | + let html = openapi::swagger_ui_html(); | |
1318 | + Ok(Response::builder() | |
1319 | + .status(StatusCode::OK) | |
1320 | + .header("content-type", "text/html; charset=utf-8") | |
1321 | + .header("content-length", html.len().to_string()) | |
1322 | + .body(crate::UnsyncBoxBody::new( | |
1323 | + http_body_util::Full::new(Bytes::from(html)) | |
1324 | + .map_err(|e: std::convert::Infallible| -> anyhow::Error { match e {} }), | |
1325 | + )) | |
1326 | + .unwrap()) | |
1327 | + } | |
1328 | + Handler::Redoc => { | |
1329 | + let html = openapi::redoc_html(); | |
1330 | + Ok(Response::builder() | |
1331 | + .status(StatusCode::OK) | |
1332 | + .header("content-type", "text/html; charset=utf-8") | |
1333 | + .header("content-length", html.len().to_string()) | |
1334 | + .body(crate::UnsyncBoxBody::new( | |
1335 | + http_body_util::Full::new(Bytes::from(html)) | |
1336 | + .map_err(|e: std::convert::Infallible| -> anyhow::Error { match e {} }), | |
1337 | + )) | |
1338 | + .unwrap()) | |
1339 | + } | |
1340 | + Handler::GraphQL => { | |
1341 | + if let Some(schema) = graphql_schema { | |
1342 | + crate::graphql::handle_graphql(schema, req) | |
1343 | + .await | |
1344 | + .or_else(|e| { | |
1345 | + Ok(Response::builder() | |
1346 | + .status(StatusCode::INTERNAL_SERVER_ERROR) | |
1347 | + .body(crate::UnsyncBoxBody::new( | |
1348 | + http_body_util::Full::new(Bytes::from(format!( | |
1349 | + "GraphQL Error: {}", | |
1350 | + e | |
1351 | + ))) | |
1352 | + .map_err( | |
1353 | + |e: std::convert::Infallible| -> anyhow::Error { match e {} }, | |
1354 | + ), | |
1355 | + )) | |
1356 | + .unwrap()) | |
1357 | + }) | |
1358 | + } else { | |
1359 | + Ok(Response::builder() | |
1360 | + .status(StatusCode::INTERNAL_SERVER_ERROR) | |
1361 | + .body(crate::UnsyncBoxBody::new( | |
1362 | + http_body_util::Full::new(Bytes::from("GraphQL schema not initialized")) | |
1363 | + .map_err(|e: std::convert::Infallible| -> anyhow::Error { match e {} }), | |
1364 | + )) | |
1365 | + .unwrap()) | |
1366 | + } | |
1367 | + } | |
1368 | + Handler::Custom(f) => f(req).await, | |
1369 | + } | |
1370 | +} | |
1371 | + | |
1372 | +// --------------------------------------------------------------------------- | |
1373 | +// SSE | |
1374 | +// --------------------------------------------------------------------------- | |
1375 | + | |
1376 | +fn sse_response() -> Response<ResponseBody> { | |
1377 | + let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes>>(16); | |
1378 | + | |
1379 | + tokio::spawn(async move { | |
1380 | + for i in 1..=10 { | |
1381 | + let msg = format!("data: {{\"count\":{}}}\n\n", i); | |
1382 | + if tx.send(Ok(Bytes::from(msg))).await.is_err() { | |
1383 | + break; | |
1384 | + } | |
1385 | + tokio::time::sleep(std::time::Duration::from_millis(100)).await; | |
1386 | + } | |
1387 | + }); | |
1388 | + | |
1389 | + let stream = tokio_stream::wrappers::ReceiverStream::new(rx); | |
1390 | + streaming_response(StatusCode::OK, "text/event-stream", stream) | |
1391 | +} | |
1392 | + | |
1393 | +/// Dispatch an accepted WebSocket connection to a registered handler. | |
1394 | +#[cfg(feature = "ws")] | |
1395 | +async fn dispatch_ws(read: WsRead, write: WsWrite, path: String, handler: &WsHandler) { | |
1396 | + if let Err(e) = handler(path, read, write).await { | |
1397 | + tracing::warn!("WebSocket handler error: {}", e); | |
1398 | + } | |
1399 | +} | |
1400 | + | |
1401 | +// --------------------------------------------------------------------------- | |
1402 | +// TLS | |
1403 | +// --------------------------------------------------------------------------- | |
1404 | + | |
1405 | +#[cfg(feature = "tls")] | |
1406 | +fn init_tls_provider() { | |
1407 | + rustls::crypto::ring::default_provider() | |
1408 | + .install_default() | |
1409 | + .expect("ring crypto provider should install"); | |
1410 | +} | |
1411 | + | |
1412 | +#[cfg(feature = "tls")] | |
1413 | +async fn serve_with_tls( | |
1414 | + listener: TcpListener, | |
1415 | + chain: Arc<MiddlewareChain>, | |
1416 | + config: TlsConfig, | |
1417 | + static_dir: Option<StaticDir>, | |
1418 | + metrics: Metrics, | |
1419 | + shutdown: Option<CancellationToken>, | |
1420 | + wasm_middleware: Option<Arc<crate::wasm::WasmEngine>>, | |
1421 | + #[cfg(feature = "ws")] ws_handler: Option<WsHandler>, | |
1422 | +) -> Result<()> { | |
1423 | + use std::fs::File; | |
1424 | + use std::io::BufReader; | |
1425 | + use std::sync::Arc as StdArc; | |
1426 | + | |
1427 | + init_tls_provider(); | |
1428 | + | |
1429 | + let certs = { | |
1430 | + let mut reader = BufReader::new(File::open(&config.cert_path)?); | |
1431 | + rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()? | |
1432 | + }; | |
1433 | + let key = { | |
1434 | + let mut reader = BufReader::new(File::open(&config.key_path)?); | |
1435 | + rustls_pemfile::private_key(&mut reader)? | |
1436 | + .ok_or_else(|| anyhow::anyhow!("no private key in {}", config.key_path))? | |
1437 | + }; | |
1438 | + | |
1439 | + let mut server_config = rustls::ServerConfig::builder() | |
1440 | + .with_no_client_auth() | |
1441 | + .with_single_cert(certs, key)?; | |
1442 | + | |
1443 | + server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; | |
1444 | + | |
1445 | + let acceptor = tokio_rustls::TlsAcceptor::from(StdArc::new(server_config)); | |
1446 | + | |
1447 | + let local_addr = listener.local_addr()?; | |
1448 | + tracing::info!( | |
1449 | + "Listening on {} (TLS, HTTP/1.1 + HTTP/2 via ALPN)", | |
1450 | + local_addr | |
1451 | + ); | |
1452 | + | |
1453 | + let mut connections = tokio::task::JoinSet::new(); | |
1454 | + | |
1455 | + loop { | |
1456 | + let token = shutdown.as_ref().cloned(); | |
1457 | + | |
1458 | + let (stream, peer) = if let Some(token) = &token { | |
1459 | + tokio::select! { | |
1460 | + result = listener.accept() => result?, | |
1461 | + _ = token.cancelled() => { | |
1462 | + tracing::info!("Shutdown signal received, stopping TLS accept loop"); | |
1463 | + break; | |
1464 | + } | |
1465 | + } | |
1466 | + } else { | |
1467 | + listener.accept().await? | |
1468 | + }; | |
1469 | + let acceptor = acceptor.clone(); | |
1470 | + let chain = chain.clone(); | |
1471 | + let static_dir = static_dir.clone(); | |
1472 | + let conn_metrics = metrics.clone(); | |
1473 | + let wasm_middleware = wasm_middleware.clone(); | |
1474 | + | |
1475 | + conn_metrics.connection_opened(); | |
1476 | + let spawn_metrics = conn_metrics.clone(); | |
1477 | + let peer_addr = peer; | |
1478 | + #[cfg(feature = "ws")] | |
1479 | + let ws_handler = ws_handler.clone(); | |
1480 | + let token_clone = shutdown.clone(); | |
1481 | + connections.spawn(async move { | |
1482 | + match acceptor.accept(stream).await { | |
1483 | + Ok(tls_stream) => { | |
1484 | + let io = TokioIo::new(tls_stream); | |
1485 | + let arena = Arc::new(SharedArena::new()); | |
1486 | + let svc = service_fn(move |mut req| { | |
1487 | + arena.reset(); | |
1488 | + let chain = chain.clone(); | |
1489 | + let arena = arena.clone(); | |
1490 | + let static_dir = static_dir.clone(); | |
1491 | + let metrics = spawn_metrics.clone(); | |
1492 | + let wasm_middleware = wasm_middleware.clone(); | |
1493 | + #[cfg(feature = "ws")] | |
1494 | + let ws_handler = ws_handler.clone(); | |
1495 | + async move { | |
1496 | + let path = req.uri().path().to_string(); | |
1497 | + let _arena_path = arena.alloc_str(&path); | |
1498 | + let start = std::time::Instant::now(); | |
1499 | + req.extensions_mut().insert(peer_addr); | |
1500 | + | |
1501 | + metrics.record_request(); | |
1502 | + | |
1503 | + #[cfg(feature = "ws")] | |
1504 | + if let Some(ref handler) = ws_handler { | |
1505 | + let is_ws = req.headers().get(hyper::header::UPGRADE) | |
1506 | + .and_then(|v| v.to_str().ok()) | |
1507 | + .map(|s| s.eq_ignore_ascii_case("websocket")) | |
1508 | + .unwrap_or(false); | |
1509 | + | |
1510 | + if is_ws { | |
1511 | + if let Some(key) = req.headers().get("sec-websocket-key") { | |
1512 | + let key_str = key.to_str().unwrap_or(""); | |
1513 | + let mut sha1 = sha1::Sha1::new(); | |
1514 | + use sha1::Digest; | |
1515 | + sha1.update(key_str.as_bytes()); | |
1516 | + sha1.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); | |
1517 | + use base64::Engine; | |
1518 | + let accept_key = base64::engine::general_purpose::STANDARD.encode(sha1.finalize()); | |
1519 | + | |
1520 | + let mut res = crate::json_response(StatusCode::SWITCHING_PROTOCOLS, ""); | |
1521 | + res.headers_mut().insert(hyper::header::UPGRADE, hyper::header::HeaderValue::from_static("websocket")); | |
1522 | + res.headers_mut().insert(hyper::header::CONNECTION, hyper::header::HeaderValue::from_static("upgrade")); | |
1523 | + res.headers_mut().insert("sec-websocket-accept", hyper::header::HeaderValue::from_str(&accept_key).unwrap()); | |
1524 | + | |
1525 | + let handler = handler.clone(); | |
1526 | + let path_clone = path.clone(); | |
1527 | + | |
1528 | + tokio::task::spawn(async move { | |
1529 | + match hyper::upgrade::on(&mut req).await { | |
1530 | + Ok(upgraded) => { | |
1531 | + let upgraded_io = hyper_util::rt::TokioIo::new(upgraded); | |
1532 | + let ws_stream = tokio_tungstenite::WebSocketStream::from_raw_socket( | |
1533 | + upgraded_io, | |
1534 | + tokio_tungstenite::tungstenite::protocol::Role::Server, | |
1535 | + None | |
1536 | + ).await; | |
1537 | + let (write, read) = ws_stream.split(); | |
1538 | + dispatch_ws(Box::new(read), Box::new(write), path_clone, &handler).await; | |
1539 | + } | |
1540 | + Err(e) => tracing::error!("WebSocket upgrade error: {}", e), | |
1541 | + } | |
1542 | + }); | |
1543 | + return Ok(res); | |
1544 | + } | |
1545 | + } | |
1546 | + } | |
1547 | + | |
1548 | + if let Some(ref wasm) = wasm_middleware { | |
1549 | + let mut hdrs = serde_json::Map::new(); | |
1550 | + for (k, v) in req.headers() { | |
1551 | + if let Ok(v_str) = v.to_str() { | |
1552 | + hdrs.insert( | |
1553 | + k.as_str().to_string(), | |
1554 | + serde_json::Value::String(v_str.to_string()), | |
1555 | + ); | |
1556 | + } | |
1557 | + } | |
1558 | + let req_json = serde_json::json!({ | |
1559 | + "path": path, | |
1560 | + "method": req.method().as_str(), | |
1561 | + "headers": hdrs, | |
1562 | + }) | |
1563 | + .to_string(); | |
1564 | + | |
1565 | + match wasm.execute_middleware(&req_json).await { | |
1566 | + Ok(res_json) => { | |
1567 | + if let Ok(parsed) = | |
1568 | + serde_json::from_str::<serde_json::Value>(&res_json) | |
1569 | + { | |
1570 | + if let Some(status) = parsed.get("status") { | |
1571 | + if let Some(status_code) = status.as_u64() { | |
1572 | + if status_code != 200 { | |
1573 | + let body = parsed | |
1574 | + .get("body") | |
1575 | + .and_then(|v| v.as_str()) | |
1576 | + .unwrap_or("") | |
1577 | + .to_string(); | |
1578 | + metrics.record_status( | |
1579 | + StatusCode::from_u16( | |
1580 | + status_code as u16, | |
1581 | + ) | |
1582 | + .unwrap_or(StatusCode::FORBIDDEN), | |
1583 | + ); | |
1584 | + metrics.record_latency( | |
1585 | + start.elapsed().as_secs_f64() * 1000.0, | |
1586 | + ); | |
1587 | + return Ok(json_response( | |
1588 | + StatusCode::from_u16( | |
1589 | + status_code as u16, | |
1590 | + ) | |
1591 | + .unwrap_or(StatusCode::FORBIDDEN), | |
1592 | + &body, | |
1593 | + )); | |
1594 | + } | |
1595 | + } | |
1596 | + } | |
1597 | + } | |
1598 | + } | |
1599 | + Err(e) => { | |
1600 | + tracing::error!("WASM execution error: {:?}", e); | |
1601 | + } | |
1602 | + } | |
1603 | + } | |
1604 | + | |
1605 | + let resp = chain.run(req).await; | |
1606 | + | |
1607 | + match resp { | |
1608 | + Ok(r) => { | |
1609 | + let status = r.status(); | |
1610 | + metrics.record_status(status); | |
1611 | + metrics.record_latency(start.elapsed().as_secs_f64() * 1000.0); | |
1612 | + | |
1613 | + if status == StatusCode::NOT_FOUND { | |
1614 | + if let Some(ref sd) = static_dir { | |
1615 | + if let Some(file_path) = sd.resolve(&path) { | |
1616 | + if tokio::fs::metadata(&file_path) | |
1617 | + .await | |
1618 | + .map(|m| m.is_file()) | |
1619 | + .unwrap_or(false) | |
1620 | + { | |
1621 | + return sd.serve_file(&file_path).await; | |
1622 | + } | |
1623 | + } | |
1624 | + } | |
1625 | + } | |
1626 | + Ok(r) | |
1627 | + } | |
1628 | + Err(_) => { | |
1629 | + metrics.record_status(StatusCode::NOT_FOUND); | |
1630 | + metrics.record_latency(start.elapsed().as_secs_f64() * 1000.0); | |
1631 | + | |
1632 | + if let Some(ref sd) = static_dir { | |
1633 | + if let Some(file_path) = sd.resolve(&path) { | |
1634 | + if tokio::fs::metadata(&file_path) | |
1635 | + .await | |
1636 | + .map(|m| m.is_file()) | |
1637 | + .unwrap_or(false) | |
1638 | + { | |
1639 | + return sd.serve_file(&file_path).await; | |
1640 | + } | |
1641 | + } | |
1642 | + } | |
1643 | + Ok(json_response( | |
1644 | + StatusCode::NOT_FOUND, | |
1645 | + r#"{"error":"not found"}"#, | |
1646 | + )) | |
1647 | + } | |
1648 | + } | |
1649 | + } | |
1650 | + }); | |
1651 | + | |
1652 | + let builder = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new()); | |
1653 | + let mut conn = std::pin::pin!(builder.serve_connection_with_upgrades(io, svc)); | |
1654 | + | |
1655 | + if let Some(token) = token_clone { | |
1656 | + tokio::select! { | |
1657 | + res = &mut conn => { | |
1658 | + if let Err(e) = res { | |
1659 | + tracing::warn!("TLS connection error from {}: {}", peer, e); | |
1660 | + } | |
1661 | + } | |
1662 | + _ = token.cancelled() => { | |
1663 | + conn.as_mut().graceful_shutdown(); | |
1664 | + if let Err(e) = conn.await { | |
1665 | + tracing::warn!("TLS connection error during shutdown: {}", e); | |
1666 | + } | |
1667 | + } | |
1668 | + } | |
1669 | + } else { | |
1670 | + if let Err(e) = conn.await { | |
1671 | + tracing::warn!("TLS connection error from {}: {}", peer, e); | |
1672 | + } | |
1673 | + } | |
1674 | + } | |
1675 | + Err(e) => { | |
1676 | + tracing::warn!("TLS handshake failed from {}: {}", peer, e); | |
1677 | + } | |
1678 | + } | |
1679 | + conn_metrics.connection_closed(); | |
1680 | + }); | |
1681 | + } | |
1682 | + | |
1683 | + if let Some(_) = shutdown { | |
1684 | + tracing::info!( | |
1685 | + "Waiting for {} active TLS connections to drain...", | |
1686 | + connections.len() | |
1687 | + ); | |
1688 | + tokio::select! { | |
1689 | + _ = async { while connections.join_next().await.is_some() {} } => { | |
1690 | + tracing::debug!("all connections closed after fallback drain"); | |
1691 | + } | |
1692 | + _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => { | |
1693 | + tracing::warn!("Graceful shutdown timeout exceeded. Dropping remaining TLS connections."); | |
1694 | + } | |
1695 | + } | |
1696 | + } else { | |
1697 | + while connections.join_next().await.is_some() {} | |
1698 | + } | |
1699 | + | |
1700 | + Ok(()) | |
1701 | +} | |
1702 | + | |
1703 | +// --------------------------------------------------------------------------- | |
1704 | +// Tests | |
1705 | +// --------------------------------------------------------------------------- | |
1706 | + | |
1707 | +#[cfg(test)] | |
1708 | +mod tests { | |
1709 | + use super::*; | |
1710 | + use crate::router::Router; | |
1711 | + | |
1712 | + #[tokio::test] | |
1713 | + async fn test_hello_route_via_router() { | |
1714 | + let mut router: Router<Handler> = Router::new(); | |
1715 | + router | |
1716 | + .insert( | |
1717 | + Method::GET, | |
1718 | + "/hello", | |
1719 | + Handler::Static { | |
1720 | + status: StatusCode::OK, | |
1721 | + body: r#"{"message":"hello"}"#, | |
1722 | + }, | |
1723 | + ) | |
1724 | + .unwrap(); | |
1725 | + let m = router.at(&Method::GET, "/hello").unwrap(); | |
1726 | + assert!(matches!(m.handler, Handler::Static { .. })); | |
1727 | + } | |
1728 | + | |
1729 | + #[tokio::test] | |
1730 | + async fn test_unknown_route_fallback() { | |
1731 | + let mut router: Router<Handler> = Router::new(); | |
1732 | + router.set_fallback(Handler::Static { | |
1733 | + status: StatusCode::NOT_FOUND, | |
1734 | + body: r#"{"error":"not found"}"#, | |
1735 | + }); | |
1736 | + let fb = router.fallback().unwrap(); | |
1737 | + assert!(matches!(fb, Handler::Static { .. })); | |
1738 | + } | |
1739 | + | |
1740 | + #[tokio::test] | |
1741 | + async fn test_users_param_route() { | |
1742 | + let mut router: Router<Handler> = Router::new(); | |
1743 | + router | |
1744 | + .insert( | |
1745 | + Method::GET, | |
1746 | + "/users/{id}", | |
1747 | + Handler::Static { | |
1748 | + status: StatusCode::OK, | |
1749 | + body: r#"{"message":"user lookup"}"#, | |
1750 | + }, | |
1751 | + ) | |
1752 | + .unwrap(); | |
1753 | + let m = router.at(&Method::GET, "/users/42").unwrap(); | |
1754 | + assert_eq!(m.params.get("id"), Some("42")); | |
1755 | + } | |
1756 | + | |
1757 | + #[tokio::test] | |
1758 | + async fn test_arena_reuse() { | |
1759 | + let arena = SharedArena::new(); | |
1760 | + arena.alloc_str("/hello").unwrap(); | |
1761 | + arena.reset(); | |
1762 | + let s = arena.alloc_str("/users/42").unwrap(); | |
1763 | + assert_eq!(s, "/users/42"); | |
1764 | + } | |
1765 | + | |
1766 | + #[tokio::test] | |
1767 | + async fn test_sse_route_exists() { | |
1768 | + let mut router: Router<Handler> = Router::new(); | |
1769 | + router.insert(Method::GET, "/events", Handler::Sse).unwrap(); | |
1770 | + let m = router.at(&Method::GET, "/events").unwrap(); | |
1771 | + assert!(matches!(m.handler, Handler::Sse)); | |
1772 | + } | |
1773 | + | |
1774 | + #[tokio::test] | |
1775 | + async fn test_health_route_exists() { | |
1776 | + let mut router: Router<Handler> = Router::new(); | |
1777 | + router | |
1778 | + .insert(Method::GET, "/health", Handler::Health) | |
1779 | + .unwrap(); | |
1780 | + let m = router.at(&Method::GET, "/health").unwrap(); | |
1781 | + assert!(matches!(m.handler, Handler::Health)); | |
1782 | + } | |
1783 | + | |
1784 | + #[tokio::test] | |
1785 | + async fn test_prometheus_route_exists() { | |
1786 | + let mut router: Router<Handler> = Router::new(); | |
1787 | + router | |
1788 | + .insert(Method::GET, "/metrics", Handler::Prometheus) | |
1789 | + .unwrap(); | |
1790 | + let m = router.at(&Method::GET, "/metrics").unwrap(); | |
1791 | + assert!(matches!(m.handler, Handler::Prometheus)); | |
1792 | + } | |
1793 | +} | |
1794 | + | |
1795 | +// --- gRPC stub placeholder --- | |
1796 | +// (We'll fully integrate the gRPC side-by-side loop next) | |
0 | 1797 | |