feat: add tls to endpoints (#11)

common/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "common"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.75"
+serde = "1.0.92"
+serde_derive = "1.0.92"
+tokio = { version = "1", features = ["full"] }
+tonic = { version = "0.8.1", features = ["tls", "transport"] }

common/src/lib.rs

@@ -0,0 +1 @@
+pub mod tls;

common/src/tls.rs

@@ -0,0 +1,48 @@
+use anyhow::bail;
+use std::path::Path;
+use tonic::transport::{Certificate, Identity};
+
+#[derive(Clone)]
+pub struct Config {
+ pub ca_cert: Certificate,
+ pub identity: Identity,
+}
+
+impl Config {
+ pub async fn new(
+ ca_cert_path: String,
+ cert_path: String,
+ key_path: String,
+ ) -> anyhow::Result<Self> {
+ let (ca_cert, identity) = get_cert_and_identity(ca_cert_path, cert_path, key_path).await?;
+ Ok(Config { ca_cert, identity })
+ }
+}
+
+async fn get_cert_and_identity(
+ ca_cert_path: String,
+ cert_path: String,
+ key_path: String,
+) -> anyhow::Result<(Certificate, Identity)> {
+ let ca_cert_path = Path::new(&ca_cert_path);
+ let cert_path = Path::new(&cert_path);
+ let key_path = Path::new(&key_path);
+ if !ca_cert_path.is_file() || !cert_path.is_file() || !key_path.is_file() {
+ bail!("both ca_cert_path, cert_path and key_path should be valid file")
+ }
+
+ let ca_cert = tokio::fs::read(ca_cert_path)
+ .await
+ .unwrap_or_else(|err| panic!("Failed to read {:?}, err: {:?}", ca_cert_path, err));
+ let ca_cert = Certificate::from_pem(ca_cert);
+
+ let cert = tokio::fs::read(cert_path)
+ .await
+ .unwrap_or_else(|err| panic!("Failed to read {:?}, err: {:?}", cert_path, err));
+ let key = tokio::fs::read(key_path)
+ .await
+ .unwrap_or_else(|err| panic!("Failed to read {:?}, err: {:?}", key_path, err));
+ let identity = Identity::from_pem(cert, key);
+
+ Ok((ca_cert, identity))
+}

service/Cargo.toml

@@ -9,6 +9,7 @@ edition = "2021"
 prover = { path = "../prover" }
 stage = {path = "../stage"}
 executor = {path = "../executor"}
+common = {path = "../common"}
 tonic = "0.8.1"
 prost = "0.11.0"
 tokio = { version = "1.21.0", features = ["macros", "rt-multi-thread", "signal"] }
@@ -21,5 +22,6 @@ env_logger = "0.10"
 toml = "0.5.1"
 lazy_static = "1.4"
 clap = "4.5.2"
+anyhow = "1.0.75"
 [build-dependencies]
 tonic-build = "0.8.0"

service/config/README.md

@@ -0,0 +1,19 @@
+# README
+
+## Description
+
+The script file `gen_config.sh` allow you generate multi prover toml in a easy way.
+
+First, you should set these variables according to your environment.
+
+- provers
+- stage
+- snarks
+- tls
+- base_dir
+
+Then you can run this script in below way.
+
+```bash
+bash gen_config.sh
+```

service/config/gen_config.sh

@@ -0,0 +1,96 @@
+#!/bin/bash
+
+# You should provide some variable to use this config bash
+provers=("localhost:50001" "localhost:50002")
+stage="localhost:50000"
+snarks=("localhost:50051")
+tls=false
+base_dir="/tmp/zkm/test/test_proof"
+
+# Generate tls certs
+if [ "$tls" = true ]; then
+ IFS=':' read -r host port <<< "$stage"
+ cd ./../../tools/certs
+ bash certgen.sh --cn stage --ssl-dns $host
+ rm -rf stage.csr
+ id=1
+ for prover in "${provers[@]}"; do
+ prover_name="prover${id}"
+ IFS=':' read -r host port <<< "$prover"
+ bash certgen.sh --cn $prover_name --ssl-dns ${host}
+ rm -rf ${prover_name}.csr
+ ((id++))
+ done
+ bash certgen.sh --cn client --ssl-dns localhost
+ rm -rf client.csr
+ rm -rf ca.srl
+ rm -rf openssl.cnf
+ cd -
+fi
+
+# Generate stage toml
+# Read templeta content first
+if [ "$tls" = true ]; then
+ stage_template_content=$(cat stage_tls.toml.template)
+else
+ stage_template_content=$(cat stage.toml.template)
+fi
+stage_config="$stage_template_content"
+IFS=':' read -r host port <<< "$stage"
+stage_config="${stage_config//\{\{addr\}\}/0.0.0.0:${port}}"
+# generate prover addrs
+prover_addrs=""
+for prover in "${provers[@]}"; do
+ if [ -z "$prover_addrs" ]; then
+ prover_addrs="$prover\""
+ else
+ prover_addrs="$prover_addrs, \"$prover"
+ fi
+done
+stage_config="${stage_config//\{\{prover_addrs\}\}/\"${prover_addrs}\"}"
+# generate snark addrs
+snark_addrs=""
+for snark in "${snarks[@]}"; do
+ if [ -z "$snark_addrs" ]; then
+ snark_addrs="$snark\""
+ else
+ snark_addrs="$snark_addrs, \"$snark"
+ fi
+done
+stage_config="${stage_config//\{\{snark_addrs\}\}/\"${snark_addrs}\"}"
+stage_config="${stage_config//\{\{base_dir\}\}/${base_dir}}"
+if [ "$tls" = true ]; then
+ echo "$stage_config" > stage_tls.toml 
+else
+ echo "$stage_config" > stage.toml 
+fi
+
+# Generate provers toml
+# Read templeta content first
+if [ "$tls" = true ]; then
+ prover_template_content=$(cat prover_tls.toml.template)
+else
+ prover_template_content=$(cat prover.toml.template)
+fi
+
+id=1
+for prover in "${provers[@]}"; do
+ if [ "$tls" = true ]; then
+ prover_path="prover${id}_tls.toml"
+ else
+ prover_path="prover${id}.toml"
+ fi
+ IFS=':' read -r host port <<< "$prover"
+ prover_config="$prover_template_content"
+ addr="0.0.0.0:${port}"
+ prover_config="${prover_config//\{\{addr\}\}/${addr}}"
+ prover_config="${prover_config//\{\{prover_addrs\}\}/\"${addr}\"}"
+ prover_config="${prover_config//\{\{base_dir\}\}/${base_dir}}"
+ prover_config="${prover_config//\{\{prover_name\}\}/prover${id}}"
+ if [ "$tls" = true ]; then
+ echo "$prover_config" > "prover${id}_tls.toml"
+ else
+ echo "$prover_config" > "prover${id}.toml"
+ fi
+ ((id++))
+done

service/config/prover.toml.template

@@ -0,0 +1,4 @@
+addr = "{{addr}}"
+prover_addrs = [{{prover_addrs}}]
+snark_addrs = []
+base_dir = "{{base_dir}}"

service/config/prover_tls.toml.template

@@ -0,0 +1,7 @@
+addr = "{{addr}}"
+prover_addrs = [{{prover_addrs}}]
+snark_addrs = []
+base_dir = "{{base_dir}}"
+ca_cert_path = "tools/certs/ca.pem"
+cert_path = "tools/certs/{{prover_name}}.pem"
+key_path = "tools/certs/{{prover_name}}.key"

service/config/stage.toml.template

@@ -0,0 +1,4 @@
+addr = "{{addr}}"
+prover_addrs = [{{prover_addrs}}]
+snark_addrs = [{{snark_addrs}}]
+base_dir = "{{base_dir}}"

service/config/stage_tls.toml.template

@@ -0,0 +1,7 @@
+addr = "{{addr}}"
+prover_addrs = [{{prover_addrs}}]
+snark_addrs = [{{snark_addrs}}]
+base_dir = "{{base_dir}}"
+ca_cert_path = "tools/certs/ca.pem"
+cert_path = "tools/certs/stage.pem"
+key_path = "tools/certs/stage.key"

service/examples/README.md

@@ -28,19 +28,22 @@ cargo build --release
 * Start prover_server.
 
 ```
+# use prover1_tls.toml and prover2_tls.toml instead if tls is enabled
 $ ./target/release/service --config ./service/config/prover1.toml
 $ ./target/release/service --config ./service/config/prover2.toml
 ```
 
 * Start stage_server.
 
 ```
+# use stage_tls.toml instead if tls is enabled
 ./target/release/service --config ./service/config/stage.toml
 ```
 
 * Start example stage
 
 ```
+# set CA_CERT_PATH, CERT_PATH and KEY_PATH if tls is enabled
 cargo run --release --example stage
 ```
 

service/examples/stage.rs

@@ -1,11 +1,12 @@
+use common::tls::Config;
 use stage_service::stage_service_client::StageServiceClient;
 use stage_service::{BlockFileItem, GenerateProofRequest};
-
 use std::env;
 use std::fs;
 use std::path::Path;
-
 use std::time::Instant;
+use tonic::transport::ClientTlsConfig;
+use tonic::transport::Endpoint;
 
 pub mod stage_service {
  tonic::include_proto!("stage.v1");
@@ -17,8 +18,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
  let block_path = env::var("BLOCK_PATH").unwrap_or("/tmp/zkm/test/0_13284491".to_string());
  let block_no = env::var("BLOCK_NO").unwrap_or("13284491".to_string());
  let block_no = block_no.parse::<_>().unwrap_or(13284491);
- let seg_size = env::var("SEG_SIZE").unwrap_or("262144".to_string());
- let seg_size = seg_size.parse::<_>().unwrap_or(262144);
+ let seg_size = env::var("SEG_SIZE").unwrap_or("16384".to_string());
+ let seg_size = seg_size.parse::<_>().unwrap_or(16384);
+ let endpoint = env::var("ENDPOINT").unwrap_or("http://127.0.0.1:50000".to_string());
+ let ca_cert_path = env::var("CA_CERT_PATH").unwrap_or("".to_string());
+ let cert_path = env::var("CERT_PATH").unwrap_or("".to_string());
+ let key_path = env::var("KEY_PATH").unwrap_or("".to_string());
+ let ssl_config = if ca_cert_path.is_empty() {
+ None
+ } else {
+ Some(Config::new(ca_cert_path, cert_path, key_path).await?)
+ };
 
  let elf_data = prover::provers::read_file_bin(&elf_path).unwrap();
  let mut block_data = Vec::new();
@@ -47,7 +57,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
  };
  println!("request: {:?}", request.proof_id.clone());
  let start = Instant::now();
- let mut stage_client = StageServiceClient::connect("http://127.0.0.1:50000").await?;
+ let endpoint = match ssl_config {
+ Some(config) => {
+ let tls_config = ClientTlsConfig::new()
+ .ca_certificate(config.ca_cert)
+ .identity(config.identity);
+ Endpoint::new(endpoint)?.tls_config(tls_config)?
+ }
+ None => Endpoint::new(endpoint)?,
+ };
+ let mut stage_client = StageServiceClient::connect(endpoint).await?;
  let response = stage_client.generate_proof(request).await?.into_inner();
  println!("response: {:?}", response);
  let end = Instant::now();

service/src/config.rs

@@ -11,12 +11,15 @@ pub fn instance() -> &'static Mutex<RuntimeConfig> {
  INSTANCE.get_or_init(|| Mutex::new(RuntimeConfig::new()))
 }
 
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, Clone)]
 pub struct RuntimeConfig {
  pub addr: String,
  pub prover_addrs: Vec<String>,
  pub snark_addrs: Vec<String>,
  pub base_dir: String,
+ pub ca_cert_path: Option<String>,
+ pub cert_path: Option<String>,
+ pub key_path: Option<String>,
 }
 
 impl RuntimeConfig {
@@ -26,6 +29,9 @@ impl RuntimeConfig {
  prover_addrs: ["0.0.0.0:50000".to_string()].to_vec(),
  snark_addrs: ["0.0.0.0:50000".to_string()].to_vec(),
  base_dir: "/tmp".to_string(),
+ ca_cert_path: None,
+ cert_path: None,
+ key_path: None,
  }
  }
 
@@ -50,6 +56,17 @@ impl RuntimeConfig {
  return None;
  }
  };
+ // both of ca_cert_path, cert_path, key_path should be some or none
+ if (config.ca_cert_path.is_some()
+ || config.cert_path.is_some()
+ || config.key_path.is_some())
+ && (config.ca_cert_path.is_none()
+ || config.cert_path.is_none()
+ || config.key_path.is_none())
+ {
+ error!("both of ca_cert_path, cert_path, key_path should be some or none");
+ return None;
+ }
  instance().lock().unwrap().addr.clone_from(&config.addr);
  instance()
  .lock()
@@ -66,6 +83,21 @@ impl RuntimeConfig {
  .unwrap()
  .snark_addrs
  .clone_from(&config.snark_addrs);
+ instance()
+ .lock()
+ .unwrap()
+ .ca_cert_path
+ .clone_from(&config.ca_cert_path);
+ instance()
+ .lock()
+ .unwrap()
+ .cert_path
+ .clone_from(&config.cert_path);
+ instance()
+ .lock()
+ .unwrap()
+ .key_path
+ .clone_from(&config.key_path);
  Some(config)
  }
 }