Skip to content

Commit

Permalink
refactor: working aptos proof_server
Browse files Browse the repository at this point in the history
  • Loading branch information
@tchataigner
tchataigner committed Aug 29, 2024
1 parent e25b2d3 commit 397da74
Show file tree
Hide file tree
Showing 8 changed files with 447 additions and 355 deletions.
78 changes: 51 additions & 27 deletions aptos/proof-server/benches/proof_server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@

use anyhow::anyhow;
use bcs::from_bytes;
use proof_server::error::ClientError;
use proof_server::types::aptos::{AccountInclusionProofResponse, EpochChangeProofResponse};
use proof_server::types::proof_server::{EpochChangeData, InclusionData, Request};
use proof_server::utils::{read_bytes, write_bytes};
use proof_server::types::proof_server::{EpochChangeData, InclusionData, ProvingMode, Request};
use serde::Serialize;
use sphinx_sdk::artifacts::try_install_plonk_bn254_artifacts;
use std::env;
Expand Down Expand Up @@ -201,9 +201,7 @@ async fn bench_proving_inclusion(final_snark: bool) -> Result<ProofData, anyhow:
// Connect to primary server
let primary_address =
env::var("PRIMARY_ADDR").map_err(|_| anyhow::anyhow!("PRIMARY_ADDR not set"))?;
let mut tcp_stream = TcpStream::connect(primary_address)
.await
.map_err(|e| anyhow!(e))?;
let client = reqwest::Client::new();

// Read the binary file
let mut file = File::open(ACCOUNT_INCLUSION_DATA_PATH).map_err(|e| anyhow!(e))?;
Expand All @@ -218,21 +216,35 @@ async fn bench_proving_inclusion(final_snark: bool) -> Result<ProofData, anyhow:
let inclusion_data: InclusionData = account_inclusion_proof_response.into();

