From 6e178d26ff17c0b8ff4c161205042e94e5ab78e2 Mon Sep 17 00:00:00 2001
From: Steven Gu <88301195+stevenatcrypto@users.noreply.github.com>
Date: Sun, 4 Dec 2022 21:43:55 +0800
Subject: [PATCH] Implement recovering account from derive path. (#4374)

---
 Cargo.lock       | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml       |  2 ++
 sdk/Cargo.toml   |  2 ++
 sdk/src/types.rs | 52 ++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 106 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index ccbce25ce59a7..49242e742cf5a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1202,11 +1202,13 @@ dependencies = [
  "aptos-types",
  "bcs 0.1.3 (git+https://github.com/aptos-labs/bcs?rev=2cde3e8446c460cb17b0c1d6bac7e27e964ac169)",
  "cached-packages",
+ "ed25519-dalek-bip32",
  "move-core-types",
  "once_cell",
  "rand 0.7.3",
  "rand_core 0.5.1",
  "serde 1.0.144",
+ "tiny-bip39",
  "tokio",
  "url",
 ]
@@ -3290,6 +3292,12 @@ dependencies = [
  "serde 1.0.144",
 ]
 
+[[package]]
+name = "derivation-path"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e5c37193a1db1d8ed868c03ec7b152175f26160a5b740e5e484143877e0adf0"
+
 [[package]]
 name = "derive_arbitrary"
 version = "1.1.6"
@@ -3525,6 +3533,18 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "ed25519-dalek-bip32"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d2be62a4061b872c8c0873ee4fc6f101ce7b889d039f019c5fa2af471a59908"
+dependencies = [
+ "derivation-path",
+ "ed25519-dalek",
+ "hmac 0.12.1",
+ "sha2 0.10.2",
+]
+
 [[package]]
 name = "ed25519-dalek-fiat"
 version = "0.1.0"
@@ -4268,8 +4288,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce"
 dependencies = [
  "cfg-if",
+ "js-sys",
  "libc",
  "wasi 0.9.0+wasi-snapshot-preview1",
+ "wasm-bindgen",
 ]
 
 [[package]]
@@ -6929,6 +6951,15 @@ dependencies = [
  "serde 1.0.144",
 ]
 
+[[package]]
+name = "pbkdf2"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "216eaa586a190f0a738f2f918511eecfa90f13295abec0e457cdebcceda80cbd"
+dependencies = [
+ "crypto-mac 0.8.0",
+]
+
 [[package]]
 name = "peeking_take_while"
 version = "0.1.2"
@@ -9366,6 +9397,25 @@ dependencies = [
  "lazy_static 0.2.11",
 ]
 
+[[package]]
+name = "tiny-bip39"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ffc59cb9dfc85bb312c3a78fd6aa8a8582e310b0fa885d5bb877f6dcc601839d"
+dependencies = [
+ "anyhow",
+ "hmac 0.8.1",
+ "once_cell",
+ "pbkdf2",
+ "rand 0.7.3",
+ "rustc-hash",
+ "sha2 0.9.9",
+ "thiserror",
+ "unicode-normalization",
+ "wasm-bindgen",
+ "zeroize",
+]
+
 [[package]]
 name = "tiny-keccak"
 version = "2.0.2"
diff --git a/Cargo.toml b/Cargo.toml
index 05bb5531a011d..ee4bb2f5d5639 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -322,6 +322,7 @@ digest = "0.9.0"
 dir-diff = "0.3.2"
 dirs = "4.0.0"
 ed25519-dalek = { version = "1.0.1", features = ["std", "serde"] }
+ed25519-dalek-bip32 = "0.2.0"
 either = "1.6.1"
 enum_dispatch = "0.3.8"
 env_logger = "0.9.0"
@@ -428,6 +429,7 @@ tempfile = "3.3.0"
 termcolor = "1.1.2"
 textwrap = "0.15.0"
 thiserror = "1.0.31"
+tiny-bip39 = "0.8.2"
 tiny-keccak = { version = "2.0.2", features = ["keccak", "sha3"] }
 tracing = "0.1.34"
 tracing-subscriber = "0.3.11"
diff --git a/sdk/Cargo.toml b/sdk/Cargo.toml
index 37a268bd6063c..ce6e1a31ea8c5 100644
--- a/sdk/Cargo.toml
+++ b/sdk/Cargo.toml
@@ -20,9 +20,11 @@ aptos-rest-client = { workspace = true }
 aptos-types = { workspace = true }
 bcs = { workspace = true }
 cached-packages = { workspace = true }
+ed25519-dalek-bip32 = { workspace = true }
 move-core-types = { workspace = true }
 rand_core = { workspace = true }
 serde = { workspace = true }
+tiny-bip39 = { workspace = true }
 
 [dev-dependencies]
 once_cell = { workspace = true }
diff --git a/sdk/src/types.rs b/sdk/src/types.rs
index f28c669e44dd7..ae72e4a0985d7 100644
--- a/sdk/src/types.rs
+++ b/sdk/src/types.rs
@@ -13,8 +13,12 @@ use crate::{
     },
 };
 
