shack-payment-gateway 0.1.0

use clap::Parser;
use rand::Rng;
use serde_json::{json, Value};
use sha3::{Digest, Keccak256};
use k256::ecdsa::{RecoveryId, Signature, VerifyingKey};
use std::time::SystemTime;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tracing::{info, warn, Level};
use tracing_subscriber::FmtSubscriber;

#[derive(Parser, Debug)]
#[command(author, version, about = "Shack MCP Payment Gateway Server")]
struct Args {
    /// Expected Ethereum wallet address for authorization.
    /// If not specified, the server operates in permissive mode: any valid signature
    /// recovers its signer address, and that address is returned as authorized.
    #[arg(short, long, env = "SHACK_AUTHORIZED_WALLET")]
    authorized_wallet: Option<String>,
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // All diagnostics go to STDERR; STDOUT carries only JSON-RPC messages.
    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .with_writer(std::io::stderr)
        .finish();
    tracing::subscriber::set_global_default(subscriber)
        .expect("setting default subscriber failed");

    let args = Args::parse();

    if let Some(ref wallet) = args.authorized_wallet {
        // Validate the configured wallet at startup so misconfiguration is caught early.
        if let Err(e) = validate_eth_address(wallet) {
            eprintln!("Invalid --authorized-wallet address: {}", e);
            std::process::exit(1);
        }
        info!("Payment gateway running. Authorized wallet: {}", wallet);
    } else {
        info!("Payment gateway running in permissive mode (no wallet restriction).");
    }

    let stdin = tokio::io::stdin();
    let mut stdout = tokio::io::stdout();
    let mut reader = BufReader::new(stdin).lines();

    while let Ok(Some(line)) = reader.next_line().await {
        if line.trim().is_empty() {
            continue;
        }

        let req: Value = match serde_json::from_str(&line) {
            Ok(v) => v,
            Err(e) => {
                let err_resp = json!({
                    "jsonrpc": "2.0",
                    "error": { "code": -32700, "message": format!("Parse error: {}", e) },
                    "id": null
                });
                stdout
                    .write_all(format!("{}\n", serde_json::to_string(&err_resp)?).as_bytes())
                    .await?;
                stdout.flush().await?;
                continue;
            }
        };

        let response = handle_json_rpc(req, &args).await;
        if let Some(resp_val) = response {
            let payload = serde_json::to_string(&resp_val)? + "\n";
            stdout.write_all(payload.as_bytes()).await?;
            stdout.flush().await?;
        }
    }

    Ok(())
}

async fn handle_json_rpc(req: Value, args: &Args) -> Option<Value> {
    let id = req.get("id").cloned().unwrap_or(Value::Null);
    let method = match req.get("method").and_then(|m| m.as_str()) {
        Some(m) => m,
        None => {
            return Some(json!({
                "jsonrpc": "2.0",
                "error": { "code": -32600, "message": "Invalid Request: missing method" },
                "id": id
            }));
        }
    };

    match method {
        "initialize" => Some(json!({
            "jsonrpc": "2.0",
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": { "tools": {} },
                "serverInfo": {
                    "name": "shack-payment-gateway",
                    "version": "0.1.0"
                }
            },
            "id": id
        })),

        // Client sends this notification after the initialize handshake; no response needed.
        "notifications/initialized" => None,

        "tools/list" => Some(json!({
            "jsonrpc": "2.0",
            "result": {
                "tools": [
                    {
                        "name": "request_authorization",
                        "description": "Generates a cryptographic EIP-191 challenge string that the client must sign with their Ethereum private key before calling a paid tool.",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "tool_name": {
                                    "type": "string",
                                    "description": "Name of the tool the client wants to call"
                                },
                                "cost": {
                                    "type": "string",
                                    "description": "Cost in Shack-Credits (e.g. \"1.5\")"
                                }
                            },
                            "required": ["tool_name", "cost"]
                        }
                    },
                    {
                        "name": "verify_payment",
                        "description": "Verifies an EIP-191 signature against the challenge returned by request_authorization. On success, returns the recovered signer address and whether it matches the authorized wallet.",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "challenge": {
                                    "type": "string",
                                    "description": "The exact challenge string returned by request_authorization"
                                },
                                "signature": {
                                    "type": "string",
                                    "description": "Hex-encoded 65-byte EIP-191 signature (with or without 0x prefix)"
                                }
                            },
                            "required": ["challenge", "signature"]
                        }
                    }
                ]
            },
            "id": id
        })),

        "tools/call" => {
            let params = req.get("params").cloned().unwrap_or(Value::Null);
            let name = match params.get("name").and_then(|n| n.as_str()) {
                Some(n) => n,
                None => {
                    return Some(json!({
                        "jsonrpc": "2.0",
                        "error": { "code": -32602, "message": "Invalid params: missing tool name" },
                        "id": id
                    }));
                }
            };
            let tool_args = params.get("arguments").cloned().unwrap_or(Value::Null);

            match handle_tool_call(name, tool_args, args).await {
                Ok(res_val) => Some(json!({
                    "jsonrpc": "2.0",
                    "result": res_val,
                    "id": id
                })),
                Err(err_msg) => {
                    warn!("Tool '{}' error: {}", name, err_msg);
                    Some(json!({
                        "jsonrpc": "2.0",
                        "error": { "code": -32603, "message": err_msg },
                        "id": id
                    }))
                }
            }
        }

        _ => Some(json!({
            "jsonrpc": "2.0",
            "error": { "code": -32601, "message": format!("Method not found: {}", method) },
            "id": id
        })),
    }
}

