rostrum 14.0.1

An efficient implementation of Electrum Server with token support
Documentation
use crate::config::Config;
use crate::doslimit::ConnectionLimits;
use crate::edition::Edition;
use crate::errors::*;
use crate::query::Query;
use crate::rpc::blockchain::BlockchainRpc;
use crate::rpc::mempool::MempoolRPC;
use crate::rpc::rpcstats::RpcStats;
use crate::rpc::server::{
    server_banner, server_features, server_peers_subscribe, server_version, ServerRPC,
};
use crate::rpc::token::TokenRPC;
use crate::signal::{NetworkNotifier, Waiter};
use crate::thread::run_with_timeout;
use crate::timeout::LeakyBucket;
use crate::util::Channel;
use anyhow::{Context, Result};
use serde_json::{from_str, Value};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;

use super::communication::Communcation;
use super::server::ConnectionId;
use super::{to_rpc_error, Message};
use crate::query::scripthash_subscriptions::ScripthashSubscriptions;

pub struct Connection {
    query: Arc<Query>,
    addr: SocketAddr,
    stats: Arc<RpcStats>,
    blockchainrpc: BlockchainRpc,
    server_rpc: ServerRPC,
    token_rpc: TokenRPC,
    mempool_rpc: MempoolRPC,
    idle_timeout: u64,
    rpc_timeout: Mutex<LeakyBucket>,
    communication: Box<dyn Communcation + Send + Sync>,
    #[allow(dead_code)] // Used by BlockchainRpc, cleanup handled in communication layer
    conn_id: ConnectionId,
    #[allow(dead_code)] // Used by BlockchainRpc, cleanup handled in communication layer
    subscription_manager: Arc<ScripthashSubscriptions>,
}

