feat: unlock, lock and withdraw funds in sdk and cli

crates/cli/src/main.rs

@@ -4,6 +4,9 @@ use std::io::BufReader;
 use std::io::Write;
 use std::path::PathBuf;
 use std::str::FromStr;
+use std::time::Duration;
+use std::time::SystemTime;
+use std::time::UNIX_EPOCH;
 
 use aligned_sdk::aggregation_layer;
 use aligned_sdk::aggregation_layer::AggregationModeVerificationData;
@@ -19,13 +22,18 @@ use aligned_sdk::verification_layer::estimate_fee;
 use aligned_sdk::verification_layer::get_chain_id;
 use aligned_sdk::verification_layer::get_nonce_from_batcher;
 use aligned_sdk::verification_layer::get_nonce_from_ethereum;
+use aligned_sdk::verification_layer::get_unlock_block_time;
+use aligned_sdk::verification_layer::lock_balance_in_aligned;
+use aligned_sdk::verification_layer::unlock_balance_in_aligned;
+use aligned_sdk::verification_layer::withdraw_balance_from_aligned;
 use aligned_sdk::verification_layer::{deposit_to_aligned, get_balance_in_aligned};
 use aligned_sdk::verification_layer::{get_vk_commitment, save_response, submit_multiple};
 use clap::Args;
 use clap::Parser;
 use clap::Subcommand;
 use clap::ValueEnum;
 use env_logger::Env;
+use ethers::core::k256::pkcs8::der::asn1::UtcTime;
 use ethers::prelude::*;
 use ethers::utils::format_ether;
 use ethers::utils::hex;
@@ -41,8 +49,11 @@ use crate::AlignedCommands::GetUserBalance;
 use crate::AlignedCommands::GetUserNonce;
 use crate::AlignedCommands::GetUserNonceFromEthereum;
 use crate::AlignedCommands::GetVkCommitment;
+use crate::AlignedCommands::LockFunds;
 use crate::AlignedCommands::Submit;
+use crate::AlignedCommands::UnlockFunds;
 use crate::AlignedCommands::VerifyProofOnchain;
+use crate::AlignedCommands::WithdrawFunds;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
@@ -65,6 +76,12 @@ pub enum AlignedCommands {
         name = "deposit-to-batcher"
     )]
     DepositToBatcher(DepositToBatcherArgs),
+    #[clap(about = "Unlocks funds from the batcher", name = "unlock-funds")]
+    UnlockFunds(LockUnlockFundsArgs),
+    #[clap(about = "Lock funds in the batcher", name = "lock-funds")]
+    LockFunds(LockUnlockFundsArgs),
+    #[clap(about = "Withdraw funds from the batcher", name = "withdraw-funds")]
+    WithdrawFunds(WithdrawFundsArgs),
     #[clap(about = "Get user balance from the batcher", name = "get-user-balance")]
     GetUserBalance(GetUserBalanceArgs),
     #[clap(
@@ -208,6 +225,38 @@ pub struct DepositToBatcherArgs {
     amount: String,
 }
 
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+pub struct LockUnlockFundsArgs {
+    #[command(flatten)]
+    private_key_type: PrivateKeyType,
+    #[arg(
+        name = "Ethereum RPC provider address",
+        long = "rpc_url",
+        default_value = "http://localhost:8545"
+    )]
+    eth_rpc_url: String,
+    #[clap(flatten)]
+    network: NetworkArg,
+}
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+pub struct WithdrawFundsArgs {
+    #[command(flatten)]
+    private_key_type: PrivateKeyType,
+    #[arg(
+        name = "Ethereum RPC provider address",
+        long = "rpc_url",
+        default_value = "http://localhost:8545"
+    )]
+    eth_rpc_url: String,
+    #[clap(flatten)]
+    network: NetworkArg,
+    #[arg(name = "Amount to withdraw", long = "amount", required = true)]
+    amount: String,
+}
+
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 pub struct VerifyProofOnchainArgs {
@@ -752,6 +801,197 @@ async fn main() -> Result<(), AlignedError> {
                 }
             }
         }
