Add AccountRestorationProcessor for Improved Key Rotation and Multi-Key Account Discoverability

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_multi_key_after_rotation/auth_key_account_addresses.json

@@ -0,0 +1,8 @@
+[
+  {
+    "auth_key": "0x36f31938e6995bc4ef7d91ff0278520e4d5f66719d25288240783701f9fc2b8d",
+    "address": "0x21fccd4b01606fa6d094f3b99d6c4134107bf3bd72ef386316554356953e478a",
+    "verified": true,
+    "last_transaction_version": 2200077800
+  }
+]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_multi_key_after_rotation/auth_key_multikey_layout.json

@@ -0,0 +1,13 @@
+[
+  {
+    "auth_key": "0x36f31938e6995bc4ef7d91ff0278520e4d5f66719d25288240783701f9fc2b8d",
+    "signatures_required": 2,
+    "multikey_layout_with_prefixes": [
+      "0x001ea7665d2fe8fb111bf76df50ca1256d28513b51d32ff4644703cf5d3d3c7c1e",
+      "0x00eaba369227ebddfc4495a2cd741f5c04c7ed09da47acaf29ec50a806a3e882d3",
+      "0x01042f9041a53413f15777a369d4751e20d5009ef32b6625ab2cd41f97f2da49d84270a054f1b9c84d6961cc1f3d47d461c50005bfb26a82c358caa967000b81104c"
+    ],
+    "multikey_type": "multi_key",
+    "last_transaction_version": 2200077800
+  }
+]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_multi_key_after_rotation/public_key_auth_keys.json

@@ -0,0 +1,23 @@
+[
+  {
+    "public_key": "0x042f9041a53413f15777a369d4751e20d5009ef32b6625ab2cd41f97f2da49d84270a054f1b9c84d6961cc1f3d47d461c50005bfb26a82c358caa967000b81104c",
+    "public_key_type": "secp256k1_ecdsa",
+    "auth_key": "0x36f31938e6995bc4ef7d91ff0278520e4d5f66719d25288240783701f9fc2b8d",
+    "verified": true,
+    "last_transaction_version": 2200077800
+  },
+  {
+    "public_key": "0x1ea7665d2fe8fb111bf76df50ca1256d28513b51d32ff4644703cf5d3d3c7c1e",
+    "public_key_type": "ed25519",
+    "auth_key": "0x36f31938e6995bc4ef7d91ff0278520e4d5f66719d25288240783701f9fc2b8d",
+    "verified": true,
+    "last_transaction_version": 2200077800
+  },
+  {
+    "public_key": "0xeaba369227ebddfc4495a2cd741f5c04c7ed09da47acaf29ec50a806a3e882d3",
+    "public_key_type": "ed25519",
+    "auth_key": "0x36f31938e6995bc4ef7d91ff0278520e4d5f66719d25288240783701f9fc2b8d",
+    "verified": false,
+    "last_transaction_version": 2200077800
+  }
+]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_single_key_ed25519/auth_key_account_addresses.json

@@ -0,0 +1,8 @@
+[
+  {
+    "auth_key": "0xc9dd6c2111c805cfa3baeb8f964bca007318a3b2738253e1501c5da4025d6385",
+    "address": "0x21fccd4b01606fa6d094f3b99d6c4134107bf3bd72ef386316554356953e478a",
+    "verified": true,
+    "last_transaction_version": 2200077591
+  }
+]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_single_key_ed25519/auth_key_multikey_layout.json

@@ -0,0 +1 @@
+[]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_single_key_ed25519/public_key_auth_keys.json

@@ -0,0 +1 @@
+[]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_single_key_secp256k1_after_rotation/auth_key_account_addresses.json

@@ -0,0 +1,8 @@
+[
+  {
+    "auth_key": "0xff8deae4f70d192454ace42147b1f4ee5dc0f281c9248b845a8d8e9090a0ea16",
+    "address": "0x19169c4c4e275aba9f20de1ee740dfc4914989e8ead71c6604cb832c99f26626",
+    "verified": true,
+    "last_transaction_version": 2200077877
+  }
+]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_single_key_secp256k1_after_rotation/auth_key_multikey_layout.json

