Use jsonrpc for simple balance check

src/bin/index_server.rs

@@ -3,7 +3,6 @@ extern crate simplelog;
 extern crate argparse;
 extern crate crossbeam;
 extern crate indexrs;
-extern crate zmq;
 
 use argparse::{ArgumentParser, StoreFalse};
 use std::fs::OpenOptions;
@@ -28,27 +27,11 @@ fn setup_logging() {
 }
 
 const DB_PATH: &str = "./db/mainnet";
-const RPC_ADDRESS: &str = "ipc:///tmp/indexrs.rpc";
 
 struct Config {
     enable_indexing: bool,
 }
 
-fn handle_queries(store: &store::Store, daemon: &daemon::Daemon) {
-    let query = index::Query::new(&store, &daemon);
-
-    let ctx = zmq::Context::new();
-    let sock = ctx.socket(zmq::SocketType::REP).unwrap();
-    sock.bind(RPC_ADDRESS).unwrap();
-
-    loop {
-        let script_hash = sock.recv_bytes(0).unwrap();
-        let balance = query.balance(&script_hash);
-        let reply = format!("{}", balance);
-        sock.send(&reply.into_bytes(), 0).unwrap();
-    }
-}
-
 fn run_server(config: Config) {
     let mut index = index::Index::new();
     let waiter = waiter::Waiter::new("tcp://localhost:28332");
@@ -68,10 +51,10 @@ fn run_server(config: Config) {
     }
 
     let store = store::Store::open(DB_PATH, store::StoreOptions { auto_compact: true });
+    let query = index::Query::new(&store, &daemon);
 
     crossbeam::scope(|scope| {
-        scope.spawn(|| handle_queries(&store, &daemon));
-        scope.spawn(|| rpc::serve());
+        scope.spawn(|| rpc::serve("localhost:50001", &query));
         loop {
             waiter.wait();
             if config.enable_indexing {

src/rpc.rs

+use bitcoin::util::hash::Sha256dHash;
 use serde_json::{from_str, Number, Value};
-
+use std::io::{BufRead, BufReader, Write};
 use std::net::{SocketAddr, TcpListener, TcpStream};
 
-use std::io::{BufRead, BufReader, Write};
+use index;
 
 error_chain!{}
 
-fn blockchain_headers_subscribe() -> Result<Value> {
-    Ok(json!({}))
+struct Handler<'a> {
+    query: &'a index::Query<'a>,
 }
 
-fn server_version() -> Result<Value> {
-    Ok(json!(["LES 0.1.0", "1.2"]))
-}
+impl<'a> Handler<'a> {
+    fn blockchain_headers_subscribe(&self) -> Result<Value> {
+        Ok(json!({}))
+    }
 
-fn server_banner() -> Result<Value> {
-    Ok(json!("Welcome to Local Electrum Server!\n"))
-}
+    fn server_version(&self) -> Result<Value> {
+        Ok(json!(["LES 0.1.0", "1.2"]))
+    }
 
-fn server_donation_address() -> Result<Value> {
-    Ok(json!("No, thanks :)\n"))
-}
+    fn server_banner(&self) -> Result<Value> {
+        Ok(json!("Welcome to Local Electrum Server!\n"))
+    }
 
-fn server_peers_subscribe() -> Result<Value> {
-    Ok(json!([]))
-}
+    fn server_donation_address(&self) -> Result<Value> {
+        Ok(json!("No, thanks :)\n"))
+    }
 
-fn mempool_get_fee_histogram() -> Result<Value> {
-    Ok(json!([])) // TODO: consult with actual mempool
-}
+    fn server_peers_subscribe(&self) -> Result<Value> {
+        Ok(json!([]))
+    }
 
-fn blockchain_estimatefee(_params: &[&str]) -> Result<Value> {
-    Ok(json!(1e-5)) // TODO: consult with actual mempool
-}
+    fn mempool_get_fee_histogram(&self) -> Result<Value> {
+        Ok(json!([])) // TODO: consult with actual mempool
+    }
 
-fn blockchain_scripthash_subscribe(_params: &[&str]) -> Result<Value> {
-    Ok(json!("HEX_STATUS"))
-}
+    fn blockchain_estimatefee(&self, _params: &[&str]) -> Result<Value> {
+        Ok(json!(1e-5)) // TODO: consult with actual mempool
+    }
 
-fn blockchain_scripthash_get_history(_params: &[&str]) -> Result<Value> {
-    Ok(json!([])) // TODO: list of {tx_hash: "ABC", height: 123}
-}
+    fn blockchain_scripthash_subscribe(&self, _params: &[&str]) -> Result<Value> {
+        Ok(json!("HEX_STATUS"))
+    }
 
-fn blockchain_transaction_get(_params: &[&str]) -> Result<Value> {
-    Ok(json!("HEX_TX")) // TODO: list of {tx_hash: "ABC", height: 123}
-}
+    fn blockchain_scripthash_get_balance(&self, params: &[&str]) -> Result<Value> {
+        let script_hash_hex = params.get(0).chain_err(|| "missing parameter")?;
+        let script_hash =
+            Sha256dHash::from_hex(script_hash_hex).chain_err(|| "invalid scripthash")?;
+        let confirmed = self.query.balance(&script_hash[..]);
+        Ok(json!({ "confirmed": confirmed })) // TODO: "unconfirmed"
+    }
 
-fn blockchain_transaction_get_merkle(_params: &[&str]) -> Result<Value> {
-    Ok(json!({"block_height": 123, "merkle": ["A", "B", "C"], "pos": 45}))
-}
+    fn blockchain_scripthash_get_history(&self, _params: &[&str]) -> Result<Value> {
+        Ok(json!([])) // TODO: list of {tx_hash: "ABC", height: 123}
+    }
 
-fn handle_command(method: &str, params_values: &[Value], id: &Number) -> Result<Value> {
-    let mut params = Vec::<&str>::new();
-    for value in params_values {
-        if let Some(s) = value.as_str() {
-            params.push(s);
-        } else {
-            bail!("invalid param: {:?}", value);
-        }
+    fn blockchain_transaction_get(&self, _params: &[&str]) -> Result<Value> {
+        Ok(json!("HEX_TX")) // TODO: list of {tx_hash: "ABC", height: 123}
     }
-    let result = match method {
-        "blockchain.headers.subscribe" => blockchain_headers_subscribe(),
-        "server.version" => server_version(),
-        "server.banner" => server_banner(),
-        "server.donation_address" => server_donation_address(),
-        "server.peers.subscribe" => server_peers_subscribe(),
-        "mempool.get_fee_histogram" => mempool_get_fee_histogram(),
-        "blockchain.estimatefee" => blockchain_estimatefee(&params),
-        "blockchain.scripthash.subscribe" => blockchain_scripthash_subscribe(&params),
-        "blockchain.scripthash.get_history" => blockchain_scripthash_get_history(&params),
-        "blockchain.transaction.get" => blockchain_transaction_get(&params),
-        "blockchain.transaction.get_merkle" => blockchain_transaction_get_merkle(&params),
-        &_ => bail!("unknown method {} {:?}", method, params),
-    }?;
-    let reply = json!({"jsonrpc": "2.0", "id": id, "result": result});
-    Ok(reply)
-}
 
-fn handle_client(mut stream: TcpStream, addr: SocketAddr) -> Result<()> {
-    let mut reader = BufReader::new(stream
-        .try_clone()
-        .chain_err(|| "failed to clone TcpStream")?);
-    let mut line = String::new();
+    fn blockchain_transaction_get_merkle(&self, _params: &[&str]) -> Result<Value> {
+        Ok(json!({"block_height": 123, "merkle": ["A", "B", "C"], "pos": 45}))
+    }
 
-    loop {
-        line.clear();
-        reader
-            .read_line(&mut line)
-            .chain_err(|| "failed to read a request")?;
-        if line.is_empty() {
-            break;
+    fn handle_command(&self, method: &str, params_values: &[Value], id: &Number) -> Result<Value> {
+        let mut params = Vec::<&str>::new();
+        for value in params_values {
+            if let Some(s) = value.as_str() {
+                params.push(s);
+            } else {
+                bail!("invalid param: {:?}", value);
+            }
+        }
+        let result = match method {
+            "blockchain.headers.subscribe" => self.blockchain_headers_subscribe(),
+            "server.version" => self.server_version(),
+            "server.banner" => self.server_banner(),
+            "server.donation_address" => self.server_donation_address(),
+            "server.peers.subscribe" => self.server_peers_subscribe(),
+            "mempool.get_fee_histogram" => self.mempool_get_fee_histogram(),
+            "blockchain.estimatefee" => self.blockchain_estimatefee(&params),
+            "blockchain.scripthash.subscribe" => self.blockchain_scripthash_subscribe(&params),
+            "blockchain.scripthash.get_balance" => self.blockchain_scripthash_get_balance(&params),
+            "blockchain.scripthash.get_history" => self.blockchain_scripthash_get_history(&params),
+            "blockchain.transaction.get" => self.blockchain_transaction_get(&params),
+            "blockchain.transaction.get_merkle" => self.blockchain_transaction_get_merkle(&params),
+            &_ => bail!("unknown method {} {:?}", method, params),
+        }?;
+        let reply = json!({"jsonrpc": "2.0", "id": id, "result": result});
+        Ok(reply)
+    }
+
+    pub fn run(self, mut stream: TcpStream, addr: SocketAddr) -> Result<()> {
+        let mut reader = BufReader::new(stream
+            .try_clone()
+            .chain_err(|| "failed to clone TcpStream")?);
+        let mut line = String::new();
+
+        loop {
+            line.clear();
+            reader
+                .read_line(&mut line)
+                .chain_err(|| "failed to read a request")?;
+            if line.is_empty() {
+                break;
+            }
+            let line = line.trim_right();
+            let cmd: Value = from_str(line).chain_err(|| "invalid JSON format")?;
+
+            let reply = match (cmd.get("method"), cmd.get("params"), cmd.get("id")) {
+                (
+                    Some(&Value::String(ref method)),
+                    Some(&Value::Array(ref params)),
+                    Some(&Value::Number(ref id)),
+                ) => self.handle_command(method, params, id)?,
+                _ => bail!("invalid command: {}", cmd),
+            };
+
+            debug!("[{}] {} -> {}", addr, cmd, reply);
+            let mut line = reply.to_string();
+            line.push_str("\n");
+            stream
+                .write_all(line.as_bytes())
+                .chain_err(|| "failed to send response")?;
         }
-        let line = line.trim_right();
-        let cmd: Value = from_str(line).chain_err(|| "invalid JSON format")?;
-
-        let reply = match (cmd.get("method"), cmd.get("params"), cmd.get("id")) {
-            (
-                Some(&Value::String(ref method)),
-                Some(&Value::Array(ref params)),
-                Some(&Value::Number(ref id)),
-            ) => handle_command(method, params, id)?,
-            _ => bail!("invalid command: {}", cmd),
-        };
-
-        info!("[{}] {} -> {}", addr, cmd, reply);
-        let mut line = reply.to_string();
-        line.push_str("\n");
-        stream
-            .write_all(line.as_bytes())
-            .chain_err(|| "failed to send response")?;
+        Ok(())
     }
-    Ok(())
 }
 
-pub fn serve() {
-    let listener = TcpListener::bind("127.0.0.1:50001").unwrap();
+pub fn serve(addr: &str, query: &index::Query) {
+    let listener = TcpListener::bind(addr).unwrap();
+    info!("RPC server running on {}", addr);
     loop {
         let (stream, addr) = listener.accept().unwrap();
         info!("[{}] connected peer", addr);
-        match handle_client(stream, addr) {
+        let handler = Handler { query };
+        match handler.run(stream, addr) {
             Ok(()) => info!("[{}] disconnected peer", addr),
             Err(ref e) => {
                 error!("[{}] {}", addr, e);

tools/client.py

 import hashlib
+import json
 import sys
 
 from logbook import Logger, StreamHandler
@@ -7,7 +8,7 @@ from pycoin.coins.bitcoin.networks import BitcoinMainnet
 import pycoin.ui.key_from_text
 import pycoin.key
 
-import zmq
+import socket
 
 script_for_address = BitcoinMainnet.ui.script_for_address
 
@@ -15,9 +16,8 @@ log = Logger(__name__)
 
 
 def main():
-    c = zmq.Context()
-    s = c.socket(zmq.REQ)
-    s.connect('ipc:///tmp/indexrs.rpc')
+    s = socket.create_connection(('localhost', 50001))
+    f = s.makefile('r')
 
     xpub, = sys.argv[1:]
     total = 0
@@ -28,13 +28,18 @@ def main():
             address = k.subkey(change).subkey(n).address()
             script = script_for_address(address)
             script_hash = hashlib.sha256(script).digest()
-            s.send(script_hash)
-            res = s.recv()
-            b = float(res)
-            total += b
-            if b:
+            req = {
+                'id': 1,
+                'method': 'blockchain.scripthash.get_balance',
+                'params': [script_hash[::-1].hex()]
+            }
+            msg = json.dumps(req) + '\n'
+            s.sendall(msg.encode('ascii'))
+            res = json.loads(f.readline())['result']
+            total += res['confirmed']
+            if res['confirmed']:
                 log.info('{}/{} => {} has {:11.8f} BTC',
-                         change, n, address, b)
+                         change, n, address, res['confirmed'])
                 empty = 0
             else:
                 empty += 1