+        UnlockFunds(args) => {
+            let eth_rpc_url = args.eth_rpc_url;
+            let eth_rpc_provider =
+                Provider::<Http>::try_from(eth_rpc_url.clone()).map_err(|e| {
+                    SubmitError::EthereumProviderError(format!(
+                        "Error while connecting to Ethereum: {}",
+                        e
+                    ))
+                })?;
+
+            let keystore_path = &args.private_key_type.keystore_path;
+            let private_key = &args.private_key_type.private_key;
+
+            let mut wallet = if let Some(keystore_path) = keystore_path {
+                let password = rpassword::prompt_password("Please enter your keystore password:")
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?;
+                Wallet::decrypt_keystore(keystore_path, password)
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?
+            } else if let Some(private_key) = private_key {
+                private_key
+                    .parse::<LocalWallet>()
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?
+            } else {
+                warn!("Missing keystore or private key used for payment.");
+                return Ok(());
+            };
+
+            let chain_id = get_chain_id(eth_rpc_url.as_str()).await?;
+            wallet = wallet.with_chain_id(chain_id);
+
+            let client = SignerMiddleware::new(eth_rpc_provider, wallet);
+
+            match unlock_balance_in_aligned(&client, args.network.into()).await {
+                Ok(receipt) => {
+                    info!(
+                        "Funds in batcher unlocked successfully. Receipt: 0x{:x}",
+                        receipt.transaction_hash
+                    );
+                }
+                Err(e) => {
+                    error!("Transaction failed: {:?}", e);
+                }
+            }
+        }
+        LockFunds(args) => {
+            let eth_rpc_url = args.eth_rpc_url;
+            let eth_rpc_provider =
+                Provider::<Http>::try_from(eth_rpc_url.clone()).map_err(|e| {
+                    SubmitError::EthereumProviderError(format!(
+                        "Error while connecting to Ethereum: {}",
+                        e
+                    ))
+                })?;
+
+            let keystore_path = &args.private_key_type.keystore_path;
+            let private_key = &args.private_key_type.private_key;
+
+            let mut wallet = if let Some(keystore_path) = keystore_path {
+                let password = rpassword::prompt_password("Please enter your keystore password:")
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?;
+                Wallet::decrypt_keystore(keystore_path, password)
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?
+            } else if let Some(private_key) = private_key {
+                private_key
+                    .parse::<LocalWallet>()
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?
+            } else {
+                warn!("Missing keystore or private key used for payment.");
+                return Ok(());
+            };
+
+            let chain_id = get_chain_id(eth_rpc_url.as_str()).await?;
+            wallet = wallet.with_chain_id(chain_id);
+
+            let client = SignerMiddleware::new(eth_rpc_provider, wallet);
+
+            match lock_balance_in_aligned(&client, args.network.into()).await {
+                Ok(receipt) => {
+                    info!(
+                        "Funds in batcher locked successfully. Receipt: 0x{:x}",
+                        receipt.transaction_hash
+                    );
+                }
+                Err(e) => {
+                    error!("Transaction failed: {:?}", e);
+                }
+            }
+        }
+        WithdrawFunds(args) => {
+            if !args.amount.ends_with("ether") {
+                error!("Amount should be in the format XX.XXether");
+                return Ok(());
+            }
+
+            let amount_ether = args.amount.replace("ether", "");
+
+            let amount_wei = parse_ether(&amount_ether).map_err(|e| {
+                SubmitError::EthereumProviderError(format!("Error while parsing amount: {}", e))
+            })?;
+
+            let eth_rpc_url = args.eth_rpc_url;
+            let eth_rpc_provider =
+                Provider::<Http>::try_from(eth_rpc_url.clone()).map_err(|e| {
+                    SubmitError::EthereumProviderError(format!(
+                        "Error while connecting to Ethereum: {}",
+                        e
+                    ))
+                })?;
+
+            let keystore_path = &args.private_key_type.keystore_path;
+            let private_key = &args.private_key_type.private_key;
+
+            let mut wallet = if let Some(keystore_path) = keystore_path {
+                let password = rpassword::prompt_password("Please enter your keystore password:")
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?;
+                Wallet::decrypt_keystore(keystore_path, password)
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?
+            } else if let Some(private_key) = private_key {
+                private_key
+                    .parse::<LocalWallet>()
+                    .map_err(|e| SubmitError::GenericError(e.to_string()))?
+            } else {
+                warn!("Missing keystore or private key used for payment.");
+                return Ok(());
+            };
+
+            let unlock_block_time = match get_unlock_block_time(
+                wallet.address(),
+                &eth_rpc_url,
+                args.network.clone().into(),
+            )
+            .await
+            {
+                Ok(value) => value,
+                Err(e) => {
+                    error!("Failed to get unlock time: {:?}", e);
+                    return Ok(());
+                }
+            };
+
+            let current_timestamp = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .expect("Time went backwards")
+                .as_secs();
+
+            let unlock_time = UtcTime::from_unix_duration(Duration::from_secs(unlock_block_time))
+                .expect("invalid unlock time");
+            let now_time =
+                UtcTime::from_system_time(SystemTime::now()).expect("invalid system time");
+
+            let retry_after_minutes = if unlock_block_time > current_timestamp {
+                (unlock_block_time - current_timestamp) / 60
+            } else {
+                0
+            };
+
+            if unlock_block_time == 0 {
+                error!("Funds are locked, you need to unlock them first.");
+                return Ok(());
+            }
+
+            if unlock_block_time > current_timestamp {
+                warn!(
+                    "Funds are still locked. You need to wait {} minutes before being able to withdraw after unlocking the funds.\n\
+                    Unlocks in block time: {}\n\
+                    Current time: {}",
+                    retry_after_minutes,
+                    format_utc_time(unlock_time),
+                    format_utc_time(now_time)
+                );
+
+                return Ok(());
+            }
+
+            let chain_id = get_chain_id(eth_rpc_url.as_str()).await?;
+            wallet = wallet.with_chain_id(chain_id);
+
+            let client = SignerMiddleware::new(eth_rpc_provider, wallet);
+
+            match withdraw_balance_from_aligned(&client, args.network.into(), amount_wei).await {
+                Ok(receipt) => {
+                    info!(
+                        "Balance withdraw from batcher successfully. Receipt: 0x{:x}",
+                        receipt.transaction_hash
+                    );
+                }
+                Err(e) => {
+                    error!("Transaction failed: {:?}", e);
+                }
+            }
+        }
         GetUserBalance(get_user_balance_args) => {
             let user_address = H160::from_str(&get_user_balance_args.user_address).unwrap();
             match get_balance_in_aligned(
@@ -1021,3 +1261,17 @@ pub async fn get_user_balance(
         ))
     }
 }
+
+fn format_utc_time(date: UtcTime) -> String {
+    let dt = date.to_date_time();
+
+    format!(
+        "{:04}-{:02}-{:02} {:02}:{:02}:{:02}Z",
+        dt.year(),
+        dt.month(),
+        dt.day(),
+        dt.hour(),
+        dt.minutes(),
+        dt.seconds()
+    )
+}

crates/sdk/src/eth/batcher_payment_service.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use ethers::prelude::*;
+use ethers::{core::k256::ecdsa::SigningKey, prelude::*};
 
 use crate::common::errors::VerificationError;
 
@@ -11,6 +11,9 @@ abigen!(
 
 pub type BatcherPaymentService = BatcherPaymentServiceContract<Provider<Http>>;
 
+pub type SignerMiddlewareT = SignerMiddleware<Provider<Http>, Wallet<SigningKey>>;
+pub type BatcherPaymentServiceWithSigner = BatcherPaymentServiceContract<SignerMiddlewareT>;
+
 pub async fn batcher_payment_service(
     provider: Provider<Http>,
     contract_address: H160,