fix(coprocessor): wip

Author: rudy-6-4

coprocessor/fhevm-engine/Cargo.lock

@@ -32,6 +32,7 @@ axum = "0.7"
 tower-http = { version = "0.5", features = ["trace"] }
 anyhow = "1.0.98"
 aws-config = "1.8.5"
+aws-credential-types = "1.2.6"
 aws-sdk-kms = { version = "1.68.0", default-features = false }
 aws-sdk-s3 = { version = "1.103.0", features = ["test-util"] }
 bigdecimal = "0.4.8"

coprocessor/fhevm-engine/Cargo.toml

@@ -11,6 +11,7 @@ alloy = { workspace = true }
 anyhow = { workspace = true }
 axum = { workspace = true }
 aws-config = { workspace = true }
+aws-credential-types = { workspace = true }
 aws-sdk-s3 = { workspace = true }
 clap = { workspace = true }
 futures-util = { workspace = true }
@@ -26,6 +27,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

@@ -1,9 +1,10 @@
 use std::time::Duration;
 
 use aws_config::{retry::RetryConfig, timeout::TimeoutConfig, BehaviorVersion};
-use aws_sdk_s3::{config::Builder, Client};
+use aws_sdk_s3::config::{Builder, ProvideCredentials};
+use aws_sdk_s3::Client;
 use tokio_util::bytes;
-use tracing::{error, info};
+use tracing::{error, info, warn};
 
 #[derive(Clone, Debug, Default)]
 pub struct S3Policy {
@@ -26,8 +27,22 @@ 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,
+) -> anyhow::Result<aws_sdk_s3::Client> {
     let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
+
+    let credentials = sdk_config
+        .credentials_provider()
+        .ok_or(anyhow::Error::msg("s3 client: no credential provider"))?
+        .provide_credentials()
+        .await?;
+    info!(access_key = %credentials.access_key_id(), "Loaded AWS credentials");
+
+    let region = sdk_config.region();
+    info!(region = ?region, "Using AWS region");
+
     let timeout_config = TimeoutConfig::builder()
         .connect_timeout(retry_policy.connect_timeout)
         .operation_attempt_timeout(retry_policy.max_retries_timeout)
@@ -40,43 +55,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)
+    Ok(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> {
+        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 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 bucket = bucket.trim_start_matches('/');
+    info!(s3_bucket_url, host, bucket, "Parsed S3 url");
+    Ok((host.to_owned(), bucket.to_owned()))
+}
+
 pub async fn download_key_from_s3<A: AwsS3Interface>(
     s3_client: &A,
     s3_bucket_urls: &[String],
@@ -92,7 +129,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

@@ -38,11 +38,11 @@ pub async fn update_tenant_key(
     key_bytes: &[u8],
     reduced_key_bytes: Option<Vec<u8>>,
     tenant_id: TenantId,
-    chain_id: ChainId,
+    host_chain_id: ChainId,
 ) -> anyhow::Result<()> {
     let query = match key_type {
         KeyType::ServerKey => {
-            info!(tenant_id, chain_id, key_id, "Updating server key");
+            info!(tenant_id, host_chain_id, key_id, "Updating server key");
             let Some(reduced_key_bytes) = reduced_key_bytes else {
                 anyhow::bail!("Reduced key bytes must be provided for server key");
             };
@@ -58,11 +58,11 @@ pub async fn update_tenant_key(
                 reduced_key_bytes,
                 key_id,
                 tenant_id as i32,
-                chain_id as i64,
+                host_chain_id as i64,
             )
         }
         KeyType::PublicKey => {
-            info!(tenant_id, chain_id, key_id, "Updating public key");
+            info!(tenant_id, host_chain_id, key_id, "Updating public key");
             sqlx::query!(
                 "UPDATE tenants
                 SET
@@ -72,7 +72,7 @@ pub async fn update_tenant_key(
                 key_bytes,
                 key_id,
                 tenant_id as i32,
-                chain_id as i64,
+                host_chain_id as i64,
             )
         }
     };
@@ -81,7 +81,7 @@ pub async fn update_tenant_key(
         anyhow::bail!(
             "No tenant found for tenant_id {} and chain_id {}",
             tenant_id,
-            chain_id
+            host_chain_id
         );
     }
     Ok(())

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

@@ -13,7 +13,7 @@ use crate::aws_s3::{download_key_from_s3, AwsS3Interface};
 use crate::database::{tenant_id, update_tenant_crs, update_tenant_key};
 use crate::digest::{digest_crs, digest_key};
 use crate::sks_key::extract_server_key_without_ns;
-use crate::{ChainId, ConfigSettings, HealthStatus, KeyId, KeyType, TenantId};
+use crate::{ChainId, ConfigSettings, HealthStatus, KeyId, KeyType};
 
 sol!(
     #[sol(rpc)]
@@ -27,11 +27,6 @@ sol!(
     "./../../../gateway-contracts/artifacts/contracts/KMSManagement.sol/KMSManagement.json"
 );
 
-struct TenantInfo {
-    tenant_id: Option<TenantId>,
-    chain_id: ChainId,
-}
-
 pub struct GatewayListener<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> {
     input_verification_address: Address,
     kms_management_address: Address,
@@ -105,12 +100,8 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
         let kms_management = KMSManagement::new(self.kms_management_address, &self.provider);
 
         let mut from_block = self.get_last_block_num(db_pool).await?;
-        let chain_id = self.provider.get_chain_id().await?;
-        let tenant_id = tenant_id(db_pool, chain_id).await?;
-        let tenant_info = TenantInfo {
-            tenant_id,
-            chain_id,
-        };
+        // TODO: make CHAIN_ID a command line parameter
+        let host_chain_id = std::env::var("CHAIN_ID")?.parse::<i64>()? as ChainId;
 
         // We call `from_block` here, but we expect that most nodes will not honour it. That doesn't lead to issues as described below.
         //
@@ -163,7 +154,7 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
                         return Err(anyhow::anyhow!("Block stream closed"));
                     };
                     let (request, _log) = item?;
-                    self.activate_key(db_pool, request, &self.aws_s3_client, &tenant_info).await?;
+                    self.activate_key(db_pool, request, &self.aws_s3_client, host_chain_id).await?;
                     info!("ActivateKey event successful");
                 }
                 item = activate_crs.next() => {
@@ -172,7 +163,7 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
                         return Err(anyhow::anyhow!("Block stream closed"));
                     };
                     let (request, _log) = item?;
-                    self.activate_crs(db_pool, request, &self.aws_s3_client, &tenant_info).await?;
+                    self.activate_crs(db_pool, request, &self.aws_s3_client, host_chain_id).await?;
                     info!("ActivateCrs event successful");
                 }
             }
@@ -218,7 +209,7 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
         db_pool: &Pool<Postgres>,
         request: KMSManagement::ActivateKey,
         s3_client: &A,
-        tenant_info: &TenantInfo,
+        host_chain_id: ChainId,
     ) -> anyhow::Result<()> {
         let key_id: KeyId = request.keyId;
         let s3_bucket_urls = request.kmsNodeStorageUrls;
@@ -255,12 +246,12 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
             }
             keys_bytes.push(bytes);
         }
