Add SDK support for prover and pepper use with keyless (#15191)

* Add sdk support for prover and pepper use with keyless

* update

* update

* fix

crates/aptos-rest-client/src/client_builder.rs

@@ -23,13 +23,9 @@ pub enum AptosBaseUrl {
 impl AptosBaseUrl {
     pub fn to_url(&self) -> Url {
         match self {
-            AptosBaseUrl::Mainnet => {
-                Url::from_str("https://fullnode.mainnet.aptoslabs.com").unwrap()
-            },
-            AptosBaseUrl::Devnet => Url::from_str("https://fullnode.devnet.aptoslabs.com").unwrap(),
-            AptosBaseUrl::Testnet => {
-                Url::from_str("https://fullnode.testnet.aptoslabs.com").unwrap()
-            },
+            AptosBaseUrl::Mainnet => Url::from_str("https://api.mainnet.aptoslabs.com").unwrap(),
+            AptosBaseUrl::Devnet => Url::from_str("https://api.devnet.aptoslabs.com").unwrap(),
+            AptosBaseUrl::Testnet => Url::from_str("https://api.testnet.aptoslabs.com").unwrap(),
             AptosBaseUrl::Custom(url) => url.to_owned(),
         }
     }

crates/aptos-rest-client/src/lib.rs

@@ -7,6 +7,7 @@ extern crate core;
 pub mod aptos;
 pub mod error;
 pub mod faucet;
+use error::AptosErrorResponse;
 pub use faucet::FaucetClient;
 pub mod response;
 pub use response::Response;
@@ -23,18 +24,19 @@ pub use aptos_api_types::{
 use aptos_api_types::{
     deserialize_from_string,
     mime_types::{BCS, BCS_SIGNED_TRANSACTION, BCS_VIEW_FUNCTION, JSON},
-    AptosError, BcsBlock, Block, GasEstimation, HexEncodedBytes, IndexResponse, MoveModuleId,
-    TransactionData, TransactionOnChainData, TransactionsBatchSubmissionResult, UserTransaction,
-    VersionedEvent, ViewFunction, ViewRequest,
+    AptosError, AptosErrorCode, BcsBlock, Block, GasEstimation, HexEncodedBytes, IndexResponse,
+    MoveModuleId, TransactionData, TransactionOnChainData, TransactionsBatchSubmissionResult,
+    UserTransaction, VersionedEvent, ViewFunction, ViewRequest,
 };
 use aptos_crypto::HashValue;
 use aptos_logger::{debug, info, sample, sample::SampleRate};
 use aptos_types::{
     account_address::AccountAddress,
     account_config::{AccountResource, NewBlockEvent, CORE_CODE_ADDRESS},
     contract_event::EventWithVersion,
+    keyless::{Groth16Proof, Pepper, ZeroKnowledgeSig, ZKP},
     state_store::state_key::StateKey,
-    transaction::SignedTransaction,
+    transaction::{authenticator::EphemeralSignature, SignedTransaction},
 };
 use move_core_types::{
     ident_str,
@@ -64,13 +66,65 @@ const X_APTOS_SDK_HEADER_VALUE: &str = concat!("aptos-rust-sdk/", env!("CARGO_PK
 
 type AptosResult<T> = Result<T, RestError>;
 
+#[derive(Deserialize)]
+pub struct Table {
+    pub handle: AccountAddress,
+}
+
+#[derive(Deserialize)]
+pub struct OriginatingAddress {
+    pub address_map: Table,
+}
+
 #[derive(Clone, Debug)]
 pub struct Client {
     inner: ReqwestClient,
     base_url: Url,
     version_path_base: String,
 }
 
+// TODO: Dedupe the pepper/prover request/response types with the ones defined in the service.
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct PepperRequest {
+    pub jwt_b64: String,
+    #[serde(with = "hex")]
+    pub epk: Vec<u8>,
+    #[serde(with = "hex")]
+    pub epk_blinder: Vec<u8>,
+    pub exp_date_secs: u64,
+    pub uid_key: String,
+}
+
+#[derive(Debug, serde::Deserialize)]
+struct PepperResponse {
+    #[serde(with = "hex")]
+    pub pepper: Vec<u8>,
+}
+
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct ProverRequest {
+    pub jwt_b64: String,
+    #[serde(with = "hex")]
+    pub epk: Vec<u8>,
+    #[serde(with = "hex")]
+    pub epk_blinder: Vec<u8>,
+    pub exp_date_secs: u64,
+    pub exp_horizon_secs: u64,
+    #[serde(with = "hex")]
+    pub pepper: Vec<u8>,
+    pub uid_key: String,
+}
+
+#[derive(Debug, serde::Deserialize)]
+struct ProverResponse {
+    proof: Groth16Proof,
+    #[serde(with = "hex")]
+    #[allow(dead_code)]
+    public_inputs_hash: [u8; 32],
+    #[serde(with = "hex")]
+    training_wheels_signature: Vec<u8>,
+}
+
 impl Client {
     pub fn builder(aptos_base_url: AptosBaseUrl) -> ClientBuilder {
         ClientBuilder::new(aptos_base_url)
@@ -101,6 +155,14 @@ impl Client {
         Ok(self.base_url.join(&self.version_path_base)?.join(path)?)
     }
 
+    pub fn get_prover_url(&self) -> Url {
+        self.base_url.join("keyless/prover/v0/prove").unwrap()
+    }
+
+    pub fn get_pepper_url(&self) -> Url {
+        self.base_url.join("keyless/pepper/v0/fetch").unwrap()
+    }
+
     pub async fn get_aptos_version(&self) -> AptosResult<Response<AptosVersion>> {
         self.get_resource::<AptosVersion>(CORE_CODE_ADDRESS, "0x1::version::Version")
             .await
@@ -204,6 +266,54 @@ impl Client {
         Ok(response.and_then(|inner| bcs::from_bytes(&inner))?)
     }
 
+    pub async fn lookup_address(
+        &self,
+        address_key: AccountAddress,
+        must_exist: bool,
+    ) -> AptosResult<Response<AccountAddress>> {
+        let originating_address_table: Response<OriginatingAddress> = self
+            .get_account_resource_bcs(CORE_CODE_ADDRESS, "0x1::account::OriginatingAddress")
+            .await?;
+
+        let table_handle = originating_address_table.inner().address_map.handle;
+
+        // The derived address that can be used to look up the original address
+        match self
+            .get_table_item_bcs(
+                table_handle,
+                "address",
+                "address",
+                address_key.to_hex_literal(),
+            )
+            .await
+        {
+            Ok(inner) => Ok(inner),
+            Err(RestError::Api(AptosErrorResponse {
+                error:
+                    AptosError {
+                        error_code: AptosErrorCode::TableItemNotFound,
+                        ..
+                    },
+                ..
+            })) => {
+                // If the table item wasn't found, we may check if the account exists
+                if !must_exist {
+                    Ok(Response::new(
+                        address_key,
+                        originating_address_table.state().clone(),
+                    ))
+                } else {
+                    self.get_account_bcs(address_key)
+                        .await
+                        .map(|account_resource| {
+                            Response::new(address_key, account_resource.state().clone())
+                        })
+                }
+            },
+            Err(err) => Err(err),
+        }
+    }
+
     async fn view_account_balance_bcs_impl(
         &self,
         address: AccountAddress,
@@ -454,7 +564,10 @@ impl Client {
         self.json::<PendingTransaction>(response).await
     }
 
-    pub async fn submit_without_serializing_response(&self, txn: &SignedTransaction) -> Result<()> {
+    pub async fn submit_without_deserializing_response(
+        &self,
+        txn: &SignedTransaction,
+    ) -> Result<()> {
         let txn_payload = bcs::to_bytes(txn)?;
         let url = self.build_path("transactions")?;
 
@@ -1227,6 +1340,32 @@ impl Client {
         self.json(response).await
     }
 
+    pub async fn get_account_sequence_number(
+        &self,
+        address: AccountAddress,
+    ) -> AptosResult<Response<u64>> {
+        let res = self.get_account_bcs(address).await;
+
+        match res {
+            Ok(account) => account.and_then(|account| Ok(account.sequence_number())),
+            Err(error) => match error {
+                RestError::Api(error) => {
+                    if matches!(error.error.error_code, AptosErrorCode::AccountNotFound) {
+                        if let Some(state) = error.state {
+                            Ok(Response::new(0, state))
+                        } else {
+                            let ledger_info = self.get_ledger_information().await?;
+                            Ok(Response::new(0, ledger_info.state().clone()))
+                        }
+                    } else {
+                        Err(error::RestError::Api(error))
+                    }
+                },
+                _ => Err(error),
+            },
+        }
+    }
+
     pub async fn get_account_events_bcs(
         &self,
         address: AccountAddress,
@@ -1499,6 +1638,53 @@ impl Client {
         self.check_and_parse_bcs_response(response).await
     }
 
+    pub async fn make_prover_request(&self, req: ProverRequest) -> AptosResult<ZeroKnowledgeSig> {
+        let response: ProverResponse = self
+            .post_json_no_state(self.get_prover_url(), serde_json::to_value(req.clone())?)
+            .await?;
+        let proof = response.proof;
+        let ephem_sig = Some(
+            EphemeralSignature::try_from(response.training_wheels_signature.as_slice())
+                .map_err(anyhow::Error::from)?,
+        );
+        Ok(ZeroKnowledgeSig {
+            proof: ZKP::Groth16(proof),
+            exp_horizon_secs: req.exp_horizon_secs,
+            extra_field: None,
+            override_aud_val: None,
+            training_wheels_signature: ephem_sig,
+        })
+    }
+
+    pub async fn make_pepper_request(&self, req: PepperRequest) -> AptosResult<Pepper> {
+        let response: PepperResponse = self
+            .post_json_no_state(self.get_pepper_url(), serde_json::to_value(req.clone())?)
+            .await?;
+        let pepper = response.pepper;
+        Ok(Pepper::new(
+            pepper.as_slice().try_into().map_err(anyhow::Error::from)?,
+        ))
+    }
+
+    async fn post_json_no_state<T: serde::de::DeserializeOwned>(
+        &self,
+        url: Url,
+        data: serde_json::Value,
+    ) -> AptosResult<T> {
+        let response = self
+            .inner
+            .post(url)
+            .header(ACCEPT, JSON)
+            .json(&data)
+            .send()
+            .await?;
+        if !response.status().is_success() {
+            Err(parse_error(response).await)
+        } else {
+            Ok(response.json().await.map_err(anyhow::Error::from)?)
+        }
+    }
+
     async fn post_bcs(
         &self,
         url: Url,

crates/aptos/src/account/key_rotation.rs

@@ -14,11 +14,7 @@ use aptos_crypto::{
     PrivateKey, SigningKey,
 };
 use aptos_ledger;
-use aptos_rest_client::{
-    aptos_api_types::{AptosError, AptosErrorCode},
-    error::{AptosErrorResponse, RestError},
-    Client,
-};
+use aptos_rest_client::{error::RestError, Client};
 use aptos_types::{
     account_address::AccountAddress,
     account_config::{RotationProofChallenge, CORE_CODE_ADDRESS},
@@ -391,52 +387,8 @@ pub async fn lookup_address(
     address_key: AccountAddress,
     must_exist: bool,
 ) -> Result<AccountAddress, RestError> {
-    let originating_resource: OriginatingResource = rest_client
-        .get_account_resource_bcs(CORE_CODE_ADDRESS, "0x1::account::OriginatingAddress")
+    Ok(rest_client
+        .lookup_address(address_key, must_exist)
         .await?
-        .into_inner();
-
-    let table_handle = originating_resource.address_map.handle;
-
-    // The derived address that can be used to look up the original address
-    match rest_client
-        .get_table_item_bcs(
-            table_handle,
-            "address",
-            "address",
-            address_key.to_hex_literal(),
-        )
-        .await
-    {
-        Ok(inner) => Ok(inner.into_inner()),
-        Err(RestError::Api(AptosErrorResponse {
-            error:
-                AptosError {
-                    error_code: AptosErrorCode::TableItemNotFound,
-                    ..
-                },
-            ..
-        })) => {
-            // If the table item wasn't found, we may check if the account exists
-            if !must_exist {
-                Ok(address_key)
-            } else {
-                rest_client
-                    .get_account_bcs(address_key)
-                    .await
-                    .map(|_| address_key)
-            }
-        },
-        Err(err) => Err(err),
-    }
-}
-
-#[derive(Deserialize)]
-pub struct OriginatingResource {
-    pub address_map: Table,
-}
-
-#[derive(Deserialize)]
-pub struct Table {
-    pub handle: AccountAddress,
+        .into_inner())
 }

crates/transaction-emitter-lib/src/emitter/transaction_executor.rs

@@ -277,17 +277,11 @@ pub async fn query_sequence_number_with_client(
 ) -> Result<u64> {
     let result = FETCH_ACCOUNT_RETRY_POLICY
         .retry_if(
-            move || rest_client.get_account_bcs(account_address),
+            move || rest_client.get_account_sequence_number(account_address),
             |error: &RestError| !is_account_not_found(error),
         )
         .await;
-    match result {
-        Ok(account) => Ok(account.into_inner().sequence_number()),
-        Err(error) => match is_account_not_found(&error) {
-            true => Ok(0),
-            false => Err(error.into()),
-        },
-    }
+    Ok(*result?.inner())
 }
 
 fn is_account_not_found(error: &RestError) -> bool {

sdk/Cargo.toml

@@ -26,6 +26,7 @@ ed25519-dalek-bip32 = { workspace = true }
 hex = { workspace = true }
 move-core-types = { workspace = true }
 rand_core = { workspace = true }
+reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 tiny-bip39 = { workspace = true }