@@ -0,0 +1 @@
+[]

rust/integration-tests/sdk_expected_db_output_files/account_restoration_processor/test_single_key_secp256k1_after_rotation/public_key_auth_keys.json

@@ -0,0 +1 @@
+[]

rust/integration-tests/src/diff_test_helper/account_restoration_processor.rs

@@ -0,0 +1,52 @@
+use crate::models::account_restoration_models::{
+    AuthKeyAccountAddress, AuthKeyMultikeyLayout, PublicKeyAuthKey,
+};
+use anyhow::Result;
+use diesel::{
+    query_dsl::methods::{FilterDsl, ThenOrderDsl},
+    ExpressionMethods, PgConnection, RunQueryDsl,
+};
+use processor::schema::{
+    auth_key_account_addresses::dsl as aa_dsl, auth_key_multikey_layout::dsl as am_dsl,
+    public_key_auth_keys::dsl as pa_dsl,
+};
+use serde_json::Value;
+use std::collections::HashMap;
+
+#[allow(dead_code)]
+pub fn load_data(
+    conn: &mut PgConnection,
+    txn_versions: Vec<i64>,
+) -> Result<HashMap<String, Value>> {
+    let mut result_map: HashMap<String, Value> = HashMap::new();
+
+    let aa_result = aa_dsl::auth_key_account_addresses
+        .filter(aa_dsl::last_transaction_version.eq_any(&txn_versions))
+        .then_order_by(aa_dsl::last_transaction_version.asc())
+        .load::<AuthKeyAccountAddress>(conn)?;
+    result_map.insert(
+        "auth_key_account_addresses".to_string(),
+        serde_json::to_value(&aa_result)?,
+    );
+
+    let am_result = am_dsl::auth_key_multikey_layout
+        .filter(am_dsl::last_transaction_version.eq_any(&txn_versions))
+        .then_order_by(am_dsl::last_transaction_version.asc())
+        .load::<AuthKeyMultikeyLayout>(conn)?;
+    result_map.insert(
+        "auth_key_multikey_layout".to_string(),
+        serde_json::to_value(&am_result)?,
+    );
+
+    let pa_result = pa_dsl::public_key_auth_keys
+        .filter(pa_dsl::last_transaction_version.eq_any(&txn_versions))
+        .then_order_by(pa_dsl::last_transaction_version.asc())
+        .then_order_by(pa_dsl::public_key.asc())
+        .load::<PublicKeyAuthKey>(conn)?;
+    result_map.insert(
+        "public_key_auth_keys".to_string(),
+        serde_json::to_value(&pa_result)?,
+    );
+
+    Ok(result_map)
+}

rust/integration-tests/src/diff_test_helper/mod.rs

@@ -1,3 +1,4 @@
+pub mod account_restoration_processor;
 pub mod account_transaction_processor;
 pub mod ans_processor;
 pub mod default_processor;

rust/integration-tests/src/models/account_restoration_models.rs

@@ -0,0 +1,45 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use diesel::{Identifiable, Insertable, Queryable};
+use field_count::FieldCount;
+use processor::schema::{
+    auth_key_account_addresses, auth_key_multikey_layout, public_key_auth_keys,
+};
+use serde::{Deserialize, Serialize};
+
+#[derive(
+    Clone, Debug, Default, Deserialize, FieldCount, Identifiable, Insertable, Serialize, Queryable,
+)]
+#[diesel(primary_key(address))]
+#[diesel(table_name = auth_key_account_addresses)]
+pub struct AuthKeyAccountAddress {
+    pub auth_key: String,
+    pub address: String,
+    pub verified: bool,
+    pub last_transaction_version: i64,
+}
+
+#[derive(Clone, Debug, Deserialize, FieldCount, Identifiable, Insertable, Serialize, Queryable)]
+#[diesel(primary_key(auth_key))]
+#[diesel(table_name = auth_key_multikey_layout)]
+pub struct AuthKeyMultikeyLayout {
+    pub auth_key: String,
+    pub signatures_required: i64,
+    pub multikey_layout_with_prefixes: serde_json::Value,
+    pub multikey_type: String,
+    pub last_transaction_version: i64,
+}
+
+#[derive(
+    Clone, Debug, Default, Deserialize, FieldCount, Identifiable, Insertable, Serialize, Queryable,
+)]
+#[diesel(primary_key(public_key, public_key_type, auth_key))]
+#[diesel(table_name = public_key_auth_keys)]
+pub struct PublicKeyAuthKey {
+    pub public_key: String,
+    pub public_key_type: String,
+    pub auth_key: String,
+    pub verified: bool,
+    pub last_transaction_version: i64,
+}

