Skip to content

Commit

Permalink
funds-manager: middleware: Add symmetric key auth for withdrawal
Browse files Browse the repository at this point in the history
  • Loading branch information
joeykraut committed Jul 24, 2024
1 parent 06a6747 commit 99180ed
Show file tree
Hide file tree
Showing 5 changed files with 172 additions and 12 deletions.
2 changes: 1 addition & 1 deletion funds-manager/funds-manager-api/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ pub struct DepositAddressResponse {
}

/// The request body for withdrawing funds from custody
#[derive(Debug, Serialize, Deserialize)]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct WithdrawFundsRequest {
/// The mint of the asset to withdraw
pub mint: String,
Expand Down
4 changes: 4 additions & 0 deletions funds-manager/funds-manager-server/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,10 @@ edition = "2021"
# === CLI + Server === #
clap = { version = "4.5.3", features = ["derive", "env"] }
funds-manager-api = { path = "../funds-manager-api" }
hex = "0.4.3"
hmac = "0.12.1"
http-body-util = "0.1.0"
sha2 = "0.10.6"
tokio = { version = "1.10", features = ["full"] }
warp = "0.3"

Expand Down Expand Up @@ -39,6 +42,7 @@ renegade-util = { package = "util", workspace = true }
# === Misc Dependencies === #
base64 = "0.22"
bigdecimal = { version = "0.4", features = ["serde"] }
bytes = "1.5.0"
futures = "0.3"
http = "1.1"
num-bigint = "0.4"
Expand Down
3 changes: 3 additions & 0 deletions funds-manager/funds-manager-server/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,8 @@ pub enum ApiError {
InternalError(String),
/// Bad request error
BadRequest(String),
/// Unauthenticated error
Unauthenticated(String),
}

impl Reject for ApiError {}
Expand All @@ -113,6 +115,7 @@ impl Display for ApiError {
ApiError::RedemptionError(e) => write!(f, "Redemption error: {}", e),
ApiError::InternalError(e) => write!(f, "Internal error: {}", e),
ApiError::BadRequest(e) => write!(f, "Bad request: {}", e),
ApiError::Unauthenticated(e) => write!(f, "Unauthenticated: {}", e),
}
}
}
Expand Down
71 changes: 60 additions & 11 deletions funds-manager/funds-manager-server/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ pub mod db;
pub mod error;
pub mod fee_indexer;
pub mod handlers;
pub mod middleware;
pub mod relayer_client;

