cipher-gate 0.3.0

mod decode;
mod log;
mod proxy;
mod rpc;
mod ws;

use std::sync::Arc;

use actix_cors::Cors;
use actix_web::{App, HttpRequest, HttpResponse, HttpServer, web};
use clap::Parser;
use rust_embed::Embed;
use serde_json::Value;

use rpc::{JsonRpcRequest, JsonRpcResponse, MethodType, classify_method};
use ws::AppState;

#[derive(Embed)]
#[folder = "frontend/out"]
struct FrontendAssets;

/// Baked-in Reown (WalletConnect) project ID so the tool works with zero config.
/// It is not a secret — it ships in the frontend regardless. Override with
/// --reown-project-id / REOWN_PROJECT_ID if you'd rather use your own quota.
/// NOTE: replace this with a real project ID (with localhost origins allowed)
/// before publishing.
const DEFAULT_REOWN_PROJECT_ID: &str = "15362e9c1f56f36ebec18ad143adf101";

#[derive(Parser)]
#[command(
    name = "cipher-gate",
    about = "Cipher Gate — proxy RPC that routes signing requests to a browser wallet"
)]
struct Cli {
    /// Upstream RPC URL to forward read calls to
    #[arg(long)]
    rpc_url: String,

    /// Port for JSON-RPC proxy + frontend UI
    #[arg(long, default_value = "8545")]
    port: u16,

    /// Port for the frontend UI (if you want it on a separate port)
    #[arg(long)]
    ui_port: Option<u16>,

    /// Address to bind to. Defaults to localhost only; set to 0.0.0.0 to
    /// expose the proxy + signing UI to your whole network (use with care —
    /// anyone who can reach it can drive your wallet).
    #[arg(long, default_value = "127.0.0.1")]
    host: String,

    /// Extra browser Origin(s) allowed to open the signing WebSocket, e.g.
    /// http://localhost:3000 for `npm run dev`. Repeatable. The bound host/ports
    /// are always allowed.
    #[arg(long = "allowed-origin")]
    allowed_origins: Vec<String>,

    /// Reown (WalletConnect) project ID from https://cloud.reown.com.
    /// Optional — falls back to a built-in default if unset.
    #[arg(long, env = "REOWN_PROJECT_ID", default_value = DEFAULT_REOWN_PROJECT_ID)]
    reown_project_id: String,
}

/// Runtime config shared across handlers.
pub struct Config {
    pub reown_project_id: String,
    pub ws_url: String,
}

async fn handle_rpc(state: web::Data<Arc<AppState>>, body: web::Json<Value>) -> HttpResponse {
    let raw = body.into_inner();

    // Handle batch requests
    if let Value::Array(batch) = &raw {
        let requests: Vec<JsonRpcRequest> = match batch
            .iter()
            .map(|v| serde_json::from_value(v.clone()))
            .collect::<Result<Vec<_>, _>>()
        {
            Ok(r) => r,
            Err(e) => {
                return HttpResponse::Ok().json(JsonRpcResponse::error(
                    Value::Null,
                    -32700,
                    format!("Parse error: {e}"),
                ));
            }
        };

        let mut responses = Vec::with_capacity(requests.len());
        for req in &requests {
            responses.push(handle_single_request(&state, req).await);
        }
        return HttpResponse::Ok().json(responses);
    }

    // Single request
    let request: JsonRpcRequest = match serde_json::from_value(raw) {
        Ok(r) => r,
        Err(e) => {
            return HttpResponse::Ok().json(JsonRpcResponse::error(
                Value::Null,
                -32700,
                format!("Parse error: {e}"),
            ));
        }
    };

    let response = handle_single_request(&state, &request).await;
    HttpResponse::Ok().json(response)
}

