swadhinbiswas

[email protected]
2026-07-07 08:30:00 +0600 committed NaN years ago
0dbd480
commit 0dbd480
8 files changed
+3046 -0
▼crates/justapi-core/Cargo.toml
+79-0
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
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
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
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
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
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
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
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