fix(coprocessor): wip

Author: rudy-6-4

coprocessor/fhevm-engine/Cargo.lock

@@ -26,6 +26,7 @@ tokio-util = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 tower-http = { workspace = true }
+url = "2.5.7"
 
 fhevm-engine-common = { path = "../fhevm-engine-common" }
 

coprocessor/fhevm-engine/gw-listener/Cargo.toml

@@ -3,7 +3,7 @@ use std::time::Duration;
 use aws_config::{retry::RetryConfig, timeout::TimeoutConfig, BehaviorVersion};
 use aws_sdk_s3::{config::Builder, Client};
 use tokio_util::bytes;
-use tracing::{error, info};
+use tracing::{error, info, warn};
 
 #[derive(Clone, Debug, Default)]
 pub struct S3Policy {
@@ -26,7 +26,7 @@ impl S3Policy {
     };
 }
 
-pub async fn create_s3_client(retry_policy: &S3Policy) -> aws_sdk_s3::Client {
+pub async fn create_s3_client(retry_policy: &S3Policy, url: &str) -> aws_sdk_s3::Client {
     let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
     let timeout_config = TimeoutConfig::builder()
         .connect_timeout(retry_policy.connect_timeout)
@@ -40,43 +40,65 @@ pub async fn create_s3_client(retry_policy: &S3Policy) -> aws_sdk_s3::Client {
     let config = Builder::from(&sdk_config)
         .timeout_config(timeout_config)
         .retry_config(retry_config)
+        .endpoint_url(url)
         .build();
 
     Client::from_conf(config)
 }
 
-pub async fn default_aws_s3_client() -> AwsS3Client {
-    let s3_client = create_s3_client(&S3Policy::DEFAULT).await;
-    AwsS3Client { s3_client }
-}
-
-// Let's wrap Aws Client to have an interface for it so we can mock it.
+// Let's wrap Aws access to have an interface for it so we can mock it.
 #[derive(Clone)]
-pub struct AwsS3Client {
-    pub s3_client: Client,
-}
+pub struct AwsS3Client {}
 
 impl AwsS3Interface for AwsS3Client {
-    async fn get_bucket_key(&self, bucket: &str, key: &str) -> anyhow::Result<bytes::Bytes> {
-        let result = self
-            .s3_client
+    async fn get_bucket_key(
+        &self,
+        url: &str,
+        bucket: &str,
+        key: &str,
+    ) -> anyhow::Result<bytes::Bytes> {
+        // let parsed_url_and_bucket = url::Url::parse(_url_and_bucket)?;
+        // let url = parsed_url_and_bucket.
+        Ok(create_s3_client(&S3Policy::DEFAULT, url)
+            .await
             .get_object()
             .bucket(bucket)
             .key(key)
             .send()
-            .await?;
-        Ok(result.body.collect().await?.into_bytes())
+            .await?
+            .body
+            .collect()
+            .await?
+            .into_bytes())
     }
 }
 
 pub trait AwsS3Interface {
     fn get_bucket_key(
         &self,
+        url: &str,
         bucket: &str,
         key: &str,
     ) -> impl std::future::Future<Output = anyhow::Result<bytes::Bytes>>;
 }
 
+fn split_url(s3_bucket_url: &String) -> anyhow::Result<(String, String)> {
+    let parsed_url_and_bucket = url::Url::parse(s3_bucket_url)?;
+    let bucket = parsed_url_and_bucket.path();
+    let scheme = parsed_url_and_bucket.scheme();
+    let host = s3_bucket_url.replace(bucket, "").trim_end_matches('/').to_owned();
+    let host = if host.contains("minio:9000") {
+        // TODO: replace by docker configuration
+        warn!(s3_bucket_url, "Using localhost for minio access");
+        host.replace("minio:9000", "172.17.0.1:9000")
+    } else {
+        host.to_owned()
+    };
+    let url = format!("{scheme}://{host}",);
+    info!(s3_bucket_url, url, bucket, "Parsed S3 url");
+    Ok((url.to_owned(), bucket.to_owned()))
+}
+
 pub async fn download_key_from_s3<A: AwsS3Interface>(
     s3_client: &A,
     s3_bucket_urls: &[String],
@@ -92,7 +114,11 @@ pub async fn download_key_from_s3<A: AwsS3Interface>(
             key_path,
             s3_bucket_url, i_s3_bucket_url, nb_urls, url_index, "Try downloading"
         );
-        let result = s3_client.get_bucket_key(s3_bucket_url, &key_path).await;
+        let Ok((url, bucket)) = split_url(s3_bucket_url) else {
+            error!(s3_bucket_url, "Failed to parse S3 url");
+            continue;
+        };
+        let result = s3_client.get_bucket_key(&url, &bucket, &key_path).await;
         let Ok(result) = result else {
             error!(s3_bucket_url, key_path, result = ?result, "Downloading failed");
             continue;

coprocessor/fhevm-engine/gw-listener/src/aws_s3.rs

@@ -3,7 +3,7 @@ use std::time::Duration;
 use alloy::providers::{ProviderBuilder, WsConnect};
 use alloy::{primitives::Address, transports::http::reqwest::Url};
 use clap::Parser;
-use gw_listener::aws_s3::default_aws_s3_client;
+use gw_listener::aws_s3::AwsS3Client;
 use gw_listener::gw_listener::GatewayListener;
 use gw_listener::http_server::HttpServer;
 use gw_listener::ConfigSettings;
@@ -116,7 +116,7 @@ async fn main() -> anyhow::Result<()> {
         }
     };
 
-    let aws_s3_client = default_aws_s3_client().await;
+    let aws_s3_client = AwsS3Client {};
 
     let cancel_token = CancellationToken::new();
 

coprocessor/fhevm-engine/gw-listener/src/bin/gw_listener.rs

@@ -9,17 +9,17 @@ use alloy::{
     sol,
 };
 
-use aws_sdk_s3::operation::get_object::GetObjectError;
-use gw_listener::gw_listener::GatewayListener;
-use gw_listener::ConfigSettings;
+use aws_sdk_s3::{operation::get_object::GetObjectError, Client};
 use gw_listener::{
-    aws_s3::{default_aws_s3_client, AwsS3Client},
-    gw_listener::{key_id_to_key_bucket, to_bucket_key_prefix},
+    aws_s3::{AwsS3Client, AwsS3Interface},
+    gw_listener::{key_id_to_key_bucket, to_bucket_key_prefix, GatewayListener},
+    ConfigSettings,
 };
 use serial_test::serial;
 use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
 use test_harness::instance::ImportMode;
 use tokio::time::sleep;
+use tokio_util::bytes;
 use tokio_util::sync::CancellationToken;
 use tracing::Level;
 
@@ -105,7 +105,7 @@ async fn verify_proof_request_inserted_into_db() -> anyhow::Result<()> {
         .wallet(env.wallet)
         .connect_ws(WsConnect::new(env.anvil.ws_endpoint_url()))
         .await?;
-    let aws_s3_client = default_aws_s3_client().await;
+    let aws_s3_client = AwsS3Client {};
     let input_verification = InputVerification::deploy(&provider).await?;
     let kms_management = KMSManagement::deploy(&provider).await?;
     let gw_listener = GatewayListener::new(
@@ -213,6 +213,30 @@ async fn has_crs(db_pool: &Pool<Postgres>) -> anyhow::Result<bool> {
     Ok(false)
 }
 
+#[derive(Clone)]
+pub struct AwsS3ClientMocked(Client);
+
+impl AwsS3Interface for AwsS3ClientMocked {
+    async fn get_bucket_key(
+        &self,
+        url: &str,
+        bucket: &str,
+        key: &str,
+    ) -> anyhow::Result<bytes::Bytes> {
+        Ok(self
+            .0
+            .get_object()
+            .bucket(bucket)
+            .key(key)
+            .send()
+            .await?
+            .body
+            .collect()
+            .await?
+            .into_bytes())
+    }
+}
+
 // test bad bucket
 // test bad key
 #[tokio::test]
@@ -233,10 +257,10 @@ async fn keygen_ok() -> anyhow::Result<()> {
 
     // see ../contracts/KMSManagement.sol
     let buckets = [
-        "https://s3.amazonaws.com/test-bucket1/PUB-P1",
-        "https://s3.amazonaws.com/test-bucket2/PUB-P2",
-        "https://s3.amazonaws.com/test-bucket3/PUB-P3",
-        "https://s3.amazonaws.com/test-bucket4/PUB-P4",
+        "/test-bucket1/PUB-P1",
+        "/test-bucket2/PUB-P2",
+        "/test-bucket3/PUB-P3",
+        "/test-bucket4/PUB-P4",
     ];
 
     let keys_digests = [KeyType::PublicKey, KeyType::ServerKey];
@@ -283,7 +307,7 @@ async fn keygen_ok() -> anyhow::Result<()> {
         .wallet(env.wallet)
         .connect_ws(WsConnect::new(env.anvil.ws_endpoint_url()))
         .await?;
-    let aws_s3_client = AwsS3Client { s3_client: s3 };
+    let aws_s3_client = AwsS3ClientMocked(s3);
     let input_verification = InputVerification::deploy(&provider).await?;
     let kms_management = KMSManagement::deploy(&provider).await?;
     let gw_listener = GatewayListener::new(
@@ -383,7 +407,7 @@ async fn keygen_compromised_key() -> anyhow::Result<()> {
         .wallet(env.wallet)
         .connect_ws(WsConnect::new(env.anvil.ws_endpoint_url()))
         .await?;
-    let aws_s3_client = AwsS3Client { s3_client: s3 };
+    let aws_s3_client = AwsS3ClientMocked(s3);
     let input_verification = InputVerification::deploy(&provider).await?;
     let kms_management = KMSManagement::deploy(&provider).await?;
     let gw_listener = GatewayListener::new(
@@ -489,7 +513,7 @@ async fn keygen_bad_key_or_bucket() -> anyhow::Result<()> {
         .wallet(env.wallet)
         .connect_ws(WsConnect::new(env.anvil.ws_endpoint_url()))
         .await?;
-    let aws_s3_client = AwsS3Client { s3_client: s3 };
+    let aws_s3_client = AwsS3ClientMocked(s3);
     let input_verification = InputVerification::deploy(&provider).await?;
     let kms_management = KMSManagement::deploy(&provider).await?;
     let gw_listener = GatewayListener::new(

coprocessor/fhevm-engine/gw-listener/tests/gw_listener_tests.rs

@@ -154,9 +154,10 @@ impl SwitchNSquashService {
     }
 }
 
-async fn get_keyset(pool: PgPool,
+async fn get_keyset(
+    pool: PgPool,
     keys_cache: Arc<RwLock<lru::LruCache<String, KeySet>>>,
-    tenant_api_key: &String
+    tenant_api_key: &String,
 ) -> Result<Option<KeySet>, ExecutionError> {
     let t = telemetry::tracer("worker_loop_init");
     let s = t.child_span("fetch_keyset");
@@ -191,7 +192,10 @@ pub(crate) async fn run_loop(
             Some(keys)
         }
         None => {
-            warn!(tenant_api_key = tenant_api_key, "No keys found for the given tenant_api_key");
+            warn!(
+                tenant_api_key = tenant_api_key,
+                "No keys found for the given tenant_api_key"
+            );
             None
         }
     };
@@ -205,7 +209,10 @@ pub(crate) async fn run_loop(
         update_last_active(last_active_at.clone()).await;
 
         let Some(keys) = keys.as_ref() else {
-            warn!(tenant_api_key = tenant_api_key, "No keys available, retrying in 5 seconds");
+            warn!(
+                tenant_api_key = tenant_api_key,
+                "No keys available, retrying in 5 seconds"
+            );
             tokio::time::sleep(Duration::from_secs(5)).await;
             if token.is_cancelled() {
                 return Ok(());
@@ -214,7 +221,7 @@ pub(crate) async fn run_loop(
             continue;
         };
 
-        let maybe_remaining = fetch_and_execute_sns_tasks(&pool, &tx, &keys, &conf, &token).await?;
+        let maybe_remaining = fetch_and_execute_sns_tasks(&pool, &tx, keys, &conf, &token).await?;
         if maybe_remaining {
             if token.is_cancelled() {
                 return Ok(());