feat: use axum for ohttp gateway middleware

Author: Harshdev098

Cargo-minimal.lock

@@ -2411,6 +2411,8 @@ dependencies = [
 name = "ohttp-relay"
 version = "0.0.11"
 dependencies = [
+ "bhttp",
+ "bitcoin-ohttp",
  "byteorder",
  "bytes",
  "futures",

Cargo-recent.lock

@@ -2411,6 +2411,8 @@ dependencies = [
 name = "ohttp-relay"
 version = "0.0.11"
 dependencies = [
+ "bhttp",
+ "bitcoin-ohttp",
  "byteorder",
  "bytes",
  "futures",

ohttp-relay/Cargo.toml

@@ -47,6 +47,8 @@ tokio-util = { version = "0.7.16", features = ["net", "codec"] }
 tower = "0.5"
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
+ohttp = { package = "bitcoin-ohttp", version = "0.6" }
+bhttp = { version = "0.6.1", features = ["http"] }
 
 [dev-dependencies]
 mockito = "1.7.0"

ohttp-relay/src/gateway_helpers.rs

@@ -0,0 +1,126 @@
+use http_body_util::combinators::BoxBody;
+use http_body_util::{BodyExt, Full};
+use hyper::body::Bytes;
+use hyper::{Request, Response, Uri};
+
+pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
+
+pub const CHACHA20_POLY1305_NONCE_LEN: usize = 32;
+pub const POLY1305_TAG_SIZE: usize = 16;
+pub const ENCAPSULATED_MESSAGE_BYTES: usize = 65536;
+pub const BHTTP_REQ_BYTES: usize =
+    ENCAPSULATED_MESSAGE_BYTES - (CHACHA20_POLY1305_NONCE_LEN + POLY1305_TAG_SIZE);
+
+#[derive(Debug)]
+pub enum GatewayError {
+    BadRequest(String),
+    OhttpKeyRejection(String),
+    InternalServerError(String),
+}
+
+impl std::fmt::Display for GatewayError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            GatewayError::BadRequest(msg) => write!(f, "Bad request: {}", msg),
+            GatewayError::OhttpKeyRejection(msg) => write!(f, "OHTTP key rejection: {}", msg),
+            GatewayError::InternalServerError(msg) => write!(f, "Internal server error: {}", msg),
+        }
+    }
+}
+
+impl std::error::Error for GatewayError {}
+
+pub async fn decapsulate_ohttp_request<B>(
+    req: Request<B>,
+    ohttp_server: &ohttp::Server,
+) -> Result<(Request<BoxBody<Bytes, hyper::Error>>, ohttp::ServerResponse), GatewayError>
+where
+    B: hyper::body::Body<Data = Bytes> + Send + 'static,
+    B::Error: Into<BoxError>,
+{
+    // Collect OHTTP body
+    let ohttp_body = req
+        .into_body()
+        .collect()
+        .await
+        .map_err(|e| GatewayError::BadRequest(format!("Failed to read body: {}", e.into())))?
+        .to_bytes();
+
+    // Decapsulate using OHTTP server
+    let (bhttp_req, res_ctx) = ohttp_server.decapsulate(&ohttp_body).map_err(|e| {
+        GatewayError::OhttpKeyRejection(format!("OHTTP decapsulation failed: {}", e))
+    })?;
+
+    // Parse BHTTP message
+    let mut cursor = std::io::Cursor::new(bhttp_req);
+    let bhttp_msg = bhttp::Message::read_bhttp(&mut cursor)
+        .map_err(|e| GatewayError::BadRequest(format!("Invalid BHTTP: {}", e)))?;
+
+    let uri = Uri::builder()
+        .scheme(bhttp_msg.control().scheme().unwrap_or_default())
+        .authority(bhttp_msg.control().authority().unwrap_or_default())
+        .path_and_query(bhttp_msg.control().path().unwrap_or_default())
+        .build()
+        .map_err(|e| GatewayError::BadRequest(format!("Invalid URI: {}", e)))?;
+
+    let body = bhttp_msg.content().to_vec();
+    let mut http_req =
+        Request::builder().uri(uri).method(bhttp_msg.control().method().unwrap_or_default());
+
+    for header in bhttp_msg.header().fields() {
+        http_req = http_req.header(header.name(), header.value());
+    }
+
+    let request = http_req.body(full(body)).map_err(|e| {
+        GatewayError::InternalServerError(format!("Failed to build request: {}", e))
+    })?;
+
+    Ok((request, res_ctx))
+}
+
+pub async fn encapsulate_ohttp_response(
+    response: Response<BoxBody<Bytes, hyper::Error>>,
+    res_ctx: ohttp::ServerResponse,
+) -> Result<Response<BoxBody<Bytes, hyper::Error>>, GatewayError> {
+    let (parts, body) = response.into_parts();
+
+    let mut bhttp_res =
+        bhttp::Message::response(bhttp::StatusCode::try_from(parts.status.as_u16()).map_err(
+            |e| GatewayError::InternalServerError(format!("Invalid status code: {}", e)),
+        )?);
+
+    for (name, value) in parts.headers.iter() {
+        bhttp_res.put_header(name.as_str(), value.to_str().unwrap_or_default());
+    }
+
+    let full_body = body
+        .collect()
+        .await
+        .map_err(|e| GatewayError::InternalServerError(format!("Failed to collect body: {}", e)))?
+        .to_bytes();
+    bhttp_res.write_content(&full_body);
+
+    let mut bhttp_bytes = Vec::new();
+    bhttp_res.write_bhttp(bhttp::Mode::KnownLength, &mut bhttp_bytes).map_err(|e| {
+        GatewayError::InternalServerError(format!("BHTTP serialization failed: {}", e))
+    })?;
+
+    bhttp_bytes.resize(BHTTP_REQ_BYTES, 0);
+
+    let ohttp_res = res_ctx.encapsulate(&bhttp_bytes).map_err(|e| {
+        GatewayError::InternalServerError(format!("OHTTP encapsulation failed: {}", e))
+    })?;
+
+    assert!(
+        ohttp_res.len() == ENCAPSULATED_MESSAGE_BYTES,
+        "Unexpected OHTTP response size: {} != {}",
+        ohttp_res.len(),
+        ENCAPSULATED_MESSAGE_BYTES
+    );
+
+    Ok(Response::new(full(ohttp_res)))
+}
+
+fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
+    Full::new(chunk.into()).map_err(|never| match never {}).boxed()
+}