// Send the InclusionData as a request payload to the primary server
let request_bytes = if final_snark {
bcs::to_bytes(&Request::SnarkProveInclusion(inclusion_data)).map_err(|e| anyhow!(e))?
let proving_type = if final_snark {
ProvingMode::SNARK
} else {
bcs::to_bytes(&Request::ProveInclusion(inclusion_data)).map_err(|e| anyhow!(e))?
ProvingMode::STARK
};

write_bytes(&mut tcp_stream, &request_bytes)
.await
.map_err(|e| anyhow!(e))?;
let request_bytes = bcs::to_bytes(&Request::ProveInclusion(Box::new((
proving_type,
inclusion_data,
))))
.map_err(|e| anyhow!(e))?;

// Start measuring proving time
let start = Instant::now();

// Measure the time taken to get a response and the size of the response payload
let response_bytes = read_bytes(&mut tcp_stream).await.map_err(|e| anyhow!(e))?;
let response = client
.post(format!("http://{primary_address}/inclusion/proof"))
.header("Accept", "application/octet-stream")
.body(request_bytes)
.send()
.await
.map_err(|err| ClientError::Request {
endpoint: primary_address,
source: err.into(),
})?;

let response_bytes = response
.bytes()
.await
.map_err(|err| ClientError::Internal { source: err.into() })?;

Ok(ProofData {
proving_time: start.elapsed().as_millis(),
Expand All @@ -244,9 +256,7 @@ async fn bench_proving_epoch_change(final_snark: bool) -> Result<ProofData, anyh
// Connect to primary server
let primary_address =
env::var("PRIMARY_ADDR").map_err(|_| anyhow::anyhow!("PRIMARY_ADDR not set"))?;
let mut tcp_stream = TcpStream::connect(primary_address)
.await
.map_err(|e| anyhow!(e))?;
let client = reqwest::Client::new();

// Read the binary file
let mut file = File::open(EPOCH_CHANGE_DATA_PATH).map_err(|e| anyhow!(e))?;
Expand All @@ -258,24 +268,38 @@ async fn bench_proving_epoch_change(final_snark: bool) -> Result<ProofData, anyh
from_bytes(&buffer).map_err(|e| anyhow!(e))?;

// Convert the EpochChangeProofResponse structure into an EpochChangeData structure
let inclusion_data: EpochChangeData = account_inclusion_proof_response.into();
let epoch_change_data: EpochChangeData = account_inclusion_proof_response.into();

// Send the InclusionData as a request payload to the primary server
let request_bytes = if final_snark {
bcs::to_bytes(&Request::SnarkProveEpochChange(inclusion_data)).map_err(|e| anyhow!(e))?
let proving_type = if final_snark {
ProvingMode::SNARK
} else {
bcs::to_bytes(&Request::ProveEpochChange(inclusion_data)).map_err(|e| anyhow!(e))?
ProvingMode::STARK
};

write_bytes(&mut tcp_stream, &request_bytes)
.await
.map_err(|e| anyhow!(e))?;
let request_bytes = bcs::to_bytes(&Request::ProveEpochChange(Box::new((
proving_type,
epoch_change_data,
))))
.map_err(|e| anyhow!(e))?;

// Start measuring proving time
let start = Instant::now();

// Measure the time taken to get a response and the size of the response payload
let response_bytes = read_bytes(&mut tcp_stream).await.map_err(|e| anyhow!(e))?;
let response = client
.post(format!("http://{primary_address}/epoch/proof"))
.header("Accept", "application/octet-stream")
.body(request_bytes)
.send()
.await
.map_err(|err| ClientError::Request {
endpoint: primary_address,
source: err.into(),
})?;

let response_bytes = response
.bytes()
.await
.map_err(|err| ClientError::Internal { source: err.into() })?;

Ok(ProofData {
proving_time: start.elapsed().as_millis(),
Expand Down
87 changes: 57 additions & 30 deletions aptos/proof-server/src/bin/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,14 @@ use proof_server::error::ClientError;
use proof_server::types::aptos::{
AccountInclusionProofResponse, EpochChangeProofResponse, LedgerInfoResponse,
};
use proof_server::types::proof_server::ProvingMode;
use proof_server::utils::validate_and_format_url;
use proof_server::{
aptos_inclusion_proof_endpoint,
types::proof_server::Request,
utils::{read_bytes, write_bytes},
aptos_inclusion_proof_endpoint, types::proof_server::Request,
APTOS_EPOCH_CHANGE_PROOF_ENDPOINT, APTOS_LEDGER_INFO_ENDPOINT,
};
use sphinx_sdk::SphinxProofWithPublicValues;
use std::env;
use std::fmt::Display;
use std::str::FromStr;
use std::sync::Arc;
Expand Down Expand Up @@ -438,28 +438,31 @@ async fn request_prover(
request: &Request,
) -> Result<Vec<u8>, ClientError> {
debug!("Connecting to the proof server at {}", proof_server_address);
let mut stream = TcpStream::connect(&proof_server_address)
.await
.map_err(|err| ClientError::Internal {
source: format!("Error while connecting to proof server: {err}").into(),
})?;
let client = reqwest::Client::new();
debug!("Successfully connected to the proof server");

info!("Sending request to prover: {}", request);

let request_bytes =
bcs::to_bytes(request).map_err(|err| ClientError::Internal { source: err.into() })?;

write_bytes(&mut stream, &request_bytes)
let response = client
.post(proof_server_address)
.header("Accept", "application/octet-stream")
.body(request_bytes)
.send()
.await
.map_err(|err| ClientError::Request {
endpoint: "prover".into(),
endpoint: proof_server_address.into(),
source: err.into(),
})?;

read_bytes(&mut stream)
let response_bytes = response
.bytes()
.await
.map_err(|err| ClientError::Internal { source: err.into() })
.map_err(|err| ClientError::Internal { source: err.into() })?;

Ok(response_bytes.to_vec())
}

/// This method verifies the validator verifier predicate, ie: that the validator committee that
Expand Down Expand Up @@ -534,10 +537,17 @@ async fn epoch_change_proving_task(
// Request a proof generation for the latest epoch change.
debug!("Sending epoch change proof request to the prover");

let request = Request::ProveEpochChange(epoch_change_proof_data.clone().into());
let request = Request::ProveEpochChange(Box::new((
get_proving_mode(),
epoch_change_proof_data.clone().into(),
)));

let epoch_change_proof: SphinxProofWithPublicValues = bcs::from_bytes(
&request_prover(&proof_server_address, &request).await?,
&request_prover(
&format!("http://{}/epoch/proof", proof_server_address),
&request,
)
.await?,
)
.map_err(|err| ClientError::ResponsePayload {
endpoint: format!("{}", &request),
Expand Down Expand Up @@ -590,13 +600,16 @@ async fn epoch_change_verifying_task(
info!("Starting epoch change verification task");
// Verifying the received epoch change proof and the validator verifier hash.
let request = Request::VerifyEpochChange(epoch_change_proof.clone());
let epoch_change_proof_verified = *request_prover(&proof_server_address, &request)
.await?
.first()
.ok_or_else(|| ClientError::ResponsePayload {
endpoint: format!("{}", &request),
source: "No response from prover".into(),
})?;
let epoch_change_proof_verified = *request_prover(
&format!("http://{}/epoch/verify", proof_server_address),
&request,
)
.await?
.first()
.ok_or_else(|| ClientError::ResponsePayload {
endpoint: format!("{}", &request),
source: "No response from prover".into(),
})?;

if epoch_change_proof_verified != 1 {
return Err(ClientError::Verification(String::from(
Expand Down Expand Up @@ -638,9 +651,14 @@ async fn inclusion_proving_task(
let inclusion_proof_data = fetch_inclusion_proof_data(&aptos_node_url).await?;

debug!("Sending account inclusion proof request to the prover");
let request = Request::ProveInclusion(inclusion_proof_data.into());
let request =
Request::ProveInclusion(Box::new((get_proving_mode(), inclusion_proof_data.into())));
let account_inclusion_proof: SphinxProofWithPublicValues = bcs::from_bytes(
&request_prover(&proof_server_address, &request).await?,
&request_prover(
&format!("http://{}/inclusion/proof", proof_server_address),
&request,
)
.await?,
)
.map_err(|err| ClientError::ResponsePayload {
endpoint: format!("{}", &request),
Expand Down Expand Up @@ -671,13 +689,16 @@ async fn inclusion_verifying_task(
info!("Verifying account inclusion proof");
// Verifying the received account inclusion proof and the validator verifier hash.
let request = Request::VerifyInclusion(account_inclusion_proof.clone());
let inclusion_proof_verified = *request_prover(&proof_server_address, &request)
.await?
.first()
.ok_or_else(|| ClientError::ResponsePayload {
endpoint: format!("{}", &request),
source: "No response from prover".into(),
})?;
let inclusion_proof_verified = *request_prover(
&format!("http://{}/inclusion/verify", proof_server_address),
&request,
)
.await?
.first()
.ok_or_else(|| ClientError::ResponsePayload {
endpoint: format!("{}", &request),
source: "No response from prover".into(),
})?;

if inclusion_proof_verified != 1 {
return Err(ClientError::Verification(String::from(
Expand Down Expand Up @@ -783,3 +804,9 @@ async fn verifier_task(
}
}
}

fn get_proving_mode() -> ProvingMode {
// Get proving mode for the light client.
let mode_str: String = env::var("MODE").unwrap_or_else(|_| "STARK".into());
ProvingMode::try_from(mode_str.as_str()).expect("MODE should be STARK or SNARK")
}
Loading

0 comments on commit 397da74

Please sign in to comment.