async fn handle_tool_call(name: &str, args: Value, gateway_args: &Args) -> Result<Value, String> {
    match name {
        "request_authorization" => {
            let tool_name = args
                .get("tool_name")
                .and_then(|t| t.as_str())
                .ok_or("Missing parameter: tool_name")?;
            let cost = args
                .get("cost")
                .and_then(|c| c.as_str())
                .ok_or("Missing parameter: cost")?;

            if tool_name.is_empty() {
                return Err("tool_name must not be empty".to_string());
            }
            if cost.is_empty() {
                return Err("cost must not be empty".to_string());
            }
            // Reject obviously invalid cost strings early.
            if cost.parse::<f64>().is_err() {
                return Err(format!("cost is not a valid number: {}", cost));
            }

            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .map(|d| d.as_secs())
                .unwrap_or(0);

            // A 16-byte random nonce makes each challenge globally unique.
            let nonce: u128 = rand::thread_rng().r#gen();

            let challenge = format!(
                "Authorize execution of tool: {} costing {} Shack-Credits. Timestamp: {}. Nonce: {:032x}",
                tool_name, cost, now, nonce
            );

            info!("Challenge issued for tool '{}' at cost {}", tool_name, cost);

            Ok(json!({
                "content": [{
                    "type": "text",
                    "text": serde_json::to_string(&json!({
                        "challenge": challenge,
                        "tool_name": tool_name,
                        "cost": cost
                    })).unwrap()
                }]
            }))
        }

        "verify_payment" => {
            let challenge = args
                .get("challenge")
                .and_then(|c| c.as_str())
                .ok_or("Missing parameter: challenge")?;
            let signature = args
                .get("signature")
                .and_then(|s| s.as_str())
                .ok_or("Missing parameter: signature")?;

            if challenge.is_empty() {
                return Err("challenge must not be empty".to_string());
            }
            if signature.is_empty() {
                return Err("signature must not be empty".to_string());
            }

            let msg_hash = eip191_hash(challenge);
            let signer_addr = recover_eth_address(&msg_hash, signature)
                .map_err(|e| format!("Signature recovery failed: {}", e))?;

            let verified = match &gateway_args.authorized_wallet {
                Some(expected) => expected.to_lowercase() == signer_addr.to_lowercase(),
                // No wallet configured: accept any valid signature.
                None => true,
            };

            info!(
                "Verification result: verified={}, recovered={}",
                verified, signer_addr
            );

            Ok(json!({
                "content": [{
                    "type": "text",
                    "text": serde_json::to_string(&json!({
                        "verified": verified,
                        "recovered_address": signer_addr,
                        "authorized_address": gateway_args.authorized_wallet
                    })).unwrap()
                }]
            }))
        }

        _ => Err(format!("Unknown tool: {}", name)),
    }
}