ohttp-relay/src/lib.rs

@@ -36,6 +36,12 @@ pub mod gateway_prober;
 mod gateway_uri;
 pub mod sentinel;
 pub use sentinel::SentinelTag;
+pub mod gateway_helpers;
+
+pub use gateway_helpers::{
+    decapsulate_ohttp_request, encapsulate_ohttp_response, BHTTP_REQ_BYTES,
+    ENCAPSULATED_MESSAGE_BYTES,
+};
 
 use crate::error::{BoxError, Error};
 

payjoin-service/src/lib.rs

@@ -6,34 +6,29 @@ use config::Config;
 use ohttp_relay::SentinelTag;
 use rand::Rng;
 use tokio_listener::{Listener, SystemOptions, UserOptions};
-use tower::Service;
+use tower::{Service, ServiceExt};
 use tracing::info;
-pub mod ohttp;
-
-use http_body_util::combinators::BoxBody;
-use hyper::body::Bytes;
-use hyper::{Request, StatusCode};
-use ohttp::{OhttpGatewayConfig, OhttpGatewayLayer};
-use tower::{ServiceBuilder, ServiceExt};
 
 pub mod cli;
 pub mod config;
+pub mod ohttp;
+
+use crate::ohttp::OhttpGatewayConfig;
 
 #[derive(Clone)]
 struct Services {
     directory: payjoin_directory::Service<payjoin_directory::FilesDb>,
     relay: ohttp_relay::Service,
-    sentinel_tag: SentinelTag,
+    ohttp_config: OhttpGatewayConfig,
 }
 
 pub async fn serve(config: Config) -> anyhow::Result<()> {
     let sentinel_tag = generate_sentinel_tag();
+    let directory = init_directory(&config, sentinel_tag).await?;
+    let ohttp_config = OhttpGatewayConfig::new(directory.ohttp.clone(), sentinel_tag);
 
-    let services = Services {
-        directory: init_directory(&config, sentinel_tag).await?,
-        relay: ohttp_relay::Service::new(sentinel_tag).await,
-        sentinel_tag,
-    };
+    let services =
+        Services { directory, relay: ohttp_relay::Service::new(sentinel_tag).await, ohttp_config };
     let app = Router::new().fallback(route_request).with_state(services);
 
     let listener =
@@ -61,12 +56,15 @@ pub async fn serve_manual_tls(
     use std::net::SocketAddr;
 
     let sentinel_tag = generate_sentinel_tag();
+    let directory = init_directory(&config, sentinel_tag).await?;
+    let ohttp_config = OhttpGatewayConfig::new(directory.ohttp.clone(), sentinel_tag);
 
     let services = Services {
-        directory: init_directory(&config, sentinel_tag).await?,
+        directory,
         relay: ohttp_relay::Service::new_with_roots(root_store, sentinel_tag).await,
-        sentinel_tag,
+        ohttp_config,
     };
+
     let app = Router::new().fallback(route_request).with_state(services);
 
     let addr: SocketAddr = config
@@ -134,63 +132,49 @@ async fn route_request(State(services): State<Services>, req: axum::extract::Req
         let mut relay = services.relay.clone();
         match relay.call(req).await {
             Ok(res) => res.into_response(),
-            Err(e) => (StatusCode::BAD_GATEWAY, e.to_string()).into_response(),
+            Err(e) => (axum::http::StatusCode::BAD_GATEWAY, e.to_string()).into_response(),
         }
     } else {
         handle_directory_request(services, req).await
     }
 }
 
 async fn handle_directory_request(services: Services, req: axum::extract::Request) -> Response {
-    let ohttp_server = services.directory.ohttp.clone();
-
-    let ohttp_config = OhttpGatewayConfig::new(ohttp_server, services.sentinel_tag);
-
-    let (parts, body) = req.into_parts();
-
-    use http_body_util::BodyExt as _;
-
-    let body_bytes = body
-        .collect()
-        .await
-        .map_err(|_| "Failed to collect body")
-        .expect("Failed to collect body")
-        .to_bytes();
-
-    let boxed_body = BoxBody::new(http_body_util::Full::new(body_bytes));
-
-    let hyper_req = Request::from_parts(parts, boxed_body);
-
-    let directory_service = tower::service_fn({
-        let directory = services.directory.clone();
-        move |req: Request<BoxBody<Bytes, hyper::Error>>| {
-            let mut dir = directory.clone();
-            async move {
-                dir.call(req).await.map_err(|e| {
-                    Box::new(std::io::Error::other(e.to_string()))
-                        as Box<dyn std::error::Error + Send + Sync>
-                })
-            }
-        }
-    });
-
-    let mut service_with_ohttp = ServiceBuilder::new()
-        .layer(OhttpGatewayLayer::new(ohttp_config))
-        .service(directory_service)
-        .boxed_clone();
-
-    match service_with_ohttp.ready().await {
-        Ok(ready_service) => match ready_service.call(hyper_req).await {
-            Ok(response) => {
-                let (parts, body) = response.into_parts();
-                let axum_body = axum::body::Body::new(body);
-                Response::from_parts(parts, axum_body).into_response()
-            }
+    let is_ohttp_request = matches!(
+        (req.method(), req.uri().path()),
+        (&Method::POST, "/.well-known/ohttp-gateway") | (&Method::POST, "/")
+    );
+
+    if is_ohttp_request {
+        let app = Router::new()
+            .fallback(directory_handler)
+            .layer(axum::middleware::from_fn_with_state(
+                services.ohttp_config.clone(),
+                crate::ohttp::ohttp_gateway,
+            ))
+            .with_state(services.directory.clone());
+
+        match app.oneshot(req).await {
+            Ok(response) => response,
             Err(e) =>
-                (StatusCode::INTERNAL_SERVER_ERROR, format!("Service error: {}", e)).into_response(),
-        },
+                (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Service error: {}", e))
+                    .into_response(),
+        }
+    } else {
+        directory_handler(State(services.directory), req).await
+    }
+}
+
+async fn directory_handler(
+    State(directory): State<payjoin_directory::Service<payjoin_directory::FilesDb>>,
+    req: axum::extract::Request,
+) -> Response {
+    let mut dir = directory.clone();
+    match dir.call(req).await {
+        Ok(response) => response.into_response(),
         Err(e) =>
-            (StatusCode::INTERNAL_SERVER_ERROR, format!("Service not ready: {}", e)).into_response(),
+            (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Directory error: {}", e))
+                .into_response(),
     }
 }