Skip to content

Commit

Permalink
Add sdk support for prover and pepper use with keyless
Browse files Browse the repository at this point in the history
  • Loading branch information
heliuchuan committed Nov 5, 2024
1 parent 7ea4594 commit 75b5bd6
Show file tree
Hide file tree
Showing 9 changed files with 434 additions and 99 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 3 additions & 7 deletions crates/aptos-rest-client/src/client_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
Expand Down
195 changes: 190 additions & 5 deletions crates/aptos-rest-client/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand Down Expand Up @@ -64,13 +66,64 @@ 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 OriginatingResource {
pub address_map: Table,
}

#[derive(Clone, Debug)]
pub struct Client {
inner: ReqwestClient,
base_url: Url,
version_path_base: String,
}

#[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)
Expand Down Expand Up @@ -101,6 +154,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
Expand Down Expand Up @@ -204,6 +265,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_resource: Response<OriginatingResource> = self
.get_account_resource_bcs(CORE_CODE_ADDRESS, "0x1::account::OriginatingAddress")
.await?;

let table_handle = originating_resource.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_resource.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,
Expand Down Expand Up @@ -454,7 +563,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")?;

Expand Down Expand Up @@ -1227,6 +1339,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,
Expand Down Expand Up @@ -1499,6 +1637,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,
Expand Down
46 changes: 4 additions & 42 deletions crates/aptos/src/account/key_rotation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -391,44 +387,10 @@ 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),
}
.into_inner())
}

#[derive(Deserialize)]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions sdk/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
Loading

0 comments on commit 75b5bd6

Please sign in to comment.