diff --git a/crates/starknet_transaction_prover/src/server.rs b/crates/starknet_transaction_prover/src/server.rs index 7a7d447b084..577e63cb368 100644 --- a/crates/starknet_transaction_prover/src/server.rs +++ b/crates/starknet_transaction_prover/src/server.rs @@ -29,12 +29,15 @@ pub const OHTTP_JSONRPSEE_BODY_BUILDER: fn(Full) -> HttpBody = HttpBody:: pub mod config; pub mod cors; pub mod errors; +pub mod health; #[cfg(test)] pub mod mock_rpc; pub mod rpc_api; pub mod rpc_impl; pub mod tls; +pub use health::{HealthLayer, HEALTH_PATH}; + #[cfg(test)] mod rpc_spec_test; @@ -75,7 +78,10 @@ pub async fn start_server( // type it expects. `HttpBody::new` is a zero-cost wrapper, so // non-OHTTP requests still stream through unbuffered. .set_http_middleware( + // `HealthLayer` sits outermost so `GET /health` is answered + // before any other middleware runs. ServiceBuilder::new() + .layer(HealthLayer) .option_layer(cors_layer) .layer(MapRequestBodyLayer::new(HttpBody::new)) .option_layer(ohttp_layer) diff --git a/crates/starknet_transaction_prover/src/server/health.rs b/crates/starknet_transaction_prover/src/server/health.rs new file mode 100644 index 00000000000..ec49cf8a59c --- /dev/null +++ b/crates/starknet_transaction_prover/src/server/health.rs @@ -0,0 +1,64 @@ +//! HTTP `/health` endpoint as a tower middleware layer. +//! +//! Short-circuits `GET /health` before the jsonrpsee service sees the request +//! (which would 405 a GET). Any other request passes through unchanged. + +use std::task::{Context, Poll}; + +use bytes::Bytes; +use futures::future::{ready, Either, Ready}; +use http::{header, Method, Request, Response, StatusCode}; +use http_body_util::Full; +use jsonrpsee::server::HttpBody; +use tower::{Layer, Service}; + +#[cfg(test)] +#[path = "health_test.rs"] +mod health_test; + +pub const HEALTH_PATH: &str = "/health"; + +const HEALTHY_BODY: &[u8] = br#"{"status":"ok"}"#; + +#[derive(Clone, Copy, Default)] +pub struct HealthLayer; + +impl Layer for HealthLayer { + type Service = HealthService; + + fn layer(&self, inner: S) -> Self::Service { + HealthService { inner } + } +} + +#[derive(Clone)] +pub struct HealthService { + inner: S, +} + +impl Service> for HealthService +where + S: Service, Response = Response>, +{ + type Response = Response; + type Error = S::Error; + type Future = Either>, S::Future>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + // Always ready: the health fast-path doesn't need the inner service. + // Inner backpressure is driven on demand by `inner.call` below. + Poll::Ready(Ok(())) + } + + fn call(&mut self, request: Request) -> Self::Future { + if request.method() == Method::GET && request.uri().path() == HEALTH_PATH { + let response = Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "application/json") + .body(HttpBody::new(Full::new(Bytes::from_static(HEALTHY_BODY)))) + .expect("response build with a static body is infallible"); + return Either::Left(ready(Ok(response))); + } + Either::Right(self.inner.call(request)) + } +} diff --git a/crates/starknet_transaction_prover/src/server/health_test.rs b/crates/starknet_transaction_prover/src/server/health_test.rs new file mode 100644 index 00000000000..288855aa25e --- /dev/null +++ b/crates/starknet_transaction_prover/src/server/health_test.rs @@ -0,0 +1,69 @@ +use bytes::Bytes; +use http::{Method, Request, Response, StatusCode}; +use http_body_util::{BodyExt, Full}; +use jsonrpsee::server::HttpBody; +use tower::{Layer, ServiceExt}; + +use crate::server::health::{HealthLayer, HEALTH_PATH}; + +/// Inner stub returning 418 so we can tell whether `HealthLayer` short-circuited. +fn fallthrough_service() -> impl tower::Service< + Request, + Response = Response, + Error = std::convert::Infallible, + Future = futures::future::Ready, std::convert::Infallible>>, +> + Clone { + tower::service_fn(|_req: Request| { + let response = Response::builder() + .status(StatusCode::IM_A_TEAPOT) + .body(HttpBody::new(Full::new(Bytes::from_static(b"fallthrough")))) + .expect("static body is infallible"); + futures::future::ready(Ok::<_, std::convert::Infallible>(response)) + }) +} + +fn empty_request(method: Method, path: &str) -> Request { + Request::builder() + .method(method) + .uri(path) + .body(HttpBody::new(Full::new(Bytes::new()))) + .expect("static body is infallible") +} + +async fn read_body(response: Response) -> (StatusCode, Vec, http::HeaderMap) { + let (parts, body) = response.into_parts(); + let bytes = body.collect().await.expect("body collect").to_bytes().to_vec(); + (parts.status, bytes, parts.headers) +} + +#[tokio::test] +async fn get_health_returns_200_with_json_body() { + let svc = HealthLayer.layer(fallthrough_service()); + + let response = svc.oneshot(empty_request(Method::GET, HEALTH_PATH)).await.unwrap(); + + let (status, body, headers) = read_body(response).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body, br#"{"status":"ok"}"#); + assert_eq!(headers.get(http::header::CONTENT_TYPE).unwrap(), "application/json"); +} + +#[tokio::test] +async fn non_get_health_falls_through() { + let svc = HealthLayer.layer(fallthrough_service()); + + let response = svc.oneshot(empty_request(Method::POST, HEALTH_PATH)).await.unwrap(); + + let (status, _body, _) = read_body(response).await; + assert_eq!(status, StatusCode::IM_A_TEAPOT); +} + +#[tokio::test] +async fn get_other_path_falls_through() { + let svc = HealthLayer.layer(fallthrough_service()); + + let response = svc.oneshot(empty_request(Method::GET, "/")).await.unwrap(); + + let (status, _body, _) = read_body(response).await; + assert_eq!(status, StatusCode::IM_A_TEAPOT); +} diff --git a/crates/starknet_transaction_prover/src/server/tls.rs b/crates/starknet_transaction_prover/src/server/tls.rs index f41ea44cb28..d03e93650b0 100644 --- a/crates/starknet_transaction_prover/src/server/tls.rs +++ b/crates/starknet_transaction_prover/src/server/tls.rs @@ -27,7 +27,7 @@ use tower_http::map_request_body::MapRequestBodyLayer; use tower_http::map_response_body::MapResponseBodyLayer; use tracing::warn; -use super::OhttpJsonrpseeLayer; +use crate::server::{HealthLayer, OhttpJsonrpseeLayer}; /// Maximum time allowed for a TLS handshake before the connection is dropped. const TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); @@ -59,7 +59,10 @@ pub async fn start_tls_server( let svc_builder = ServerBuilder::default() .set_config(server_config) .set_http_middleware( + // `HealthLayer` sits outermost so `GET /health` is answered before + // any other middleware runs. ServiceBuilder::new() + .layer(HealthLayer) .option_layer(cors_layer) .layer(MapRequestBodyLayer::new(HttpBody::new)) .option_layer(ohttp_layer)