use crate::address::Address;
use crate::client::{ClientBuilder, Response};
use crate::dht::{Dht, DhtMetrics};
use crate::error::{Error, Result};
use crate::identity::NodeId;
use crate::poll::{wait_readable, wait_readable2, Waker};
use crate::server::ServerBuilder;
use crate::wire::Method;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
use std::thread::JoinHandle;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{mpsc, oneshot};
const POLL_CAP_MS: u32 = 200;
const ANNOUNCE_INTERVAL_SECS: u64 = 30;
fn now_ms() -> i64 {
static START: OnceLock<Instant> = OnceLock::new();
START.get_or_init(Instant::now).elapsed().as_millis() as i64
}
fn now_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
}
enum ServerCommand {
Resolve {
target: NodeId,
timeout: Duration,
reply: oneshot::Sender<Result<Address>>,
},
DhtMetrics {
reply: oneshot::Sender<DhtMetrics>,
},
}
pub struct RunningServer {
node_id: NodeId,
port: u16,
stop: Arc<AtomicBool>,
join: Option<JoinHandle<()>>,
cmd: Option<mpsc::Sender<ServerCommand>>,
}
impl RunningServer {
pub fn node_id(&self) -> NodeId {
self.node_id
}
pub fn local_port(&self) -> u16 {
self.port
}
pub async fn resolve(&self, target: &NodeId, timeout: Duration) -> Result<Address> {
let cmd = self.cmd.as_ref().ok_or(Error::ConfigMissing)?;
let (reply, rx) = oneshot::channel();
cmd.send(ServerCommand::Resolve {
target: *target,
timeout,
reply,
})
.await
.map_err(|_| Error::NetworkClosed)?;
rx.await.map_err(|_| Error::NetworkClosed)?
}
pub async fn dht_metrics(&self) -> Option<DhtMetrics> {
let cmd = self.cmd.as_ref()?;
let (reply, rx) = oneshot::channel();
cmd.send(ServerCommand::DhtMetrics { reply }).await.ok()?;
rx.await.ok()
}
pub fn shutdown(mut self) {
self.stop.store(true, Ordering::Relaxed);
if let Some(join) = self.join.take() {
let _ = join.join();
}
}
}
impl Drop for RunningServer {
fn drop(&mut self) {
self.stop.store(true, Ordering::Relaxed);
}
}
impl ServerBuilder {
pub async fn serve(mut self) -> Result<RunningServer> {
let (ready_tx, ready_rx) = oneshot::channel::<Result<(NodeId, u16)>>();
let stop = Arc::new(AtomicBool::new(false));
let stop_for_thread = stop.clone();
let dht_cfg = self.take_managed_dht();
let has_dht = dht_cfg.is_some();
let (cmd_tx, cmd_rx) = mpsc::channel::<ServerCommand>(32);
let join = std::thread::Builder::new()
.name("nwep-server".into())
.spawn(move || {
let server = match self.build() {
Ok(s) => s,
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
};
let mut announce = None;
let dht = match dht_cfg {
Some((contacts, ann, seq)) => {
announce = ann;
match attach_managed_dht(&server, contacts, seq) {
Ok(d) => Some(d),
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
}
}
None => None,
};
let info = server.node_id().map(|n| (n, server.local_port()));
let info = match info {
Ok(i) => i,
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
};
if ready_tx.send(Ok(info)).is_err() {
return; }
run_server_loop(&server, dht.as_ref(), &announce, cmd_rx, &stop_for_thread);
})
.map_err(|_| Error::Internal)?;
match ready_rx.await {
Ok(Ok((node_id, port))) => Ok(RunningServer {
node_id,
port,
stop,
join: Some(join),
cmd: has_dht.then_some(cmd_tx),
}),
Ok(Err(e)) => {
let _ = join.join();
Err(e)
}
Err(_) => {
let _ = join.join();
Err(Error::Internal)
}
}
}
}
fn attach_managed_dht(
server: &crate::Server,
contacts: Vec<crate::Bootstrap>,
initial_seq: u64,
) -> Result<Dht<'_>> {
let dht = Dht::builder(server)
.bootstraps(contacts)
.initial_seq(initial_seq)
.attach()?;
dht.join(now_secs())?;
Ok(dht)
}
fn run_server_loop(
server: &crate::Server,
dht: Option<&Dht<'_>>,
announce: &Option<Address>,
mut cmd_rx: mpsc::Receiver<ServerCommand>,
stop: &AtomicBool,
) {
let mut announced_at = 0u64;
while !stop.load(Ordering::Relaxed) {
if server.tick(now_ms()).is_err() {
break;
}
if let Some(dht) = dht {
let secs = now_secs();
let _ = dht.tick(secs);
if let Some(addr) = announce {
if announced_at == 0 || secs - announced_at >= ANNOUNCE_INTERVAL_SECS {
let _ = dht.announce(addr, secs);
announced_at = secs;
}
}
while let Ok(cmd) = cmd_rx.try_recv() {
match cmd {
ServerCommand::Resolve {
target,
timeout,
reply,
} => {
let _ = reply.send(resolve_on_thread(server, dht, &target, timeout));
}
ServerCommand::DhtMetrics { reply } => {
let _ = reply.send(dht.metrics());
}
}
}
}
let mut wait = server.next_timeout(now_ms()).unwrap_or(POLL_CAP_MS);
if let Some(dht) = dht {
if let Some(d) = dht.next_timeout(now_secs()) {
wait = wait.min(d);
}
}
wait_readable(server.fd(), wait.min(POLL_CAP_MS));
}
}
fn resolve_on_thread(
server: &crate::Server,
dht: &Dht<'_>,
target: &NodeId,
timeout: Duration,
) -> Result<Address> {
if let Some(rec) = dht.lookup_result(target) {
return Ok(rec.address());
}
dht.start_lookup(target, now_secs())?;
let deadline = Instant::now() + timeout;
while Instant::now() < deadline {
server.tick(now_ms())?;
let secs = now_secs();
dht.tick(secs)?;
if let Some(rec) = dht.lookup_result(target) {
return Ok(rec.address());
}
let wait = dht.next_timeout(secs).unwrap_or(20).min(20);
wait_readable(server.fd(), wait);
}
Err(Error::IdentityNotFound)
}
enum ClientCommand {
Send {
method: Method,
path: String,
headers: Vec<(String, String)>,
body: Vec<u8>,
reply: oneshot::Sender<Result<Response>>,
},
}
pub struct AsyncClient {
cmd: Option<mpsc::Sender<ClientCommand>>,
wake: Arc<Waker>,
_join: JoinHandle<()>,
}
impl ClientBuilder {
pub async fn connect_async(self, target: &NodeId, addr: &Address) -> Result<AsyncClient> {
let target = *target;
let addr = *addr;
let wake = Arc::new(Waker::new().map_err(|_| Error::Internal)?);
let wake_owner = wake.clone();
let (ready_tx, ready_rx) = oneshot::channel::<Result<()>>();
let (cmd_tx, cmd_rx) = mpsc::channel::<ClientCommand>(64);
let join = std::thread::Builder::new()
.name("nwep-client".into())
.spawn(move || {
let client = match self.connect(&target, &addr) {
Ok(c) => c,
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
};
if ready_tx.send(Ok(())).is_err() {
return;
}
run_client_loop(&client, cmd_rx, &wake_owner);
})
.map_err(|_| Error::Internal)?;
match ready_rx.await {
Ok(Ok(())) => Ok(AsyncClient {
cmd: Some(cmd_tx),
wake,
_join: join,
}),
Ok(Err(e)) => {
let _ = join.join();
Err(e)
}
Err(_) => {
let _ = join.join();
Err(Error::Internal)
}
}
}
}
impl Drop for AsyncClient {
fn drop(&mut self) {
self.cmd.take();
self.wake.wake();
}
}
fn run_client_loop(
client: &crate::Client,
mut cmd_rx: mpsc::Receiver<ClientCommand>,
wake: &Waker,
) {
use tokio::sync::mpsc::error::TryRecvError;
let cfd = client.fd();
let wfd = wake.raw();
let mut pending: Vec<(crate::RequestId, oneshot::Sender<Result<Response>>)> = Vec::new();
loop {
loop {
match cmd_rx.try_recv() {
Ok(ClientCommand::Send {
method,
path,
headers,
body,
reply,
}) => match client.submit_request(method, &path, &headers, &body) {
Ok(id) => pending.push((id, reply)),
Err(e) => {
let _ = reply.send(Err(e));
}
},
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => {
for (_, reply) in pending.drain(..) {
let _ = reply.send(Err(Error::NetworkClosed));
}
return;
}
}
}
if client.tick(now_ms()).is_err() || !client.is_alive() {
for (_, reply) in pending.drain(..) {
let _ = reply.send(Err(Error::NetworkClosed));
}
return;
}
let mut i = 0;
while i < pending.len() {
match client.poll_request(pending[i].0) {
Ok(Some(resp)) => {
let (_, reply) = pending.swap_remove(i);
let _ = reply.send(Ok(resp));
}
Ok(None) => i += 1,
Err(e) => {
let (_, reply) = pending.swap_remove(i);
let _ = reply.send(Err(e));
}
}
}
let wait = if pending.is_empty() {
POLL_CAP_MS
} else {
client
.next_timeout(now_ms())
.unwrap_or(POLL_CAP_MS)
.min(POLL_CAP_MS)
};
wait_readable2(cfd, wfd, wait);
wake.drain();
}
}
impl AsyncClient {
pub fn request(&self, method: Method, path: &str) -> AsyncRequestBuilder<'_> {
AsyncRequestBuilder {
client: self,
method,
path: path.to_owned(),
headers: Vec::new(),
body: Vec::new(),
}
}
pub async fn send(&self, method: Method, path: &str, body: &[u8]) -> Result<Response> {
self.request(method, path).body(body).send().await
}
async fn dispatch(
&self,
method: Method,
path: String,
headers: Vec<(String, String)>,
body: Vec<u8>,
) -> Result<Response> {
let cmd = self.cmd.as_ref().ok_or(Error::NetworkClosed)?;
let (reply_tx, reply_rx) = oneshot::channel();
cmd.send(ClientCommand::Send {
method,
path,
headers,
body,
reply: reply_tx,
})
.await
.map_err(|_| Error::NetworkClosed)?;
self.wake.wake();
reply_rx.await.map_err(|_| Error::NetworkClosed)?
}
}
pub struct AsyncRequestBuilder<'c> {
client: &'c AsyncClient,
method: Method,
path: String,
headers: Vec<(String, String)>,
body: Vec<u8>,
}
impl AsyncRequestBuilder<'_> {
pub fn header(mut self, name: &str, value: &str) -> Self {
self.headers.push((name.to_owned(), value.to_owned()));
self
}
pub fn body(mut self, body: impl Into<Vec<u8>>) -> Self {
self.body = body.into();
self
}
pub async fn send(self) -> Result<Response> {
self.client
.dispatch(self.method, self.path, self.headers, self.body)
.await
}
}
type StreamItem = Result<Option<Vec<u8>>>;
const STREAM_CHUNK: usize = 64 * 1024;
pub struct AsyncStream {
status: crate::Status,
headers: Vec<(String, String)>,
chunks: mpsc::Receiver<StreamItem>,
_join: JoinHandle<()>,
}
impl AsyncStream {
pub fn status(&self) -> crate::Status {
self.status
}
pub fn header(&self, name: &str) -> Option<&str> {
self.headers
.iter()
.find(|(n, _)| n == name)
.map(|(_, v)| v.as_str())
}
pub fn headers(&self) -> impl Iterator<Item = (&str, &str)> {
self.headers.iter().map(|(n, v)| (n.as_str(), v.as_str()))
}
pub async fn recv(&mut self) -> Result<Option<Vec<u8>>> {
match self.chunks.recv().await {
Some(item) => item,
None => Ok(None), }
}
}
impl ClientBuilder {
pub async fn stream(
self,
method: Method,
path: &str,
target: &NodeId,
addr: &Address,
) -> Result<AsyncStream> {
let target = *target;
let addr = *addr;
let path = path.to_owned();
let (ready_tx, ready_rx) =
oneshot::channel::<Result<(crate::Status, Vec<(String, String)>)>>();
let (chunk_tx, chunk_rx) = mpsc::channel::<StreamItem>(8);
let join = std::thread::Builder::new()
.name("nwep-stream".into())
.spawn(move || {
let client = match self.connect(&target, &addr) {
Ok(c) => c,
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
};
let stream = match client.open_stream(method, &path) {
Ok(s) => s,
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
};
let meta = match stream.response() {
Ok(m) => m,
Err(e) => {
let _ = ready_tx.send(Err(e));
return;
}
};
let status = meta.status().unwrap_or(crate::Status::Error);
let headers: Vec<(String, String)> = meta
.headers()
.map(|(n, v)| (n.to_owned(), v.to_owned()))
.collect();
if ready_tx.send(Ok((status, headers))).is_err() {
return; }
let peer = client.peer_pubkey().unwrap_or([0u8; 32]);
let mut buf = vec![0u8; STREAM_CHUNK];
loop {
match stream.recv(&mut buf) {
Ok((n, ended)) => {
if n > 0 && chunk_tx.blocking_send(Ok(Some(buf[..n].to_vec()))).is_err()
{
return; }
if ended {
let end = match stream.verify(&peer) {
Ok(()) => Ok(None),
Err(e) => Err(e),
};
let _ = chunk_tx.blocking_send(end);
return;
}
}
Err(e) => {
let _ = chunk_tx.blocking_send(Err(e));
return;
}
}
}
})
.map_err(|_| Error::Internal)?;
match ready_rx.await {
Ok(Ok((status, headers))) => Ok(AsyncStream {
status,
headers,
chunks: chunk_rx,
_join: join,
}),
Ok(Err(e)) => {
let _ = join.join();
Err(e)
}
Err(_) => {
let _ = join.join();
Err(Error::Internal)
}
}
}
}