fix: Handle routing using worker header

Author: gonzalezzfelipe

Cargo.lock

@@ -7,6 +7,7 @@ edition = "2021"
 balius-runtime = { git = "https://github.com/gonzalezzfelipe/balius.git", branch="feat/support-object-urls", features = ["http", "aws"] }
 config = { version = "0.15.9", default-features = false, features = ["toml", "json"] }
 futures-util = "0.3.30"
+hex = "0.4.3"
 k8s-openapi = { version = "0.24.0", features = ["latest"] }
 kube = { version = "0.98.0", features = ["runtime", "client", "derive"] }
 lazy_static = "1.4.0"
@@ -21,3 +22,4 @@ tokio-util = "0.7.13"
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"
 url = "2.5.4"
+warp = { version = "0.3.7" }

instance/Cargo.toml

@@ -5,6 +5,7 @@ use tracing::{debug, warn, Level};
 
 mod config;
 mod runtime;
+mod server;
 
 async fn wait_for_exit_signal() {
     let mut sigterm =
@@ -59,7 +60,7 @@ async fn main() -> miette::Result<()> {
     let cancel = hook_exit_token();
 
     let jsonrpc_server = async {
-        balius_runtime::drivers::jsonrpc::serve(config.rpc.clone(), runtime.clone(), cancel.clone())
+        server::serve(config.rpc.clone(), runtime.clone(), cancel.clone())
             .await
             .into_diagnostic()
             .context("Running JsonRPC server")

instance/src/main.rs

@@ -0,0 +1,101 @@
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use std::net::SocketAddr;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, error};
+use warp::Filter as _;
+
+use balius_runtime::{wit, Error, Runtime};
+
+#[derive(Deserialize)]
+struct Request {
+    pub id: Option<String>,
+    pub method: String,
+    pub params: serde_json::Value,
+}
+
+#[derive(Serialize)]
+struct ErrorResponse {
+    error: String,
+}
+
+fn parse_request(body: serde_json::Value) -> Result<Request, ErrorResponse> {
+    match serde_json::from_value(body) {
+        Ok(x) => Ok(x),
+        Err(x) => Err(ErrorResponse {
+            error: x.to_string(),
+        }),
+    }
+}
+
+pub async fn handle_request(
+    runtime: Runtime,
+    worker: String,
+    body: serde_json::Value,
+) -> warp::reply::Json {
+    let request = match parse_request(body) {
+        Ok(x) => x,
+        Err(err) => return warp::reply::json(&err),
+    };
+
+    debug!(
+        worker,
+        id = request.id,
+        method = request.method,
+        "handling request"
+    );
+
+    let params = serde_json::to_vec(&request.params).unwrap();
+
+    let reply = runtime
+        .handle_request(&worker, &request.method, params)
+        .await;
+
+    match reply {
+        Ok(x) => {
+            debug!(worker, id = request.id, "request successful");
+
+            let x = match x {
+                wit::Response::Acknowledge => json!({}),
+                wit::Response::Json(x) => serde_json::from_slice(&x).unwrap(),
+                wit::Response::Cbor(x) => json!({ "cbor": hex::encode(x) }),
+                wit::Response::PartialTx(x) => json!({ "tx": hex::encode(x) }),
+            };
+
+            warp::reply::json(&x)
+        }
+        Err(err) => {
+            error!(worker, id = request.id, "request failed");
+            warp::reply::json(&ErrorResponse {
+                error: err.to_string(),
+            })
+        }
+    }
+}
+
+pub async fn serve(
+    config: balius_runtime::drivers::jsonrpc::Config,
+    runtime: Runtime,
+    cancel: CancellationToken,
+) -> Result<(), Error> {
+    let filter = warp::any()
+        .map(move || runtime.clone())
+        .and(warp::header::<String>("worker"))
+        .and(warp::post())
+        .and(warp::body::json())
+        .then(handle_request);
+
+    let address: SocketAddr = config
+        .listen_address
+        .parse()
+        .map_err(|x: std::net::AddrParseError| Error::Config(x.to_string()))?;
+
+    let (addr, server) =
+        warp::serve(filter).bind_with_graceful_shutdown(address, cancel.cancelled_owned());
+
+    tracing::info!(%addr, "Json-RPC server listening");
+
+    server.await;
+
+    Ok(())
+}

instance/src/server.rs

@@ -142,6 +142,12 @@ impl ProxyHttp for BaliusProxy {
             ctx.consumer.network, self.config.balius_dns, self.config.balius_port
         );
 
+        // replace existing header if any
+        session
+            .req_header_mut()
+            .insert_header("worker", &ctx.consumer.port_name)
+            .unwrap();
+
         if self.limiter(&ctx.consumer).await? {
             session.respond_error(429).await?;
             return Ok(true);