netapp/examples/basalt.rs

use std::io::Write;
use std::sync::Arc;
use std::time::Duration;

use log::{debug, info, warn};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use structopt::StructOpt;

use sodiumoxide::crypto::auth;
use sodiumoxide::crypto::sign::{ed25519, PublicKey};

use tokio::sync::watch;

use tokio_unix_tcp::NamedSocketAddr;

use netapp::endpoint::*;
use netapp::message::*;
use netapp::peering::basalt::*;
use netapp::util::parse_peer_addr;
use netapp::{NetApp, NodeID};

#[derive(StructOpt, Debug)]
#[structopt(name = "netapp")]
pub struct Opt {
	#[structopt(long = "network-key", short = "n")]
	network_key: Option<String>,

	#[structopt(long = "private-key", short = "p")]
	private_key: Option<String>,

	#[structopt(long = "bootstrap-peer", short = "b")]
	bootstrap_peers: Vec<String>,

	#[structopt(long = "listen-addr", short = "l", default_value = "127.0.0.1:1980")]
	listen_addr: String,

	#[structopt(long = "public-addr", short = "a")]
	public_addr: Option<String>,

	#[structopt(long = "view-size", short = "v", default_value = "100")]
	view_size: usize,

	#[structopt(long = "cache-size", short = "c", default_value = "1000")]
	cache_size: usize,

	#[structopt(long = "exchange-interval-secs", short = "x", default_value = "1")]
	exchange_interval: u64,

	#[structopt(long = "reset-interval-secs", short = "r", default_value = "10")]
	reset_interval: u64,

	#[structopt(long = "reset-count", short = "k", default_value = "20")]
	reset_count: usize,
}

struct Example {
	netapp: Arc<NetApp>,
	basalt: Arc<Basalt>,
	example_endpoint: Arc<Endpoint<ExampleMessage, Self>>,
}

#[tokio::main]
async fn main() {
	env_logger::Builder::new()
		.parse_env("RUST_LOG")
		.format(|buf, record| {
			writeln!(
				buf,
				"{} {} {} {}",
				chrono::Local::now().format("%s%.6f"),
				record.module_path().unwrap_or("_"),
				record.level(),
				record.args()
			)
		})
		.init();

	let opt = Opt::from_args();

	let netid = match &opt.network_key {
		Some(k) => auth::Key::from_slice(&hex::decode(k).unwrap()).unwrap(),
		None => auth::gen_key(),
	};
	info!("KYEV NK {}", hex::encode(&netid));

	let privkey = match &opt.private_key {
		Some(k) => ed25519::SecretKey::from_slice(&hex::decode(k).unwrap()).unwrap(),
		None => {
			let (_pk, sk) = ed25519::gen_keypair();
			sk
		}
	};

	info!("KYEV SK {}", hex::encode(&privkey));
	info!("KYEV PK {}", hex::encode(&privkey.public_key()));

	let netapp = NetApp::new(0u64, netid, privkey);

	let mut bootstrap_peers: Vec<(PublicKey, NamedSocketAddr)> = vec![];
	for peer in opt.bootstrap_peers.iter() {
		bootstrap_peers.push(
			parse_peer_addr(peer)
				.map(|(node_id, socket_bind_addr)| (node_id, socket_bind_addr.into()))
				.expect("Invalid peer address"),
		);
	}

	let basalt_params = BasaltParams {
		view_size: opt.view_size,
		cache_size: opt.cache_size,
		exchange_interval: Duration::from_secs(opt.exchange_interval),
		reset_interval: Duration::from_secs(opt.reset_interval),
		reset_count: opt.reset_count,
	};
	let basalt = Basalt::new(netapp.clone(), bootstrap_peers, basalt_params);

	let example = Arc::new(Example {
		netapp: netapp.clone(),
		basalt,
		example_endpoint: netapp.endpoint("__netapp/examples/basalt.rs/Example".into()),
	});
	example.example_endpoint.set_handler(example.clone());

	let listen_addr = opt.listen_addr.parse().unwrap();
	let public_addr = opt.public_addr.map(|x| x.parse().unwrap());

	let watch_cancel = netapp::util::watch_ctrl_c();

	tokio::join!(
		example.clone().sampling_loop(watch_cancel.clone()),
		example
			.netapp
			.clone()
			.listen(listen_addr, public_addr, watch_cancel.clone()),
		example.basalt.clone().run(watch_cancel.clone()),
	);
}

impl Example {
	async fn sampling_loop(self: Arc<Self>, must_exit: watch::Receiver<bool>) {
		while !*must_exit.borrow() {
			tokio::time::sleep(Duration::from_secs(10)).await;

			let peers = self.basalt.sample(10);
			for p in peers {
				debug!("kyev S {}", hex::encode(p));

				let self2 = self.clone();
				tokio::spawn(async move {
					match self2
						.example_endpoint
						.call(&p, ExampleMessage { example_field: 42 }, PRIO_NORMAL)
						.await
					{
						Ok(resp) => debug!("Got example response: {:?}", resp),
						Err(e) => warn!("Error with example request: {}", e),
					}
				});
			}
		}
	}
}

#[async_trait]
impl EndpointHandler<ExampleMessage> for Example {
	async fn handle(self: &Arc<Self>, msg: &ExampleMessage, _from: NodeID) -> ExampleResponse {
		debug!("Got example message: {:?}, sending example response", msg);
		ExampleResponse {
			example_field: false,
		}
	}
}

#[derive(Serialize, Deserialize, Debug)]
struct ExampleMessage {
	example_field: usize,
}

#[derive(Serialize, Deserialize, Debug)]
struct ExampleResponse {
	example_field: bool,
}

impl Message for ExampleMessage {
	type Response = ExampleResponse;
}