Saving
This commit is contained in:
Igor Katson
parent
f6656841c0
commit
950d47ab31
7 changed files with 451 additions and 167 deletions
|
|
@ -1,143 +1,297 @@
|
|||
use std::{collections::HashMap, net::SocketAddrV4};
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
net::{SocketAddr, SocketAddrV4},
|
||||
time::Instant,
|
||||
};
|
||||
|
||||
use crate::bprotocol::MessageKind;
|
||||
use bencode::ByteString;
|
||||
use dht::{
|
||||
bprotocol::{
|
||||
self, CompactNodeInfo, CompactPeerInfo, FindNodeRequest, GetPeersRequest, Message,
|
||||
MessageKind,
|
||||
},
|
||||
id20::Id20,
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use librqbit_core::peer_id::generate_peer_id;
|
||||
use log::debug;
|
||||
use parking_lot::Mutex;
|
||||
use tokio::{
|
||||
net::UdpSocket,
|
||||
sync::mpsc::{channel, Receiver, Sender},
|
||||
};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
use crate::bprotocol::Message;
|
||||
|
||||
mod bprotocol;
|
||||
|
||||
struct SocketManager {
|
||||
socket: tokio::net::UdpSocket,
|
||||
rx: tokio::sync::mpsc::Receiver<(
|
||||
SocketAddrV4,
|
||||
MessageKind<ByteString>,
|
||||
tokio::sync::oneshot::Sender<Message<ByteString>>,
|
||||
)>,
|
||||
struct OutstandingRequest {
|
||||
transaction_id: u16,
|
||||
addr: SocketAddr,
|
||||
request: Request,
|
||||
time: Instant,
|
||||
}
|
||||
|
||||
impl SocketManager {
|
||||
pub async fn spawn() -> anyhow::Result<SocketManagerHandle> {
|
||||
let socket = tokio::net::UdpSocket::bind("0.0.0.0:0").await?;
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(1);
|
||||
let mgr = SocketManager { socket, rx };
|
||||
tokio::spawn(mgr.run());
|
||||
Ok(SocketManagerHandle { tx })
|
||||
struct DhtState {
|
||||
id: Id20,
|
||||
next_transaction_id: u16,
|
||||
outstanding_requests: Vec<OutstandingRequest>,
|
||||
searching_for_peers: Vec<Id20>,
|
||||
}
|
||||
|
||||
enum PeersOrNodes {
|
||||
Nodes(CompactNodeInfo),
|
||||
Peers(Vec<CompactPeerInfo>),
|
||||
}
|
||||
|
||||
impl DhtState {
|
||||
fn add_searching_for_peers(&mut self, info_hash: Id20) {
|
||||
self.searching_for_peers.push(info_hash)
|
||||
}
|
||||
pub async fn run(self) -> anyhow::Result<()> {
|
||||
let Self { socket, mut rx } = self;
|
||||
|
||||
let mut transaction_id = 0u16;
|
||||
let mut next_transaction_id = move || {
|
||||
let next = transaction_id;
|
||||
transaction_id = next + 1;
|
||||
next
|
||||
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];
|
||||
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,
|
||||
}),
|
||||
},
|
||||
};
|
||||
|
||||
let outstanding = Mutex::new(HashMap::<
|
||||
u16,
|
||||
tokio::sync::oneshot::Sender<Message<ByteString>>,
|
||||
>::new());
|
||||
|
||||
let writer = async {
|
||||
let mut buf = Vec::new();
|
||||
while let Some((addr, msg, tx)) = rx.recv().await {
|
||||
let transaction_id = next_transaction_id();
|
||||
let transaction_id_buf =
|
||||
[(transaction_id >> 8) as u8, (transaction_id & 0xff) as u8];
|
||||
buf.clear();
|
||||
bprotocol::serialize_message(
|
||||
&mut buf,
|
||||
// this is bad, allocates
|
||||
ByteString::from(transaction_id_buf.as_ref()),
|
||||
None,
|
||||
None,
|
||||
msg,
|
||||
self.outstanding_requests.push(OutstandingRequest {
|
||||
transaction_id,
|
||||
addr,
|
||||
request,
|
||||
time: Instant::now(),
|
||||
});
|
||||
message
|
||||
}
|
||||
fn on_incoming_from_remote(
|
||||
&mut self,
|
||||
msg: Message<ByteString>,
|
||||
addr: SocketAddr,
|
||||
) -> anyhow::Result<()> {
|
||||
match msg.kind {
|
||||
MessageKind::Error(_) | MessageKind::Response(_) => {}
|
||||
other => anyhow::bail!("requests from DHT not supported, but got {:?}", other),
|
||||
};
|
||||
if msg.transaction_id.len() != 2 {
|
||||
anyhow::bail!("transaction id unrecognized")
|
||||
}
|
||||
let tid = ((msg.transaction_id[0] as u16) << 8) + (msg.transaction_id[1] as u16);
|
||||
// O(n) but whatever
|
||||
let outstanding_id = self
|
||||
.outstanding_requests
|
||||
.iter()
|
||||
.position(|req| req.transaction_id == tid && req.addr == addr)
|
||||
.ok_or_else(|| anyhow::anyhow!("outstanding request not found"))?;
|
||||
let outstanding = self.outstanding_requests.remove(outstanding_id);
|
||||
let response = match msg.kind {
|
||||
MessageKind::Error(e) => {
|
||||
anyhow::bail!(
|
||||
"request {:?} received error response {:?}",
|
||||
outstanding.request,
|
||||
e
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
debug!("inserting transaction id {}", transaction_id);
|
||||
assert!(outstanding.lock().insert(transaction_id, tx).is_none());
|
||||
debug!("sending msg to {}", addr);
|
||||
socket.send_to(&buf, addr).await.unwrap();
|
||||
}
|
||||
MessageKind::Response(r) => r,
|
||||
_ => unreachable!(),
|
||||
};
|
||||
match outstanding.request {
|
||||
Request::FindNode(id) => {
|
||||
if response.id != id {
|
||||
anyhow::bail!(
|
||||
"response id does not match: expected {:?}, received {:?}",
|
||||
id,
|
||||
response.id
|
||||
)
|
||||
};
|
||||
let nodes = response
|
||||
.nodes
|
||||
.ok_or_else(|| anyhow::anyhow!("expected nodes for find_node requests"))?;
|
||||
self.on_found_nodes(id, nodes)
|
||||
}
|
||||
Request::GetPeers(id) => {
|
||||
if response.id != id {
|
||||
anyhow::bail!(
|
||||
"response id does not match: expected {:?}, received {:?}",
|
||||
id,
|
||||
response.id
|
||||
)
|
||||
};
|
||||
let nodes = response
|
||||
.nodes
|
||||
.ok_or_else(|| anyhow::anyhow!("expected nodes for find_node requests"))?;
|
||||
let pn = match (response.nodes, response.values) {
|
||||
(Some(nodes), None) => PeersOrNodes::Nodes(nodes),
|
||||
(None, Some(peers)) => PeersOrNodes::Peers(peers),
|
||||
_ => anyhow::bail!("expected nodes or values to be set in find_peers response"),
|
||||
};
|
||||
self.on_found_peers_or_nodes(id, pn)
|
||||
}
|
||||
};
|
||||
Ok(())
|
||||
}
|
||||
fn on_found_nodes(&mut self, target: Id20, nodes: CompactNodeInfo) {
|
||||
todo!("on_found_nodes not implemented")
|
||||
}
|
||||
|
||||
let reader = async {
|
||||
let mut buf = vec![0u8; 16384];
|
||||
while let Ok(size) = socket.recv(&mut buf).await {
|
||||
debug!("received {}", size);
|
||||
let msg = match bprotocol::deserialize_message::<ByteString>(&buf[..size]) {
|
||||
Ok(msg) => msg,
|
||||
// todo handle errors
|
||||
Err(e) => panic!("{}", e),
|
||||
};
|
||||
assert!(msg.transaction_id.len() == 2);
|
||||
let b0 = msg.transaction_id[0];
|
||||
let b1 = msg.transaction_id[1];
|
||||
let tid = ((b0 as u16) << 8) + b1 as u16;
|
||||
let tx = outstanding.lock().remove(&tid).unwrap();
|
||||
debug!("sending oneshot result, tid {}", tid);
|
||||
tx.send(msg).unwrap();
|
||||
fn on_found_peers_or_nodes(&mut self, target: Id20, data: PeersOrNodes) {
|
||||
todo!("on_found_nodes not implemented")
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_framer(
|
||||
socket: &UdpSocket,
|
||||
mut input_rx: Receiver<(Message<ByteString>, SocketAddr)>,
|
||||
output_tx: Sender<Message<ByteString>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let writer = async {
|
||||
let mut buf = Vec::new();
|
||||
while let Some((msg, addr)) = input_rx.recv().await {
|
||||
buf.clear();
|
||||
bprotocol::serialize_message(
|
||||
&mut buf,
|
||||
msg.transaction_id,
|
||||
msg.version,
|
||||
msg.ip,
|
||||
msg.kind,
|
||||
)
|
||||
.unwrap();
|
||||
socket.send_to(&buf, addr).await.unwrap();
|
||||
}
|
||||
};
|
||||
let reader = async {
|
||||
let mut buf = vec![0u8; 16384];
|
||||
while let Ok((size, addr)) = socket.recv_from(&mut buf).await {
|
||||
match bprotocol::deserialize_message::<ByteString>(&buf[..size]) {
|
||||
Ok(msg) => match output_tx.send(msg).await {
|
||||
Ok(_) => {}
|
||||
Err(_) => break,
|
||||
},
|
||||
Err(e) => log::warn!("error deseriaizing msg: {}", e),
|
||||
}
|
||||
}
|
||||
};
|
||||
tokio::select! {
|
||||
_ = writer => {},
|
||||
_ = reader => {},
|
||||
};
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum Request {
|
||||
GetPeers(Id20),
|
||||
FindNode(Id20),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum Response {
|
||||
Peer(SocketAddr),
|
||||
}
|
||||
|
||||
struct Dht {
|
||||
request_tx: Sender<(Request, Sender<Response>)>,
|
||||
}
|
||||
|
||||
struct DhtWorker {
|
||||
socket: UdpSocket,
|
||||
request_rx: Receiver<(Request, Sender<Response>)>,
|
||||
next_transaction_id: u16,
|
||||
peer_id: Id20,
|
||||
}
|
||||
|
||||
impl DhtWorker {
|
||||
fn on_request(&self, request: Request, sender: Sender<Response>) {}
|
||||
|
||||
async fn start(&mut self, bootstrap_addrs: Vec<String>) -> anyhow::Result<()> {
|
||||
let (in_tx, in_rx) = channel(1);
|
||||
let (out_tx, out_rx) = channel(1);
|
||||
let framer = run_framer(&self.socket, in_rx, out_tx);
|
||||
|
||||
let bootstrap = async {
|
||||
// bootstrap
|
||||
for addr in bootstrap_addrs {
|
||||
for addr in tokio::net::lookup_host(addr).await.unwrap() {
|
||||
let msg = MessageKind::FindNodeRequest(FindNodeRequest {
|
||||
id: self.peer_id,
|
||||
target: self.peer_id,
|
||||
});
|
||||
in_tx.send((msg, addr)).await.unwrap();
|
||||
}
|
||||
}
|
||||
};
|
||||
let mut bootstrap_done = false;
|
||||
|
||||
let request_reader = async {
|
||||
while let Some((request, sender)) = self.request_rx.recv().await {
|
||||
self.on_request(request, sender)
|
||||
}
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
_ = writer => {},
|
||||
_ = reader => {}
|
||||
_ = framer => {
|
||||
anyhow::bail!("framer quit")
|
||||
},
|
||||
_ = bootstrap, if !bootstrap_done => {
|
||||
bootstrap_done = true
|
||||
},
|
||||
_ = request_reader => {}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SocketManagerHandle {
|
||||
tx: tokio::sync::mpsc::Sender<(
|
||||
SocketAddrV4,
|
||||
MessageKind<ByteString>,
|
||||
tokio::sync::oneshot::Sender<Message<ByteString>>,
|
||||
)>,
|
||||
}
|
||||
|
||||
impl SocketManagerHandle {
|
||||
async fn request(
|
||||
&self,
|
||||
addr: SocketAddrV4,
|
||||
kind: MessageKind<ByteString>,
|
||||
) -> anyhow::Result<bprotocol::Message<ByteString>> {
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
self.tx.send((addr, kind, tx)).await?;
|
||||
let msg = rx.await?;
|
||||
Ok(msg)
|
||||
impl Dht {
|
||||
pub async fn new(bootstrap_addrs: &[&str]) -> anyhow::Result<Self> {
|
||||
let (request_tx, request_rx) = channel(1);
|
||||
let socket = UdpSocket::bind("0.0.0.0:0").await?;
|
||||
let mut worker = DhtWorker {
|
||||
socket,
|
||||
request_rx,
|
||||
next_transaction_id: 0,
|
||||
peer_id: Id20(generate_peer_id()),
|
||||
};
|
||||
let bootstrap_addrs = bootstrap_addrs.iter().map(|s| s.to_string()).collect();
|
||||
tokio::spawn(async move { worker.start(bootstrap_addrs).await });
|
||||
Ok(Dht { request_tx })
|
||||
}
|
||||
pub async fn get_peers(&self, info_hash: Id20) -> impl StreamExt<Item = SocketAddr> {
|
||||
let (tx, rx) = channel::<Response>(1);
|
||||
self.request_tx
|
||||
.send((Request::GetPeers(info_hash), tx))
|
||||
.await
|
||||
.unwrap();
|
||||
ReceiverStream::new(rx).map(|r| match r {
|
||||
Response::Peer(addr) => addr,
|
||||
_ => panic!("programming error"),
|
||||
})
|
||||
}
|
||||
// async fn run(self) -> anyhow::Result<Self> {
|
||||
// let socket = UdpSocket::bind("0.0.0.0:0").await?;
|
||||
// let (in_tx, in_rx) = channel(1);
|
||||
// let (out_tx, out_rx) = channel(1);
|
||||
// let framer = run_framer(socket, in_rx, out_tx);
|
||||
// }
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
std::env::set_var("RUST_LOG", "trace");
|
||||
pretty_env_logger::init();
|
||||
|
||||
let mgr = SocketManager::spawn().await.unwrap();
|
||||
|
||||
let peer_id = bprotocol::Id20(generate_peer_id());
|
||||
for first_addr in tokio::net::lookup_host("dht.transmissionbt.com:6881")
|
||||
.await
|
||||
.unwrap()
|
||||
.filter_map(|a| match a {
|
||||
std::net::SocketAddr::V4(v4) => Some(v4),
|
||||
std::net::SocketAddr::V6(_) => None,
|
||||
})
|
||||
.skip(1)
|
||||
{
|
||||
let msg = bprotocol::MessageKind::FindNodeRequest(bprotocol::FindNodeRequest {
|
||||
id: peer_id,
|
||||
target: peer_id,
|
||||
});
|
||||
|
||||
dbg!(mgr.request(first_addr, msg).await.unwrap());
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let info_hash = Id20([0u8; 20]);
|
||||
let dht = Dht::new(&["dht.transmissionbt.com:6881"]).await.unwrap();
|
||||
let mut stream = dht.get_peers(info_hash).await;
|
||||
while let Some(peer) = stream.next().await {
|
||||
log::info!("peer found: {}", peer)
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue