Kademlia

extern crate clap;
extern crate libp2p;
extern crate libp2p_swarm;
extern crate tokio;
extern crate colored;

use clap::{Arg, App};
use libp2p::{Multiaddr, PeerId};
use libp2p::kad::Kademlia;
use libp2p::kad::record::store::MemoryStore;
use libp2p::identity::{ed25519, Keypair, PublicKey};
use libp2p::Swarm;
use tokio::signal;
use std::path::Path;
use std::fs;
use colored::*;
use chrono::Local;
use std::io::{stdout, Write};

fn main() {
    // コマンドラインパラメータの解析
    let matches = App::new("Kadmlia-based Simple KVS").version("1.0").author("TAKAMI Torao <koiroha@gmail.com>")
        .about("Simple key-value store implementation using the Kademlia DHT of libp2p.")
        .arg(Arg::with_name("ID")
            .help("numeric node id")
            .required(true))
        .arg(Arg::with_name("daemon")
            .short("d")
            .long("daemon")
            .help("use when no input from stdin is accepted.")
            .required(false)
            .takes_value(false))
        .arg(Arg::with_name("bootstrap")
            .short("b")
            .long("bootstrap")
            .value_name("PEER_ID=NODE_ADDRESS")
            .help("bootstrap node address (e.g., 12D3KooWCzsfEkYrFf2djHem3qdagCjV4M6EXkcPUQzP9Xq76EPV=/ip4/127.0.0.1/tcp/6200)")
            .takes_value(true))
        .get_matches();

    // ノード ID を取得し公開鍵ペアをロード
    let id:u32 = matches.value_of("ID").map(|s| s.parse().unwrap()).unwrap_or(0);
    let key_pair = load_keypair(id);

    let peer_id = PeerId::from(PublicKey::Ed25519(key_pair.public()));
    let kv_store = MemoryStore::new(peer_id.clone());
    let mut kademlia = Kademlia::new(peer_id.clone(), kv_store);
    log(id, format!("Peer ID: {}", peer_id));
    log(id, format!("Protocol Name: {}", String::from_utf8(kademlia.protocol_name().to_vec()).unwrap()));

    let bootstrap:Option<u32> = matches.value_of("bootstrap").map(|s| s.parse().unwrap());
    if let Some(bs_id) = bootstrap {
        let bs_key_pair = load_keypair(bs_id);
        let bs_peer_id = PeerId::from(PublicKey::Ed25519(bs_key_pair.public()));
        kademlia.add_address(&bs_peer_id, address(bs_id));
        log(id, format!("Bootstrap: {}", bs_id));
    }

    let bind_address = address(id);
    log(id, format!("Bind Address: {}", bind_address.to_string()));
    let transport = libp2p::build_development_transport(Keypair::Ed25519(key_pair)).unwrap();
    let mut swarm = Swarm::new(transport, kademlia, peer_id);
    Swarm::listen_on(&mut swarm, bind_address).unwrap();

    if matches.is_present("daemon") {
        let ctrl_c = signal::ctrl_c();
        let mut rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            log(id, "To exit, press CTRL+C.".to_string());
            ctrl_c.await.unwrap();
            log(id, "\nshutting-down".to_string());
            0
        });
    } else {
        loop {
            let cmd:String = prompt(id, "");
            log(id, cmd);
        }
    }
}

fn load_keypair(id:u32) -> ed25519::Keypair {
    let path = format!("node/{}/", id);
    let dir = Path::new(&path);
    if !dir.is_dir() {
        fs::create_dir_all(dir).unwrap();
    }
    let file = dir.join("key");
    if file.is_file() {
        let mut binary = fs::read(file).unwrap();
        ed25519::Keypair::decode(&mut binary).unwrap()
    } else {
        let key_pair = ed25519::Keypair::generate();
        let binary = key_pair.encode();
        fs::write(file, binary.to_vec()).unwrap();
        key_pair
    }
}

fn address(id:u32) -> Multiaddr {
    let bind_address = format!("/ip4/127.0.0.1/tcp/{}", 6200 + id);
    bind_address.parse().unwrap()
}

fn prompt<T: std::str::FromStr>(id:u32, msg:&str) -> T {
    let color = color(id);
    print!("{}> ", msg.color(color));
    stdout().flush().unwrap();
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).unwrap();
    s.trim().parse().ok().unwrap()
}

fn color(id:u32) -> Color {
    let colors = [Color::BrightRed, Color::BrightGreen, Color::BrightYellow, Color::BrightBlue, Color::BrightMagenta,
        Color::BrightCyan];
    colors[id as usize % colors.len()]
}

fn log(id:u32, msg:String) {
    let color = color(id);
    let tm = Local::now().format("%Y-%m-%d %H:%M:%S%.3f").to_string();
    println!("{} {}", tm.bright_black(), format!("Node{} {}", id, msg).color(color));
}