-        let Some(tenant_id) = tenant_info.tenant_id else {
+        let Some(tenant_id) = tenant_id(db_pool, host_chain_id).await? else {
             error!(
-                chain_id = tenant_info.chain_id,
+                host_chain_id,
                 "No tenant found for chain id, stopping"
             );
-            anyhow::bail!("No tenant found for chain id {}", tenant_info.chain_id);
+            anyhow::bail!("No tenant found for chain id {}", host_chain_id);
         };
         let key_id = key_id_to_bytes(key_id);
         let mut tx = db_pool.begin().await?;
@@ -276,7 +267,7 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
                 &key_bytes,
                 reduced_key_bytes,
                 tenant_id,
-                tenant_info.chain_id,
+                host_chain_id,
             )
             .await?;
         }
@@ -289,7 +280,7 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
         db_pool: &Pool<Postgres>,
         request: KMSManagement::ActivateCrs,
         s3_client: &A,
-        tenant_info: &TenantInfo,
+        host_chain_id: ChainId,
     ) -> anyhow::Result<()> {
         let crs_id: KeyId = request.crsId;
         let s3_bucket_urls = request.kmsNodeStorageUrls;
@@ -312,15 +303,15 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface> GatewayListener
             error!(download_digest = ?download_digest, expected_digest = ?expected_digest, "Key digest mismatch, stopping");
             anyhow::bail!("Invalid Key digest for key id:{crs_id}");
         }
-        let Some(tenant_id) = tenant_info.tenant_id else {
+        let Some(tenant_id) = tenant_id(db_pool, host_chain_id).await? else {
             error!(
-                chain_id = tenant_info.chain_id,
+                host_chain_id,
                 "No tenant found for chain id, stopping"
             );
-            anyhow::bail!("No tenant found for chain id {}", tenant_info.chain_id);
+            anyhow::bail!("No tenant found for chain id {}", host_chain_id);
         };
         let mut tx = db_pool.begin().await?;
-        update_tenant_crs(&mut tx, &bytes, tenant_id, tenant_info.chain_id).await?;
+        update_tenant_crs(&mut tx, &bytes, tenant_id, host_chain_id).await?;
         tx.commit().await?;
         Ok(())
     }