Skip to content

Commit

Permalink
fix: delegate rotation cap before calling rotate
Browse files Browse the repository at this point in the history
  • Loading branch information
@0xmovses
0xmovses committed Jan 24, 2025
1 parent a63d41c commit 376382f
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 104 deletions.
157 changes: 59 additions & 98 deletions networks/movement/movement-client/src/bin/e2e/key_rotation.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,19 @@
use anyhow::Context;
use aptos_sdk::rest_client::{
aptos_api_types::{Address, EntryFunctionId, IdentifierWrapper, MoveModuleId, ViewRequest},
Account, Response,
use aptos_sdk::{
rest_client::{Client, FaucetClient, Response},
types::{
account_address::AccountAddress,
transaction::{Script, TransactionArgument, TransactionPayload},
},
};
use aptos_sdk::types::account_address::AccountAddress;
use ed25519_dalek::{Keypair, PublicKey, SecretKey, Signer};
use movement_client::{
coin_client::CoinClient,
rest_client::{Client, FaucetClient},
types::LocalAccount,
types::{LocalAccount, RotationProofChallenge},
};
use once_cell::sync::Lazy;
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH};
use std::{fs, str::FromStr};
use tracing;
use url::Url;

Expand All @@ -35,7 +38,7 @@ static NODE_URL: Lazy<Url> = Lazy::new(|| {
.clone();
let node_connection_url =
format!("http://{}:{}", node_connection_address, node_connection_port);
Url::from_str(node_connection_url.as_str()).unwrap()
Url::from_str(&node_connection_url).unwrap()
});