// Returns the EIP-191 Keccak-256 digest of a personal_sign message.
// EIP-191 format: "\x19Ethereum Signed Message:\n" + len(msg) + msg
fn eip191_hash(message: &str) -> [u8; 32] {
    let prefix = format!("\x19Ethereum Signed Message:\n{}", message.len());
    let mut hasher = Keccak256::new();
    hasher.update(prefix.as_bytes());
    hasher.update(message.as_bytes());
    let result = hasher.finalize();
    let mut out = [0u8; 32];
    out.copy_from_slice(&result);
    out
}

// Recovers the Ethereum address that produced the given EIP-191 signature.
// `signature_hex` is a 65-byte hex string (with or without leading 0x).
// The last byte is the recovery identifier v (0, 1, 27, or 28; EIP-155 variants also accepted).
fn recover_eth_address(message_hash: &[u8; 32], signature_hex: &str) -> Result<String, anyhow::Error> {
    let stripped = signature_hex.trim_start_matches("0x");

    // Reject non-hex characters before attempting decode.
    if !stripped.chars().all(|c| c.is_ascii_hexdigit()) {
        return Err(anyhow::anyhow!("Signature contains non-hex characters"));
    }

    let sig_bytes = hex::decode(stripped)?;
    if sig_bytes.len() != 65 {
        return Err(anyhow::anyhow!(
            "Signature must be exactly 65 bytes (130 hex chars); got {}",
            sig_bytes.len()
        ));
    }

    let v_byte = sig_bytes[64];

    // Normalize v to a 0/1 recovery bit.
    let rec_id_byte: u8 = match v_byte {
        0 | 1 => v_byte,
        27 => 0,
        28 => 1,
        // EIP-155 replay-protection encoding: v = 35 + chain_id*2 + recovery_bit
        v if v >= 35 => (v - 35) & 1,
        other => return Err(anyhow::anyhow!("Invalid recovery byte v={}", other)),
    };

    let signature = Signature::from_slice(&sig_bytes[0..64])
        .map_err(|e| anyhow::anyhow!("Malformed signature bytes: {:?}", e))?;
    let recovery_id = RecoveryId::try_from(rec_id_byte)
        .map_err(|e| anyhow::anyhow!("Invalid recovery id: {:?}", e))?;

    let public_key = VerifyingKey::recover_from_prehash(message_hash, &signature, recovery_id)
        .map_err(|e| anyhow::anyhow!("Key recovery failed: {:?}", e))?;

    // Ethereum address = last 20 bytes of Keccak-256(uncompressed public key without 0x04 prefix)
    let encoded_point = public_key.to_encoded_point(false);
    let pk_bytes = encoded_point.as_bytes(); // [0x04, x(32), y(32)]
    if pk_bytes.first() != Some(&0x04) || pk_bytes.len() != 65 {
        return Err(anyhow::anyhow!("Unexpected public key encoding"));
    }

    let pk_hash = Keccak256::digest(&pk_bytes[1..]);
    let eth_address_bytes = &pk_hash[12..]; // take the last 20 bytes

    Ok(format!("0x{}", hex::encode(eth_address_bytes)))
}