async fn handle_single_request(state: &AppState, request: &JsonRpcRequest) -> JsonRpcResponse {
    let method_type = classify_method(&request.method);

    match method_type {
        MethodType::Read => {
            proxy::forward_to_upstream(&state.http_client, &state.rpc_url, request).await
        }
        MethodType::Account => {
            let addr = state.connected_address.read().await;
            match addr.as_ref() {
                Some(address) => JsonRpcResponse::success(
                    request.id.clone(),
                    Value::Array(vec![Value::String(address.clone())]),
                ),
                None => JsonRpcResponse::success(request.id.clone(), Value::Array(vec![])),
            }
        }
        MethodType::Write => {
            log::intercepted(&request.method);

            // For eth_sendTransaction: run simulation + calldata decode concurrently
            let (simulation, decoded_calldata) = if request.method == "eth_sendTransaction" {
                // Extract calldata hex + target address from tx params
                let tx_obj = if let Value::Array(arr) = &request.params {
                    arr.first()
                } else {
                    Some(&request.params)
                };
                let calldata_hex = tx_obj.and_then(|obj| {
                    obj.get("data")
                        .or(obj.get("input"))
                        .and_then(|v| v.as_str())
                        .filter(|s| s.len() > 2 && *s != "0x")
                        .map(String::from)
                });
                let to_addr = tx_obj
                    .and_then(|obj| obj.get("to"))
                    .and_then(|v| v.as_str())
                    .map(String::from);
                // Snapshot chain id without holding the lock across the await below.
                let chain_id = state.chain_id.read().await.clone();

                let (sim, decoded) = tokio::join!(
                    async {
                        let sim = proxy::simulate_transaction(
                            &state.http_client,
                            &state.rpc_url,
                            &request.params,
                        )
                        .await;
                        if sim.success {
                            log::simulation_passed(sim.gas_estimate.as_deref().unwrap_or("?"));
                        } else {
                            log::simulation_failed(
                                sim.revert_reason.as_deref().unwrap_or("unknown"),
                            );
                        }
                        Some(sim)
                    },
                    async {
                        match calldata_hex {
                            Some(ref hex) => {
                                let result = decode::decode_calldata(
                                    &state.http_client,
                                    &state.selector_cache,
                                    &state.rpc_url,
                                    chain_id.as_deref(),
                                    to_addr.as_deref(),
                                    hex,
                                )
                                .await;
                                if let Some(ref d) = result {
                                    log::decoded(&d.signature);
                                    for w in &d.warnings {
                                        log::calldata_warning(&format!("{:?}", w));
                                    }
                                }
                                result
                            }
                            None => None,
                        }
                    }
                );
                (sim, decoded)
            } else {
                (None, None)
            };

            let request_id = uuid::Uuid::new_v4().to_string();
            let method_for_log = request.method.clone();

            // Overall budget for the whole interaction (queue wait + signing).
            let deadline = std::time::Instant::now() + std::time::Duration::from_secs(600);
            let mut logged_waiting = false;

            // Send to the frontend. If no UI is connected yet, wait for one to
            // connect (bounded by the deadline) instead of failing immediately.
            let rx = loop {
                // Arm the connect notification *before* attempting to send, so a
                // frontend that connects mid-attempt is not missed.
                let connected = state.frontend_connected.notified();

                if let Some(rx) = state
                    .send_signing_request(
                        request_id.clone(),
                        request.method.clone(),
                        request.params.clone(),
                        simulation.clone(),
                        decoded_calldata.clone(),
                    )
                    .await
                {
                    break rx;
                }

                let remaining = deadline.saturating_duration_since(std::time::Instant::now());
                if remaining.is_zero() {
                    log::timeout(&method_for_log);
                    return JsonRpcResponse::error(
                        request.id.clone(),
                        -32603,
                        "No frontend connected before timeout. Open the signing UI and connect a wallet.",
                    );
                }

                if !logged_waiting {
                    log::waiting_for_frontend();
                    logged_waiting = true;
                }

                tokio::select! {
                    _ = connected => {}
                    _ = tokio::time::sleep(remaining) => {
                        log::timeout(&method_for_log);
                        return JsonRpcResponse::error(
                            request.id.clone(),
                            -32603,
                            "No frontend connected before timeout. Open the signing UI and connect a wallet.",
                        );
                    }
                }
            };

            // Wait for the frontend to respond, using the remaining time budget.
            let remaining = deadline.saturating_duration_since(std::time::Instant::now());
            match tokio::time::timeout(remaining, rx).await {
                Ok(Ok(mut resp)) => {
                    resp.id = request.id.clone();
                    if resp.error.is_some() {
                        log::rejected(&method_for_log);
                    } else {
                        log::signed(&method_for_log);
                    }
                    resp
                }
                Ok(Err(_)) => {
                    log::rejected(&method_for_log);
                    JsonRpcResponse::error(
                        request.id.clone(),
                        -32603,
                        "Frontend connection lost while waiting for signature",
                    )
                }
                Err(_) => {
                    state.pending.remove(&request_id);
                    log::timeout(&method_for_log);
                    JsonRpcResponse::error(
                        request.id.clone(),
                        -32603,
                        "Signing request timed out (600s)",
                    )
                }
            }
        }
    }
}

async fn ws_handler(
    req: HttpRequest,
    stream: web::Payload,
    state: web::Data<Arc<AppState>>,
) -> Result<HttpResponse, actix_web::Error> {
    // Require the secret token (passed as ?token=… in the WS URL). This keeps
    // other local processes from driving the wallet.
    let token_ok = req.query_string().split('&').any(|kv| {
        let mut it = kv.splitn(2, '=');
        it.next() == Some("token") && it.next() == Some(state.auth_token.as_str())
    });
    if !token_ok {
        return Ok(HttpResponse::Unauthorized().body("invalid or missing token"));
    }

    // Reject cross-origin browser connections. The Origin header is set by the
    // browser and can't be forged from JS, so this blocks malicious local pages.
    if let Some(origin) = req.headers().get("origin").and_then(|v| v.to_str().ok())
        && !state.allowed_origins.iter().any(|o| o == origin)
    {
        return Ok(HttpResponse::Forbidden().body("origin not allowed"));
    }

    let (response, session, msg_stream) = actix_ws::handle(&req, stream)?;

    let state = state.get_ref().clone();
    actix_web::rt::spawn(ws::handle_ws_connection(state, session, msg_stream));

    Ok(response)
}