rust/integration-tests/src/models/mod.rs

@@ -1,3 +1,4 @@
+pub mod account_restoration_models;
 pub mod account_transaction_models;
 pub mod ans_models;
 pub mod default_models;

rust/integration-tests/src/sdk_tests/account_restoration_processor_tests.rs

@@ -0,0 +1,131 @@
+use ahash::AHashMap;
+use aptos_indexer_testing_framework::sdk_test_context::SdkTestContext;
+use sdk_processor::config::{
+    db_config::{DbConfig, PostgresConfig},
+    indexer_processor_config::IndexerProcessorConfig,
+    processor_config::{DefaultProcessorConfig, ProcessorConfig},
+};
+use std::collections::HashSet;
+
+pub fn setup_account_restoration_processor_config(
+    test_context: &SdkTestContext,
+    db_url: &str,
+) -> (IndexerProcessorConfig, &'static str) {
+    let transaction_stream_config = test_context.create_transaction_stream_config();
+    let postgres_config = PostgresConfig {
+        connection_string: db_url.to_string(),
+        db_pool_size: 100,
+    };
+
+    let db_config = DbConfig::PostgresConfig(postgres_config);
+    let default_processor_config = DefaultProcessorConfig {
+        per_table_chunk_sizes: AHashMap::new(),
+        channel_size: 100,
+        deprecated_tables: HashSet::new(),
+    };
+
+    let processor_config = ProcessorConfig::AccountRestorationProcessor(default_processor_config);
+
+    let processor_name = processor_config.name();
+    (
+        IndexerProcessorConfig {
+            processor_config,
+            transaction_stream_config,
+            db_config,
+            backfill_config: None,
+        },
+        processor_name,
+    )
+}
+
+#[allow(clippy::needless_return)]
+#[cfg(test)]
+mod sdk_account_restoration_processor_tests {
+    use super::setup_account_restoration_processor_config;
+    use crate::{
+        diff_test_helper::account_restoration_processor::load_data,
+        sdk_tests::{
+            run_processor_test, setup_test_environment, validate_json, DEFAULT_OUTPUT_FOLDER,
+        },
+    };
+    use aptos_indexer_test_transactions::json_transactions::generated_transactions::{
+        IMPORTED_MAINNET_TXNS_2200077591_ACCOUNT_RESTORATION_SINGLE_ED25519,
+        IMPORTED_MAINNET_TXNS_2200077800_ACCOUNT_RESTORATION_ROTATED_TO_MULTI_KEY,
+        IMPORTED_MAINNET_TXNS_2200077877_ACCOUNT_RESTORATION_ROTATED_TO_SINGLE_SECP256K1,
+    };
+    use aptos_indexer_testing_framework::{cli_parser::get_test_config, database::TestDatabase};
+    use sdk_processor::processors::account_restoration_processor::AccountRestorationProcessor;
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn test_single_key_ed25519() {
+        process_single_transaction(
+            IMPORTED_MAINNET_TXNS_2200077591_ACCOUNT_RESTORATION_SINGLE_ED25519,
+            Some("test_single_key_ed25519".to_string()),
+        )
+        .await;
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn test_multi_key_after_rotation() {
+        process_single_transaction(
+            IMPORTED_MAINNET_TXNS_2200077800_ACCOUNT_RESTORATION_ROTATED_TO_MULTI_KEY,
+            Some("test_multi_key_after_rotation".to_string()),
+        )
+        .await;
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn test_single_key_secp256k1_after_rotation() {
+        process_single_transaction(
+            IMPORTED_MAINNET_TXNS_2200077877_ACCOUNT_RESTORATION_ROTATED_TO_SINGLE_SECP256K1,
+            Some("test_single_key_secp256k1_after_rotation".to_string()),
+        )
+        .await;
+    }
+
+    // Helper function to abstract out the transaction processing
+    async fn process_single_transaction(txn: &[u8], test_case_name: Option<String>) {
+        let (diff_flag, custom_output_path) = get_test_config();
+        let output_path = custom_output_path.unwrap_or_else(|| DEFAULT_OUTPUT_FOLDER.to_string());
+
+        let (db, mut test_context) = setup_test_environment(&[txn]).await;
+
+        let db_url = db.get_db_url();
+        let (indexer_processor_config, processor_name) =
+            setup_account_restoration_processor_config(&test_context, &db_url);
+
+        let account_restoration_processor =
+            AccountRestorationProcessor::new(indexer_processor_config)
+                .await
+                .expect("Failed to create AccountRestorationProcessor");
+
+        match run_processor_test(
+            &mut test_context,
+            account_restoration_processor,
+            load_data,
+            db_url,
+            diff_flag,
+            output_path.clone(),
+            test_case_name.clone(),
+        )
+        .await
+        {
+            Ok(mut db_value) => {
+                let _ = validate_json(
+                    &mut db_value,
+                    test_context.get_request_start_version(),
+                    processor_name,
+                    output_path.clone(),
+                    test_case_name,
+                );
+            },
+            Err(e) => {
+                panic!(
+                    "Test failed on transactions {:?} due to processor error: {}",
+                    test_context.get_test_transaction_versions(),
+                    e
+                );
+            },
+        }
+    }
+}

rust/integration-tests/src/sdk_tests/mod.rs

@@ -34,6 +34,8 @@ pub mod objects_processor_tests;
 #[cfg(test)]
 pub mod stake_processor_tests;
 
+#[cfg(test)]
+pub mod account_restoration_processor_tests;
 #[cfg(test)]
 pub mod user_transaction_processor_tests;
 

rust/processor/src/db/postgres/migrations/2024-12-03-064752_add_account_restoration_tables/down.sql

@@ -0,0 +1,4 @@
+-- This file should undo anything in `up.sql`
+DROP TABLE IF EXISTS public_key_auth_keys;
+DROP TABLE IF EXISTS auth_key_multikey_layout;
+DROP TABLE IF EXISTS auth_key_account_addresses;

rust/processor/src/db/postgres/migrations/2024-12-03-064752_add_account_restoration_tables/up.sql

@@ -0,0 +1,31 @@
+-- Public key to associated multikey auth keys
+CREATE TABLE public_key_auth_keys (
+  public_key VARCHAR(200) NOT NULL,
+  public_key_type VARCHAR(50) NOT NULL,
+  auth_key VARCHAR(66) NOT NULL,
+  verified BOOLEAN NOT NULL,
+  last_transaction_version BIGINT NOT NULL,
+  -- Constraints
+  PRIMARY KEY (
+    public_key,
+    public_key_type,
+    auth_key
+  )
+);
+
+-- Auth key to its corresponding multikey layout
+CREATE TABLE auth_key_multikey_layout (
+  auth_key VARCHAR(66) PRIMARY KEY NOT NULL,
+  signatures_required BIGINT NOT NULL,
+  multikey_layout_with_prefixes jsonb NOT NULL,
+  multikey_type VARCHAR(50) NOT NULL,
+  last_transaction_version BIGINT NOT NULL
+);
+
+-- Auth key to account addresses
+CREATE TABLE auth_key_account_addresses (
+  auth_key VARCHAR(66) NOT NULL,
+  address VARCHAR(66) PRIMARY KEY NOT NULL,
+  verified BOOLEAN NOT NULL,
+  last_transaction_version BIGINT NOT NULL
+);