// Validates that an address string looks like a well-formed Ethereum address.
// Accepts both checksummed and lowercased variants.
fn validate_eth_address(addr: &str) -> Result<(), String> {
    let stripped = addr.strip_prefix("0x").unwrap_or(addr);
    if stripped.len() != 40 {
        return Err(format!(
            "Expected 40 hex chars after 0x prefix, got {}",
            stripped.len()
        ));
    }
    if !stripped.chars().all(|c| c.is_ascii_hexdigit()) {
        return Err("Address contains non-hex characters".to_string());
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    // Known-good test vector: message "Hello World" signed by the key whose corresponding
    // Ethereum address is 0x1a642f0e3c3af545e7acbd38b07251b3990914f1.
    #[test]
    fn test_recover_known_vector() {
        let message = "Hello World";
        let signature_hex = "e9d4cfddf4da3a1b1800c7c178456cf0e1d285417b0204c7cd038ad450fe541464548e164dfb8acc8340cfff0fbed9aa5585f1c78f7241aa5a8434b3e39ec2b61b";
        let expected_address = "0x1a642f0e3c3af545e7acbd38b07251b3990914f1";

        let hash = eip191_hash(message);
        let recovered = recover_eth_address(&hash, signature_hex).expect("should recover address");

        assert_eq!(recovered.to_lowercase(), expected_address.to_lowercase());
    }

    // Same vector but with 0x prefix on the signature.
    #[test]
    fn test_recover_with_0x_prefix() {
        let message = "Hello World";
        let signature_hex = "0xe9d4cfddf4da3a1b1800c7c178456cf0e1d285417b0204c7cd038ad450fe541464548e164dfb8acc8340cfff0fbed9aa5585f1c78f7241aa5a8434b3e39ec2b61b";
        let expected_address = "0x1a642f0e3c3af545e7acbd38b07251b3990914f1";

        let hash = eip191_hash(message);
        let recovered = recover_eth_address(&hash, signature_hex).expect("should recover address");

        assert_eq!(recovered.to_lowercase(), expected_address.to_lowercase());
    }

    #[test]
    fn test_recover_bad_length_fails() {
        let hash = eip191_hash("test");
        // 64 bytes (128 hex chars) — missing the recovery byte
        let short_sig = "a".repeat(128);
        let err = recover_eth_address(&hash, &short_sig).unwrap_err();
        assert!(err.to_string().contains("65 bytes"), "expected length error, got: {}", err);
    }

    #[test]
    fn test_recover_non_hex_fails() {
        let hash = eip191_hash("test");
        let bad_sig = "zz".repeat(65); // 130 chars but not hex
        let err = recover_eth_address(&hash, &bad_sig).unwrap_err();
        assert!(err.to_string().contains("non-hex"), "expected non-hex error, got: {}", err);
    }

    #[test]
    fn test_validate_eth_address_valid() {
        assert!(validate_eth_address("0x1a642f0e3c3af545e7acbd38b07251b3990914f1").is_ok());
        assert!(validate_eth_address("0x1A642F0E3C3AF545E7ACBD38B07251B3990914F1").is_ok());
    }

    #[test]
    fn test_validate_eth_address_no_prefix() {
        // Without 0x prefix is also accepted (40 raw hex chars)
        assert!(validate_eth_address("1a642f0e3c3af545e7acbd38b07251b3990914f1").is_ok());
    }

    #[test]
    fn test_validate_eth_address_too_short() {
        let err = validate_eth_address("0x1a642f").unwrap_err();
        assert!(err.contains("40 hex chars"), "got: {}", err);
    }

    #[test]
    fn test_validate_eth_address_non_hex() {
        // 40 chars after 0x, but containing non-hex characters.
        let addr = format!("0x{}", "G".repeat(40));
        let err = validate_eth_address(&addr).unwrap_err();
        assert!(err.contains("non-hex"), "got: {}", err);
    }

    #[test]
    fn test_eip191_hash_deterministic() {
        // Same input must always produce the same digest.
        let h1 = eip191_hash("deterministic input");
        let h2 = eip191_hash("deterministic input");
        assert_eq!(h1, h2);
    }

    #[test]
    fn test_eip191_hash_different_messages() {
        let h1 = eip191_hash("message A");
        let h2 = eip191_hash("message B");
        assert_ne!(h1, h2);
    }

    // v=27 and v=28 normalization mirrors the common wallet encoding.
    #[test]
    fn test_v27_v28_normalization() {
        // Build a signature with v=27 from the known vector (which already has v=0x1b = 27).
        let message = "Hello World";
        let sig_hex = "e9d4cfddf4da3a1b1800c7c178456cf0e1d285417b0204c7cd038ad450fe541464548e164dfb8acc8340cfff0fbed9aa5585f1c78f7241aa5a8434b3e39ec2b61b";
        let expected = "0x1a642f0e3c3af545e7acbd38b07251b3990914f1";

        let hash = eip191_hash(message);
        let recovered = recover_eth_address(&hash, sig_hex).expect("v=27 should work");
        assert_eq!(recovered.to_lowercase(), expected);
    }
}