/// Serves /api/config — runtime config for the frontend.
async fn api_config(config: web::Data<Arc<Config>>) -> HttpResponse {
    HttpResponse::Ok().json(serde_json::json!({
        "projectId": config.reown_project_id,
        "wsUrl": config.ws_url,
    }))
}

/// Serves embedded frontend static files.
async fn frontend_handler(req: HttpRequest) -> HttpResponse {
    let path = req.path().trim_start_matches('/');
    let path = if path.is_empty() { "index.html" } else { path };

    // Try exact path first
    if let Some(file) = FrontendAssets::get(path) {
        let mime = mime_guess::from_path(path).first_or_octet_stream();
        return HttpResponse::Ok()
            .content_type(mime.as_ref())
            .body(file.data.into_owned());
    }

    // For paths without extension (SPA routes), serve index.html
    if !path.contains('.')
        && let Some(file) = FrontendAssets::get("index.html")
    {
        return HttpResponse::Ok()
            .content_type("text/html")
            .body(file.data.into_owned());
    }

    HttpResponse::NotFound().body("Not found")
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    // Suppress actix/framework noise — our log module handles user-facing output
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("warn")).init();

    let cli = Cli::parse();
    let ui_port = cli.ui_port.unwrap_or(cli.port);
    let separate_ui = cli.ui_port.is_some();
    let exposed = !matches!(cli.host.as_str(), "127.0.0.1" | "localhost" | "::1");

    // Per-run secret the signing UI must present on the WebSocket.
    let auth_token = uuid::Uuid::new_v4().to_string();

    // Browser Origins allowed to open the WS: loopback on both ports, plus any
    // the user explicitly allowed (e.g. http://localhost:3000 for `npm run dev`).
    let mut allowed_origins: Vec<String> = Vec::new();
    for hostname in ["localhost", "127.0.0.1"] {
        for p in [cli.port, ui_port] {
            allowed_origins.push(format!("http://{hostname}:{p}"));
        }
    }
    allowed_origins.extend(cli.allowed_origins.iter().cloned());
    allowed_origins.sort();
    allowed_origins.dedup();

    log::banner(
        &cli.rpc_url,
        &cli.host,
        cli.port,
        ui_port,
        separate_ui,
        exposed,
    );

    let state = Arc::new(AppState::new(
        cli.rpc_url,
        auth_token.clone(),
        allowed_origins,
    ));

    // Fetch chain ID from upstream
    if let Some(chain_id_val) = proxy::fetch_chain_id(&state.http_client, &state.rpc_url).await {
        if let Some(cid) = chain_id_val.as_str() {
            log::chain_id(cid);
            let mut lock = state.chain_id.write().await;
            *lock = Some(cid.to_string());
        }
    } else {
        log::chain_id_failed();
    }

    log::ready();

    // Always return explicit WS URL so it works regardless of how the frontend is
    // accessed. The per-run token is embedded so the UI authenticates automatically.
    let ws_url = format!("ws://localhost:{}/ws?token={}", cli.port, auth_token);

    let config = Arc::new(Config {
        reown_project_id: cli.reown_project_id,
        ws_url,
    });

    if separate_ui {
        // Two separate servers: proxy on --port, UI on --ui-port
        let state2 = state.clone();
        let config2 = config.clone();

        let proxy_server = HttpServer::new(move || {
            App::new()
                .wrap(Cors::permissive())
                .app_data(web::Data::new(state.clone()))
                .app_data(web::Data::new(config.clone()))
                .route("/", web::post().to(handle_rpc))
                .route("/ws", web::get().to(ws_handler))
                .route("/api/config", web::get().to(api_config))
        })
        .bind((cli.host.as_str(), cli.port))?
        .run();

        let ui_server = HttpServer::new(move || {
            App::new()
                .wrap(Cors::permissive())
                .app_data(web::Data::new(state2.clone()))
                .app_data(web::Data::new(config2.clone()))
                .route("/api/config", web::get().to(api_config))
                .default_service(web::to(frontend_handler))
        })
        .bind((cli.host.as_str(), ui_port))?
        .run();

        tokio::try_join!(proxy_server, ui_server)?;
    } else {
        // Single server: everything on one port
        // POST / = JSON-RPC, GET / = frontend UI, GET /ws = WebSocket
        HttpServer::new(move || {
            App::new()
                .wrap(Cors::permissive())
                .app_data(web::Data::new(state.clone()))
                .app_data(web::Data::new(config.clone()))
                .route("/", web::post().to(handle_rpc))
                .route("/ws", web::get().to(ws_handler))
                .route("/api/config", web::get().to(api_config))
                // All other GET requests serve frontend static files
                .default_service(web::get().to(frontend_handler))
        })
        .bind((cli.host.as_str(), cli.port))?
        .run()
        .await?;
    }

    Ok(())
}