rqbit/crates/dht/src/dht.rs
2023-04-23 17:55:52 +02:00

702 lines
24 KiB
Rust

use std::{
collections::{hash_map::Entry, HashMap, HashSet},
net::SocketAddr,
sync::Arc,
task::Poll,
};
use crate::{
bprotocol::{
self, CompactNodeInfo, CompactPeerInfo, FindNodeRequest, GetPeersRequest, Message,
MessageKind, Node,
},
routing_table::{InsertResult, RoutingTable},
};
use anyhow::Context;
use bencode::ByteString;
use futures::{stream::FuturesUnordered, Stream, StreamExt};
use indexmap::IndexSet;
use librqbit_core::{id20::Id20, peer_id::generate_peer_id};
use log::{debug, info, trace, warn};
use parking_lot::RwLock;
use rand::Rng;
use serde::Serialize;
use tokio::{
net::UdpSocket,
sync::mpsc::{channel, unbounded_channel, Sender, UnboundedReceiver, UnboundedSender},
};
use tokio_stream::wrappers::{errors::BroadcastStreamRecvError, BroadcastStream};
#[derive(Debug, Serialize)]
pub struct DhtStats {
#[serde(serialize_with = "crate::utils::serialize_id20")]
pub id: Id20,
pub outstanding_requests: usize,
pub seen_peers: usize,
pub made_requests: usize,
pub routing_table_size: usize,
}
struct DhtState {
id: Id20,
next_transaction_id: u16,
outstanding_requests: HashMap<(u16, SocketAddr), Request>,
routing_table: RoutingTable,
listen_addr: SocketAddr,
// This sender sends requests to the worker.
// It is unbounded so that the methods on Dht state don't need to be async.
// If the methods on Dht state were async, we would have a problem, as it's behind
// a lock.
// Alternatively, we can lock only the parts that change, and use that internally inside DhtState...
sender: UnboundedSender<(Message<ByteString>, SocketAddr)>,
seen_peers: HashMap<Id20, IndexSet<SocketAddr>>,
get_peers_subscribers: HashMap<Id20, tokio::sync::broadcast::Sender<SocketAddr>>,
made_requests: HashSet<(Request, SocketAddr)>,
}
impl DhtState {
fn new(
id: Id20,
sender: UnboundedSender<(Message<ByteString>, SocketAddr)>,
routing_table: Option<RoutingTable>,
listen_addr: SocketAddr,
) -> Self {
let routing_table = routing_table.unwrap_or_else(|| RoutingTable::new(id));
Self {
id,
next_transaction_id: 0,
outstanding_requests: Default::default(),
routing_table,
sender,
listen_addr,
seen_peers: Default::default(),
get_peers_subscribers: Default::default(),
made_requests: Default::default(),
}
}
fn create_request(&mut self, request: Request, addr: SocketAddr) -> Message<ByteString> {
let transaction_id = self.next_transaction_id;
let transaction_id_buf = [(transaction_id >> 8) as u8, (transaction_id & 0xff) as u8];
self.next_transaction_id = if transaction_id == u16::MAX {
0
} else {
transaction_id + 1
};
let message = match request {
Request::GetPeers(info_hash) => Message {
transaction_id: ByteString::from(transaction_id_buf.as_ref()),
version: None,
ip: None,
kind: MessageKind::GetPeersRequest(GetPeersRequest {
id: self.id,
info_hash,
}),
},
Request::FindNode(target) => Message {
transaction_id: ByteString::from(transaction_id_buf.as_ref()),
version: None,
ip: None,
kind: MessageKind::FindNodeRequest(FindNodeRequest {
id: self.id,
target,
}),
},
};
self.outstanding_requests
.insert((transaction_id, addr), request);
message
}
fn on_incoming_from_remote(
&mut self,
msg: Message<ByteString>,
addr: SocketAddr,
) -> anyhow::Result<()> {
let generate_compact_nodes = |target| {
let nodes = self
.routing_table
.sorted_by_distance_from(target)
.into_iter()
.filter_map(|r| {
Some(Node {
id: r.id(),
addr: match r.addr() {
SocketAddr::V4(v4) => v4,
SocketAddr::V6(_) => return None,
},
})
})
.take(8)
.collect::<Vec<_>>();
CompactNodeInfo { nodes }
};
match &msg.kind {
MessageKind::Error(_) | MessageKind::Response(_) => {
if msg.transaction_id.len() != 2 {
anyhow::bail!(
"{}: transaction id unrecognized, expected its length == 2. Message: {:?}",
addr,
msg
)
}
let tid = ((msg.transaction_id[0] as u16) << 8) + (msg.transaction_id[1] as u16);
let request = match self.outstanding_requests.remove(&(tid, addr)) {
Some(req) => req,
None => anyhow::bail!("outstanding request not found. Message: {:?}", msg),
};
let response = match msg.kind {
MessageKind::Error(e) => {
anyhow::bail!("request {:?} received error response {:?}", request, e)
}
MessageKind::Response(r) => r,
_ => unreachable!(),
};
self.routing_table.mark_response(&response.id);
match request {
Request::FindNode(id) => {
let nodes = response.nodes.ok_or_else(|| {
anyhow::anyhow!("expected nodes for find_node requests")
})?;
self.on_found_nodes(response.id, addr, id, nodes)
}
Request::GetPeers(id) => {
self.on_found_peers_or_nodes(response.id, addr, id, response)
}
}
}
MessageKind::PingRequest(_) => {
let message = Message {
transaction_id: msg.transaction_id,
version: None,
ip: None,
kind: MessageKind::Response(bprotocol::Response {
id: self.id,
..Default::default()
}),
};
self.sender.send((message, addr))?;
Ok(())
}
MessageKind::GetPeersRequest(req) => {
let peers = self.seen_peers.get(&req.info_hash).map(|peers| {
peers
.iter()
.copied()
.filter_map(|a| match a {
SocketAddr::V4(v4) => Some(CompactPeerInfo { addr: v4 }),
// this should never happen in practice
SocketAddr::V6(_) => None,
})
.take(50)
.collect::<Vec<_>>()
});
let token = if peers.is_some() {
let mut token = [0u8; 20];
rand::thread_rng().fill(&mut token);
Some(ByteString::from(token.as_ref()))
} else {
None
};
let compact_node_info = generate_compact_nodes(req.info_hash);
let message = Message {
transaction_id: msg.transaction_id,
version: None,
ip: None,
kind: MessageKind::Response(bprotocol::Response {
id: self.id,
nodes: Some(compact_node_info),
values: peers,
token,
}),
};
self.sender.send((message, addr))?;
Ok(())
}
MessageKind::FindNodeRequest(req) => {
let compact_node_info = generate_compact_nodes(req.target);
let message = Message {
transaction_id: msg.transaction_id,
version: None,
ip: None,
kind: MessageKind::Response(bprotocol::Response {
id: self.id,
nodes: Some(compact_node_info),
..Default::default()
}),
};
self.sender.send((message, addr))?;
Ok(())
}
}
}
pub fn get_stats(&self) -> DhtStats {
DhtStats {
id: self.id,
outstanding_requests: self.outstanding_requests.len(),
seen_peers: self.seen_peers.values().map(|v| v.len()).sum(),
made_requests: self.made_requests.len(),
routing_table_size: self.routing_table.len(),
}
}
#[allow(clippy::type_complexity)]
fn get_peers(
&mut self,
info_hash: Id20,
) -> anyhow::Result<(
Option<(usize, usize)>,
tokio::sync::broadcast::Receiver<SocketAddr>,
)> {
match self.get_peers_subscribers.entry(info_hash) {
Entry::Occupied(o) => {
let pos = self.seen_peers.get(&info_hash).and_then(|p| {
if p.is_empty() {
None
} else {
Some((0, p.len()))
}
});
let rx = o.get().subscribe();
Ok((pos, rx))
}
Entry::Vacant(v) => {
// DHT sends peers REALLY fast, so ideally the consumer of this broadcast should not lag behind.
// In case it does though we have PeerStream to replay.
let (tx, rx) = tokio::sync::broadcast::channel(100);
v.insert(tx);
// We don't need to allocate/collect here, but the borrow checker is not happy otherwise.
let nodes_to_query = self
.routing_table
.sorted_by_distance_from(info_hash)
.iter()
.map(|n| (n.id(), n.addr()))
.take(8)
.collect::<Vec<_>>();
for (id, addr) in nodes_to_query {
self.send_find_peers_if_not_yet(info_hash, id, addr)?;
}
Ok((None, rx))
}
}
}
fn send_find_peers_if_not_yet(
&mut self,
info_hash: Id20,
target_node: Id20,
addr: SocketAddr,
) -> anyhow::Result<()> {
let request = Request::GetPeers(info_hash);
if self.made_requests.insert((request, addr)) {
self.routing_table.mark_outgoing_request(&target_node);
let msg = self.create_request(request, addr);
self.sender.send((msg, addr))?;
}
Ok(())
}
fn send_find_node_if_not_yet(
&mut self,
search_id: Id20,
target_node: Id20,
addr: SocketAddr,
) -> anyhow::Result<()> {
let request = Request::FindNode(search_id);
if self.made_requests.insert((request, addr)) {
self.routing_table.mark_outgoing_request(&target_node);
let msg = self.create_request(request, addr);
self.sender.send((msg, addr))?;
}
Ok(())
}
fn on_found_nodes(
&mut self,
source: Id20,
source_addr: SocketAddr,
target: Id20,
nodes: CompactNodeInfo,
) -> anyhow::Result<()> {
// We don't need to allocate/collect here, but the borrow checker is not happy
// otherwise when we iterate self.searching_for_peers and mutating self in the loop.
let searching_for_peers = self
.get_peers_subscribers
.keys()
.copied()
.collect::<Vec<_>>();
// On newly discovered nodes, ask them for peers that we are interested in.
match self.routing_table.add_node(source, source_addr) {
InsertResult::ReplacedBad(_) | InsertResult::Added => {
for info_hash in &searching_for_peers {
self.send_find_peers_if_not_yet(*info_hash, source, source_addr)?;
}
}
_ => {}
};
for node in nodes.nodes {
match self.routing_table.add_node(node.id, node.addr.into()) {
InsertResult::ReplacedBad(_) | InsertResult::Added => {
for info_hash in &searching_for_peers {
self.send_find_peers_if_not_yet(*info_hash, node.id, node.addr.into())?;
}
// recursively find nodes closest to us until we can't find more.
self.send_find_node_if_not_yet(target, source, source_addr)?;
}
_ => {}
};
}
Ok(())
}
fn on_found_peers_or_nodes(
&mut self,
source: Id20,
source_addr: SocketAddr,
target: Id20,
data: bprotocol::Response<ByteString>,
) -> anyhow::Result<()> {
self.routing_table.add_node(source, source_addr);
self.routing_table.mark_response(&source);
let bsender = match self.get_peers_subscribers.get(&target) {
Some(s) => s,
None => {
warn!(
"ignoring get_peers response, no subscribers for {:?}",
target
);
return Ok(());
}
};
if let Some(peers) = data.values {
let seen = self.seen_peers.entry(target).or_default();
for peer in peers.iter() {
if peer.addr.port() < 1024 {
debug!("bad peer port, ignoring: {}", peer.addr);
continue;
}
let addr = SocketAddr::V4(peer.addr);
if seen.insert(addr) {
bsender
.send(addr)
.context("error sending peers to subscribers")?;
}
}
};
if let Some(nodes) = data.nodes {
for node in nodes.nodes {
self.routing_table.add_node(node.id, node.addr.into());
self.send_find_peers_if_not_yet(target, node.id, node.addr.into())?;
}
};
Ok(())
}
}
async fn run_framer(
socket: &UdpSocket,
mut input_rx: UnboundedReceiver<(Message<ByteString>, SocketAddr)>,
output_tx: Sender<(Message<ByteString>, SocketAddr)>,
) -> anyhow::Result<()> {
let writer = async {
let mut buf = Vec::new();
while let Some((msg, addr)) = input_rx.recv().await {
let addr = match addr {
SocketAddr::V4(v4) => v4,
SocketAddr::V6(_) => continue,
};
trace!("{}: sending {:?}", addr, &msg);
buf.clear();
bprotocol::serialize_message(
&mut buf,
msg.transaction_id,
msg.version,
msg.ip,
msg.kind,
)
.unwrap();
if let Err(e) = socket.send_to(&buf, addr).await {
warn!("could not send to {:?}: {}", addr, e)
}
}
Err::<(), _>(anyhow::anyhow!(
"DHT UDP socket writer over, nowhere to read messages from"
))
};
let reader = async {
let mut buf = vec![0u8; 16384];
loop {
let (size, addr) = socket
.recv_from(&mut buf)
.await
.context("error reading from UDP socket")?;
match bprotocol::deserialize_message::<ByteString>(&buf[..size]) {
Ok(msg) => {
trace!("{}: received {:?}", addr, &msg);
match output_tx.send((msg, addr)).await {
Ok(_) => {}
Err(_) => break,
}
}
Err(e) => log::debug!("{}: error deserializing incoming message: {}", addr, e),
}
}
Err::<(), _>(anyhow::anyhow!(
"DHT UDP socket reader over, nowhere to send responses to"
))
};
let result = tokio::select! {
err = writer => err,
err = reader => err,
};
result.context("DHT UDP framer closed")
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Request {
GetPeers(Id20),
FindNode(Id20),
}
#[derive(Clone)]
pub struct Dht {
state: Arc<RwLock<DhtState>>,
}
struct DhtWorker {
socket: UdpSocket,
peer_id: Id20,
state: Arc<RwLock<DhtState>>,
}
impl DhtWorker {
fn on_response(&self, msg: Message<ByteString>, addr: SocketAddr) -> anyhow::Result<()> {
self.state.write().on_incoming_from_remote(msg, addr)
}
async fn start(
self,
in_tx: UnboundedSender<(Message<ByteString>, SocketAddr)>,
in_rx: UnboundedReceiver<(Message<ByteString>, SocketAddr)>,
bootstrap_addrs: &[String],
) -> anyhow::Result<()> {
let (out_tx, mut out_rx) = channel(1);
let framer = run_framer(&self.socket, in_rx, out_tx);
let bootstrap = async {
let mut futs = FuturesUnordered::new();
// bootstrap
for addr in bootstrap_addrs.iter() {
let addr = addr;
let this = &self;
let in_tx = &in_tx;
futs.push(async move {
match tokio::net::lookup_host(addr).await {
Ok(addrs) => {
for addr in addrs {
let request = this
.state
.write()
.create_request(Request::FindNode(this.peer_id), addr);
in_tx.send((request, addr))?;
}
}
Err(e) => warn!("error looking up {}: {}", addr, e),
}
Ok::<_, anyhow::Error>(())
});
}
let mut successes = 0;
while let Some(resp) = futs.next().await {
match resp {
Ok(_) => successes += 1,
Err(e) => warn!("error in one of the bootstrappers: {}", e),
}
}
if successes == 0 {
anyhow::bail!("bootstrapping did not succeed")
}
Ok(())
};
let mut bootstrap_done = false;
let response_reader = {
let this = &self;
async move {
while let Some((response, addr)) = out_rx.recv().await {
if let Err(e) = this.on_response(response, addr) {
debug!("error in on_response, addr={:?}: {}", addr, e)
}
}
Err::<(), _>(anyhow::anyhow!(
"closed response reader, nowhere to send results to, DHT closed"
))
}
};
tokio::pin!(framer);
tokio::pin!(bootstrap);
tokio::pin!(response_reader);
loop {
tokio::select! {
err = &mut framer => {
anyhow::bail!("framer quit: {:?}", err)
},
result = &mut bootstrap, if !bootstrap_done => {
bootstrap_done = true;
result?;
},
err = &mut response_reader => {anyhow::bail!("response reader quit: {:?}", err)}
}
}
}
}
struct PeerStream {
info_hash: Id20,
state: Arc<RwLock<DhtState>>,
absolute_stream_pos: usize,
initial_peers_pos: Option<(usize, usize)>,
broadcast_rx: BroadcastStream<SocketAddr>,
}
impl Stream for PeerStream {
type Item = SocketAddr;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
loop {
if let Some((pos, end)) = self.initial_peers_pos.take() {
let addr = *self
.state
.read()
.seen_peers
.get(&self.info_hash)
.unwrap()
.get_index(pos)
.unwrap();
if pos + 1 < end {
self.initial_peers_pos = Some((pos + 1, end));
}
self.absolute_stream_pos += 1;
return Poll::Ready(Some(addr));
}
match self.broadcast_rx.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(v))) => {
self.absolute_stream_pos += 1;
return Poll::Ready(Some(v));
}
Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(lagged_by)))) => {
debug!("peer stream is lagged by {}", lagged_by);
let s = self.absolute_stream_pos;
let e = s + lagged_by as usize;
self.initial_peers_pos = Some((s, e));
continue;
}
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => return Poll::Pending,
};
}
}
}
#[derive(Default)]
pub struct DhtConfig {
pub peer_id: Option<Id20>,
pub bootstrap_addrs: Option<Vec<String>>,
pub routing_table: Option<RoutingTable>,
pub listen_addr: Option<SocketAddr>,
}
impl Dht {
pub async fn new() -> anyhow::Result<Self> {
Self::with_config(DhtConfig::default()).await
}
pub async fn with_config(config: DhtConfig) -> anyhow::Result<Self> {
let socket = match config.listen_addr {
Some(addr) => UdpSocket::bind(addr)
.await
.with_context(|| format!("error binding socket, address {addr}")),
None => UdpSocket::bind("0.0.0.0:0")
.await
.context("error binding socket, address 0.0.0.0:0"),
}?;
let listen_addr = socket
.local_addr()
.context("cannot determine UDP listen addr")?;
info!("DHT listening on {:?}", listen_addr);
let peer_id = config.peer_id.unwrap_or_else(generate_peer_id);
info!("starting up DHT with peer id {:?}", peer_id);
let bootstrap_addrs = config
.bootstrap_addrs
.unwrap_or_else(|| crate::DHT_BOOTSTRAP.iter().map(|v| v.to_string()).collect());
let (in_tx, in_rx) = unbounded_channel();
let state = Arc::new(RwLock::new(DhtState::new(
peer_id,
in_tx.clone(),
config.routing_table,
listen_addr,
)));
tokio::spawn({
let state = state.clone();
async move {
let worker = DhtWorker {
socket,
peer_id,
state,
};
let result = worker.start(in_tx, in_rx, &bootstrap_addrs).await;
warn!("DHT worker finished with {:?}", result);
}
});
Ok(Dht { state })
}
pub async fn get_peers(
&self,
info_hash: Id20,
) -> anyhow::Result<impl Stream<Item = SocketAddr> + Unpin> {
let (pos, rx) = self.state.write().get_peers(info_hash)?;
Ok(PeerStream {
info_hash,
state: self.state.clone(),
absolute_stream_pos: 0,
initial_peers_pos: pos,
broadcast_rx: BroadcastStream::new(rx),
})
}
pub fn listen_addr(&self) -> SocketAddr {
self.state.read().listen_addr
}
pub fn stats(&self) -> DhtStats {
self.state.read().get_stats()
}
pub fn with_routing_table<R, F: FnOnce(&RoutingTable) -> R>(&self, f: F) -> R {
f(&self.state.read().routing_table)
}
pub fn clone_routing_table(&self) -> RoutingTable {
self.state.read().routing_table.clone()
}
}