starknet_transaction_prover: add /health endpoint via HealthLayer

crates/starknet_transaction_prover/src/server.rs

@@ -29,12 +29,15 @@ pub const OHTTP_JSONRPSEE_BODY_BUILDER: fn(Full<Bytes>) -> 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)

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<S> Layer<S> for HealthLayer {
+    type Service = HealthService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        HealthService { inner }
+    }
+}
+
+#[derive(Clone)]
+pub struct HealthService<S> {
+    inner: S,
+}
+
+impl<S, ReqB> Service<Request<ReqB>> for HealthService<S>
+where
+    S: Service<Request<ReqB>, Response = Response<HttpBody>>,
+{
+    type Response = Response<HttpBody>;
+    type Error = S::Error;
+    type Future = Either<Ready<Result<Self::Response, Self::Error>>, S::Future>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        // 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<ReqB>) -> 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))
+    }
+}

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<HttpBody>,
+    Response = Response<HttpBody>,
+    Error = std::convert::Infallible,
+    Future = futures::future::Ready<Result<Response<HttpBody>, std::convert::Infallible>>,
+> + Clone {
+    tower::service_fn(|_req: Request<HttpBody>| {
+        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<HttpBody> {
+    Request::builder()
+        .method(method)
+        .uri(path)
+        .body(HttpBody::new(Full::new(Bytes::new())))
+        .expect("static body is infallible")
+}
+
+async fn read_body(response: Response<HttpBody>) -> (StatusCode, Vec<u8>, 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);
+}

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)