PIE bootloader

Dockerfile

@@ -3,7 +3,7 @@ RUN rustup install 1.79.0
 RUN rustup component add cargo clippy rust-docs rust-std rustc rustfmt
 
 # Use apk for package management in Alpine
-RUN apk add --no-cache build-base libressl-dev
+RUN apk add --no-cache build-base libressl-dev protobuf
 RUN cargo install cargo-chef
 
 FROM chef AS planner
@@ -32,7 +32,7 @@ FROM python:3.9.18-slim-bookworm AS final
 
 WORKDIR /
 
-RUN apt update && apt install -y build-essential libgmp-dev elfutils jq git
+RUN apt update && apt install -y build-essential libgmp-dev elfutils jq git protobuf-compiler
 RUN pip install --upgrade pip
 
 RUN git clone --depth=1 -b v2.7.0-rc.3 https://github.com/starkware-libs/cairo.git

bin/cairo-prove/src/lib.rs

@@ -13,6 +13,7 @@ pub mod prove;
 pub enum CairoVersion {
     V0,
     V1,
+    PIE,
 }
 
 impl FromStr for CairoVersion {

bin/cairo-prove/src/main.rs

@@ -20,7 +20,7 @@ pub async fn main() -> Result<(), ProveErrors> {
         } else {
             fetch_job_polling(sdk, job).await?
         };
-        let path: std::path::PathBuf = args.program_output;
+        let path = args.program_output;
         std::fs::write(path, serde_json::to_string_pretty(&job)?)?;
     }
     Ok(())

bin/cairo-prove/src/prove.rs

@@ -3,15 +3,16 @@ use crate::validate_input;
 use crate::Args;
 use crate::CairoVersion;
 use prover_sdk::sdk::ProverSDK;