+use anyhow::Result;
 use aptos_types::event::EventKey;
 pub use aptos_types::*;
+use bip39::{Language, Mnemonic, Seed};
+use ed25519_dalek_bip32::{DerivationPath, ExtendedSecretKey};
+use std::str::FromStr;
 
 /// LocalAccount represents an account on the Aptos blockchain. Internally it
 /// holds the private / public key pair and the address of the account. You can
@@ -42,6 +46,29 @@ impl LocalAccount {
         }
     }
 
+    /// Recover an account from derive path (e.g. m/44'/637'/0'/0'/0') and mnemonic phrase,
+    pub fn from_derive_path(
+        derive_path: &str,
+        mnemonic_phrase: &str,
+        sequence_number: u64,
+    ) -> Result<Self> {
+        let derive_path = DerivationPath::from_str(derive_path)?;
+        let mnemonic = Mnemonic::from_phrase(mnemonic_phrase, Language::English)?;
+        // TODO: Make `password` as an optional argument.
+        let seed = Seed::new(&mnemonic, "");
+        let key = ExtendedSecretKey::from_seed(seed.as_bytes())?
+            .derive(&derive_path)?
+            .secret_key;
+        let key = AccountKey::from(Ed25519PrivateKey::try_from(key.as_bytes().as_ref())?);
+        let address = key.authentication_key().derived_address();
+
+        Ok(Self {
+            address,
+            key,
+            sequence_number,
+        })
+    }
+
     /// Generate a new account locally. Note: This function does not actually
     /// create an account on the Aptos blockchain, it just generates a new
     /// account locally.
@@ -183,3 +210,28 @@ impl From<Ed25519PrivateKey> for AccountKey {
         Self::from_private_key(private_key)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_recover_account_from_derive_path() {
+        // Same constants in test cases of TypeScript
+        // https://github.com/aptos-labs/aptos-core/blob/main/ecosystem/typescript/sdk/src/aptos_account.test.ts
+        let derive_path = "m/44'/637'/0'/0'/0'";
+        let mnemonic_phrase =
+            "shoot island position soft burden budget tooth cruel issue economy destroy above";
+        let expected_address = "0x7968dab936c1bad187c60ce4082f307d030d780e91e694ae03aef16aba73f30";
+
+        // Validate if the expected address.
+        let account = LocalAccount::from_derive_path(derive_path, mnemonic_phrase, 0).unwrap();
+        assert_eq!(account.address().to_hex_literal(), expected_address);
+
+        // Return an error for empty derive path.
+        assert!(LocalAccount::from_derive_path("", mnemonic_phrase, 0).is_err());
+
+        // Return an error for empty mnemonic phrase.
+        assert!(LocalAccount::from_derive_path(derive_path, "", 0).is_err());
+    }
+}