starknet_os_runner: add optional TLS termination (#12918)

Author: avi-starkware

Cargo.lock

@@ -392,6 +392,7 @@ time = "0.3.37"
 tokio = "1.47.1"
 tokio-metrics = "0.4.4"
 tokio-retry = "0.3"
+tokio-rustls = "0.26.4"
 tokio-stream = "0.1.8"
 tokio-test = "0.4.4"
 tokio-util = "0.7.13"

Cargo.toml

@@ -45,6 +45,7 @@ stwo_run_and_prove_lib = { workspace = true, optional = true }
 tempfile.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["macros", "process", "rt-multi-thread", "time"] }
+tokio-rustls.workspace = true
 tower = { workspace = true, features = ["util"] }
 tower-http = { workspace = true, features = ["cors"] }
 tracing.workspace = true

crates/starknet_os_runner/Cargo.toml

@@ -11,14 +11,12 @@ fn main() {
 async fn main() -> anyhow::Result<()> {
     use std::net::SocketAddr;
 
-    use anyhow::Context;
     use clap::Parser;
-    use jsonrpsee::server::{ServerBuilder, ServerConfig};
-    use starknet_os_runner::server::config::{CliArgs, ServiceConfig};
+    use starknet_os_runner::server::config::{CliArgs, ServiceConfig, TransportMode};
     use starknet_os_runner::server::cors::{build_cors_layer, cors_mode};
     use starknet_os_runner::server::rpc_impl::ProvingRpcServerImpl;
     use starknet_os_runner::server::rpc_trait::ProvingRpcServer;
-    use tower::ServiceBuilder;
+    use starknet_os_runner::server::start_server;
     use tracing::info;
     use tracing_subscriber::prelude::*;
     use tracing_subscriber::{fmt, EnvFilter};
@@ -36,27 +34,32 @@ async fn main() -> anyhow::Result<()> {
     // Build and start the JSON-RPC server.
     let rpc_impl = ProvingRpcServerImpl::from_config(&config);
     let addr = SocketAddr::new(config.ip, config.port);
-
     let cors_layer = build_cors_layer(&config.cors_allow_origin)?;
 
-    let server_config = ServerConfig::builder().max_connections(config.max_connections).build();
-    let server = ServerBuilder::default()
-        .set_config(server_config)
-        .set_http_middleware(ServiceBuilder::new().option_layer(cors_layer))
-        .build(&addr)
-        .await
-        .context(format!("Failed to bind JSON-RPC server to {addr}"))?;
+    let scheme = match &config.transport {
+        TransportMode::Http => "http",
+        TransportMode::Https { .. } => "https",
+    };
+
+    let (local_addr, server_handle) = start_server(
+        addr,
+        &config.transport,
+        rpc_impl.into_rpc().into(),
+        config.max_connections,
+        cors_layer,
+    )
+    .await?;
 
-    let handle = server.start(rpc_impl.into_rpc());
     info!(
-        local_address = %addr,
+        local_address = %local_addr,
+        scheme,
         max_concurrent_requests = config.max_concurrent_requests,
         max_connections = config.max_connections,
         cors_mode = cors_mode(&config.cors_allow_origin),
         cors_allow_origin = ?config.cors_allow_origin,
         "JSON-RPC proving server is running."
     );
 
-    handle.stopped().await;
+    server_handle.stopped().await;
     Ok(())
 }

crates/starknet_os_runner/src/main.rs

@@ -1,6 +1,51 @@
+use std::net::SocketAddr;
+
+use anyhow::Context;
+use jsonrpsee::server::{Methods, ServerBuilder, ServerConfig, ServerHandle};
+use tower::ServiceBuilder;
+use tower_http::cors::CorsLayer;
+
+use self::config::TransportMode;
+
 pub mod config;
 pub mod cors;
 pub mod error;
 pub mod mock_rpc;
 pub mod rpc_impl;
 pub mod rpc_trait;
+pub mod tls;
+
+/// Starts the JSON-RPC server in either HTTP or HTTPS mode depending on the transport.
+pub async fn start_server(
+    addr: SocketAddr,
+    transport: &TransportMode,
+    methods: Methods,
+    max_connections: u32,
+    cors_layer: Option<CorsLayer>,
+) -> anyhow::Result<(SocketAddr, ServerHandle)> {
+    match transport {
+        TransportMode::Http => {
+            let server_config = ServerConfig::builder().max_connections(max_connections).build();
+            let server = ServerBuilder::default()
+                .set_config(server_config)
+                .set_http_middleware(ServiceBuilder::new().option_layer(cors_layer))
+                .build(&addr)
+                .await
+                .context(format!("Failed to bind JSON-RPC server to {addr}"))?;
+            let local_addr = server.local_addr()?;
+            let server_handle = server.start(methods);
+            Ok((local_addr, server_handle))
+        }
+        TransportMode::Https { tls_cert_file, tls_key_file } => {
+            tls::start_tls_server(
+                addr,
+                tls_cert_file,
+                tls_key_file,
+                methods,
+                max_connections,
+                cors_layer,
+            )
+            .await
+        }
+    }
+}

crates/starknet_os_runner/src/server.rs

@@ -1,4 +1,4 @@
-//! Configuration for the HTTP proving service.
+//! Configuration for the proving service.
 
 use std::net::{IpAddr, Ipv4Addr};
 use std::path::PathBuf;
@@ -21,27 +21,54 @@ const DEFAULT_PORT: u16 = 3000;
 const DEFAULT_MAX_CONCURRENT_REQUESTS: usize = 2;
 const DEFAULT_MAX_CONNECTIONS: u32 = 10;
 
-/// Configuration for the HTTP proving service.
+/// Transport mode for the JSON-RPC server.
+#[derive(Clone, Debug)]
+pub enum TransportMode {
+    Http,
+    Https { tls_cert_file: PathBuf, tls_key_file: PathBuf },
+}
+
+impl TransportMode {
+    /// Constructs a `TransportMode` from optional cert and key paths.
+    ///
+    /// Returns `Http` when both are `None`, `Https` when both are `Some`, or an error when only
+    /// one is provided.
+    pub fn new(
+        tls_cert_file: Option<PathBuf>,
+        tls_key_file: Option<PathBuf>,
+    ) -> Result<Self, ConfigError> {
+        match (tls_cert_file, tls_key_file) {
+            (None, None) => Ok(Self::Http),
+            (Some(tls_cert_file), Some(tls_key_file)) => {
+                Ok(Self::Https { tls_cert_file, tls_key_file })
+            }
+            (Some(_), None) => Err(ConfigError::IncompleteTlsConfig(
+                "tls_cert_file is set but tls_key_file is missing".to_string(),
+            )),
+            (None, Some(_)) => Err(ConfigError::IncompleteTlsConfig(
+                "tls_key_file is set but tls_cert_file is missing".to_string(),
+            )),
+        }
+    }
+}
+
+/// Raw configuration as deserialized from JSON. TLS fields are optional and validated after CLI
+/// overrides are applied.
 #[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(default)]
-pub struct ServiceConfig {
-    /// Configuration for the prover.
+struct RawServiceConfig {
     #[serde(flatten)]
-    pub prover_config: ProverConfig,
-    /// IP address to bind the server to.
-    pub ip: IpAddr,
-    /// Port to bind the server to.
-    pub port: u16,
-    /// Maximum number of concurrent proving requests.
-    pub max_concurrent_requests: usize,
-    /// Maximum number of simultaneous JSON-RPC connections (safety net).
-    pub max_connections: u32,
-    /// List of allowed web origins (domains) that may call this HTTP service from a browser
-    /// (CORS). Examples: `http://localhost:5173`, `https://app.example.com`, or `*` to allow any origin.
-    pub cors_allow_origin: Vec<String>,
+    prover_config: ProverConfig,
+    ip: IpAddr,
+    port: u16,
+    max_concurrent_requests: usize,
+    max_connections: u32,
+    cors_allow_origin: Vec<String>,
+    tls_cert_file: Option<PathBuf>,
+    tls_key_file: Option<PathBuf>,
 }
 
-impl Default for ServiceConfig {
+impl Default for RawServiceConfig {
     fn default() -> Self {
         Self {
             prover_config: ProverConfig::default(),
@@ -50,10 +77,33 @@ impl Default for ServiceConfig {
             max_concurrent_requests: DEFAULT_MAX_CONCURRENT_REQUESTS,
             max_connections: DEFAULT_MAX_CONNECTIONS,
             cors_allow_origin: Vec::new(),
+            tls_cert_file: None,
+            tls_key_file: None,
         }
     }
 }
 
+/// Configuration for the proving service.
+#[derive(Clone, Debug)]
+pub struct ServiceConfig {
+    /// Configuration for the prover.
+    pub prover_config: ProverConfig,
+    /// IP address to bind the server to.
+    pub ip: IpAddr,
+    /// Port to bind the server to.
+    pub port: u16,
+    /// Maximum number of concurrent proving requests.
+    pub max_concurrent_requests: usize,
+    /// Maximum number of simultaneous JSON-RPC connections (safety net).
+    pub max_connections: u32,
+    /// List of allowed web origins (domains) that may call this HTTP service from a browser
+    /// (CORS). Examples: `http://localhost:5173`, `https://app.example.com`, or `*` to allow any
+    /// origin.
+    pub cors_allow_origin: Vec<String>,
+    /// Transport mode (HTTP or HTTPS with TLS).
+    pub transport: TransportMode,
+}
+
 impl ServiceConfig {
     /// Creates a ServiceConfig from CLI arguments.
     pub fn from_args(args: CliArgs) -> Result<Self, ConfigError> {
@@ -65,15 +115,15 @@ impl ServiceConfig {
                     e
                 ))
             })?;
-            serde_json::from_str(&contents).map_err(|e| {
+            serde_json::from_str::<RawServiceConfig>(&contents).map_err(|e| {
                 ConfigError::ConfigFileError(format!(
                     "Failed to parse config file {}: {}",
                     config_file.display(),
                     e
                 ))
             })?
         } else {
-            ServiceConfig::default()
+            RawServiceConfig::default()
         };
 
         // Override with CLI arguments if provided.
@@ -126,6 +176,24 @@ impl ServiceConfig {
                 config.max_connections = max;
             }
         }
+        if let Some(tls_cert_file) = args.tls_cert_file {
+            if Some(&tls_cert_file) != config.tls_cert_file.as_ref() {
+                info!(
+                    "CLI override: tls_cert_file: {:?} -> {:?}",
+                    config.tls_cert_file, tls_cert_file
+                );
+                config.tls_cert_file = Some(tls_cert_file);
+            }
+        }
+        if let Some(tls_key_file) = args.tls_key_file {
+            if Some(&tls_key_file) != config.tls_key_file.as_ref() {
+                info!(
+                    "CLI override: tls_key_file: {:?} -> {:?}",
+                    config.tls_key_file, tls_key_file
+                );
+                config.tls_key_file = Some(tls_key_file);
+            }
+        }
 
         if args.no_cors && !args.cors_allow_origin.is_empty() {
             return Err(ConfigError::InvalidArgument(
@@ -194,19 +262,28 @@ impl ServiceConfig {
                 "max_connections must be at least 1".to_string(),
             ));
         }
-        config.cors_allow_origin = normalize_cors_allow_origins(config.cors_allow_origin)?;
-        if config.cors_allow_origin == ["*"] {
+        let transport = TransportMode::new(config.tls_cert_file, config.tls_key_file)?;
+        let cors_allow_origin = normalize_cors_allow_origins(config.cors_allow_origin)?;
+        if cors_allow_origin == ["*"] {
             info!("CORS allow-origin configured as wildcard '*'.");
         }
 
-        Ok(config)
+        Ok(ServiceConfig {
+            prover_config: config.prover_config,
+            ip: config.ip,
+            port: config.port,
+            max_concurrent_requests: config.max_concurrent_requests,
+            max_connections: config.max_connections,
+            cors_allow_origin,
+            transport,
+        })
     }
 }
 
 /// CLI arguments for the proving service.
 #[derive(Parser, Debug)]
 #[command(name = "starknet-os-runner")]
-#[command(about = "HTTP service for generating Starknet OS proofs", long_about = None)]
+#[command(about = "HTTP/HTTPS service for generating Starknet OS proofs", long_about = None)]
 pub struct CliArgs {
     /// Path to JSON configuration file.
     #[arg(long, value_name = "FILE")]
@@ -236,6 +313,14 @@ pub struct CliArgs {
     #[arg(long, value_name = "N")]
     pub max_connections: Option<u32>,
 
+    /// Path to TLS certificate chain PEM file. Requires --tls-key-file.
+    #[arg(long, value_name = "FILE")]
+    pub tls_cert_file: Option<PathBuf>,
+
+    /// Path to TLS private key PEM file. Requires --tls-cert-file.
+    #[arg(long, value_name = "FILE")]
+    pub tls_key_file: Option<PathBuf>,
+
     /// Override STRK fee token address (hex, e.g. for custom environments that share a chain ID).
     #[arg(long, value_name = "ADDRESS")]
     pub strk_fee_token_address: Option<String>,
@@ -279,4 +364,6 @@ pub enum ConfigError {
     InvalidArgument(String),
     #[error("Missing required field: {0}")]
     MissingRequiredField(String),
+    #[error("Incomplete TLS configuration: {0}")]
+    IncompleteTlsConfig(String),
 }