impl Connection {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        config: Arc<Config>,
        query: Arc<Query>,
        addr: SocketAddr,
        stats: Arc<RpcStats>,
        doslimits: ConnectionLimits,
        communication: Box<dyn Communcation + Send + Sync>,
        network_notifier: Arc<NetworkNotifier>,
        conn_id: ConnectionId,
        subscription_manager: Arc<ScripthashSubscriptions>,
    ) -> Connection {
        let rpc_timeout = Mutex::new(LeakyBucket::new(0.5, doslimits.rpc_timeout as u32));
        Connection {
            query: query.clone(),
            addr,
            stats: stats.clone(),
            blockchainrpc: BlockchainRpc::new(
                Arc::clone(&query),
                stats,
                doslimits,
                config.network_type,
                network_notifier,
                conn_id,
                subscription_manager.clone(),
                Edition::Original,
            ),
            server_rpc: ServerRPC::new(config.donation_address.clone(), config.network_type),
            token_rpc: TokenRPC::new(query.clone(), config.network_type),
            mempool_rpc: MempoolRPC::new(query.unconfirmed_index_cloned()),
            idle_timeout: config.rpc_idle_timeout,
            rpc_timeout,
            communication,
            conn_id,
            subscription_manager,
        }
    }

    pub async fn single_req(&mut self, method: &str, params: &[Value], id: &Value) -> Value {
        // server.peers.subscribe is not really a subscription call
        if method != "server.peers.subscribe" && method.contains("subscribe") {
            return to_rpc_error(
                rpc_invalid_request(
                    "Subscriptions not available for this transport protocol".into(),
                ),
                id,
            );
        }

        let timeout = self.rpc_timeout.lock().await.check_out();
        let resp = timeout.timeout(self.handle_command(method, params, id));
        let resp = match resp.await {
            Ok(r) => r,
            Err(e) => to_rpc_error(e, id),
        };
        self.rpc_timeout.lock().await.check_in(&timeout);
        resp
    }

    async fn handle_command(&mut self, method: &str, params: &[Value], id: &Value) -> Value {
        let timer = self
            .stats
            .latency
            .with_label_values(&[method])
            .start_timer();
        let result = match method {
            "blockchain.address.decode" => self.blockchainrpc.address_decode(params),
            "blockchain.address.get_balance" => {
                self.blockchainrpc.address_get_balance(params).await
            }
            "blockchain.address.get_first_use" => {
                self.blockchainrpc.address_get_first_use(params).await
            }
            "blockchain.address.get_first_spend" => {
                self.blockchainrpc.address_get_first_spend(params).await
            }
            "blockchain.address.get_history" => {
                self.blockchainrpc.address_get_history(params).await
            }
            "blockchain.address.get_mempool" => {
                self.blockchainrpc.address_get_mempool(params).await
            }
            "blockchain.address.get_scripthash" => {
                self.blockchainrpc.address_get_scripthash(params)
            }
            "blockchain.address.subscribe" => self.blockchainrpc.address_subscribe(params).await,
            "blockchain.address.listunspent" => {
                self.blockchainrpc.address_listunspent(params).await
            }
            "blockchain.address.unsubscribe" => {
                self.blockchainrpc.address_unsubscribe(params).await
            }
            "blockchain.address.get_volume" => self.blockchainrpc.address_get_volume(params).await,
            "blockchain.block.get" => self.blockchainrpc.block_get(params).await,
            "blockchain.block.header" => self.blockchainrpc.block_header(params).await,
            "blockchain.block.header_verbose" => {
                self.blockchainrpc.block_header_verbose(params).await
            }
            "blockchain.block.headers" => self.blockchainrpc.block_headers(params).await,
            "blockchain.estimatefee" => self.blockchainrpc.estimatefee(params).await,
            "blockchain.headers.subscribe" => self.blockchainrpc.headers_subscribe().await,
            "blockchain.headers.tip" => self.blockchainrpc.headers_tip().await,
            "blockchain.relayfee" => self.blockchainrpc.relayfee().await,
            "blockchain.scripthash.get_balance" => {
                self.blockchainrpc.scripthash_get_balance(params).await
            }
            "blockchain.scripthash.get_first_use" => {
                self.blockchainrpc.scripthash_get_first_use(params).await
            }
            "blockchain.scripthash.get_first_spend" => {
                self.blockchainrpc.scripthash_get_first_spend(params).await
            }
            "blockchain.scripthash.get_history" => {
                self.blockchainrpc.scripthash_get_history(params).await
            }
            "blockchain.scripthash.get_mempool" => {
                self.blockchainrpc.scripthash_get_mempool(params).await
            }
            "blockchain.scripthash.listunspent" => {
                self.blockchainrpc.scripthash_listunspent(params).await
            }
            "blockchain.scripthash.subscribe" => {
                self.blockchainrpc.scripthash_subscribe(params).await
            }
            "blockchain.scripthash.unsubscribe" => {
                self.blockchainrpc.scripthash_unsubscribe(params).await
            }
            "blockchain.scripthash.get_volume" => {
                self.blockchainrpc.scripthash_get_volume(params).await
            }
            "blockchain.transaction.broadcast" => {
                self.blockchainrpc.transaction_broadcast(params).await
            }
            "blockchain.transaction.get" => self.blockchainrpc.transaction_get(params).await,
            "blockchain.transaction.get_confirmed_blockhash" => {
                self.blockchainrpc
                    .transaction_get_confirmed_blockhash(params)
                    .await
            }
            "blockchain.transaction.get_merkle" => {
                self.blockchainrpc.transaction_get_merkle(params).await
            }
            "blockchain.transaction.id_from_pos" => {
                self.blockchainrpc.transaction_id_from_pos(params).await
            }
            "blockchain.utxo.get" => self.blockchainrpc.utxo_get(params).await,
            "mempool.get_fee_histogram" => Ok(self.mempool_rpc.mempool_get_fee_histogram().await),
            "mempool.count" => self.mempool_rpc.mempool_count().await,
            "mempool.get" => self.mempool_rpc.mempool_get(params).await,
            "server.add_peer" => {
                self.server_rpc
                    .server_add_peer(
                        params,
                        &self.addr,
                        &self.query.chain().genesis_hash().await,
                        self.query.discoverer(),
                    )
                    .await
            }
            "server.banner" => server_banner(&self.query).await,
            "server.donation_address" => self.server_rpc.server_donation_address(),
            "server.features" => server_features(&self.query),
            "server.peers.subscribe" => Ok(server_peers_subscribe(self.query.discoverer()).await),
            "server.ping" => Ok(Value::Null),
            "server.version" => {
                let (edition, result) =
                    server_version(params, self.query.announcer().server_version());
                if result.is_ok() {
                    self.blockchainrpc.set_edition(edition);
                }
                result
            }
            "token.address.get_balance" => self.token_rpc.address_get_balance(params).await,
            "token.address.get_history" => self.token_rpc.address_get_history(params).await,
            "token.address.get_mempool" => self.token_rpc.address_get_mempool(params).await,
            "token.address.listunspent" => self.token_rpc.address_listunspent(params).await,
            "token.genesis.info" => self.token_rpc.genesis_info(params).await,
            "token.transaction.get_history" => {
                self.token_rpc.token_transaction_history(params).await
            }
            "token.scripthash.get_balance" => self.token_rpc.scripthash_get_balance(params).await,
            "token.scripthash.get_history" => self.token_rpc.scripthash_get_history(params).await,
            "token.scripthash.get_mempool" => self.token_rpc.scripthash_get_mempool(params).await,
            "token.scripthash.listunspent" => self.token_rpc.scripthash_listunspent(params).await,
            "token.nft.list" => self.token_rpc.list_nft(params).await,
            &_ => Err(rpc_method_not_found(method)),
        };
        timer.observe_duration();

        match result {
            Ok(value) => {
                json!({"jsonrpc": "2.0", "id": id, "result": value})
            }
            Err(e) => {
                if let Some(rpc_e) = e.downcast_ref::<RpcError>() {
                    // Use (at most) two errors from the error chain to produce
                    // an error descrption.
                    let errmsgs: Vec<String> = e.chain().take(2).map(|x| x.to_string()).collect();
                    let errmsgs = errmsgs.join("; ");
                    json!({"jsonrpc": "2.0",
                    "id": id,
                    "error": {
                        "code": rpc_e.code as i32,
                        "message": errmsgs,
                    }})
                } else {
                    trace!(
                        "rpc #{} {} {:?} failed: {:?}",
                        id,
                        method,
                        params,
                        e.to_string()
                    );

                    json!({"jsonrpc": "2.0",
                    "id": id,
                    "error": {
                        "code": RpcErrorCode::InternalError as i32,
                        "message": e.to_string()
                    }})
                }
            }
        }
    }

    async fn handle_replies(&mut self, rpc_queue: Channel<Message>) -> Result<()> {
        let empty_params = json!([]);

        loop {
            let mut receiver = rpc_queue.receiver().await;
            let timeout_duration = Duration::from_secs(self.idle_timeout);
            let recv_result = run_with_timeout(timeout_duration, receiver.recv()).await;

            let msg = match recv_result {
                Ok(Some(message)) => message,
                Ok(None) => return Err(anyhow!("Channel closed")),
                Err(_) => return Err(anyhow!("Idle timeout {} exceeded", self.idle_timeout)),
            };

            let timeout = self.rpc_timeout.lock().await.check_out();

            match msg {
                Message::Request(line) => {
                    let cmd: Value = match from_str(&line).context("invalid JSON format") {
                        Ok(v) => v,
                        Err(err) => {
                            return self
                                .communication
                                .send_values(&[json!({"jsonrpc": "2.0",
                                     "id": -1,
                                    "error": {
                                    "code": RpcErrorCode::InvalidRequest as i32,
                                    "message": err.to_string(),
                                }})])
                                .await
                        }
                    };
                    let reply = match (
                        cmd.get("method"),
                        cmd.get("params").unwrap_or(&empty_params),
                        cmd.get("id"),
                    ) {
                        (Some(Value::String(method)), Value::Array(params), Some(id)) => {
                            let timeout = self.rpc_timeout.lock().await.check_out();
                            let resp = timeout.timeout(self.handle_command(method, params, id));
                            match resp.await {
                                Ok(r) => r,
                                Err(e) => to_rpc_error(e, id),
                            }
                        }
                        _ => bail!("invalid command: {}", cmd),
                    };
                    self.communication.send_values(&[reply]).await?
                }
                Message::ScriptHashChange(updates) => {
                    let notifications = timeout
                        .timeout(self.blockchainrpc.on_scripthash_change(updates))
                        .await;
                    self.rpc_timeout.lock().await.check_in(&timeout);

                    let notifications = match notifications {
                        Ok(n) => n,
                        Err(e) => {
                            trace!("scripthash notification: {}", e);
                            return Err(e);
                        }
                    };

                    self.communication.send_values(&notifications).await?;
                }
                Message::ChainTipChange(tip) => {
                    let notification = self.blockchainrpc.on_chaintip_change(*tip).await?;
                    if let Some(n) = notification {
                        self.communication.send_values(&[n]).await?;
                    }
                }
                Message::Done => return Ok(()),
            }
            Waiter::shutdown_check()?;
        }
    }

    pub async fn run(mut self, rpc_queue: Channel<Message>) {
        self.communication
            .start_request_receiver(rpc_queue.sender(), self.idle_timeout);

        if let Err(e) = self.handle_replies(rpc_queue).await {
            error!(
                "[{}] connection handling failed: {:?}",
                self.addr,
                e.chain().map(|x| x.to_string()).collect::<Vec<String>>()
            );
        }

        self.stats
            .subscriptions
            .sub(self.blockchainrpc.get_num_subscriptions().await);
        trace!("[{}] shutting down connection", self.addr);

        // Shutdown the stream
        let _ = self.communication.shutdown().await;
    }
}