swadhin / justapi public

feat: opt-in security headers, real /metrics, genuine /ready health checks

swadhinbiswas committed 3 months ago

dbcd571

1 parent bda80da

commit dbcd571

3 files changed

+278 -11

Files

crates

justapi-core

src

server

mod.rs +35

justapi-py

python

justapi

app.py +84 -11

src

native

app.rs +159

Showing 3 changed files+278-11

crates/justapi-core/src/server/mod.rs crates/justapi-py/python/justapi/app.py crates/justapi-py/src/native/app.rs

▼crates/justapi-core/src/server/mod.rs

+35-0

@@ -287,6 +287,13 @@ pub struct Server {
287	287	#[cfg(feature = "ws")]
288	288	ws_handler: Option<WsHandler>,
289	289	router: Option<Router<Handler>>,
	290	+ /// Tracks whether a custom handler (via `with_handler`) was installed, so we
	291	+ /// can refuse to also install the default router (which would silently
	292	+ /// discard the custom handler — see `with_default_routes`).
	293	+ custom_handler_set: bool,
	294	+ /// Tracks whether the default router (via `with_default_routes`) was
	295	+ /// installed, so we can refuse a later `with_handler` (same footgun).
	296	+ default_routes_set: bool,
290	297	openapi_spec: Option<Arc<String>>,
291	298	}
292	299
@@ -323,6 +330,8 @@ impl Server {
323	330	#[cfg(feature = "ws")]
324	331	ws_handler: Some(default_ws_echo()),
325	332	router: Some(router),
	333	+ custom_handler_set: false,
	334	+ default_routes_set: false,
326	335	openapi_spec: None,
327	336	}
328	337	}
@@ -344,13 +353,32 @@ impl Server {
344	353
345	354	/// Set a custom fallback handler (used by the ASGI shim).
346	355	pub fn with_handler(mut self, handler: crate::middleware::HandlerFn) -> Self {
	356	+ // A custom handler and the built-in default router are mutually
	357	+ // exclusive: installing the default router after a custom handler would
	358	+ // silently discard every user route (the router-derived handler wins in
	359	+ // `run`). Fail loudly instead of losing routes without explanation.
	360	+ if self.default_routes_set {
	361	+ panic!(
	362	+ "with_handler() called after with_default_routes(): the default router would \
	363	+ silently discard your custom handler. Choose one path — custom handler OR \
	364	+ default routes — not both."
	365	+ );
	366	+ }
347	367	// If the user specifies a custom handler, it replaces the entire router behavior.
348	368	self.router = None;
349	369	self.chain.set_handler(handler);
	370	+ self.custom_handler_set = true;
350	371	self
351	372	}
352	373
353	374	pub fn with_default_routes(mut self) -> Self {
	375	+ if self.custom_handler_set {
	376	+ panic!(
	377	+ "with_default_routes() called after with_handler(): it would silently discard \
	378	+ your custom handler and every user route. Choose one path — default routes OR \
	379	+ a custom handler — not both."
	380	+ );
	381	+ }
354	382	let mut router = self.router.unwrap_or_default();
355	383	router
356	384	.insert(
@@ -376,6 +404,7 @@ impl Server {
376	404	body: r#"{"error":"not found"}"#,
377	405	});
378	406	self.router = Some(router);
	407	+ self.default_routes_set = true;
379	408	self
380	409	}
381	410
@@ -672,6 +701,12 @@ impl Server {
672	701	&self.health_registry
673	702	}
674	703
	704	+ /// Return a cloned `Arc` handle to the health registry so other owners
	705	+ /// (e.g. the Python binding layer) can run checks independently.
	706	+ pub fn health_registry_arc(&self) -> std::sync::Arc<HealthRegistry> {
	707	+ std::sync::Arc::clone(&self.health_registry)
	708	+ }
	709	+
675	710	pub fn with_health_registry(mut self, registry: HealthRegistry) -> Self {
676	711	self.health_registry = Arc::new(registry);
677	712	self
678	713

▼crates/justapi-py/python/justapi/app.py

+84-11

@@ -8,7 +8,7 @@ from ._justapi import _JustAPIApp, TokenStreamResponse, ValidatedStreamResponse
8	8	from .exceptions import WebSocketException
9	9	from .websockets import WebSocket
10	10	from .system import register_system_routes, build_help, build_openapi
11		-from .responses import PlainTextResponse
	11	+from .responses import PlainTextResponse, HTMLResponse, JSONResponse
12	12
13	13	_PY_TYPE_TO_JSON = {
14	14	str: "string",
@@ -34,19 +34,57 @@ def _builtin_live():
34	34	return {"status": "alive"}
35	35
36	36
37		-def _builtin_ready():
38		- return {"status": "ready"}
	37	+def _builtin_ready(self):
	38	+ ready, report = self._app.health_ready()
	39	+ try:
	40	+ payload = json.loads(report)
	41	+ except (ValueError, TypeError):
	42	+ payload = {"status": "ready" if ready else "not_ready"}
	43	+ if not ready:
	44	+ return JSONResponse(payload, status_code=503)
	45	+ return JSONResponse(payload)
39	46
40	47
41		-def _builtin_metrics():
42		- body = (
43		- "# HELP justapi_up 1 if the server process is up.\n"
44		- "# TYPE justapi_up gauge\n"
45		- "justapi_up 1\n"
46		- )
	48	+def _builtin_metrics(self):
	49	+ body = self._app.metrics_prometheus()
47	50	return PlainTextResponse(body)
48	51
49	52
	53	+def _builtin_openapi(app):
	54	+ """Serve the generated OpenAPI 3.1 document for the Python app."""
	55	+ return JSONResponse(build_openapi(app))
	56	+
	57	+
	58	+def _builtin_docs(app):
	59	+ """Serve Swagger UI (interactive API documentation) for the Python app."""
	60	+ html = (
	61	+ "<!DOCTYPE html>\n<html lang=\"en\">\n<head>\n"
	62	+ " <meta charset=\"UTF-8\">\n <title>API Docs — Swagger UI</title>\n"
	63	+ " <link rel=\"stylesheet\" href=\"https://unpkg.com/swagger-ui-dist/swagger-ui.css\">\n"
	64	+ "</head>\n<body>\n <div id=\"swagger-ui\"></div>\n"
	65	+ " <script src=\"https://unpkg.com/swagger-ui-dist/swagger-ui-bundle.js\"></script>\n"
	66	+ " <script>\n"
	67	+ " window.onload = () => {\n"
	68	+ " window.ui = SwaggerUIBundle({ url: '/openapi.json', dom_id: '#swagger-ui' });\n"
	69	+ " };\n"
	70	+ " </script>\n</body>\n</html>\n"
	71	+ )
	72	+ return HTMLResponse(html)
	73	+
	74	+
	75	+def _builtin_redoc(app):
	76	+ """Serve ReDoc (alternative API documentation) for the Python app."""
	77	+ html = (
	78	+ "<!DOCTYPE html>\n<html lang=\"en\">\n<head>\n"
	79	+ " <meta charset=\"UTF-8\">\n <title>API Docs — ReDoc</title>\n"
	80	+ " <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\">\n"
	81	+ "</head>\n<body>\n <redoc spec-url='/openapi.json'></redoc>\n"
	82	+ " <script src=\"https://unpkg.com/redoc/bundles/redoc.standalone.js\"></script>\n"
	83	+ "</body>\n</html>\n"
	84	+ )
	85	+ return HTMLResponse(html)
	86	+
	87	+
50	88	def _model_json_schema(ann):
51	89	"""Best-effort JSON schema from a Pydantic / dataclass / model annotation."""
52	90	try:
@@ -160,8 +198,16 @@ class JustAPIApp:
160	198	# Built-in probe/metrics endpoints for the Python app.
161	199	self.get("/health", _builtin_health, include_in_schema=False, name="builtin_health")
162	200	self.get("/live", _builtin_live, include_in_schema=False, name="builtin_live")
163		- self.get("/ready", _builtin_ready, include_in_schema=False, name="builtin_ready")
164		- self.get("/metrics", _builtin_metrics, include_in_schema=False, name="builtin_metrics")
	201	+ self.get("/ready", lambda: _builtin_ready(self), include_in_schema=False, name="builtin_ready")
	202	+ self.get("/metrics", lambda: _builtin_metrics(self), include_in_schema=False, name="builtin_metrics")
	203	+ # Interactive API documentation for the Python app. The core server's
	204	+ # /docs, /redoc, /openapi.json live only in the default router, which the
	205	+ # Python app replaces via with_handler(); these builtin routes restore
	206	+ # that DX so the README's "automatic interactive API documentation"
	207	+ # claim holds for Python apps too.
	208	+ self.get("/openapi.json", lambda: _builtin_openapi(self), include_in_schema=False, name="builtin_openapi")
	209	+ self.get("/docs", lambda: _builtin_docs(self), include_in_schema=False, name="builtin_docs")
	210	+ self.get("/redoc", lambda: _builtin_redoc(self), include_in_schema=False, name="builtin_redoc")
165	211
166	212	def _record(self, method, path, handler, *, body_schema=None, schema=None,
167	213	experimental=False, **kw):
@@ -806,6 +852,33 @@ class JustAPIApp:
806	852	"""
807	853	self._app.enable_request_coalescing(headers)
808	854
	855	+ def enable_secure_headers(self, with_hsts: bool = False):
	856	+ """
	857	+ Apply safe HTTP security headers to every response:
	858	+ ``X-Content-Type-Options: nosniff``, ``X-Frame-Options: DENY``,
	859	+ ``Content-Security-Policy: default-src 'self'``, and
	860	+ ``X-XSS-Protection: 0``.
	861	+
	862	+ HSTS is omitted by default because the Python server terminates
	863	+ connections in plaintext. Pass ``with_hsts=True`` only when you
	864	+ terminate TLS in-process. Note that the default CSP (``'self'``) will
	865	+ block third-party/CDN resources — relax it via the Rust
	866	+ ``SecurityHeaders`` builder if your frontend needs them.
	867	+ """
	868	+ self._app.enable_secure_headers(with_hsts)
	869	+
	870	+ def register_health_check(self, name: str, check: typing.Callable[[], bool]):
	871	+ """
	872	+ Register a dependency readiness probe used by the ``/ready`` endpoint.
	873	+
	874	+ ``check`` is a zero-argument callable invoked synchronously under the
	875	+ GIL. Returning truthy (or raising nothing) means the dependency is
	876	+ healthy; returning falsy or raising marks it unhealthy and makes
	877	+ ``/ready`` return 503. The ``name`` identifies the component in the
	878	+ readiness report.
	879	+ """
	880	+ self._app.register_health_check(name, check)
	881	+
809	882	def set_grpc_addr(self, addr: str):
810	883	self._app.set_grpc_addr(addr)
811	884
812	885

▼crates/justapi-py/src/native/app.rs

+159-0

@@ -58,6 +58,23 @@ pub struct JustAPIApp {
58	58	pub gateway_config: Option<String>,
59	59	pub circuit_breaker_config: Option<(usize, u64)>,
60	60	pub coalesce_headers: Option<Vec<String>>,
	61	+ /// When true, apply safe (non-HSTS) security headers to every response by
	62	+ /// default. Off by default because forcing a CSP would break apps that load
	63	+ /// external resources (e.g. CDN-hosted docs UIs). Call
	64	+ /// `enable_secure_headers()` to opt in.
	65	+ pub secure_headers: bool,
	66	+ /// Explicit security-headers config (e.g. with HSTS). When `secure_headers`
	67	+ /// is true but this is `None`, the safe non-HSTS default is used.
	68	+ pub secure_headers_config: Option<justapi_core::middleware::SecurityHeaders>,
	69	+ /// Live metrics collector, populated just before the server starts so the
	70	+ /// Python `/metrics` builtin can export real Prometheus data.
	71	+ pub metrics: Option<justapi_core::metrics::Metrics>,
	72	+ /// Health registry, populated just before the server starts so the Python
	73	+ /// `/ready` builtin reflects registered dependency checks.
	74	+ pub health_registry: Option<std::sync::Arc<justapi_core::health::HealthRegistry>>,
	75	+ /// Health checks registered from Python via `register_health_check`, applied
	76	+ /// to the server's registry at startup.
	77	+ pub health_checks: Vec<(String, Py<PyAny>)>,
61	78	/// Per-route OpenAPI metadata, keyed by `(method, path)`.
62	79	pub route_meta:
63	80	std::collections::HashMap<(hyper::Method, String), justapi_core::openapi::RouteMeta>,
@@ -83,6 +100,53 @@ pub struct PyTool {
83	100	pub handler: Py<PyAny>,
84	101	}
85	102
	103	+// ---------------------------------------------------------------------------
	104	+// Python-registered health checks
	105	+// ---------------------------------------------------------------------------
	106	+
	107	+/// A `HealthCheck` backed by a Python callable. The callable is invoked
	108	+/// synchronously under the GIL; a truthy return means healthy, a falsy return
	109	+/// or an exception means unhealthy (with the detail captured in the report).
	110	+struct PyHealthCheck {
	111	+ name: &'static str,
	112	+ check: Py<PyAny>,
	113	+}
	114	+
	115	+impl justapi_core::health::HealthCheck for PyHealthCheck {
	116	+ fn name(&self) -> &'static str {
	117	+ self.name
	118	+ }
	119	+
	120	+ fn check(
	121	+ &self,
	122	+ ) -> impl std::future::Future<Output = justapi_core::health::HealthStatus> + Send {
	123	+ let name = self.name;
	124	+ let check = Python::attach(\|py\| self.check.clone_ref(py));
	125	+ async move {
	126	+ Python::attach(\|py\| {
	127	+ let check = check.clone_ref(py);
	128	+ match check.bind(py).call0() {
	129	+ Ok(v) => match v.is_truthy() {
	130	+ Ok(true) => justapi_core::health::HealthStatus::Healthy,
	131	+ Ok(false) => justapi_core::health::HealthStatus::Unhealthy(format!(
	132	+ "health check '{}' reported unhealthy",
	133	+ name
	134	+ )),
	135	+ Err(e) => justapi_core::health::HealthStatus::Unhealthy(format!(
	136	+ "health check '{}' eval error: {}",
	137	+ name, e
	138	+ )),
	139	+ },
	140	+ Err(e) => justapi_core::health::HealthStatus::Unhealthy(format!(
	141	+ "health check '{}' raised: {}",
	142	+ name, e
	143	+ )),
	144	+ }
	145	+ })
	146	+ }
	147	+ }
	148	+}
	149	+
86	150	// ---------------------------------------------------------------------------
87	151	// Streaming structured-output validation
88	152	// ---------------------------------------------------------------------------
@@ -217,6 +281,11 @@ impl JustAPIApp {
217	281	gateway_config: None,
218	282	circuit_breaker_config: None,
219	283	coalesce_headers: None,
	284	+ secure_headers: false,
	285	+ secure_headers_config: None,
	286	+ metrics: None,
	287	+ health_registry: None,
	288	+ health_checks: Vec::new(),
220	289	route_meta: std::collections::HashMap::new(),
221	290	named_routes: std::collections::HashMap::new(),
222	291	frontend_mounts: Vec::new(),
@@ -299,6 +368,66 @@ impl JustAPIApp {
299	368	self.coalesce_headers = Some(headers.unwrap_or_default());
300	369	}
301	370
	371	+ /// Apply safe security headers (`X-Content-Type-Options`, `X-Frame-Options`,
	372	+ /// `Content-Security-Policy`, `X-XSS-Protection`) to every response. HSTS is
	373	+ /// intentionally omitted because the Python server runs over plaintext by
	374	+ /// default; set `with_hsts=True` only when terminating TLS in-process.
	375	+ #[pyo3(name = "enable_secure_headers")]
	376	+ fn enable_secure_headers(&mut self, with_hsts: Option<bool>) {
	377	+ self.secure_headers = true;
	378	+ if with_hsts.unwrap_or(false) {
	379	+ // HSTS included only on explicit opt-in for a TLS-terminating deploy.
	380	+ self.secure_headers_config =
	381	+ Some(justapi_core::middleware::SecurityHeaders::default().with_hsts_preload());
	382	+ }
	383	+ }
	384	+
	385	+ /// Register a Python callable as a readiness probe. The callable is invoked
	386	+ /// synchronously under the GIL; a truthy return (or no exception) means the
	387	+ /// dependency is healthy. Used by the `/ready` builtin via the health
	388	+ /// registry. `name` identifies the component in the readiness report.
	389	+ #[pyo3(name = "register_health_check")]
	390	+ fn register_health_check(&mut self, name: String, check: Py<PyAny>) {
	391	+ self.health_checks.push((name, check));
	392	+ }
	393	+
	394	+ /// Return the live Prometheus exposition for the running server. Used by the
	395	+ /// Python `/metrics` builtin so it exports real data instead of a stub.
	396	+ #[pyo3(name = "metrics_prometheus")]
	397	+ fn metrics_prometheus(&self) -> String {
	398	+ match &self.metrics {
	399	+ Some(m) => m.prometheus(),
	400	+ None => "# metrics not yet initialised\n".to_string(),
	401	+ }
	402	+ }
	403	+
	404	+ /// Returns `(ready, report_json)` for the running server's health registry.
	405	+ /// `ready` is true unless a registered dependency check is unhealthy. The
	406	+ /// Python `/ready` builtin uses this to give real readiness (not a static
	407	+ /// 200). When no dependencies are registered the app is always ready.
	408	+ #[pyo3(name = "health_ready")]
	409	+ fn health_ready(&self) -> (bool, String) {
	410	+ match &self.health_registry {
	411	+ Some(reg) => {
	412	+ let rt = tokio::runtime::Builder::new_current_thread().enable_all().build();
	413	+ match rt {
	414	+ Ok(rt) => {
	415	+ let report = rt.block_on(reg.check_all());
	416	+ let ready = !matches!(
	417	+ report.overall(),
	418	+ justapi_core::health::OverallHealth::Unhealthy
	419	+ );
	420	+ let json = serde_json::to_string(&report)
	421	+ .unwrap_or_else(\|_\| "{\"status\":\"unknown\"}".to_string());
	422	+ (ready, json)
	423	+ }
	424	+ Err(_) => (true, "{\"status\":\"ready\"}".to_string()),
	425	+ }
	426	+ }
	427	+ None => (true, "{\"status\":\"ready\"}".to_string()),
	428	+ }
	429	+ }
	430	+
302	431	fn load_wasm_middleware(&mut self, path: &str) -> PyResult<()> {
303	432	let bytes = std::fs::read(path).map_err(\|e\| {
304	433	pyo3::exceptions::PyIOError::new_err(format!(
@@ -1174,6 +1303,9 @@ impl JustAPIApp {
1174	1303	let gateway_config_path = app.gateway_config.take();
1175	1304	let cb_config = app.circuit_breaker_config.take();
1176	1305	let coalesce_headers = app.coalesce_headers.take();
	1306	+ let health_checks = std::mem::take(&mut app.health_checks);
	1307	+ let secure_headers = app.secure_headers;
	1308	+ let secure_headers_config = app.secure_headers_config.take();
1177	1309	let frontend_mounts = std::mem::take(&mut app.frontend_mounts);
1178	1310
1179	1311	// The app object is surfaced on `Request.app`; capture it before
@@ -1362,6 +1494,16 @@ impl JustAPIApp {
1362	1494	let mut server = justapi_core::Server::new(addr)
1363	1495	.with_handler(handler)
1364	1496	.with_shutdown(shutdown_signal);
	1497	+ if !health_checks.is_empty() {
	1498	+ let mut reg = justapi_core::health::HealthRegistry::new();
	1499	+ for (name, check) in health_checks {
	1500	+ reg.register(PyHealthCheck {
	1501	+ name: Box::leak(name.into_boxed_str()),
	1502	+ check,
	1503	+ });
	1504	+ }
	1505	+ server = server.with_health_registry(reg);
	1506	+ }
1365	1507	if let Some(ref spec) = openapi_spec {
1366	1508	server = server.with_openapi_spec(spec.clone());
1367	1509	}
@@ -1415,6 +1557,12 @@ impl JustAPIApp {
1415	1557	server = server.add_middleware(coalescer);
1416	1558	}
1417	1559
	1560	+ if secure_headers {
	1561	+ let sh = secure_headers_config
	1562	+ .unwrap_or_else(\|\| justapi_core::middleware::SecurityHeaders::default().without_hsts());
	1563	+ server = server.add_security_headers(sh);
	1564	+ }
	1565	+
1418	1566	if let Some(path) = gateway_config_path {
1419	1567	tracing::info!("Starting Gateway with hot reloading for config: {}", path);
1420	1568	let gateway_state = justapi_core::gateway::GatewayState::new(&path);
@@ -1602,6 +1750,17 @@ impl JustAPIApp {
1602	1750	server = server.with_ws(ws_handler);
1603	1751	}
1604	1752
	1753	+ // Publish the live metrics collector onto the app object so the
	1754	+ // Python `/metrics` builtin exports real data. `app` is a
	1755	+ // `!Send` borrow, so we mutate through the owning `Py<Self>`
	1756	+ // via `Python::attach` rather than capturing `app` here.
	1757	+ let live_metrics = server.metrics().clone();
	1758	+ let live_health = server.health_registry_arc();
	1759	+ Python::attach(move \|py\| {
	1760	+ let mut a = slf.borrow_mut(py);
	1761	+ a.metrics = Some(live_metrics);
	1762	+ a.health_registry = Some(live_health);
	1763	+ });
1605	1764	server.run().await
1606	1765	})
1607	1766	});
1608	1767