use aws_config::{BehaviorVersion, Region, SdkConfig};
Expand All @@ -24,6 +25,7 @@ use funds_manager_api::{
use handlers::{
get_deposit_address_handler, index_fees_handler, quoter_withdraw_handler, redeem_fees_handler,
};
use middleware::{identity, with_hmac_auth, with_json_body};
use relayer_client::RelayerClient;
use renegade_circuit_types::elgamal::DecryptionKey;
use renegade_util::{err_str, raw_err_str, telemetry::configure_telemetry};
Expand All @@ -35,7 +37,8 @@ use arbitrum_client::{
constants::Chain,
};
use clap::Parser;
use tracing::error;
use funds_manager_api::WithdrawFundsRequest;
use tracing::{error, warn};
use warp::{filters::query::query, Filter};

use crate::custody_client::CustodyClient;
Expand All @@ -62,8 +65,18 @@ const DUMMY_PRIVATE_KEY: &str =

/// The cli for the fee sweeper
#[rustfmt::skip]
#[derive(Clone, Debug, Parser)]
#[derive(Parser)]
#[clap(about = "Funds manager server")]
struct Cli {
// --- Authentication --- //

/// The HMAC key to use for authentication
#[clap(long, conflicts_with = "disable_auth", env = "HMAC_KEY")]
hmac_key: Option<String>,
/// Whether to disable authentication
#[clap(long, conflicts_with = "hmac_key")]
disable_auth: bool,

// --- Environment Configs --- //

/// The URL of the relayer to use
Expand Down Expand Up @@ -117,6 +130,30 @@ struct Cli {
datadog_logging: bool,
}

impl Cli {
/// Validate the CLI arguments
fn validate(&self) -> Result<(), String> {
if self.hmac_key.is_none() && !self.disable_auth {
Err("Either --hmac-key or --disable-auth must be provided".to_string())
} else {
Ok(())
}
}

/// Get the HMAC key as a 32-byte array
fn get_hmac_key(&self) -> Option<[u8; 32]> {
self.hmac_key.as_ref().map(|key| {
let decoded = hex::decode(key).expect("Invalid HMAC key");
if decoded.len() != 32 {
panic!("HMAC key must be 32 bytes long");
}
let mut array = [0u8; 32];
array.copy_from_slice(&decoded);
array
})
}
}

/// The server
#[derive(Clone)]
struct Server {
Expand All @@ -136,6 +173,8 @@ struct Server {
pub custody_client: CustodyClient,
/// The AWS config
pub aws_config: SdkConfig,
/// The HMAC key for custody endpoint authentication
pub hmac_key: Option<[u8; 32]>,
}

impl Server {
Expand All @@ -160,6 +199,11 @@ impl Server {
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let cli = Cli::parse();
cli.validate()?;
if cli.hmac_key.is_none() {
warn!("Authentication is disabled. This is not recommended for production use.");
}

configure_telemetry(
cli.datadog_logging, // datadog_enabled
false, // otlp_enabled
Expand All @@ -179,7 +223,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
// Build an Arbitrum client
let wallet = LocalWallet::from_str(DUMMY_PRIVATE_KEY)?;
let conf = ArbitrumClientConfig {
darkpool_addr: cli.darkpool_address,
darkpool_addr: cli.darkpool_address.clone(),
chain: cli.chain,
rpc_url: cli.rpc_url.clone(),
arb_priv_keys: vec![wallet],
Expand All @@ -190,10 +234,11 @@ async fn main() -> Result<(), Box<dyn Error>> {

// Build the indexer
let mut decryption_keys = vec![DecryptionKey::from_hex_str(&cli.relayer_decryption_key)?];
if let Some(protocol_key) = cli.protocol_decryption_key {
decryption_keys.push(DecryptionKey::from_hex_str(&protocol_key)?);
if let Some(protocol_key) = &cli.protocol_decryption_key {
decryption_keys.push(DecryptionKey::from_hex_str(protocol_key)?);
}

let hmac_key = cli.get_hmac_key();
let relayer_client = RelayerClient::new(&cli.relayer_url, &cli.usdc_mint);
let custody_client =
CustodyClient::new(cli.fireblocks_api_key, cli.fireblocks_api_secret, cli.rpc_url);
Expand All @@ -206,38 +251,41 @@ async fn main() -> Result<(), Box<dyn Error>> {
db_url: cli.db_url,
custody_client,
aws_config: config,
hmac_key,
};

// --- Routes --- //

let server = Arc::new(server);
let ping = warp::get()
.and(warp::path(PING_ROUTE))
.map(|| warp::reply::with_status("PONG", warp::http::StatusCode::OK));

let index_fees = warp::post()
.and(warp::path(INDEX_FEES_ROUTE))
.and(with_server(Arc::new(server.clone())))
.and(with_server(server.clone()))
.and_then(index_fees_handler);

let redeem_fees = warp::post()
.and(warp::path(REDEEM_FEES_ROUTE))
.and(with_server(Arc::new(server.clone())))
.and(with_server(server.clone()))
.and_then(redeem_fees_handler);

let withdraw_custody = warp::post()
.and(warp::path("custody"))
.and(warp::path("quoters"))
.and(warp::path(WITHDRAW_CUSTODY_ROUTE))
.and(warp::body::json())
.and(with_server(Arc::new(server.clone())))
.and(with_hmac_auth(server.clone()))
.map(with_json_body::<WithdrawFundsRequest>)
.and_then(identity)
.and(with_server(server.clone()))
.and_then(quoter_withdraw_handler);

let get_deposit_address = warp::get()
.and(warp::path("custody"))
.and(warp::path("quoters"))
.and(warp::path(GET_DEPOSIT_ADDRESS_ROUTE))
.and(query::<HashMap<String, String>>())
.and(with_server(Arc::new(server.clone())))
.and(with_server(server.clone()))
.and_then(get_deposit_address_handler);

let routes = ping
Expand All @@ -263,6 +311,7 @@ async fn handle_rejection(err: warp::Rejection) -> Result<impl warp::Reply, warp
ApiError::RedemptionError(msg) => (warp::http::StatusCode::BAD_REQUEST, msg),
ApiError::InternalError(msg) => (warp::http::StatusCode::INTERNAL_SERVER_ERROR, msg),
ApiError::BadRequest(msg) => (warp::http::StatusCode::BAD_REQUEST, msg),
ApiError::Unauthenticated(msg) => (warp::http::StatusCode::UNAUTHORIZED, msg),
};
error!("API Error: {:?}", api_error);
Ok(warp::reply::with_status(message.clone(), code))
Expand Down
104 changes: 104 additions & 0 deletions funds-manager/funds-manager-server/src/middleware.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
//! Middleware for the funds manager server

use crate::error::ApiError;
use crate::Server;
use bytes::Bytes;
use hmac::{Hmac, Mac};
use serde::de::DeserializeOwned;
use sha2::Sha256;
use std::sync::Arc;
use warp::Filter;

/// The header key for the HMAC signature
const X_SIGNATURE_HEADER: &str = "X-Signature";

/// Add HMAC authentication to a route
pub(crate) fn with_hmac_auth(
server: Arc<Server>,
) -> impl Filter<Extract = (Bytes,), Error = warp::Rejection> + Clone {
warp::any()
.and(warp::any().map(move || server.clone()))
.and(warp::header::optional::<String>(X_SIGNATURE_HEADER))
.and(warp::method())
.and(warp::path::full())
.and(warp::header::headers_cloned())
.and(warp::body::bytes())
.and_then(verify_hmac)
}

/// Verify the HMAC signature
async fn verify_hmac(
server: Arc<Server>,
signature: Option<String>,
method: warp::http::Method,
path: warp::path::FullPath,
headers: warp::http::HeaderMap,
body: Bytes,
) -> Result<Bytes, warp::Rejection> {
// Unwrap the key and signature
let hmac_key = match &server.hmac_key {
Some(hmac_key) => hmac_key,
None => return Ok(body), // Auth is disabled, allow the request
};

let signature = match signature {
Some(sig) => sig,
None => {
return Err(warp::reject::custom(ApiError::Unauthenticated(
"Missing signature".to_string(),
)))
},
};

// Construct the MAC
let mut mac = Hmac::<Sha256>::new_from_slice(hmac_key)
.map_err(|_| warp::reject::custom(ApiError::InternalError("HMAC error".to_string())))?;

// Update with method, path, headers and body in order
mac.update(method.as_str().as_bytes());
mac.update(path.as_str().as_bytes());
add_headers_to_hmac(&mut mac, &headers);
mac.update(&body);

// Check the signature
let expected = mac.finalize().into_bytes();
let provided = hex::decode(signature)
.map_err(|_| warp::reject::custom(ApiError::BadRequest("Invalid signature".to_string())))?;
if expected.as_slice() != provided.as_slice() {
return Err(warp::reject::custom(ApiError::Unauthenticated(
"Invalid signature".to_string(),
)));
}

Ok(body)
}

/// Hash headers into an HMAC
fn add_headers_to_hmac(mac: &mut Hmac<Sha256>, headers: &warp::http::HeaderMap) {
let mut sorted_headers: Vec<(String, String)> = headers
.iter()
.map(|(k, v)| (k.as_str().to_lowercase(), v.to_str().unwrap_or("").to_string()))
.collect();
sorted_headers.sort_by(|a, b| a.0.cmp(&b.0));

for (key, value) in sorted_headers {
// Exclude the signature header itself
if key.to_lowercase() != X_SIGNATURE_HEADER.to_lowercase() {
mac.update(key.as_bytes());
mac.update(value.as_bytes());
}
}
}

/// Extract a JSON body from a request
#[allow(clippy::needless_pass_by_value)]
pub fn with_json_body<T: DeserializeOwned + Send>(body: Bytes) -> Result<T, warp::Rejection> {
serde_json::from_slice(&body)
.map_err(|e| warp::reject::custom(ApiError::BadRequest(format!("Invalid JSON: {}", e))))
}

/// Identity map for a handler's middleware, used to chain together `map`s and
/// `and_then`s
pub async fn identity<T>(res: T) -> T {
res
}

0 comments on commit 99180ed

Please sign in to comment.