static FAUCET_URL: Lazy<Url> = Lazy::new(|| {
Expand All @@ -52,123 +55,81 @@ static FAUCET_URL: Lazy<Url> = Lazy::new(|| {
.maptos_faucet_rest_connection_port
.clone();
let faucet_listen_url = format!("http://{}:{}", faucet_listen_address, faucet_listen_port);
Url::from_str(faucet_listen_url.as_str()).unwrap()
Url::from_str(&faucet_listen_url).unwrap()
});

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
// Initialize clients
let rest_client = Client::new(NODE_URL.clone());
let faucet_client = FaucetClient::new(FAUCET_URL.clone(), NODE_URL.clone());
let coin_client = CoinClient::new(&rest_client);

// Create account for transactions and gas collection
// Generate accounts
let mut sender = LocalAccount::generate(&mut rand::rngs::OsRng);
let beneficiary = LocalAccount::generate(&mut rand::rngs::OsRng);
let delegate = LocalAccount::generate(&mut rand::rngs::OsRng);

let acc = AccountAddress::from_str(
"0xb08e0478ac871400e082f34e003145570bf4a9e4d88f17964b21fb110e93d77a",
)
.unwrap();

tracing::info!("Created test accounts");
tracing::debug!(
"Sender address: {}, Beneficiary address: {}",
sender.address(),
beneficiary.address()
);
tracing::info!("Generated sender and delegate accounts");

// Fund the sender account
faucet_client
.fund(sender.address(), 1_000_000)
.await
.context("Failed to fund sender account")?;

// Create the beneficiary account
faucet_client
.create_account(beneficiary.address())
.await
.context("Failed to create beneficiary account")?;

let view_req = ViewRequest {
function: EntryFunctionId {
module: MoveModuleId {
address: Address::from_str("0x1").unwrap(),
name: IdentifierWrapper::from_str("governed_gas_pool").unwrap(),
},
name: IdentifierWrapper::from_str("governed_gas_pool_address").unwrap(),
},
type_arguments: vec![],
arguments: vec![],
// Generate new key pair for rotation
let new_keypair = Keypair::generate(&mut rand::rngs::OsRng);
let new_public_key: PublicKey = new_keypair.public;

// Create the rotation proof challenge
let rotation_proof = RotationProofChallenge {
account_address: sender.address(),
sequence_number: sender.sequence_number(),
originator: sender.address(),
current_auth_key: sender.auth_key().to_vec(),
new_public_key: new_public_key.as_bytes().to_vec(),
};

let view_res: Response<Vec<serde_json::Value>> = rest_client
.view(&view_req, None)
.await
.context("Failed to get governed gas pool address")?;
let rotation_message = bcs::to_bytes(&rotation_proof).unwrap();

// Extract the inner field from the response
let inner_value = serde_json::to_value(view_res.inner())
.context("Failed to convert response inner to serde_json::Value")?;
// Sign the rotation proof challenge
let signature_by_new_key = new_keypair.sign(&rotation_message);

// Deserialize the inner value into your AddressResponse struct
let ggp_address: Vec<String> =
serde_json::from_value(inner_value).context("Failed to deserialize AddressResponse")?;

assert_eq!(
ggp_address,
vec!["0xb08e0478ac871400e082f34e003145570bf4a9e4d88f17964b21fb110e93d77a"],
"Governed Gas Pool Resource account is not what is expected"
);

let ggp_account_address =
AccountAddress::from_str(&ggp_address[0]).expect("Failed to parse address");

// Get initial balances
let initial_ggp_balance = coin_client
.get_account_balance(&ggp_account_address)
.await
.context("Failed to get initial framework balance")?;
// Read compiled Move script
let script_code = fs::read("path/to/compiled/script.mv").context("Failed to read script")?;
let script_payload = TransactionPayload::Script(Script::new(
script_code,
vec![],
vec![
TransactionArgument::U8(0), // Scheme for the current key (Ed25519)
TransactionArgument::U8(0), // Scheme for the new key (Ed25519)
TransactionArgument::Bytes(new_public_key.as_bytes().to_vec()),
TransactionArgument::Bytes(signature_by_new_key.to_bytes().to_vec()),
],
));

tracing::info!("Initial ggp Balance: {}", initial_ggp_balance);
// Create and submit the transaction
let expiration_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() + 60; // 60 seconds from now

// Simple transaction that will generate gas fees
tracing::info!("Executing test transaction...");
let txn_hash = coin_client
.transfer(&mut sender, beneficiary.address(), 1_000, None)
.await
.context("Failed to submit transfer transaction")?;
let txn = transaction_test_helpers::get_test_signed_transaction_with_chain_id(
sender.address(),
sender.sequence_number(),
sender.private_key(),
sender.public_key(),
Some(script_payload),
expiration_time,
100, // Max gas
None,
ChainId::new(1), // Set chain ID
);

tracing::info!("Submitting transaction for key rotation");
rest_client
.wait_for_transaction(&txn_hash)
.submit_and_wait(&txn)
.await
.context("Failed when waiting for transfer transaction")?;
tracing::info!("Test transaction completed: {:?}", txn_hash);

// Get post-transaction balance
let post_ggp_balance = coin_client
.get_account_balance(&ggp_account_address)
.await
.context("Failed to get post-transaction framework balance")?;

tracing::info!("Initial ggp Balance: {}", initial_ggp_balance);
.context("Failed to submit key rotation transaction")?;

// Verify gas fees collection
assert!(post_ggp_balance > initial_ggp_balance, "Gas fees were not collected as expected");

// Wait to verify no additional deposits
tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;

// Check final balance
let final_framework_balance = coin_client
.get_account_balance(&ggp_account_address)
.await
.context("Failed to get final framework balance")?;

// Verify no additional deposits occurred
assert_eq!(
post_ggp_balance, final_framework_balance,
"Additional unexpected deposits were detected"
);
tracing::info!("Key rotation transaction completed successfully");

Ok(())
}
23 changes: 17 additions & 6 deletions networks/movement/movement-client/src/move-modules/scripts/rotate.move
View file
Original file line number Diff line number Diff line change
@@ -1,21 +1,31 @@
script {
use aptos_framework::account;

/// Rotate the authentication key of an account, ensuring secure verification of both current and new keys.
/// This prevents malicious attacks and ensures ownership of both current and new keys.

// For full description of how rotation works see :
// https://github.com/movementlabsxyz/aptos-core/blob/ac9de113a4afec6a26fe587bb92c982532f09d3a/aptos-move/framework/aptos-framework/sources/account.move#L298
/// Perform both offering the rotation capability and rotating the authentication key of an account.
/// This ensures that the recipient has the necessary capability before performing the key rotation.
fun main(
account: &signer,
rotation_capability_sig_bytes: vector<u8>,
from_scheme: u8,
from_public_key_bytes: vector<u8>,
to_scheme: u8,
to_public_key_bytes: vector<u8>,
cap_rotate_key: vector<u8>,
cap_update_table: vector<u8>,
account_scheme: u8,
account_public_key_bytes: vector<u8>,
recipient_address: address,
) {
// Call the `rotate_authentication_key` function from the `account` module
// Step 1: Offer rotation capability to the recipient
account::offer_rotation_capability(
account,
rotation_capability_sig_bytes,
account_scheme,
account_public_key_bytes,
recipient_address,
);

// Step 2: Rotate the authentication key
account::rotate_authentication_key(
account,
from_scheme,
Expand All @@ -27,3 +37,4 @@ script {
);
}
}

0 comments on commit 376382f

Please sign in to comment.