+use prover_sdk::PieProverInput;
 use prover_sdk::{
     Cairo0CompiledProgram, Cairo0ProverInput, CairoCompiledProgram, CairoProverInput,
 };
 use serde_json::Value;
 
 pub async fn prove(args: Args, sdk: ProverSDK) -> Result<u64, ProveErrors> {
-    let program = std::fs::read_to_string(&args.program_path)?;
     let proof = match args.cairo_version {
         CairoVersion::V0 => {
+            let program = std::fs::read_to_string(&args.program_path)?;
             let input_path = args
                 .program_input_path
                 .ok_or(ProveErrors::MissingProgramInput)?;
@@ -28,6 +29,7 @@ pub async fn prove(args: Args, sdk: ProverSDK) -> Result<u64, ProveErrors> {
             sdk.prove_cairo0(data).await?
         }
         CairoVersion::V1 => {
+            let program = std::fs::read_to_string(&args.program_path)?;
             let input = match args.clone().program_input_path {
                 Some(input_path) => {
                     let input = std::fs::read_to_string(input_path)?;
@@ -45,6 +47,16 @@ pub async fn prove(args: Args, sdk: ProverSDK) -> Result<u64, ProveErrors> {
             };
             sdk.prove_cairo(data).await?
         }
+        CairoVersion::PIE => {
+            let pie = std::fs::read(&args.program_path)?;
+            let pie_input = PieProverInput {
+                pie_zip: pie,
+                layout: args.layout,
+                pow_bits: args.pow_bits,
+                n_queries: args.n_queries,
+            };
+            sdk.prove_pie(pie_input).await?
+        }
     };
     Ok(proof)
 }

common/src/prover_input/cairo.rs

@@ -12,7 +12,6 @@ pub struct CairoProverInput {
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct CairoCompiledProgram {
-    //pub version: u64,
     pub type_declarations: serde_json::Value,
     pub libfunc_declarations: serde_json::Value,
     pub statements: serde_json::Value,

common/src/prover_input/mod.rs

@@ -1,20 +1,23 @@
 mod cairo;
 mod cairo0;
-
+mod pie;
 pub use cairo::{CairoCompiledProgram, CairoProverInput};
 pub use cairo0::{Cairo0CompiledProgram, Cairo0ProverInput};
+pub use pie::PieProverInput;
 
 #[derive(Debug)]
 pub enum ProverInput {
     Cairo0(Cairo0ProverInput),
     Cairo(CairoProverInput),
+    Pie(PieProverInput),
 }
 
 impl ProverInput {
     pub fn to_json_value(&self) -> serde_json::Value {
         match self {
             ProverInput::Cairo0(input) => serde_json::to_value(input).unwrap(),
             ProverInput::Cairo(input) => serde_json::to_value(input).unwrap(),
+            ProverInput::Pie(input) => serde_json::to_value(input).unwrap(),
         }
     }
 }

common/src/prover_input/pie.rs

@@ -0,0 +1,10 @@
+// Module: prover_input/pie.rs
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct PieProverInput {
+    pub pie_zip: Vec<u8>,
+    pub layout: String,
+    pub n_queries: Option<u32>,
+    pub pow_bits: Option<u32>,
+}

prover-sdk/src/sdk.rs

@@ -1,6 +1,6 @@
 use crate::{access_key::ProverAccessKey, errors::SdkErrors, sdk_builder::ProverSDKBuilder};
 use common::{
-    prover_input::{Cairo0ProverInput, CairoProverInput, ProverInput},
+    prover_input::{Cairo0ProverInput, CairoProverInput, PieProverInput, ProverInput},
     requests::AddKeyRequest,
 };
 use ed25519_dalek::{ed25519::signature::SignerMut, VerifyingKey};
@@ -14,6 +14,7 @@ pub struct ProverSDK {
     pub client: Client,
     pub prover_cairo0: Url,
     pub prover_cairo: Url,
+    pub pie: Url,
     pub verify: Url,
     pub get_job: Url,
     pub register: Url,
@@ -51,6 +52,9 @@ impl ProverSDK {
         self.prove(ProverInput::Cairo(data), self.prover_cairo.clone())
             .await
     }
+    pub async fn prove_pie(&self, data: PieProverInput) -> Result<u64, SdkErrors> {
+        self.prove(ProverInput::Pie(data), self.pie.clone()).await
+    }
 
     async fn prove(&self, data: ProverInput, url: Url) -> Result<u64, SdkErrors> {
         let response = self

prover-sdk/src/sdk_builder.rs

@@ -157,6 +157,7 @@ impl ProverSDKBuilder {
             client,
             prover_cairo0: self.base_url.join("prove/cairo0")?,
             prover_cairo: self.base_url.join("prove/cairo")?,
+            pie: self.base_url.join("prove/pie")?,
             verify: self.base_url.join("verify")?,
             get_job: self.base_url.join("get-job")?,
             register: self.base_url.join("register")?,

prover/src/prove/mod.rs

@@ -1,13 +1,14 @@
 use axum::{routing::post, Router};
 
 use crate::server::AppState;
-mod bootloader;
 mod cairo;
 mod cairo0;
+mod pie;
 
 pub fn router(app_state: AppState) -> Router {
     Router::new()
         .route("/cairo0", post(cairo0::root))
         .route("/cairo", post(cairo::root))
+        .route("/pie", post(pie::root))
         .with_state(app_state)
 }

prover/src/prove/pie.rs

@@ -0,0 +1,35 @@
+use crate::auth::jwt::Claims;
+use crate::extractors::workdir::TempDirHandle;
+use crate::server::AppState;
+use crate::threadpool::{CairoVersionedInput, ExecuteParams};
+use axum::Json;
+use axum::{extract::State, http::StatusCode, response::IntoResponse};
+use common::prover_input::PieProverInput;
+use serde_json::json;
+
+pub async fn root(
+    State(app_state): State<AppState>,
+    TempDirHandle(dir): TempDirHandle,
+    _claims: Claims,
+    Json(program_input): Json<PieProverInput>,
+) -> impl IntoResponse {
+    let thread_pool = app_state.thread_pool.clone();
+    let job_store = app_state.job_store.clone();
+    let job_id = job_store.create_job().await;
+    let thread = thread_pool.lock().await;
+    let execution_params = ExecuteParams {
+        job_id,
+        job_store,
+        dir,
+        program_input: CairoVersionedInput::Pie(program_input.clone()),
+        sse_tx: app_state.sse_tx.clone(),
+        n_queries: program_input.clone().n_queries,
+        pow_bits: program_input.pow_bits,
+    };
+    thread.execute(execution_params).await.into_response();
+
+    let body = json!({
+        "job_id": job_id
+    });
+    (StatusCode::ACCEPTED, body.to_string())
+}

prover/src/threadpool/prove.rs

@@ -53,6 +53,7 @@ pub async fn prove(
             CairoVersionedInput::Cairo0(_cairo0_input) => {
                 prover_result(final_result, CairoVersion::Cairo0)?
             }
+            CairoVersionedInput::Pie(_pie) => prover_result(final_result, CairoVersion::Cairo0)?,
         };
         job_store
             .update_job_status(