├── .gitignore ├── Cargo.lock ├── Cargo.toml ├── README.md ├── build.rs └── src ├── args.rs ├── das_tree.rs ├── http_rpc.rs ├── libp2p.rs ├── main.rs ├── messages.rs ├── overlay.rs └── utils.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | 3 | .idea/ 4 | 5 | data 6 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "das-tests" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | eyre = "0.6" 10 | anyhow = "1.0.66" 11 | discv5 = "0.1" 12 | tokio = { version = "1.15.0", features = ["full"] } 13 | enr = { version = "0.6.2", features = ["k256", "ed25519"] } 14 | libp2p-core = { version = "0.36.0", optional = true } 15 | hex = "0.4.3" 16 | parking_lot = "0.11.2" 17 | tokio-util = { version = "0.6.9", features = ["time"] } 18 | tracing = { version = "0.1.29" } 19 | futures = "0.3.24" 20 | tokio-stream = "0.1.10" 21 | eth2_ssz = "0.4.0" 22 | eth2_ssz_derive = "0.3.0" 23 | eth2_ssz_types = "0.2.1" 24 | async-recursion = "1.0.0" 25 | warp = "0.3" 26 | strum = {version = "0.24", features = ["derive"] } 27 | itertools = "0.10" 28 | dyn-clone = "1.0" 29 | rand = "0.8" 30 | sha3 = "0.10" 31 | nanoid = "0.4.0" 32 | delay_map = "0.1.1" 33 | 34 | chrono = "0.4.22" 35 | byteorder = "1.4.3" 36 | discv5-overlay = {git = "https://github.com/timoth-y/discv5-overlay" } 37 | # discv5-overlay = {path = "../trin"} 38 | lazy_static = "1.4.0" 39 | 40 | libp2p = { version = "0.40.0" } 41 | unsigned-varint = "0.7.1" 42 | async-trait = "0.1.58" 43 | 44 | reqwest = "0.11.12" 45 | clap = { version = "4.0.18", features = ["derive"] } 46 | cli-batteries = "0.4" 47 | 48 | [build-dependencies] 49 | cli-batteries = "0.4" 50 | 51 | [patch.crates-io] 52 | discv5 = {git = "https://github.com/timoth-y/discv5"} 53 | #discv5 = {path = "../discv5"} 54 | #discv5-overlay = {path = "../trin"} 55 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DAS prototype 2 | 3 | This repo contains various prototypes of the core DAS components - dissemination and random sampling. 4 | 5 | ## Usage 6 | 7 | `cargo run -- -n -p -t [simulation-case] [ARGS]` 8 | 9 | This will spin up `num-servers` of discv5 servers starting at port `start-listen-port` all the way to `start-listen-port` + `num-servers`, and then start `simulation-case`. 10 | 11 | ### Disseminate 12 | 13 | ```bash 14 | cargo run -- -n 500 --topology uniform disseminate -n 256 --batching-strategy 'bucket-wise' --forward-mode 'FA' --replicate-mode 'RS' --redundancy 1 15 | ``` 16 | 17 | ### Sample 18 | 19 | ```bash 20 | cargo run -- -n 500 -t uniform --timeout 6 sample --validators-number 2 --samples-per-validator 75 --parallelism 30 21 | ``` 22 | 23 | ### Use snapshots 24 | Snapshots allow saving network configurations along with various RNG seeds to have more consistent measurements and for debugging. Use `--snapshot` flag with values `new`, `last`, and specific timecode eg. `2022-11-08-11:46:09`. Snapshots are saved in `--cache-dir` folder default value = `./data`. 25 | 26 | ```bash 27 | cargo run -- -n 500 --topology uniform --snapshot last disseminate -n 256 28 | ``` 29 | -------------------------------------------------------------------------------- /build.rs: -------------------------------------------------------------------------------- 1 | fn main() { 2 | cli_batteries::build_rs().unwrap() 3 | } 4 | -------------------------------------------------------------------------------- /src/args.rs: -------------------------------------------------------------------------------- 1 | use clap::{Args, Parser, Subcommand}; 2 | use strum::EnumString; 3 | 4 | #[derive(Clone, Debug, PartialEq, EnumString)] 5 | pub enum Topology { 6 | #[strum(serialize = "linear", serialize = "1")] 7 | Linear, 8 | #[strum(serialize = "uniform", serialize = "2")] 9 | Uniform, 10 | } 11 | 12 | #[derive(Clone, Debug, PartialEq, EnumString)] 13 | pub enum TalkWire { 14 | #[strum(serialize = "discv5")] 15 | Discv5, 16 | #[strum(serialize = "libp2p")] 17 | Libp2p, 18 | } 19 | 20 | // Defines settings how samples are forwarded when redundancy is enabled 21 | #[derive(Clone, Debug, PartialEq, EnumString)] 22 | pub enum ForwardPolicy { 23 | // forward only the first message of the type 24 | #[strum(serialize = "F1")] 25 | ForwardOne, 26 | // forward every time the message is received 27 | #[strum(serialize = "FA")] 28 | ForwardAll, 29 | } 30 | 31 | // Defines settings how samples are replicated when redundancy is enabled 32 | #[derive(Clone, Debug, PartialEq, EnumString)] 33 | pub enum ReplicatePolicy { 34 | // replicate only at the dispersal initiator 35 | #[strum(serialize = "R1")] 36 | ReplicateOne, 37 | // replicate at every step so a receiving node would also forward messages to certain number (based on --redundancy) of nodes in corresponding k-bucket 38 | #[strum(serialize = "RS")] 39 | ReplicateSome, 40 | // replicate at every step so a receiving node would also forward messages to every node in corresponding k-bucket 41 | #[strum(serialize = "RA")] 42 | ReplicateAll, 43 | } 44 | 45 | #[derive(Clone, Debug, PartialEq, EnumString)] 46 | pub enum BatchingStrategy { 47 | #[strum(serialize = "b", serialize = "bucket-wise")] 48 | BucketWise, 49 | /// note: must be used with `ForwardPolicy::ForwardAll` could be a bug \_(*_*)_/ 50 | #[strum(serialize = "d", serialize = "distance-wise")] 51 | DistanceWise, 52 | } 53 | 54 | #[derive(Clone, Debug, PartialEq, EnumString)] 55 | pub enum RoutingStrategy { 56 | #[strum(serialize = "i", serialize = "iterative")] 57 | Iterative, 58 | #[strum(serialize = "r", serialize = "recursive")] 59 | Recursive, 60 | } 61 | 62 | #[derive(Clone, Debug, PartialEq, EnumString)] 63 | pub enum LookupMethod { 64 | #[strum(serialize = "v", serialize = "find-value")] 65 | Discv5FindValue, 66 | #[strum(serialize = "c", serialize = "find-content")] 67 | OverlayFindContent, 68 | } 69 | 70 | #[derive(Clone, Parser)] 71 | pub struct Options { 72 | #[clap(long, short, default_value = "127.0.0.1")] 73 | pub ip_listen: String, 74 | #[clap(long, short, default_value = "9000")] 75 | pub port_udp: usize, 76 | #[clap(long, short, default_value = "10")] 77 | pub node_count: usize, 78 | #[clap(long, short, default_value = "linear")] 79 | pub topology: Topology, 80 | #[clap(long, short, default_value = "discv5")] 81 | pub wire_protocol: TalkWire, 82 | #[clap(long = "timeout", default_value = "3")] 83 | pub request_timeout: u64, 84 | #[clap(long, short, default_value = "./data")] 85 | pub cache_dir: String, 86 | #[clap(long, default_value = "new")] 87 | pub snapshot: String, 88 | 89 | #[command(subcommand)] 90 | pub simulation_case: SimulationCase, 91 | } 92 | 93 | #[derive(Clone, clap::Subcommand)] 94 | pub enum SimulationCase { 95 | Disseminate(DisseminationArgs), 96 | Sample(SamplingArgs), 97 | } 98 | 99 | #[derive(Clone, Args)] 100 | pub struct DisseminationArgs { 101 | #[clap(long, short, default_value = "256")] 102 | pub number_of_samples: usize, 103 | #[clap(long, default_value = "F1")] 104 | pub forward_mode: ForwardPolicy, 105 | #[clap(long, default_value = "R1")] 106 | pub replicate_mode: ReplicatePolicy, 107 | #[clap(long, short, default_value = "1")] 108 | pub redundancy: usize, 109 | #[clap(long, short = 'b', default_value = "b")] 110 | pub batching_strategy: BatchingStrategy, 111 | #[clap(long, short = 's', default_value = "r")] 112 | pub routing_strategy: RoutingStrategy, 113 | #[clap(long, short, default_value = "15")] 114 | pub parallelism: usize, 115 | } 116 | 117 | #[derive(Clone, Args)] 118 | pub struct SamplingArgs { 119 | #[clap(flatten)] 120 | pub dissemination_args: DisseminationArgs, 121 | 122 | #[clap(long, short = 'k', default_value = "75")] 123 | pub samples_per_validator: usize, 124 | 125 | #[clap(long, short)] 126 | pub validators_number: usize, 127 | 128 | #[clap(long, short, default_value = "find-value")] 129 | pub lookup_method: LookupMethod, 130 | } 131 | -------------------------------------------------------------------------------- /src/das_tree.rs: -------------------------------------------------------------------------------- 1 | use discv5::enr::{Enr, NodeId}; 2 | use dyn_clone::{clone_trait_object, DynClone}; 3 | use enr::EnrKey; 4 | trait TreeNode: DynClone { 5 | fn depth(&self) -> usize; 6 | fn id(&self) -> NodeId; 7 | fn score(&self) -> f64; 8 | fn sub_tree_size(&self) -> u64; 9 | // Add a node to the tree, return updated tree root, and ok == true if the node didn't already exist 10 | fn add(&self, n: Enr) -> Result>>, Box>; 11 | // Search for closest leaf nodes (log distance) and append to out, 12 | // maximum to the capacity of the out slice 13 | fn search(&self, target: NodeId, out: Vec>>) -> Vec>>; 14 | // Weakest finds the content with the weakest score at given tree depth 15 | fn weakest(&self, depth: u64) -> &dyn TreeNode; 16 | } 17 | clone_trait_object!(TreeNode); 18 | #[derive(Clone)] 19 | struct LeafNode { 20 | pub depth: usize, 21 | pub score: f64, 22 | pub _self: Enr, 23 | } 24 | 25 | #[derive(Clone)] 26 | struct PairNode { 27 | pub depth: usize, 28 | pub score: f64, 29 | pub subtree_size: u64, 30 | 31 | // Bits after depth index are zeroed 32 | pub id: NodeId, 33 | 34 | // left and right are never nil at the same time 35 | 36 | // May be nil (pair node as extension node) 37 | pub left: Option>>, 38 | // May be nil (pair node as extension node) 39 | pub right: Option>>, 40 | } 41 | 42 | impl TreeNode for PairNode 43 | where 44 | K: EnrKey + Clone, 45 | { 46 | fn depth(&self) -> usize { 47 | self.depth 48 | } 49 | 50 | fn id(&self) -> NodeId { 51 | self.id 52 | } 53 | 54 | fn score(&self) -> f64 { 55 | self.score 56 | } 57 | 58 | fn sub_tree_size(&self) -> u64 { 59 | self.subtree_size 60 | } 61 | 62 | fn add(&self, n: Enr) -> Result>>, Box> { 63 | let mut pair = self.clone(); 64 | let mut ok = false; 65 | if pair.id().raw() == n.node_id().raw() { 66 | return Ok(Some(Box::new(pair))); 67 | } 68 | if bit_check(n.node_id(), self.depth()) { 69 | if let Some(right_node) = &self.right { 70 | let right = right_node.add(n)?; 71 | pair.right = right; 72 | } else { 73 | let leaf = LeafNode { 74 | depth: pair.depth() + 1, 75 | score: 0.0, 76 | _self: n, 77 | }; 78 | pair.right = Some(Box::new(leaf)); 79 | ok = true; 80 | } 81 | } else { 82 | if let Some(left_node) = &self.left { 83 | let left = left_node.add(n)?; 84 | pair.left = left; 85 | } else { 86 | let leaf = LeafNode { 87 | depth: pair.depth() + 1, 88 | score: 0.0, 89 | _self: n, 90 | }; 91 | pair.left = Some(Box::new(leaf)); 92 | ok = true; 93 | } 94 | } 95 | if ok { 96 | pair.subtree_size += 1; 97 | pair.score = 0.0; 98 | if let Some(left) = &pair.left { 99 | pair.score += left.score(); 100 | } 101 | if let Some(right) = &pair.right { 102 | pair.score += right.score(); 103 | } 104 | } 105 | return Ok(Some(Box::new(pair))); 106 | } 107 | 108 | fn search(&self, target: NodeId, out: Vec>>) -> Vec>> { 109 | if out.len() == out.capacity() { 110 | return out; 111 | } 112 | 113 | // what happens if theyre both nil? should be unreachable, otherwise why have a pair node? 114 | match (&self.left, &self.right) { 115 | (None, None) => unreachable!(), 116 | (None, Some(right)) => return right.search(target, out), 117 | (Some(left), None) => return left.search(target, out), 118 | (Some(left), Some(right)) => { 119 | if bit_check(target, self.depth()) { 120 | let mut out = right.search(target, out); 121 | if out.len() < out.capacity() { 122 | out = left.search(target, out); 123 | } 124 | return out; 125 | } else { 126 | let mut out = left.search(target, out); 127 | if out.len() < out.capacity() { 128 | out = right.search(target, out); 129 | } 130 | return out; 131 | } 132 | } 133 | } 134 | } 135 | 136 | fn weakest(&self, depth: u64) -> &dyn TreeNode { 137 | if depth > self.depth.try_into().unwrap() { 138 | match (&self.left, &self.right) { 139 | (None, None) => return self, 140 | (None, Some(right)) => { 141 | return right.weakest(depth); 142 | } 143 | (Some(left), None) => { 144 | return left.weakest(depth); 145 | } 146 | (Some(left), Some(right)) => { 147 | if right.score() > left.score() { 148 | return right.weakest(depth); 149 | } else { 150 | return left.weakest(depth); 151 | } 152 | } 153 | } 154 | } 155 | self 156 | } 157 | } 158 | 159 | impl TreeNode for LeafNode 160 | where 161 | K: EnrKey + Clone, 162 | { 163 | fn depth(&self) -> usize { 164 | self.depth 165 | } 166 | fn score(&self) -> f64 { 167 | self.score 168 | } 169 | fn id(&self) -> NodeId { 170 | self._self.node_id() 171 | } 172 | fn sub_tree_size(&self) -> u64 { 173 | 1 174 | } 175 | /// add a node and returns a pair node 176 | fn add(&self, n: Enr) -> Result>>, Box> { 177 | if self.id().raw() == n.node_id().raw() { 178 | return Err("wtf".into()); 179 | } 180 | let pair = PairNode { 181 | depth: self.depth, 182 | score: 0.0, 183 | subtree_size: 0, 184 | id: clip(self.id(), self.depth), 185 | left: None, 186 | right: None, 187 | }; 188 | pair.add(self._self.clone())?; 189 | pair.add(n.clone())?; 190 | Ok(Some(Box::new(pair))) 191 | } 192 | 193 | fn search(&self, _target: NodeId, out: Vec>>) -> Vec>> { 194 | if out.len() == out.capacity() { 195 | return out; 196 | } 197 | let mut out = out.clone(); 198 | out.push(Box::new(self.clone())); 199 | out 200 | } 201 | 202 | fn weakest(&self, _depth: u64) -> &dyn TreeNode { 203 | self 204 | } 205 | } 206 | 207 | fn bit_check(id: NodeId, bit_index: usize) -> bool { 208 | id.raw()[bit_index >> 3] & (1 << bit_index) != 0 209 | } 210 | 211 | fn clip(id: NodeId, depth: usize) -> NodeId { 212 | let i = depth >> 3; 213 | let mut raw_id = id.raw(); 214 | // This will definitely panic lol 215 | raw_id[i] &= (1 << (depth as u8 & 7)) - 1; 216 | 217 | let mut j = i; 218 | while j < raw_id.len() { 219 | raw_id[j] = 0; 220 | j += 1; 221 | } 222 | 223 | return NodeId::new(&raw_id); 224 | } 225 | -------------------------------------------------------------------------------- /src/http_rpc.rs: -------------------------------------------------------------------------------- 1 | use discv5::Enr; 2 | use enr::k256::elliptic_curve::weierstrass::add; 3 | use enr::NodeId; 4 | use eyre::anyhow; 5 | use futures::future::BoxFuture; 6 | use futures::SinkExt; 7 | use rocket::http::Status; 8 | use rocket::response::status; 9 | use rocket::serde::{json::Json, Deserialize, Serialize}; 10 | use rocket::{routes, State}; 11 | use std::net::SocketAddr; 12 | use tokio::sync::{mpsc, oneshot}; 13 | 14 | pub enum RpcMsg { 15 | TalkReq(Vec, oneshot::Sender>), 16 | } 17 | 18 | struct Runtime { 19 | tx: mpsc::Sender, 20 | } 21 | 22 | #[post("/talk_req", data = "")] 23 | async fn api_talk_req( 24 | state: &State, 25 | req: Vec, 26 | ) -> Result, status::Custom> { 27 | let (tx, rx) = oneshot::channel(); 28 | 29 | state 30 | .tx 31 | .clone() 32 | .send(RpcMsg::TalkReq(req, tx)) 33 | .await 34 | .map_err(|e| status::Custom(Status::ServiceUnavailable, e.to_string()))?; 35 | 36 | let res = rx 37 | .await 38 | .map_err(|e| status::Custom(Status::InternalServerError, e.to_string()))?; 39 | 40 | Ok(res) 41 | } 42 | 43 | #[allow(unused_must_use)] 44 | pub async fn serve(to_runtime: mpsc::Sender, addr: impl Into) { 45 | let addr = addr.into(); 46 | let mut config = rocket::Config::default(); 47 | config.address = addr.ip(); 48 | config.port = addr.port(); 49 | config.shutdown.ctrlc = true; 50 | config.shutdown.force = true; 51 | 52 | rocket::custom(config) 53 | .manage(Runtime { tx: to_runtime }) 54 | .mount("/", routes![api_talk_req]) 55 | .launch() 56 | .await 57 | .expect("expect server to run"); 58 | } 59 | 60 | pub async fn talk_req(enr: Enr, msg: Vec) -> eyre::Result> { 61 | let addr = enr.ip4().unwrap().to_string(); 62 | let port = 3000 + (enr.udp4().unwrap() - 9000); // todo: dirty deterministic hack, should be config instead 63 | let client = reqwest::Client::new(); 64 | 65 | let mut resp = client 66 | .post(format!("http://{addr}:{port}/talk_req")) 67 | .body(msg) 68 | .send() 69 | .await 70 | .map_err(|e| anyhow!("error requesting setup: {e}"))?; 71 | 72 | if resp.status() != 200 { 73 | return Err(anyhow!("error status: {}", resp.status())); 74 | } 75 | 76 | resp.bytes() 77 | .await 78 | .map(|b| b.to_vec()) 79 | .map_err(|e| anyhow!("error decoding step1 response: {e}")) 80 | } 81 | -------------------------------------------------------------------------------- /src/libp2p.rs: -------------------------------------------------------------------------------- 1 | use discv5::Enr; 2 | use discv5_overlay::portalnet::types::messages::{ByteList, ElasticPacket}; 3 | use futures::prelude::*; 4 | use futures::stream::FuturesUnordered; 5 | use libp2p::core::transport::upgrade; 6 | use libp2p::identity::Keypair; 7 | use libp2p::noise::NoiseConfig; 8 | use libp2p::request_response::{ 9 | ProtocolSupport, RequestId, RequestResponse, RequestResponseCodec, RequestResponseConfig, 10 | RequestResponseEvent, RequestResponseMessage, ResponseChannel, 11 | }; 12 | use libp2p::swarm::{ 13 | NetworkBehaviour, NetworkBehaviourAction, NetworkBehaviourEventProcess, PollParameters, 14 | SwarmEvent, 15 | }; 16 | use libp2p::tcp::TcpConfig; 17 | use libp2p::NetworkBehaviour; 18 | use libp2p::{identity, mplex, noise, Multiaddr, PeerId, Swarm, Transport}; 19 | use parking_lot::RwLock; 20 | use ssz::{Decode, Encode}; 21 | use ssz_types::VariableList; 22 | use std::backtrace::Backtrace; 23 | use std::borrow::Cow; 24 | use std::collections::hash_map::Entry; 25 | use std::collections::{HashMap, VecDeque}; 26 | use std::pin::Pin; 27 | use std::sync::Arc; 28 | use std::task::{Context, Poll}; 29 | use std::thread::sleep; 30 | use std::time::Duration; 31 | use std::{io, iter}; 32 | use tokio::net::tcp; 33 | use tokio::select; 34 | use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; 35 | use tokio::sync::oneshot::Sender; 36 | use tokio::sync::{mpsc, oneshot}; 37 | use tokio_stream::wrappers::{ReceiverStream, UnboundedReceiverStream}; 38 | use tracing::{debug, info, log::warn}; 39 | 40 | const PROTOCOL_ID: &str = "/das/0.1.0"; 41 | 42 | /// The Libp2pService listens to events from the Libp2p swarm. 43 | pub struct Libp2pDaemon { 44 | swarm: Swarm, 45 | from_service: UnboundedReceiver, 46 | local_peer_id: PeerId, 47 | node_index: usize, 48 | } 49 | 50 | #[derive(Clone)] 51 | pub struct Libp2pService { 52 | /// Local copy of the `PeerId` of the local node. 53 | pub(crate) local_peer_id: PeerId, 54 | /// Channel for sending requests to worker. 55 | to_worker: UnboundedSender, 56 | 57 | promises: Arc>>>>, 58 | } 59 | 60 | /// Messages into the service to handle. 61 | #[derive(Debug)] 62 | pub enum NetworkMessage { 63 | RequestResponse { 64 | peer_id: PeerId, 65 | addr: Multiaddr, 66 | protocol: Vec, 67 | payload: Vec, 68 | resp_tx: oneshot::Sender>>, 69 | }, 70 | } 71 | 72 | #[derive(Debug)] 73 | pub struct TalkReqMsg { 74 | pub peer_id: PeerId, 75 | pub protocol: Vec, 76 | pub payload: Vec, 77 | pub resp_tx: oneshot::Sender>>, 78 | } 79 | 80 | impl Libp2pDaemon { 81 | pub fn new( 82 | keypair: Keypair, 83 | addr: Multiaddr, 84 | node_index: usize, 85 | ) -> ( 86 | Libp2pDaemon, 87 | mpsc::UnboundedReceiver, 88 | Libp2pService, 89 | ) { 90 | let local_peer_id = PeerId::from(keypair.public()); 91 | 92 | let transport = { 93 | let dh_keys = noise::Keypair::::new() 94 | .into_authentic(&keypair) 95 | .expect("Noise key generation failed"); 96 | 97 | TcpConfig::new() 98 | .upgrade(upgrade::Version::V1) 99 | .authenticate(NoiseConfig::xx(dh_keys).into_authenticated()) 100 | .multiplex(mplex::MplexConfig::new()) 101 | .boxed() 102 | }; 103 | 104 | let (message_sink, message_chan) = mpsc::unbounded_channel(); 105 | let mut swarm = Swarm::new( 106 | transport, 107 | Behaviour::new(message_sink, node_index), 108 | local_peer_id, 109 | ); 110 | 111 | // Listen on the addresses. 112 | if let Err(err) = swarm.listen_on(addr) { 113 | warn!(target: "sub-libp2p", "Can't listen on 'listen_address' because: {:?}", err) 114 | } 115 | 116 | let (network_sender_in, network_receiver_in) = mpsc::unbounded_channel(); 117 | 118 | let worker = Libp2pDaemon { 119 | local_peer_id, 120 | swarm, 121 | from_service: network_receiver_in, 122 | node_index, 123 | }; 124 | 125 | let service = Libp2pService { 126 | local_peer_id, 127 | to_worker: network_sender_in, 128 | promises: Arc::new(Default::default()), 129 | }; 130 | 131 | (worker, message_chan, service) 132 | } 133 | 134 | /// Starts the libp2p service networking stack. 135 | pub async fn run(mut self) { 136 | let mut swarm_stream = self.swarm.fuse(); 137 | let mut network_stream = UnboundedReceiverStream::new(self.from_service); 138 | let libp2p_node_idx = self.node_index; 139 | 140 | loop { 141 | select! { 142 | swarm_event = swarm_stream.next() => match swarm_event { 143 | // Outbound events 144 | Some(event) => match event { 145 | SwarmEvent::Behaviour(BehaviourEvent::InboundMessage{peer}) => { 146 | debug!(libp2p_node_idx=libp2p_node_idx, "Inbound message from {:?}", peer); 147 | }, 148 | SwarmEvent::NewListenAddr { address, .. } => debug!(libp2p_node_idx=libp2p_node_idx, "Listening on {:?}", address), 149 | SwarmEvent::ConnectionEstablished { peer_id, .. } => { 150 | debug!(libp2p_node_idx=libp2p_node_idx, "ConnectionEstablished with {:?}", peer_id.to_string()); 151 | }, 152 | SwarmEvent::ConnectionClosed { peer_id, .. } => { 153 | debug!(libp2p_node_idx=libp2p_node_idx, "ConnectionClosed with {:?}", peer_id.to_string()); 154 | } 155 | _ => continue 156 | } 157 | None => { break; } 158 | }, 159 | rpc_message = network_stream.next() => match rpc_message { 160 | // Inbound requests 161 | Some(request) => { 162 | let behaviour = swarm_stream.get_mut().behaviour_mut(); 163 | 164 | match request { 165 | NetworkMessage::RequestResponse { 166 | peer_id, 167 | addr, 168 | protocol, 169 | payload, 170 | resp_tx, 171 | } => { 172 | behaviour.send_request(peer_id, addr, protocol, payload, resp_tx); 173 | } 174 | } 175 | } 176 | None => { break; } 177 | } 178 | }; 179 | } 180 | } 181 | } 182 | 183 | impl Libp2pService { 184 | pub async fn talk_req( 185 | &self, 186 | peer_id: &PeerId, 187 | addr: &Multiaddr, 188 | protocol: &[u8], 189 | payload: Vec, 190 | ) -> eyre::Result> { 191 | let (tx, rx) = oneshot::channel(); 192 | 193 | let payload = 194 | ElasticPacket::Data(ByteList::from(VariableList::from(payload.clone()))).as_ssz_bytes(); 195 | 196 | self.to_worker 197 | .send(NetworkMessage::RequestResponse { 198 | peer_id: peer_id.clone(), 199 | addr: addr.clone(), 200 | protocol: protocol.to_vec(), 201 | payload, 202 | resp_tx: tx, 203 | }) 204 | .map_err(|e| eyre::eyre!("{e}"))?; 205 | 206 | let resp: ElasticPacket = ElasticPacket::from_ssz_bytes(&*rx.await.unwrap()?) 207 | .map_err(|e| eyre::eyre!("ssz decode error"))?; 208 | 209 | match resp { 210 | ElasticPacket::Data(bytes) => Ok(bytes.to_vec()), 211 | ElasticPacket::Promise(promise_id) => { 212 | let (tx, rx) = oneshot::channel(); 213 | self.promises.write().insert(promise_id, tx); 214 | rx.await 215 | .map_err(|_| eyre::eyre!("promise channel was canceled")) 216 | } 217 | _ => panic!("unsupported response"), 218 | } 219 | } 220 | 221 | pub async fn send_message( 222 | &self, 223 | peer_id: &PeerId, 224 | addr: &Multiaddr, 225 | protocol: &[u8], 226 | payload: Vec, 227 | ) -> eyre::Result> { 228 | let (tx, rx) = oneshot::channel(); 229 | 230 | self.to_worker 231 | .send(NetworkMessage::RequestResponse { 232 | peer_id: peer_id.clone(), 233 | addr: addr.clone(), 234 | protocol: protocol.to_vec(), 235 | payload, 236 | resp_tx: tx, 237 | }) 238 | .map_err(|e| eyre::eyre!("{e}"))?; 239 | 240 | rx.await.unwrap() 241 | } 242 | 243 | pub fn handle_promise_result(&self, promise_id: u16, res: Vec) { 244 | let mut promises = self.promises.write(); 245 | match promises.remove(&promise_id) { 246 | Some(tx) => tx.send(res).unwrap(), 247 | None => { 248 | drop(promises); 249 | sleep(Duration::from_millis(1)); 250 | self.handle_promise_result(promise_id, res); 251 | } 252 | } 253 | } 254 | } 255 | 256 | #[derive(NetworkBehaviour)] 257 | #[behaviour( 258 | out_event = "BehaviourEvent", 259 | poll_method = "poll", 260 | event_process = true 261 | )] 262 | pub(crate) struct Behaviour { 263 | req_resp: RequestResponse, 264 | 265 | #[behaviour(ignore)] 266 | events: VecDeque, 267 | 268 | #[behaviour(ignore)] 269 | message_sink: mpsc::UnboundedSender, 270 | 271 | #[behaviour(ignore)] 272 | pending_requests: HashMap>>>>, 273 | 274 | #[behaviour(ignore)] 275 | pending_responses: 276 | FuturesUnordered> + Send>>>, 277 | 278 | #[behaviour(ignore)] 279 | node_index: usize, 280 | } 281 | 282 | struct RequestProcessingOutcome { 283 | inner_channel: ResponseChannel, ()>>, 284 | response: Result, ()>, 285 | } 286 | 287 | pub(crate) enum BehaviourEvent { 288 | InboundMessage { 289 | /// Peer which sent us a message. 290 | peer: PeerId, 291 | }, 292 | } 293 | 294 | impl Behaviour { 295 | pub fn new(message_sink: mpsc::UnboundedSender, node_index: usize) -> Self { 296 | Self { 297 | req_resp: RequestResponse::new( 298 | GenericCodec { 299 | max_request_size: 100000, 300 | max_response_size: 100000, 301 | }, 302 | iter::once((PROTOCOL_ID.as_bytes().to_vec(), ProtocolSupport::Full)), 303 | RequestResponseConfig::default(), 304 | ), 305 | events: Default::default(), 306 | message_sink, 307 | pending_requests: Default::default(), 308 | pending_responses: Default::default(), 309 | node_index, 310 | } 311 | } 312 | 313 | pub fn send_request( 314 | &mut self, 315 | peer_id: PeerId, 316 | addr: Multiaddr, 317 | protocol: Vec, 318 | payload: Vec, 319 | resp_tx: oneshot::Sender>>, 320 | ) { 321 | self.req_resp.add_address(&peer_id, addr); 322 | let req_id = self 323 | .req_resp 324 | .send_request(&peer_id, TalkRequest { protocol, payload }); 325 | self.pending_requests.insert(req_id, Some(resp_tx)); 326 | } 327 | 328 | fn poll( 329 | &mut self, 330 | cx: &mut Context, 331 | _: &mut impl PollParameters, 332 | ) -> Poll::ProtocolsHandler>> 333 | { 334 | let node_index = self.node_index; 335 | 336 | // Poll to see if any response is ready to be sent back. 337 | while let Poll::Ready(Some(outcome)) = self.pending_responses.poll_next_unpin(cx) { 338 | let RequestProcessingOutcome { 339 | inner_channel, 340 | response, 341 | } = match outcome { 342 | Some(outcome) => outcome, 343 | // The response builder was too busy and thus the request was dropped. This is 344 | // later on reported as a `InboundFailure::Omission`. 345 | None => break, 346 | }; 347 | if let Err(_) = self.req_resp.send_response(inner_channel, response) { 348 | warn!("failed to send response"); 349 | } 350 | } 351 | 352 | if let Some(event) = self.events.pop_front() { 353 | return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event)); 354 | } 355 | 356 | Poll::Pending 357 | } 358 | } 359 | 360 | impl NetworkBehaviourEventProcess, ()>>> 361 | for Behaviour 362 | { 363 | fn inject_event(&mut self, event: RequestResponseEvent, ()>>) { 364 | let node_index = self.node_index; 365 | match event { 366 | RequestResponseEvent::Message { peer, message } => { 367 | match message { 368 | RequestResponseMessage::Request { 369 | request, channel, .. 370 | } => { 371 | let (tx, rx) = oneshot::channel(); 372 | self.message_sink 373 | .send(TalkReqMsg { 374 | peer_id: peer, 375 | protocol: request.protocol, 376 | payload: request.payload, 377 | resp_tx: tx, 378 | }) 379 | .unwrap(); 380 | 381 | self.pending_responses.push(Box::pin(async move { 382 | rx.await 383 | .map(|response| RequestProcessingOutcome { 384 | inner_channel: channel, 385 | response: response.map_err(|_| ()), 386 | }) 387 | .ok() 388 | })); 389 | } 390 | RequestResponseMessage::Response { 391 | request_id, 392 | response, 393 | } => match self.pending_requests.entry(request_id) { 394 | Entry::Occupied(e) => { 395 | if let Some(tx) = e.remove() { 396 | tx.send(response.map_err(|_| eyre::eyre!("got error resp"))) 397 | .unwrap(); 398 | } 399 | } 400 | Entry::Vacant(_) => panic!("unknown request_id"), 401 | }, 402 | } 403 | self.events 404 | .push_back(BehaviourEvent::InboundMessage { peer }); 405 | } 406 | RequestResponseEvent::OutboundFailure { peer, error, .. } => { 407 | debug!( 408 | libp2p_node_idx = node_index, 409 | "OutboundFailure {:?}: {}", peer, error 410 | ); 411 | } 412 | RequestResponseEvent::InboundFailure { peer, error, .. } => { 413 | debug!( 414 | libp2p_node_idx = node_index, 415 | "InboundFailure {:?}: {}", peer, error 416 | ); 417 | } 418 | RequestResponseEvent::ResponseSent { .. } => {} 419 | } 420 | } 421 | } 422 | 423 | #[derive(Clone, Debug)] 424 | pub struct TalkRequest { 425 | pub protocol: Vec, 426 | pub payload: Vec, 427 | } 428 | 429 | #[derive(Debug, Clone)] 430 | #[doc(hidden)] 431 | pub struct GenericCodec { 432 | pub max_request_size: u64, 433 | pub max_response_size: u64, 434 | } 435 | 436 | #[async_trait::async_trait] 437 | impl RequestResponseCodec for GenericCodec { 438 | type Protocol = Vec; 439 | type Request = TalkRequest; 440 | type Response = Result, ()>; 441 | 442 | async fn read_request( 443 | &mut self, 444 | _: &Self::Protocol, 445 | mut io: &mut T, 446 | ) -> io::Result 447 | where 448 | T: AsyncRead + Unpin + Send, 449 | { 450 | let protocol_length = unsigned_varint::aio::read_usize(&mut io) 451 | .await 452 | .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; 453 | 454 | let mut protocol = vec![0; protocol_length]; 455 | io.read_exact(&mut protocol).await?; 456 | 457 | // Read the length. 458 | let payload_length = unsigned_varint::aio::read_usize(&mut io) 459 | .await 460 | .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; 461 | 462 | if payload_length > usize::try_from(self.max_request_size).unwrap_or(usize::MAX) { 463 | return Err(io::Error::new( 464 | io::ErrorKind::InvalidInput, 465 | format!( 466 | "Request size exceeds limit: {} > {}", 467 | payload_length, self.max_request_size 468 | ), 469 | )); 470 | } 471 | 472 | // Read the payload. 473 | let mut payload = vec![0; payload_length]; 474 | io.read_exact(&mut payload).await?; 475 | 476 | Ok(TalkRequest { protocol, payload }) 477 | } 478 | 479 | async fn read_response( 480 | &mut self, 481 | _: &Self::Protocol, 482 | mut io: &mut T, 483 | ) -> io::Result 484 | where 485 | T: AsyncRead + Unpin + Send, 486 | { 487 | // Read the length. 488 | let length = match unsigned_varint::aio::read_usize(&mut io).await { 489 | Ok(l) => l, 490 | Err(unsigned_varint::io::ReadError::Io(err)) 491 | if matches!(err.kind(), io::ErrorKind::UnexpectedEof) => 492 | { 493 | return Ok(Err(())); 494 | } 495 | Err(err) => return Err(io::Error::new(io::ErrorKind::InvalidInput, err)), 496 | }; 497 | 498 | if length > usize::try_from(self.max_response_size).unwrap_or(usize::MAX) { 499 | return Err(io::Error::new( 500 | io::ErrorKind::InvalidInput, 501 | format!( 502 | "Response size exceeds limit: {} > {}", 503 | length, self.max_response_size 504 | ), 505 | )); 506 | } 507 | 508 | // Read the payload. 509 | let mut buffer = vec![0; length]; 510 | io.read_exact(&mut buffer).await?; 511 | Ok(Ok(buffer)) 512 | } 513 | 514 | async fn write_request( 515 | &mut self, 516 | _: &Self::Protocol, 517 | io: &mut T, 518 | req: Self::Request, 519 | ) -> io::Result<()> 520 | where 521 | T: AsyncWrite + Unpin + Send, 522 | { 523 | // Write the protocol_length. 524 | { 525 | let mut buffer = unsigned_varint::encode::usize_buffer(); 526 | io.write_all(unsigned_varint::encode::usize( 527 | req.protocol.len(), 528 | &mut buffer, 529 | )) 530 | .await?; 531 | } 532 | 533 | // Write protocol. 534 | io.write_all(&req.protocol).await?; 535 | 536 | // Write the payload_length. 537 | { 538 | let mut buffer = unsigned_varint::encode::usize_buffer(); 539 | io.write_all(unsigned_varint::encode::usize( 540 | req.payload.len(), 541 | &mut buffer, 542 | )) 543 | .await?; 544 | } 545 | 546 | // Write the payload. 547 | io.write_all(&req.payload).await?; 548 | 549 | io.close().await?; 550 | Ok(()) 551 | } 552 | 553 | async fn write_response( 554 | &mut self, 555 | _: &Self::Protocol, 556 | io: &mut T, 557 | res: Self::Response, 558 | ) -> io::Result<()> 559 | where 560 | T: AsyncWrite + Unpin + Send, 561 | { 562 | // If `res` is an `Err`, we jump to closing the substream without writing anything on it. 563 | if let Ok(res) = res { 564 | let mut buffer = unsigned_varint::encode::usize_buffer(); 565 | io.write_all(unsigned_varint::encode::usize(res.len(), &mut buffer)) 566 | .await?; 567 | 568 | // Write the payload. 569 | io.write_all(&res).await?; 570 | } 571 | 572 | io.close().await?; 573 | Ok(()) 574 | } 575 | } 576 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | #[macro_use] 2 | extern crate core; 3 | 4 | use crate::libp2p::Libp2pService; 5 | use crate::messages::DisseminationMsg; 6 | use crate::overlay::{DASContentKey, DASValidator}; 7 | use crate::utils::MsgCountCmd; 8 | use ::libp2p::kad::store::MemoryStore; 9 | use ::libp2p::multiaddr::Protocol::Tcp; 10 | use ::libp2p::{identity, Multiaddr, PeerId}; 11 | use args::*; 12 | use byteorder::{BigEndian, ReadBytesExt}; 13 | use chrono::{DateTime, Utc}; 14 | use clap::{Args, Parser, Subcommand}; 15 | use cli_batteries::version; 16 | use delay_map::HashMapDelay; 17 | use discv5::error::FindValueError; 18 | use discv5::kbucket::{BucketIndex, KBucketsTable, Node, NodeStatus}; 19 | use discv5::{ 20 | enr, 21 | enr::{CombinedKey, Enr, NodeId}, 22 | kbucket, ConnectionDirection, ConnectionState, Discv5, Discv5Config, Discv5ConfigBuilder, 23 | Discv5Event, Key, RequestError, TalkRequest, 24 | }; 25 | use discv5_overlay::portalnet::discovery::{Discovery, NodeAddress}; 26 | use discv5_overlay::portalnet::overlay::{OverlayConfig, OverlayProtocol}; 27 | use discv5_overlay::portalnet::overlay_service::{ 28 | OverlayCommand, OverlayRequest, OverlayRequestError, OverlayService, 29 | }; 30 | use discv5_overlay::portalnet::storage::{ 31 | ContentStore, DistanceFunction, MemoryContentStore, PortalStorage, PortalStorageConfig, 32 | }; 33 | use discv5_overlay::portalnet::types::content_key::OverlayContentKey; 34 | use discv5_overlay::portalnet::types::distance::{Distance, Metric, XorMetric}; 35 | use discv5_overlay::portalnet::types::messages::{ 36 | Content, ElasticPacket, ElasticResult, ProtocolId, SszEnr, 37 | }; 38 | use discv5_overlay::utils::bytes::hex_encode_compact; 39 | use discv5_overlay::utp::stream::{UtpListener, UtpListenerRequest}; 40 | use discv5_overlay::{portalnet, utp}; 41 | use enr::k256::elliptic_curve::bigint::Encoding; 42 | use enr::k256::elliptic_curve::weierstrass::add; 43 | use enr::k256::U256; 44 | use eyre::eyre; 45 | use futures::stream::{FuturesOrdered, FuturesUnordered}; 46 | use futures::{pin_mut, AsyncWriteExt, FutureExt, StreamExt}; 47 | use itertools::Itertools; 48 | use lazy_static::lazy_static; 49 | use nanoid::nanoid; 50 | use rand::prelude::StdRng; 51 | use rand::{thread_rng, Rng, SeedableRng}; 52 | use sha3::{Digest, Keccak256}; 53 | use ssz::{Decode, Encode}; 54 | use std::borrow::Cow; 55 | use std::collections::hash_map::Entry; 56 | use std::collections::{HashMap, HashSet, VecDeque}; 57 | use std::future::Future; 58 | use std::io::{BufReader, BufWriter, Read, Write}; 59 | use std::net::{IpAddr, Ipv4Addr, SocketAddr}; 60 | use std::ops::AddAssign; 61 | use std::path::PathBuf; 62 | use std::pin::Pin; 63 | use std::rc::Rc; 64 | use std::str::FromStr; 65 | use std::sync::Arc; 66 | use std::thread::sleep; 67 | use std::time::{Duration, SystemTime}; 68 | use std::{fs, iter}; 69 | use tokio::sync::mpsc::UnboundedReceiver; 70 | use tokio::sync::{mpsc, oneshot, RwLock}; 71 | use tokio::task::spawn_blocking; 72 | use tokio::{select, time}; 73 | use tokio_stream::wrappers::ReceiverStream; 74 | use tokio_stream::wrappers::UnboundedReceiverStream; 75 | use tracing::log::error; 76 | use tracing::{debug, info, info_span, log::warn, trace_span, Instrument}; 77 | use warp::Filter; 78 | 79 | mod args; 80 | mod libp2p; 81 | mod messages; 82 | mod overlay; 83 | mod utils; 84 | 85 | const DAS_PROTOCOL_ID: &str = "DAS"; 86 | const DISSEMINATION_PROTOCOL_ID: &[u8] = b"D2S"; 87 | 88 | #[derive(Clone)] 89 | pub struct DASNode { 90 | discovery: Arc, 91 | libp2p: Libp2pService, 92 | samples: Arc>>, 93 | handled_ids: Arc, usize>>>, 94 | overlay: Arc>, 95 | } 96 | 97 | impl DASNode { 98 | pub fn new( 99 | discovery: Arc, 100 | utp_listener_tx: mpsc::UnboundedSender, 101 | libp2p: Libp2pService, 102 | ) -> ( 103 | Self, 104 | OverlayService, 105 | ) { 106 | let config = OverlayConfig { 107 | bootnode_enrs: discovery.discv5.table_entries_enr(), 108 | // todo: setting low ping interval will hurt performance, investigate the impact of not having it 109 | ping_queue_interval: Some(Duration::from_secs(10000)), 110 | query_num_results: usize::MAX, 111 | query_timeout: Duration::from_secs(60), 112 | query_peer_timeout: Duration::from_secs(30), 113 | ..Default::default() 114 | }; 115 | let protocol = ProtocolId::Custom(DAS_PROTOCOL_ID.to_string()); 116 | let storage = { 117 | Arc::new(parking_lot::RwLock::new(MemoryContentStore::new( 118 | discovery.discv5.local_enr().node_id(), 119 | DistanceFunction::Xor, 120 | ))) 121 | }; 122 | let validator = Arc::new(DASValidator); 123 | 124 | let (overlay, service) = OverlayProtocol::new( 125 | config, 126 | discovery.clone(), 127 | utp_listener_tx, 128 | storage, 129 | Distance::MAX, 130 | protocol, 131 | validator, 132 | ); 133 | 134 | ( 135 | Self { 136 | discovery, 137 | libp2p, 138 | samples: Default::default(), 139 | handled_ids: Default::default(), 140 | overlay: Arc::new(overlay), 141 | }, 142 | service, 143 | ) 144 | } 145 | } 146 | 147 | fn main() { 148 | cli_batteries::run(version!(), app); 149 | } 150 | 151 | async fn app(options: Options) -> eyre::Result<()> { 152 | let discv5_servers = { 153 | let address = options.ip_listen.parse::().unwrap(); 154 | construct_and_start(&options, address, options.port_udp, options.node_count).await 155 | }; 156 | 157 | let node_ids = discv5_servers 158 | .iter() 159 | .map(|e| e.local_enr().node_id()) 160 | .collect::>(); 161 | 162 | let enrs = Arc::new( 163 | discv5_servers 164 | .iter() 165 | .map(|s| s.local_enr()) 166 | .collect::>(), 167 | ); 168 | 169 | let mut das_nodes = vec![]; 170 | 171 | let enr_to_libp2p = Arc::new(RwLock::new( 172 | HashMap::::default(), 173 | )); 174 | let libp2p_to_enr = Arc::new(RwLock::new(HashMap::::default())); 175 | 176 | let (msg_counter, msg_count_rx) = mpsc::unbounded_channel::(); 177 | { 178 | tokio::spawn(async move { 179 | let mut rx = UnboundedReceiverStream::new(msg_count_rx); 180 | let mut messages = 0u64; 181 | loop { 182 | if let Some(c) = rx.next().await { 183 | match c { 184 | MsgCountCmd::Increment => { 185 | messages += 1; 186 | } 187 | MsgCountCmd::Reset => { 188 | messages = 0; 189 | } 190 | MsgCountCmd::Get(tx) => tx.send(messages).unwrap(), 191 | } 192 | } 193 | } 194 | }); 195 | } 196 | 197 | for (i, discv5) in discv5_servers.into_iter().enumerate() { 198 | let mut events_str = ReceiverStream::new(discv5.event_stream().await.unwrap()); 199 | let opts = options.clone(); 200 | 201 | let (mut libp2p_worker, libp2p_msgs, libp2p_service) = { 202 | let keypair = identity::Keypair::generate_ed25519(); 203 | let peer_id = PeerId::from(keypair.public()); 204 | let mut addr = Multiaddr::from(IpAddr::from([127, 0, 0, 1])); 205 | addr.push(Tcp(4000 + i as u16)); 206 | 207 | enr_to_libp2p 208 | .write() 209 | .await 210 | .insert(discv5.local_enr().node_id(), (peer_id, addr.clone())); 211 | libp2p_to_enr 212 | .write() 213 | .await 214 | .insert(peer_id, discv5.local_enr().node_id()); 215 | 216 | libp2p::Libp2pDaemon::new(keypair, addr, i) 217 | }; 218 | let mut libp2p_msgs = UnboundedReceiverStream::new(libp2p_msgs); 219 | let discovery = Arc::new(Discovery::new_raw(discv5, Default::default())); 220 | let (utp_events_tx, utp_listener_tx, mut utp_listener_rx, mut utp_listener) = 221 | UtpListener::new(discovery.clone()); 222 | tokio::spawn(async move { utp_listener.start().await }); 223 | let (das_node, overlay_service) = DASNode::new(discovery, utp_listener_tx, libp2p_service); 224 | das_nodes.push(das_node.clone()); 225 | 226 | let talk_wire = opts.wire_protocol.clone(); 227 | if talk_wire == TalkWire::Libp2p { 228 | tokio::spawn(async move { 229 | libp2p_worker.run().await; 230 | }); 231 | } 232 | clone_all!(enr_to_libp2p, libp2p_to_enr, msg_counter, node_ids); 233 | tokio::spawn(async move { 234 | let mut overlay_service = overlay_service; 235 | let mut bucket_refresh_interval = tokio::time::interval(Duration::from_secs(60)); 236 | 237 | loop { 238 | select! { 239 | Some(e) = events_str.next() => { 240 | let chan = format!("{i} {}", das_node.discovery.discv5.local_enr().node_id().to_string()); 241 | match e { 242 | Discv5Event::Discovered(enr) => { 243 | debug!("Stream {}: Enr discovered {}", chan, enr) 244 | } 245 | Discv5Event::EnrAdded { enr, replaced: _ } => { 246 | debug!("Stream {}: Enr added {}", chan, enr) 247 | } 248 | Discv5Event::NodeInserted { 249 | node_id, 250 | replaced: _, 251 | } => debug!("Stream {}: Node inserted {}", chan, node_id), 252 | Discv5Event::SessionEstablished(enr, socket_addr) => { 253 | debug!("Stream {}: Session established {}", chan, enr); 254 | das_node.discovery.node_addr_cache 255 | .write() 256 | .put(enr.node_id(), NodeAddress { enr, socket_addr: Some(socket_addr) }); 257 | } 258 | Discv5Event::SocketUpdated(addr) => { 259 | debug!("Stream {}: Socket updated {}", chan, addr) 260 | } 261 | Discv5Event::TalkRequest(req) => { 262 | debug!("Stream {}: Talk request received", chan); 263 | msg_counter.send(MsgCountCmd::Increment); 264 | clone_all!(das_node, opts, enr_to_libp2p, node_ids, utp_events_tx); 265 | tokio::spawn(async move { 266 | let protocol = ProtocolId::from_str(&hex::encode_upper(req.protocol())).unwrap(); 267 | 268 | if protocol == ProtocolId::Utp { 269 | utp_events_tx.send(req).unwrap(); 270 | return; 271 | } 272 | 273 | if protocol == ProtocolId::Custom(DAS_PROTOCOL_ID.to_string()) { 274 | let talk_resp = match das_node.overlay.process_one_request(&req).await { 275 | Ok(response) => discv5_overlay::portalnet::types::messages::Message::from(response).into(), 276 | Err(err) => { 277 | error!("Node {chan} Error processing request: {err}"); 278 | return; 279 | }, 280 | }; 281 | 282 | if let Err(err) = req.respond(talk_resp) { 283 | error!("Unable to respond to talk request: {}", err); 284 | return; 285 | } 286 | 287 | return; 288 | } 289 | 290 | let resp = handle_talk_request(req.node_id().clone(), req.protocol(), req.body().to_vec(), das_node, opts, enr_to_libp2p, node_ids, i).await; 291 | req.respond(resp); 292 | }); 293 | }, 294 | Discv5Event::FindValue(req) => { 295 | debug!("Stream {}: FindValue request received with id {}", chan, req.id()); 296 | msg_counter.send(MsgCountCmd::Increment); 297 | clone_all!(das_node, opts, enr_to_libp2p, node_ids); 298 | tokio::spawn(async move { 299 | let resp = handle_sampling_request(req.node_id().clone(), req.key(), &das_node, &opts).await; 300 | req.respond(resp); 301 | }); 302 | }, 303 | } 304 | }, 305 | Some(crate::libp2p::TalkReqMsg{resp_tx, peer_id, payload, protocol}) = libp2p_msgs.next() => { 306 | debug!("Libp2p {i}: Talk request received"); 307 | msg_counter.send(MsgCountCmd::Increment); 308 | let from = libp2p_to_enr.read().await.get(&peer_id).unwrap().clone(); 309 | clone_all!(das_node, opts, enr_to_libp2p, node_ids); 310 | tokio::spawn(async move { 311 | resp_tx.send(Ok(handle_talk_request(from, &protocol, payload, das_node, opts, enr_to_libp2p, node_ids, i).await)); 312 | }); 313 | }, 314 | Some(command) = overlay_service.command_rx.recv() => { 315 | match command { 316 | OverlayCommand::Request(request) => overlay_service.process_request(request), 317 | OverlayCommand::FindContentQuery { target, callback } => { 318 | if let Some(query_id) = overlay_service.init_find_content_query(target.clone(), Some(callback)) { 319 | debug!( 320 | query_id=query_id.to_string(), 321 | content_id=hex_encode_compact(target.content_id()), 322 | "FindContent query initialized" 323 | ); 324 | } 325 | } 326 | } 327 | } 328 | Some(response) = overlay_service.response_rx.recv() => { 329 | // Look up active request that corresponds to the response. 330 | let optional_active_request = overlay_service.active_outgoing_requests.write().remove(&response.request_id); 331 | if let Some(active_request) = optional_active_request { 332 | 333 | // Send response to responder if present. 334 | if let Some(responder) = active_request.responder { 335 | let _ = responder.send(response.response.clone()); 336 | } 337 | 338 | // Perform background processing. 339 | match response.response { 340 | Ok(response) => overlay_service.process_response(response, active_request.destination, active_request.request, active_request.query_id), 341 | Err(error) => overlay_service.process_request_failure(response.request_id, active_request.destination, error), 342 | } 343 | 344 | } else { 345 | warn!("No request found for response"); 346 | } 347 | } 348 | Some(Ok(node_id)) = overlay_service.peers_to_ping.next() => { 349 | // If the node is in the routing table, then ping and re-queue the node. 350 | let key = discv5::kbucket::Key::from(node_id); 351 | if let discv5::kbucket::Entry::Present(ref mut entry, _) = overlay_service.kbuckets.write().entry(&key) { 352 | overlay_service.ping_node(&entry.value().enr()); 353 | overlay_service.peers_to_ping.insert(node_id); 354 | } 355 | } 356 | // todo: uncommenting next clause will servilely affect performance :/ investigate why 357 | // query_event = OverlayService::::query_event_poll(&mut overlay_service.find_node_query_pool) => { 358 | // overlay_service.handle_find_nodes_query_event(query_event); 359 | // } 360 | // Handle query events for queries in the find content query pool. 361 | query_event = OverlayService::::query_event_poll(&mut overlay_service.find_content_query_pool) => { 362 | overlay_service.handle_find_content_query_event(query_event); 363 | } 364 | _ = OverlayService::::bucket_maintenance_poll(overlay_service.protocol.clone(), &overlay_service.kbuckets) => {} 365 | _ = bucket_refresh_interval.tick() => { 366 | overlay_service.bucket_refresh_lookup(); 367 | } 368 | Some(event) = utp_listener_rx.recv() => das_node.overlay.process_utp_event(event).unwrap(), 369 | } 370 | } 371 | }); 372 | } 373 | let enrs_stats = enrs.clone(); 374 | let stats_task = tokio::spawn(async move { 375 | play_simulation( 376 | &options, 377 | &das_nodes, 378 | enr_to_libp2p.clone(), 379 | node_ids, 380 | msg_counter, 381 | ) 382 | .await; 383 | }); 384 | 385 | stats_task.await.unwrap(); 386 | 387 | tokio::signal::ctrl_c().await.unwrap(); 388 | 389 | Ok(()) 390 | } 391 | 392 | async fn construct_and_start( 393 | opts: &Options, 394 | listen_ip: Ipv4Addr, 395 | port_start: usize, 396 | node_count: usize, 397 | ) -> Vec { 398 | let mut discv5_servers = Vec::with_capacity(node_count); 399 | 400 | let snapshot_dir = match &*opts.snapshot { 401 | "new" => { 402 | let snap_time: DateTime = SystemTime::now().into(); 403 | let snap_dir = 404 | PathBuf::from(&opts.cache_dir).join(snap_time.format("%Y-%m-%d-%T").to_string()); 405 | fs::create_dir_all(&snap_dir).unwrap(); 406 | snap_dir 407 | } 408 | "last" => { 409 | let mut paths: Vec<_> = fs::read_dir(&opts.cache_dir) 410 | .unwrap() 411 | .map(|r| r.unwrap()) 412 | .collect(); 413 | paths.sort_by_key(|dir| dir.metadata().unwrap().created().unwrap()); 414 | paths.last().unwrap().path() 415 | } 416 | snap => PathBuf::from(&opts.cache_dir).join(snap), 417 | }; 418 | 419 | info!("snapshot: {}", snapshot_dir.to_str().unwrap()); 420 | 421 | for i in 0..node_count { 422 | let listen_addr = format!("{}:{}", listen_ip, port_start + i) 423 | .parse::() 424 | .unwrap(); 425 | debug!("{}", listen_addr); 426 | 427 | let enr_key = match &*opts.snapshot { 428 | "new" => { 429 | let key = CombinedKey::generate_secp256k1(); 430 | fs::write(snapshot_dir.join(format!("{i}.pem")), key.encode()).unwrap(); 431 | key 432 | } 433 | _ => { 434 | let mut key_bytes = fs::read(snapshot_dir.join(format!("{i}.pem"))).unwrap(); 435 | CombinedKey::secp256k1_from_bytes(&mut key_bytes).unwrap() 436 | } 437 | }; 438 | 439 | let enr = { 440 | let mut builder = enr::EnrBuilder::new("v4"); 441 | // TODO: Revisit this when we are not running locally 442 | // // if an IP was specified, use it 443 | // if let Some(external_address) = address { 444 | // builder.ip4(external_address); 445 | // } 446 | // // if a port was specified, use it 447 | // if std::env::args().nth(2).is_some() { 448 | // builder.udp4(port); 449 | // } 450 | builder.ip4(listen_ip); 451 | builder.udp4(port_start as u16 + i as u16); 452 | builder.build(&enr_key).unwrap() 453 | }; 454 | debug!("Node Id: {}", enr.node_id()); 455 | if enr.udp4_socket().is_some() { 456 | debug!("Base64 ENR: {}", enr.to_base64()); 457 | debug!( 458 | "IP: {}, UDP_PORT:{}", 459 | enr.ip4().unwrap(), 460 | enr.udp4().unwrap() 461 | ); 462 | } else { 463 | warn!("ENR is not printed as no IP:PORT was specified"); 464 | } 465 | 466 | // default configuration 467 | let mut config_builder = Discv5ConfigBuilder::default(); 468 | config_builder.request_retries(10); 469 | config_builder.filter_max_nodes_per_ip(None); 470 | config_builder.request_timeout(Duration::from_secs(opts.request_timeout)); 471 | config_builder.query_timeout(Duration::from_secs(60)); 472 | 473 | let config = config_builder.build(); 474 | 475 | // construct the discv5 server 476 | let discv5 = Discv5::new(enr, enr_key, config).unwrap(); 477 | discv5_servers.push(discv5); 478 | } 479 | 480 | discv5_servers = set_topology(&opts, discv5_servers); 481 | 482 | for s in discv5_servers.iter_mut() { 483 | let ip4 = s.local_enr().ip4().unwrap(); 484 | let udp4 = s.local_enr().udp4().unwrap(); 485 | s.start(format!("{}:{}", ip4, udp4).parse().unwrap()) 486 | .await 487 | .unwrap(); 488 | } 489 | discv5_servers 490 | } 491 | 492 | pub fn set_topology(opts: &Options, mut discv5_servers: Vec) -> Vec { 493 | let last_node_id = Key::from(discv5_servers.last().unwrap().local_enr().node_id()); 494 | 495 | match opts.topology { 496 | Topology::Linear => { 497 | // sort peers based on xor-distance to the latest node 498 | discv5_servers = discv5_servers 499 | .into_iter() 500 | .sorted_by_key(|s| Key::from(s.local_enr().node_id()).distance(&last_node_id)) 501 | .rev() 502 | .collect::>(); 503 | 504 | for (i, s) in discv5_servers.iter().enumerate() { 505 | if i != discv5_servers.len() - 1 { 506 | s.add_enr(discv5_servers[i + 1].local_enr().clone()) 507 | .unwrap() 508 | } 509 | } 510 | } 511 | Topology::Uniform => { 512 | let topology_seed = { 513 | let f = get_snapshot_file(&opts, "topology_seed"); 514 | let seed = fs::read(&f).map_or(rand::thread_rng().gen::(), |b| { 515 | b.as_slice().read_u64::().unwrap() 516 | }); 517 | fs::write(&f, seed.to_be_bytes()).unwrap(); 518 | seed 519 | }; 520 | 521 | let mut rng = StdRng::seed_from_u64(topology_seed); 522 | for (i, s) in discv5_servers.iter().enumerate() { 523 | let mut n = 150; 524 | while n != 0 { 525 | let i = rng.gen_range(0usize..discv5_servers.len() - 1); 526 | 527 | match s.add_enr(discv5_servers[i].local_enr().clone()) { 528 | Ok(_) => n -= 1, 529 | Err(_) => continue, 530 | } 531 | } 532 | } 533 | } 534 | } 535 | 536 | discv5_servers 537 | } 538 | 539 | pub async fn play_simulation( 540 | opts: &Options, 541 | nodes: &Vec, 542 | addr_book: Arc>>, 543 | node_ids: Vec, 544 | msg_counter: mpsc::UnboundedSender, 545 | ) { 546 | match &opts.simulation_case { 547 | SimulationCase::Disseminate(args) => { 548 | let keys = (0..args.number_of_samples).map(|i| { 549 | let mut h = Keccak256::new(); 550 | h.update(&i.to_be_bytes()); 551 | NodeId::new(&h.finalize().try_into().unwrap()) 552 | }); 553 | 554 | let (keys_per_node, nodes_per_key) = disseminate_samples( 555 | keys.clone(), 556 | opts, 557 | &args, 558 | nodes, 559 | addr_book.clone(), 560 | &node_ids, 561 | ) 562 | .await; 563 | 564 | info!("Keys per Node:"); 565 | keys_per_node 566 | .iter() 567 | .filter(|(_, keys)| **keys > 0) 568 | .for_each(|(n, keys)| { 569 | info!( 570 | "node={} ({}) keys={keys}", 571 | n.to_string(), 572 | node_ids.iter().position(|e| *e == *n).unwrap() 573 | ) 574 | }); 575 | debug!("Nodes per Key:"); 576 | nodes_per_key 577 | .iter() 578 | .for_each(|(k, nodes)| debug!("key={} nodes={}", k.to_string(), nodes.len())); 579 | debug!("Keys total: {}", nodes_per_key.len()); 580 | 581 | for k in keys { 582 | if !nodes_per_key.contains_key(&k) { 583 | warn!("missing key: {}", k.to_string()); 584 | } 585 | } 586 | 587 | let mut keys_stored_total = 0usize; 588 | keys_per_node 589 | .iter() 590 | .for_each(|(n, keys)| keys_stored_total += *keys); 591 | info!("total keys stored = {keys_stored_total} (storage overhead)"); 592 | 593 | let unique_keys_stored = nodes_per_key.len(); 594 | assert_eq!(unique_keys_stored, args.number_of_samples); 595 | 596 | let msg_count_total = { 597 | let (tx, rx) = oneshot::channel(); 598 | msg_counter.send(MsgCountCmd::Get(tx)).unwrap(); 599 | rx.await.unwrap() 600 | }; 601 | info!("total messages sent = {msg_count_total} (communication overhead)"); 602 | } 603 | SimulationCase::Sample(ref args) => { 604 | let keys = (0..args.dissemination_args.number_of_samples) 605 | .map(|i| { 606 | let mut h = Keccak256::new(); 607 | h.update(&i.to_be_bytes()); 608 | NodeId::new(&h.finalize().try_into().unwrap()) 609 | }) 610 | .collect::>(); 611 | 612 | let (keys_per_node, nodes_per_key) = disseminate_samples( 613 | keys.clone().into_iter(), 614 | opts, 615 | &args.dissemination_args, 616 | nodes, 617 | addr_book.clone(), 618 | &node_ids, 619 | ) 620 | .await; 621 | 622 | let nodes_by_node: HashMap<_, _> = nodes 623 | .iter() 624 | .map(|e| { 625 | let node_id = e.discovery.discv5.local_enr().node_id(); 626 | let known_by = nodes 627 | .iter() 628 | .filter_map(|n| { 629 | n.discovery 630 | .discv5 631 | .find_enr(&node_id) 632 | .map(|x| n.discovery.discv5.local_enr().node_id()) 633 | }) 634 | .collect_vec(); 635 | (node_id, known_by) 636 | }) 637 | .into_group_map() 638 | .into_iter() 639 | .map(|(n, ns)| { 640 | let x = ns.into_iter().flatten().collect_vec(); 641 | (n, x) 642 | }) 643 | .collect(); 644 | 645 | let overlay_nodes_by_node: HashMap<_, _> = nodes 646 | .iter() 647 | .map(|e| { 648 | let node_id = e.discovery.discv5.local_enr().node_id(); 649 | let known_by = nodes 650 | .iter() 651 | .filter_map(|n| { 652 | let key = kbucket::Key::from(n.discovery.discv5.local_enr().node_id()); 653 | if let kbucket::Entry::Present(entry, _) = 654 | e.overlay.kbuckets.write().entry(&key) 655 | { 656 | return Some(n.discovery.discv5.local_enr().node_id()); 657 | } 658 | None 659 | }) 660 | .collect_vec(); 661 | (node_id, known_by) 662 | }) 663 | .into_group_map() 664 | .into_iter() 665 | .map(|(n, ns)| { 666 | let x = ns.into_iter().flatten().collect_vec(); 667 | (n, x) 668 | }) 669 | .collect(); 670 | 671 | let mut keys_stored_total = 0usize; 672 | keys_per_node 673 | .iter() 674 | .for_each(|(n, keys)| keys_stored_total += *keys); 675 | info!("total keys stored = {keys_stored_total} (storage overhead)"); 676 | 677 | let unique_keys_stored = nodes_per_key.len(); 678 | assert_eq!( 679 | unique_keys_stored, 680 | args.dissemination_args.number_of_samples 681 | ); 682 | 683 | let msg_count_total = { 684 | let (tx, rx) = oneshot::channel(); 685 | msg_counter.send(MsgCountCmd::Get(tx)).unwrap(); 686 | rx.await.unwrap() 687 | }; 688 | info!("total messages sent = {msg_count_total} (communication overhead)"); 689 | 690 | msg_counter.send(MsgCountCmd::Reset).unwrap(); 691 | 692 | let validators_seed = { 693 | let f = get_snapshot_file(&opts, "validators_seed"); 694 | let seed = fs::read(&f).map_or(thread_rng().gen::(), |b| { 695 | b.as_slice().read_u64::().unwrap() 696 | }); 697 | fs::write(&f, seed.to_be_bytes()).unwrap(); 698 | seed 699 | }; 700 | let mut rng = StdRng::seed_from_u64(validators_seed); 701 | 702 | let validators = 703 | rand::seq::index::sample(&mut rng, nodes.len(), args.validators_number) 704 | .iter() 705 | .map(|i| nodes[i].clone()) 706 | .collect::>(); 707 | 708 | let mut futures = vec![]; 709 | 710 | for (i, validator) in validators.into_iter().enumerate() { 711 | let validator_node_id = validator.discovery.local_enr().node_id(); 712 | let samples = rand::seq::index::sample( 713 | &mut thread_rng(), 714 | keys.len(), 715 | args.samples_per_validator, 716 | ) 717 | .iter() 718 | .map(|i| (i, keys[i])) 719 | .collect_vec(); 720 | 721 | let parallelism = args.dissemination_args.parallelism; 722 | let lookup_method = args.lookup_method.clone(); 723 | 724 | clone_all!( 725 | node_ids, 726 | nodes_per_key, 727 | nodes_by_node, 728 | overlay_nodes_by_node 729 | ); 730 | 731 | futures.push(async move { 732 | let mut futures = FuturesUnordered::new(); 733 | let mut samples = samples; 734 | let mut num_waiting = 0usize; 735 | let mut num_success = 0usize; 736 | 737 | loop { 738 | if num_waiting < parallelism { 739 | if let Some((j, sample_key)) = samples.pop() { 740 | num_waiting += 1; 741 | clone_all!(validator, node_ids, nodes_per_key, lookup_method, nodes_by_node, overlay_nodes_by_node); 742 | futures.push(async move { 743 | if validator.samples.read().await.contains_key(&sample_key) { 744 | return Some((j, sample_key.clone(), b"yep".to_vec())) 745 | } 746 | 747 | info!("validator {i}: looking for a sample with key {}", NodeId::new(&DASContentKey::Sample(sample_key.raw()).content_id()).to_string()); 748 | 749 | match lookup_method { 750 | LookupMethod::Discv5FindValue => match validator.discovery.discv5.find_value(sample_key).await { 751 | Ok(res) => Some((j, sample_key.clone(), res)), 752 | Err(e) => match e { 753 | FindValueError::RequestError(e) => { 754 | error!("node {i} ({validator_node_id}) fail requesting sample {j} ({sample_key}): {e}"); 755 | None 756 | } 757 | FindValueError::RequestErrorWithEnrs((re, found_enrs)) => { 758 | error!("node {i} ({validator_node_id}) fail requesting sample {j} ({sample_key}): {re}"); 759 | 760 | let host_nodes = nodes_per_key.get(&sample_key).unwrap().clone(); 761 | let local_info = host_nodes.iter().map(|e| (e.to_string(), Key::from(e.clone()).log2_distance(&Key::from(sample_key.clone())).unwrap())).sorted_by_key(|(x, y)| *y).collect_vec(); 762 | let search_info = found_enrs.iter().map(|e| (e.node_id().to_string(), Key::from(e.node_id().clone()).log2_distance(&Key::from(sample_key.clone())).unwrap())).sorted_by_key(|(x, y)| *y).collect_vec(); 763 | info!("missing sample is stored in {:?}, visited nodes: {:?}", local_info, search_info); 764 | host_nodes.into_iter().for_each(|n| { 765 | let info = nodes_by_node.get(&n).map(|x| x.into_iter().map(|e| (e.to_string(), Key::from(e.clone()).log2_distance(&Key::from(sample_key.clone())).unwrap())).sorted_by_key(|(x, y)| *y).collect_vec()); 766 | info!("node {} that store missing samples are stored in ({:?}) {:?}", n.to_string(), info.as_ref().map(|x| x.len()), info) 767 | }); 768 | None 769 | } 770 | } 771 | } 772 | LookupMethod::OverlayFindContent => match validator.overlay.lookup_content(DASContentKey::Sample(sample_key.raw())) 773 | .await { 774 | Ok(res) => Some((j, sample_key.clone(), res)), 775 | Err(closest_nodes) => { 776 | error!("node {i} ({validator_node_id}) fail requesting sample {j} ({sample_key})"); 777 | 778 | let host_nodes = nodes_per_key.get(&sample_key).unwrap().clone(); 779 | let local_info = host_nodes.iter().map(|e| (e.to_string(), XorMetric::distance(&DASContentKey::Sample(Key::from(sample_key.clone()).hash.into()).content_id(), &e.raw()).log2())).collect_vec(); 780 | let search_info = closest_nodes.iter().map(|e| (e.to_string(), XorMetric::distance(&DASContentKey::Sample(Key::from(sample_key.clone()).hash.into()).content_id(), &e.raw()).log2().unwrap())).sorted_by_key(|(x, y)| *y).collect_vec(); 781 | info!("missing sample is stored in {:?}, visited nodes ({}): {:?}", local_info, search_info.len(), search_info); 782 | host_nodes.into_iter().for_each(|n| { 783 | let info = nodes_by_node.get(&n).map(|x| x.into_iter().map(|e| (e.to_string(), XorMetric::distance(&DASContentKey::Sample(Key::from(sample_key.clone()).hash.into()).content_id(), &e.raw()).log2().unwrap())).sorted_by_key(|(x, y)| *y).collect_vec()); 784 | info!("node {} that store missing samples are stored in ({:?}) {:?}", n.to_string(), info.as_ref().map(|x| x.len()), info) 785 | }); 786 | None 787 | } 788 | } 789 | } 790 | 791 | }); 792 | } 793 | } 794 | 795 | if let Some(res) = futures.next().await { 796 | num_waiting -= 1; 797 | if let Some((j, sample_key, value)) = res { 798 | debug!("[validator {i} ({validator_node_id})] success requesting sample {j} ({sample_key}): value='{}'", std::str::from_utf8(&value).unwrap()); 799 | 800 | num_success += 1; 801 | } 802 | } 803 | 804 | if num_waiting == 0 && samples.is_empty() { 805 | return num_success; 806 | } 807 | } 808 | }); 809 | } 810 | 811 | futures::future::join_all(futures) 812 | .instrument(info_span!("random-sampling")) 813 | .await 814 | .into_iter() 815 | .enumerate() 816 | .for_each(|(i, num_success)| { 817 | info!( 818 | "validator {i}: samples found {num_success}/{}", 819 | args.samples_per_validator 820 | ); 821 | }); 822 | 823 | let msg_count_total = { 824 | let (tx, rx) = oneshot::channel(); 825 | msg_counter.send(MsgCountCmd::Get(tx)).unwrap(); 826 | rx.await.unwrap() 827 | }; 828 | info!("total messages sent = {msg_count_total} (communication overhead)"); 829 | } 830 | _ => {} 831 | } 832 | } 833 | 834 | pub async fn handle_talk_request( 835 | from: NodeId, 836 | protocol: &[u8], 837 | message: Vec, 838 | node: DASNode, 839 | opts: Options, 840 | addr_book: Arc>>, 841 | node_ids: Vec, 842 | node_idx: usize, 843 | ) -> Vec { 844 | match protocol { 845 | DISSEMINATION_PROTOCOL_ID => match opts.simulation_case { 846 | SimulationCase::Disseminate(ref args) => handle_dissemination_request( 847 | from, 848 | message, 849 | node, 850 | opts.clone(), 851 | args.clone(), 852 | addr_book, 853 | node_ids, 854 | node_idx, 855 | ), 856 | SimulationCase::Sample(ref args) => handle_dissemination_request( 857 | from, 858 | message, 859 | node, 860 | opts.clone(), 861 | args.dissemination_args.clone(), 862 | addr_book, 863 | node_ids, 864 | node_idx, 865 | ), 866 | }, 867 | _ => panic!("unexpected protocol_id"), 868 | } 869 | } 870 | 871 | async fn disseminate_samples( 872 | keys: impl Iterator, 873 | opts: &Options, 874 | args: &DisseminationArgs, 875 | nodes: &Vec, 876 | addr_book: Arc>>, 877 | node_ids: &Vec, 878 | ) -> (HashMap, HashMap>) { 879 | match args.routing_strategy { 880 | RoutingStrategy::Recursive => { 881 | disseminate_samples_recursively(keys, opts, args, nodes, addr_book, node_ids).await 882 | } 883 | RoutingStrategy::Iterative => { 884 | disseminate_samples_iteratively(keys, opts, args, nodes, addr_book, node_ids).await 885 | } 886 | } 887 | } 888 | 889 | async fn disseminate_samples_recursively( 890 | keys: impl Iterator, 891 | opts: &Options, 892 | args: &DisseminationArgs, 893 | nodes: &Vec, 894 | addr_book: Arc>>, 895 | node_ids: &Vec, 896 | ) -> (HashMap, HashMap>) { 897 | let node = nodes[0].clone(); 898 | let local_node_id = node.discovery.discv5.local_enr().node_id(); 899 | 900 | let alloc = match args.batching_strategy { 901 | BatchingStrategy::BucketWise => { 902 | let local_view: HashMap<_, _> = node 903 | .discovery 904 | .discv5 905 | .kbuckets() 906 | .buckets_iter() 907 | .map(|kb| { 908 | kb.iter() 909 | .map(|e| e.key.preimage().clone()) 910 | .collect::>() 911 | }) 912 | .enumerate() 913 | .collect(); 914 | 915 | keys.into_iter() 916 | .flat_map(|k| { 917 | let i = 918 | BucketIndex::new(&Key::from(local_node_id.clone()).distance(&Key::from(k))) 919 | .unwrap() 920 | .get(); 921 | let local_nodes = local_view.get(&i).unwrap().clone(); 922 | /// if **replicate-all* then a receiver node applies forwards samples to more then one node in every k-bucket it handles 923 | let contacts_in_bucket = local_nodes.into_iter(); 924 | let mut forward_to: Vec<_> = match args.replicate_mode { 925 | ReplicatePolicy::ReplicateOne => contacts_in_bucket.take(1).collect(), 926 | ReplicatePolicy::ReplicateSome => { 927 | contacts_in_bucket.take(1 + &args.redundancy).collect() 928 | } 929 | ReplicatePolicy::ReplicateAll => contacts_in_bucket.collect(), 930 | }; 931 | 932 | if forward_to.is_empty() { 933 | forward_to.push(local_node_id); 934 | } 935 | 936 | forward_to 937 | .into_iter() 938 | .map(|n| (n, Key::from(k))) 939 | .collect::>() 940 | }) 941 | .into_group_map() 942 | } 943 | BatchingStrategy::DistanceWise => { 944 | let mut local_view = node 945 | .discovery 946 | .discv5 947 | .kbuckets() 948 | .buckets_iter() 949 | .flat_map(|kb| { 950 | kb.iter() 951 | .map(|e| e.key.preimage().clone()) 952 | .collect::>() 953 | }) 954 | .collect::>(); 955 | local_view.push(local_node_id); 956 | 957 | keys.into_iter() 958 | .flat_map(|k| { 959 | /// if **replicate-all* then a receiver node applies forwards samples to more then one node in every k-bucket it handles 960 | let contacts_in_bucket = local_view 961 | .clone() 962 | .into_iter() 963 | .sorted_by_key(|n| Key::from(n.clone()).distance(&Key::from(k))); 964 | let mut forward_to: Vec<_> = match args.replicate_mode { 965 | ReplicatePolicy::ReplicateOne => contacts_in_bucket.take(1).collect(), 966 | ReplicatePolicy::ReplicateSome => { 967 | contacts_in_bucket.take(1 + &args.redundancy).collect() 968 | } 969 | ReplicatePolicy::ReplicateAll => contacts_in_bucket.collect(), 970 | }; 971 | 972 | if forward_to.is_empty() { 973 | forward_to.push(local_node_id); 974 | } 975 | 976 | forward_to 977 | .into_iter() 978 | .map(|n| (n, Key::from(k))) 979 | .collect::>() 980 | }) 981 | .into_group_map() 982 | } 983 | }; 984 | 985 | let mut futures = vec![]; 986 | for (next, mut keys) in alloc.into_iter() { 987 | if next == local_node_id { 988 | debug!("no peers to forward {} keys to, saved locally", keys.len()); 989 | 990 | let mut samples = node.samples.write().await; 991 | let mut store = node.overlay.store.write(); 992 | keys.clone().into_iter().for_each(|k| { 993 | store 994 | .put(DASContentKey::Sample(k.preimage().raw()), b"yep".to_vec()) 995 | .unwrap(); 996 | 997 | match samples.entry(k.preimage().clone()) { 998 | Entry::Occupied(mut e) => e.get_mut().add_assign(1), 999 | Entry::Vacant(mut e) => { 1000 | e.insert(1); 1001 | } 1002 | } 1003 | }); 1004 | continue; 1005 | } 1006 | 1007 | let batch_id = nanoid!(8).into_bytes(); 1008 | let msg = { 1009 | let mut m = vec![]; 1010 | let mut w = BufWriter::new(&mut *m); 1011 | w.write(&*batch_id).unwrap(); 1012 | keys.iter().for_each(|k| { 1013 | let _ = w.write(&k.hash.to_vec()); 1014 | }); 1015 | w.buffer().to_vec() 1016 | }; 1017 | 1018 | let node = nodes[0].clone(); 1019 | let enr = node.discovery.find_enr(&next).unwrap(); 1020 | let addr_book = addr_book.clone(); 1021 | 1022 | { 1023 | let next_i = node_ids.iter().position(|e| *e == next).unwrap(); 1024 | debug!( 1025 | "node {0} ({}) sends {} keys for request (id={}) to {next_i} ({})", 1026 | node.discovery.local_enr().node_id(), 1027 | keys.len(), 1028 | hex::encode(&batch_id), 1029 | next 1030 | ); 1031 | } 1032 | 1033 | clone_all!(msg, keys); 1034 | futures.push(Box::pin(async move { 1035 | match opts.wire_protocol { 1036 | TalkWire::Discv5 => { 1037 | node.overlay 1038 | .send_elastic_talk_req(enr.clone(), DISSEMINATION_PROTOCOL_ID.to_vec(), msg) 1039 | .await 1040 | .unwrap(); 1041 | } 1042 | TalkWire::Libp2p => { 1043 | let (peer_id, addr) = 1044 | addr_book.read().await.get(&enr.node_id()).unwrap().clone(); 1045 | let _ = node 1046 | .libp2p 1047 | .talk_req(&peer_id, &addr, DISSEMINATION_PROTOCOL_ID, msg) 1048 | .await 1049 | .unwrap(); 1050 | } 1051 | } 1052 | })); 1053 | } 1054 | futures::future::join_all(futures) 1055 | .instrument(info_span!("dissemination")) 1056 | .await; 1057 | 1058 | let mut keys_per_node = HashMap::new(); 1059 | let mut nodes_per_key = HashMap::<_, Vec>::new(); 1060 | 1061 | for n in nodes { 1062 | let samples = n.samples.read().await; 1063 | samples.keys().for_each(|k| { 1064 | keys_per_node.insert(n.discovery.local_enr().node_id(), samples.len()); 1065 | 1066 | nodes_per_key 1067 | .entry(k.clone()) 1068 | .and_modify(|e| e.push(n.discovery.local_enr().node_id())) 1069 | .or_insert(vec![n.discovery.local_enr().node_id()]); 1070 | }) 1071 | } 1072 | 1073 | return (keys_per_node, nodes_per_key); 1074 | } 1075 | 1076 | #[derive(Clone, Debug, Eq, PartialEq, Hash)] 1077 | enum RemoteNodeId { 1078 | AskCloser(NodeId), 1079 | SendSamples(NodeId), 1080 | } 1081 | 1082 | impl RemoteNodeId { 1083 | fn unwrap(&self) -> NodeId { 1084 | match self { 1085 | RemoteNodeId::AskCloser(n) => n.clone(), 1086 | RemoteNodeId::SendSamples(n) => n.clone(), 1087 | } 1088 | } 1089 | } 1090 | 1091 | async fn disseminate_samples_iteratively( 1092 | keys: impl Iterator, 1093 | opts: &Options, 1094 | args: &DisseminationArgs, 1095 | nodes: &Vec, 1096 | addr_book: Arc>>, 1097 | node_ids: &Vec, 1098 | ) -> (HashMap, HashMap>) { 1099 | let local_node = nodes[0].clone(); 1100 | let local_node_id = local_node.discovery.discv5.local_enr().node_id(); 1101 | let local_key = Key::from(local_node_id.clone()); 1102 | 1103 | let alloc = match args.batching_strategy { 1104 | BatchingStrategy::BucketWise => { 1105 | let local_view: HashMap<_, _> = local_node 1106 | .discovery 1107 | .discv5 1108 | .kbuckets() 1109 | .buckets_iter() 1110 | .map(|kb| { 1111 | kb.iter() 1112 | .map(|e| e.key.preimage().clone()) 1113 | .collect::>() 1114 | }) 1115 | .enumerate() 1116 | .collect(); 1117 | 1118 | keys.into_group_map_by(|k| { 1119 | BucketIndex::new(&local_key.distance(&Key::from(k.clone()))) 1120 | .unwrap() 1121 | .get() 1122 | }) 1123 | .into_iter() 1124 | .map(|(i, keys)| { 1125 | let local_nodes = local_view.get(&i).unwrap().clone(); 1126 | /// if **replicate-all* then a receiver node applies forwards samples to more then one node in every k-bucket it handles 1127 | let contacts_in_bucket = local_nodes.into_iter(); 1128 | 1129 | let mut forward_to: Vec<_> = match args.replicate_mode { 1130 | ReplicatePolicy::ReplicateOne => contacts_in_bucket.take(1).collect(), 1131 | ReplicatePolicy::ReplicateSome => { 1132 | contacts_in_bucket.take(1 + &args.redundancy).collect() 1133 | } 1134 | ReplicatePolicy::ReplicateAll => contacts_in_bucket.collect(), 1135 | }; 1136 | 1137 | if forward_to.is_empty() { 1138 | forward_to.push(local_node_id); 1139 | } 1140 | 1141 | (keys, forward_to) 1142 | }) 1143 | .collect_vec() 1144 | } 1145 | BatchingStrategy::DistanceWise => unimplemented!(), 1146 | }; 1147 | 1148 | let mut futures = FuturesUnordered::< 1149 | Pin, DisseminationMsg)> + Send>>, 1150 | >::new(); 1151 | let mut pending_requests = HashMapDelay::new(Duration::from_secs(opts.request_timeout)); 1152 | let mut keys_per_batch = HashMap::new(); 1153 | let mut requests_per_batch = HashMap::>::new(); 1154 | let mut requests_per_node = HashMap::::new(); 1155 | 1156 | for (mut keys, nodes) in alloc.into_iter() { 1157 | let batch_id = utp::stream::rand(); 1158 | keys_per_batch.insert(batch_id, keys.clone()); 1159 | 1160 | let msg = DisseminationMsg::Keys((batch_id, keys.iter().map(|k| k.raw()).collect_vec())); 1161 | 1162 | for remote in nodes { 1163 | if remote == local_node_id { 1164 | let mut samples = local_node.samples.write().await; 1165 | let mut store = local_node.overlay.store.write(); 1166 | keys.clone().into_iter().for_each(|k| { 1167 | store 1168 | .put(DASContentKey::Sample(k.raw()), b"yep".to_vec()) 1169 | .unwrap(); 1170 | 1171 | match samples.entry(k.clone()) { 1172 | Entry::Occupied(mut e) => e.get_mut().add_assign(1), 1173 | Entry::Vacant(mut e) => { 1174 | e.insert(1); 1175 | } 1176 | } 1177 | }); 1178 | continue; 1179 | } 1180 | 1181 | let enr = local_node.discovery.find_enr(&remote).unwrap(); 1182 | 1183 | clone_all!(opts, local_node, msg, addr_book); 1184 | let request_id = utp::stream::rand(); 1185 | pending_requests.insert(request_id.clone(), ()); 1186 | requests_per_node.insert(remote.clone(), request_id); 1187 | requests_per_batch 1188 | .entry(batch_id) 1189 | .and_modify(|e| { 1190 | e.insert(request_id); 1191 | }) 1192 | .or_insert(HashSet::from([request_id])); 1193 | futures.push(Box::pin(async move { 1194 | let message = match opts.wire_protocol { 1195 | TalkWire::Discv5 => local_node 1196 | .overlay 1197 | .send_elastic_talk_req( 1198 | enr.clone(), 1199 | DISSEMINATION_PROTOCOL_ID.to_vec(), 1200 | msg.as_ssz_bytes(), 1201 | ) 1202 | .await 1203 | .unwrap(), 1204 | TalkWire::Libp2p => { 1205 | let (peer_id, addr) = 1206 | addr_book.read().await.get(&enr.node_id()).unwrap().clone(); 1207 | local_node 1208 | .libp2p 1209 | .talk_req( 1210 | &peer_id, 1211 | &addr, 1212 | DISSEMINATION_PROTOCOL_ID, 1213 | msg.as_ssz_bytes(), 1214 | ) 1215 | .await 1216 | .unwrap() 1217 | } 1218 | }; 1219 | 1220 | let response: DisseminationMsg = 1221 | DisseminationMsg::from_ssz_bytes(&*message).unwrap(); 1222 | 1223 | (request_id, batch_id, enr, response) 1224 | })); 1225 | } 1226 | } 1227 | 1228 | let mut responses_per_batch = HashMap::>::new(); 1229 | 1230 | clone_all!(args, opts); 1231 | let mut outbound_requests = VecDeque::new(); 1232 | let parallelism = args.parallelism; 1233 | let mut keys_per_node = HashMap::<_, HashSet>::new(); 1234 | tokio::spawn(async move { 1235 | loop { 1236 | select! { 1237 | Some((request_id, batch_id, remote, response)) = futures.next() => match response { 1238 | DisseminationMsg::CloserNodes(closer_nodes) => { 1239 | let mut pending_in_batch = requests_per_batch.get_mut(&batch_id).unwrap(); 1240 | if !pending_in_batch.remove(&request_id) { 1241 | continue 1242 | } 1243 | 1244 | let mut closer_node_ids = vec![]; 1245 | closer_nodes.into_iter().for_each(|enr| { 1246 | closer_node_ids.push(enr.node_id()); 1247 | if local_node.discovery.find_enr_or_cache(&enr.node_id()).is_none() { 1248 | local_node.discovery.node_addr_cache 1249 | .write() 1250 | .put(enr.node_id(), NodeAddress { enr: enr.into(), socket_addr: None }); 1251 | } 1252 | }); 1253 | 1254 | let remote_key = Key::from(remote.node_id()); 1255 | 1256 | let remote_view = closer_node_ids.into_iter().into_group_map_by(|closer| { 1257 | BucketIndex::new(&remote_key.distance(&Key::from(closer.clone()))) 1258 | .unwrap() 1259 | .get() 1260 | }); 1261 | 1262 | let keys = keys_per_batch.get(&batch_id).unwrap(); 1263 | 1264 | let new_alloc = keys 1265 | .clone() 1266 | .into_iter() 1267 | .flat_map(|key| { 1268 | let i = BucketIndex::new(&remote_key.distance(&Key::from(key.clone()))) 1269 | .unwrap() 1270 | .get(); 1271 | 1272 | let mut remote_nodes = remote_view.get(&i).map(|e| e.into_iter().map(|n| RemoteNodeId::AskCloser(n.clone())).collect_vec()).or_else(||Some(vec![])).unwrap(); 1273 | 1274 | if remote_nodes.is_empty() { 1275 | remote_nodes.push(RemoteNodeId::SendSamples(remote.node_id())); 1276 | } 1277 | 1278 | remote_nodes 1279 | .into_iter() 1280 | .map(|node| (key, node)) 1281 | .collect_vec() 1282 | }) 1283 | .collect_vec(); 1284 | 1285 | responses_per_batch.entry(batch_id).and_modify(|e| e.extend(new_alloc.clone())).or_insert(new_alloc); 1286 | 1287 | if !pending_in_batch.is_empty() { 1288 | requests_per_node.remove(&remote.node_id()).unwrap(); 1289 | pending_requests.remove(&request_id).unwrap(); 1290 | continue 1291 | } 1292 | 1293 | let mut new_alloc = responses_per_batch.remove(&batch_id).unwrap(); 1294 | 1295 | let new_alloc = new_alloc.into_iter() 1296 | .into_group_map() 1297 | .into_iter() 1298 | .flat_map(|(key, nodes)| { 1299 | let discovered_nodes = nodes.into_iter().unique().sorted_by_key(|node| Key::from(node.unwrap()).distance(&Key::from(key))); 1300 | let num_total = discovered_nodes.len(); 1301 | let discovered_nodes = match args.replicate_mode { 1302 | ReplicatePolicy::ReplicateOne => discovered_nodes.take(1), 1303 | ReplicatePolicy::ReplicateSome => { 1304 | discovered_nodes.take(1 + &args.redundancy) 1305 | } 1306 | ReplicatePolicy::ReplicateAll => discovered_nodes.take(num_total), 1307 | }; 1308 | 1309 | discovered_nodes 1310 | .map(|node| (node, key.clone())) 1311 | .collect_vec() 1312 | }).into_group_map(); 1313 | 1314 | for (next, mut keys) in new_alloc.into_iter() { 1315 | let request_id = utp::stream::rand(); 1316 | let batch_id = utp::stream::rand(); 1317 | keys_per_batch.insert(batch_id, keys.clone()); 1318 | 1319 | let (next, msg) = match next { 1320 | RemoteNodeId::AskCloser(next) => 1321 | (next, DisseminationMsg::Keys((request_id, keys.iter().map(|k| k.raw()).collect_vec()))), 1322 | RemoteNodeId::SendSamples(next) => { 1323 | keys = match keys_per_node.entry(next.clone()) { 1324 | Entry::Vacant(mut e) => { 1325 | e.insert(HashSet::from_iter(keys.clone())); 1326 | keys 1327 | }, 1328 | Entry::Occupied(mut e) => { 1329 | let new_keys = HashSet::from_iter(keys.clone()).difference(e.get()).map(|e| (*e).clone()).collect::>(); 1330 | e.get_mut().extend(&new_keys); 1331 | new_keys 1332 | } 1333 | }; 1334 | let samples = keys.iter().map(|key| (key.raw(), b"yep".to_vec())).collect_vec(); 1335 | (next, DisseminationMsg::Samples(samples)) 1336 | } 1337 | }; 1338 | 1339 | if keys.is_empty() { 1340 | continue 1341 | } 1342 | 1343 | if next == local_node_id { 1344 | warn!("remote send us"); 1345 | continue 1346 | } 1347 | 1348 | let enr = local_node.discovery.find_enr_or_cache(&next).unwrap(); 1349 | 1350 | if let DisseminationMsg::Keys(..) = &msg { 1351 | requests_per_batch.entry(batch_id).and_modify(|e| { e.insert(request_id); }).or_insert(HashSet::from([request_id])); 1352 | } 1353 | outbound_requests.push_back((enr, request_id, batch_id, msg.clone())) 1354 | } 1355 | 1356 | requests_per_node.remove(&remote.node_id()).unwrap(); 1357 | pending_requests.remove(&request_id).unwrap(); 1358 | if pending_requests.is_empty() && outbound_requests.is_empty() { 1359 | break 1360 | } 1361 | } 1362 | DisseminationMsg::Received(_) => { 1363 | pending_requests.remove(&request_id).unwrap(); 1364 | requests_per_node.remove(&remote.node_id()).unwrap(); 1365 | if pending_requests.is_empty() && outbound_requests.is_empty() { 1366 | break 1367 | } 1368 | } 1369 | _ => unimplemented!() 1370 | }, 1371 | Some(Ok((request_id, _))) = pending_requests.next() => { 1372 | error!("request {request_id} has timed out"); 1373 | if pending_requests.is_empty() { 1374 | break 1375 | } 1376 | } 1377 | } 1378 | 1379 | if pending_requests.len() < parallelism { 1380 | if let Some((enr, request_id, batch_id, msg)) = outbound_requests.pop_front() { 1381 | clone_all!(opts, local_node, addr_book); 1382 | let node_id = enr.node_id(); 1383 | if requests_per_node.contains_key(&node_id) { 1384 | outbound_requests.push_back((enr, request_id, batch_id, msg)); 1385 | continue 1386 | } 1387 | requests_per_node.insert(node_id.clone(), request_id); 1388 | pending_requests.insert(request_id.clone(), ()); 1389 | 1390 | futures.push(Box::pin(async move { 1391 | let message = match opts.wire_protocol { 1392 | TalkWire::Discv5 => { 1393 | local_node.overlay 1394 | .send_elastic_talk_req(enr.clone(), DISSEMINATION_PROTOCOL_ID.to_vec(), msg.as_ssz_bytes()) 1395 | .await 1396 | .unwrap() 1397 | } 1398 | TalkWire::Libp2p => { 1399 | let addr_book = addr_book.read().await; 1400 | let (peer_id, addr) = addr_book.get(&enr.node_id()).unwrap(); 1401 | 1402 | local_node 1403 | .libp2p 1404 | .talk_req(&peer_id, &addr, DISSEMINATION_PROTOCOL_ID, msg.as_ssz_bytes()) 1405 | .await 1406 | .unwrap() 1407 | } 1408 | }; 1409 | 1410 | let response: DisseminationMsg = DisseminationMsg::from_ssz_bytes(&*message).unwrap(); 1411 | 1412 | (request_id, batch_id, enr, response) 1413 | })); 1414 | } 1415 | } 1416 | 1417 | } 1418 | }).instrument(info_span!("dissemination")).await.unwrap(); 1419 | 1420 | let mut keys_per_node = HashMap::new(); 1421 | let mut nodes_per_key = HashMap::<_, Vec>::new(); 1422 | 1423 | for n in nodes { 1424 | let samples = n.samples.read().await; 1425 | samples.keys().for_each(|k| { 1426 | keys_per_node.insert(n.discovery.local_enr().node_id(), samples.len()); 1427 | 1428 | nodes_per_key 1429 | .entry(k.clone()) 1430 | .and_modify(|e| e.push(n.discovery.local_enr().node_id())) 1431 | .or_insert(vec![n.discovery.local_enr().node_id()]); 1432 | }) 1433 | } 1434 | 1435 | return (keys_per_node, nodes_per_key); 1436 | } 1437 | 1438 | fn handle_dissemination_request( 1439 | from: NodeId, 1440 | message: Vec, 1441 | node: DASNode, 1442 | opts: Options, 1443 | args: DisseminationArgs, 1444 | addr_book: Arc>>, 1445 | node_ids: Vec, 1446 | node_idx: usize, 1447 | ) -> Vec { 1448 | let promise_id = utp::stream::rand(); 1449 | 1450 | tokio::spawn(async move { 1451 | let message = { 1452 | let content: ElasticPacket = ElasticPacket::from_ssz_bytes(&*message).unwrap(); 1453 | match content { 1454 | ElasticPacket::Data(bytes) => bytes.to_vec(), 1455 | ElasticPacket::ConnectionId(conn_id) => { 1456 | let conn_id = u16::from_be(conn_id); 1457 | let enr = node.discovery.find_enr_or_cache(&from).unwrap(); 1458 | node.overlay 1459 | .init_find_content_stream(enr, conn_id) 1460 | .await 1461 | .unwrap() 1462 | } 1463 | ElasticPacket::Result((promise_id, res)) => { 1464 | let res = match res { 1465 | ElasticResult::Data(res) => res, 1466 | ElasticResult::ConnectionId(conn_id) => { 1467 | let conn_id = u16::from_be(conn_id); 1468 | let enr = node.discovery.find_enr_or_cache(&from).unwrap(); 1469 | node.overlay 1470 | .init_find_content_stream(enr, conn_id) 1471 | .await 1472 | .unwrap() 1473 | } 1474 | }; 1475 | 1476 | match opts.wire_protocol { 1477 | TalkWire::Discv5 => node.overlay.handle_promise_result(promise_id, res), 1478 | TalkWire::Libp2p => node.libp2p.handle_promise_result(promise_id, res), 1479 | } 1480 | 1481 | return; 1482 | } 1483 | _ => unreachable!(), 1484 | } 1485 | }; 1486 | 1487 | let from_i = node_ids.iter().position(|e| *e == from).unwrap(); 1488 | let local_node_id = node.discovery.local_enr().node_id(); 1489 | 1490 | let (keys, id) = match args.routing_strategy { 1491 | RoutingStrategy::Recursive => { 1492 | let mut r = BufReader::new(&*message); 1493 | let mut keys = vec![]; 1494 | 1495 | let mut id = [0; 8]; 1496 | r.read(&mut id).unwrap(); 1497 | let id = id.to_vec(); 1498 | 1499 | loop { 1500 | let mut b = [0; 32]; 1501 | if r.read(&mut b).unwrap() < 32 { 1502 | break; 1503 | } 1504 | 1505 | keys.push(NodeId::new(&b)) 1506 | } 1507 | 1508 | (keys, id) 1509 | } 1510 | RoutingStrategy::Iterative => { 1511 | let msg: DisseminationMsg = DisseminationMsg::from_ssz_bytes(&message).unwrap(); 1512 | 1513 | match msg { 1514 | DisseminationMsg::Keys((id, keys)) => ( 1515 | keys.into_iter().map(|raw| NodeId::new(&raw)).collect_vec(), 1516 | id.to_be_bytes().to_vec(), 1517 | ), 1518 | DisseminationMsg::Samples(kvp) => { 1519 | { 1520 | let mut samples = node.samples.write().await; 1521 | let mut store = node.overlay.store.write(); 1522 | for (raw_key, val) in kvp { 1523 | store 1524 | .put(DASContentKey::Sample(raw_key.clone()), val) 1525 | .unwrap(); 1526 | 1527 | match samples.entry(NodeId::new(&raw_key)) { 1528 | Entry::Occupied(mut e) => e.get_mut().add_assign(1), 1529 | Entry::Vacant(mut e) => { 1530 | e.insert(1); 1531 | } 1532 | } 1533 | } 1534 | } 1535 | 1536 | send_results( 1537 | &node, 1538 | &opts, 1539 | &from, 1540 | promise_id, 1541 | DisseminationMsg::Received(0).as_ssz_bytes(), 1542 | addr_book, 1543 | ) 1544 | .await; 1545 | return; 1546 | } 1547 | _ => unreachable!(), 1548 | } 1549 | } 1550 | }; 1551 | 1552 | { 1553 | debug!( 1554 | "node {node_idx} ({}) attempts to get lock for request (id={}) from {from_i} ({})", 1555 | node.discovery.local_enr().node_id(), 1556 | hex::encode(&id), 1557 | from 1558 | ); 1559 | let mut handled_ids = node.handled_ids.write().await; 1560 | if handled_ids.contains_key(&id) && args.forward_mode != ForwardPolicy::ForwardAll { 1561 | debug!( 1562 | "node {node_idx} ({}) skipped request (id={}) from {from_i} ({})", 1563 | node.discovery.local_enr().node_id(), 1564 | hex::encode(&id), 1565 | from 1566 | ); 1567 | 1568 | // todo: send diffrently for iter and recv d2s 1569 | send_results( 1570 | &node, 1571 | &opts, 1572 | &from, 1573 | promise_id, 1574 | DisseminationMsg::CloserNodes(vec![]).as_ssz_bytes(), 1575 | addr_book, 1576 | ) 1577 | .await; 1578 | return; 1579 | } else { 1580 | debug!( 1581 | "node {node_idx} ({}) received request (id={}) from {from_i} ({})", 1582 | node.discovery.local_enr().node_id(), 1583 | hex::encode(&id), 1584 | from 1585 | ); 1586 | match handled_ids.entry(id.clone()) { 1587 | Entry::Occupied(mut e) => e.get_mut().add_assign(1), 1588 | Entry::Vacant(mut e) => { 1589 | e.insert(1); 1590 | } 1591 | }; 1592 | drop(handled_ids); 1593 | } 1594 | } 1595 | 1596 | // debug!("node {node_idx} ({}) receives {:?} keys for request (id={}) from {from_i} ({})", node.discv5.local_enr().node_id(), keys.iter().map(|e| e.to_string()).collect_vec(), hex::encode(&id), from); 1597 | 1598 | let alloc = match args.batching_strategy { 1599 | BatchingStrategy::BucketWise => { 1600 | let local_view: HashMap<_, _> = node 1601 | .discovery 1602 | .discv5 1603 | .kbuckets() 1604 | .buckets_iter() 1605 | .map(|kb| kb.iter().map(|e| e.key.preimage().clone()).collect_vec()) 1606 | .enumerate() 1607 | .collect(); 1608 | 1609 | keys.into_iter() 1610 | .flat_map(|k| { 1611 | let i = BucketIndex::new( 1612 | &Key::from(local_node_id.clone()).distance(&Key::from(k)), 1613 | ) 1614 | .unwrap() 1615 | .get(); 1616 | let local_nodes = local_view.get(&i).unwrap().clone(); 1617 | /// if **replicate-all* then a receiver node applies forwards samples to more then one node in every k-bucket it handles 1618 | let contacts_in_bucket = local_nodes.into_iter().filter(|e| *e != from); 1619 | let mut forward_to: Vec<_> = match args.routing_strategy { 1620 | RoutingStrategy::Recursive => match args.replicate_mode { 1621 | ReplicatePolicy::ReplicateOne => { 1622 | contacts_in_bucket.take(1).collect() 1623 | } 1624 | ReplicatePolicy::ReplicateSome => { 1625 | contacts_in_bucket.take(1 + &args.redundancy).collect() 1626 | } 1627 | ReplicatePolicy::ReplicateAll => contacts_in_bucket.collect(), 1628 | }, 1629 | RoutingStrategy::Iterative => contacts_in_bucket.take(2).collect(), 1630 | }; 1631 | 1632 | if forward_to.is_empty() { 1633 | forward_to.push(local_node_id); 1634 | } 1635 | 1636 | forward_to 1637 | .into_iter() 1638 | .map(|n| (n, Key::from(k))) 1639 | .collect::>() 1640 | }) 1641 | .into_group_map() 1642 | } 1643 | BatchingStrategy::DistanceWise => { 1644 | let mut local_view = node 1645 | .discovery 1646 | .discv5 1647 | .kbuckets() 1648 | .buckets_iter() 1649 | .flat_map(|kb| { 1650 | kb.iter() 1651 | .map(|e| e.key.preimage().clone()) 1652 | .collect::>() 1653 | }) 1654 | .collect::>(); 1655 | local_view.push(local_node_id.clone()); 1656 | 1657 | keys.into_iter() 1658 | .flat_map(|k| { 1659 | let contacts_in_bucket = local_view 1660 | .clone() 1661 | .into_iter() 1662 | .filter(|e| *e != from) 1663 | .sorted_by_key(|n| Key::from(n.clone()).distance(&Key::from(k))); 1664 | let mut forward_to: Vec<_> = match args.routing_strategy { 1665 | RoutingStrategy::Recursive => match args.replicate_mode { 1666 | ReplicatePolicy::ReplicateOne => { 1667 | contacts_in_bucket.take(1).collect() 1668 | } 1669 | ReplicatePolicy::ReplicateSome => { 1670 | contacts_in_bucket.take(1 + &args.redundancy).collect() 1671 | } 1672 | ReplicatePolicy::ReplicateAll => contacts_in_bucket.collect(), 1673 | }, 1674 | RoutingStrategy::Iterative => contacts_in_bucket.take(2).collect(), 1675 | }; 1676 | 1677 | forward_to 1678 | .into_iter() 1679 | .map(|n| (n, Key::from(k))) 1680 | .collect::>() 1681 | }) 1682 | .into_group_map() 1683 | } 1684 | }; 1685 | 1686 | if args.routing_strategy == RoutingStrategy::Iterative { 1687 | let closer_nodes = alloc 1688 | .keys() 1689 | .into_iter() 1690 | .filter_map(|nid| { 1691 | node.discovery 1692 | .find_enr_or_cache(nid) 1693 | .map(|enr| SszEnr::new(enr)) 1694 | }) 1695 | .collect_vec(); 1696 | 1697 | send_results( 1698 | &node, 1699 | &opts, 1700 | &from, 1701 | promise_id, 1702 | DisseminationMsg::CloserNodes(closer_nodes).as_ssz_bytes(), 1703 | addr_book, 1704 | ) 1705 | .await; 1706 | return; 1707 | } 1708 | 1709 | let mut futures = FuturesUnordered::new(); 1710 | 1711 | for (next, keys) in alloc.into_iter() { 1712 | if next == local_node_id { 1713 | let mut samples = node.samples.write().await; 1714 | let mut store = node.overlay.store.write(); 1715 | keys.clone().into_iter().for_each(|k| { 1716 | store 1717 | .put(DASContentKey::Sample(k.preimage().raw()), b"yep".to_vec()) 1718 | .unwrap(); 1719 | 1720 | match samples.entry(k.preimage().clone()) { 1721 | Entry::Occupied(mut e) => e.get_mut().add_assign(1), 1722 | Entry::Vacant(mut e) => { 1723 | e.insert(1); 1724 | } 1725 | } 1726 | }); 1727 | 1728 | continue; 1729 | } 1730 | 1731 | let enr = node.discovery.find_enr(&next).unwrap(); 1732 | 1733 | let next_i = node_ids.iter().position(|e| *e == next).unwrap(); 1734 | debug!( 1735 | "node {node_idx} ({}) sends {:?} keys for request (id={}) to {next_i} ({})", 1736 | node.discovery.local_enr().node_id(), 1737 | keys.iter() 1738 | .map(|e| e.preimage().to_string()) 1739 | .collect::>(), 1740 | hex::encode(&id), 1741 | next 1742 | ); 1743 | 1744 | let msg = { 1745 | let mut m = vec![]; 1746 | let mut w = BufWriter::new(&mut *m); 1747 | w.write(&mut id.clone()).unwrap(); 1748 | keys.clone().into_iter().for_each(|k| { 1749 | let _ = w.write(&*k.hash.to_vec()); 1750 | }); 1751 | w.buffer().to_vec() 1752 | }; 1753 | 1754 | { 1755 | clone_all!(node, addr_book, opts, id, keys); 1756 | futures.push(async move { 1757 | match opts.wire_protocol { 1758 | TalkWire::Discv5 => node 1759 | .overlay 1760 | .send_elastic_talk_req( 1761 | enr.clone(), 1762 | DISSEMINATION_PROTOCOL_ID.to_vec(), 1763 | msg, 1764 | ) 1765 | .await 1766 | .map_err(|e| eyre::eyre!("{e}")), 1767 | TalkWire::Libp2p => { 1768 | let (peer_id, addr) = 1769 | addr_book.read().await.get(&enr.node_id()).unwrap().clone(); 1770 | node.libp2p 1771 | .talk_req(&peer_id, &addr, DISSEMINATION_PROTOCOL_ID, msg) 1772 | .await 1773 | } 1774 | } 1775 | .map_err(|e| { 1776 | eyre::eyre!( 1777 | "error making request (id={}) from {} to {}: {e}", 1778 | hex::encode(&id), 1779 | node.discovery.local_enr().node_id(), 1780 | enr.node_id(), 1781 | ) 1782 | }) 1783 | }); 1784 | } 1785 | } 1786 | 1787 | while let Some(resp) = futures.next().await { 1788 | resp.unwrap(); 1789 | } 1790 | 1791 | send_results(&node, &opts, &from, promise_id, vec![], addr_book).await; 1792 | }); 1793 | 1794 | ElasticPacket::Promise(promise_id).as_ssz_bytes() 1795 | } 1796 | 1797 | async fn send_results( 1798 | node: &DASNode, 1799 | opts: &Options, 1800 | to: &NodeId, 1801 | promise_id: u16, 1802 | msg: Vec, 1803 | addr_book: Arc>>, 1804 | ) { 1805 | match opts.wire_protocol { 1806 | TalkWire::Discv5 => { 1807 | let enr = node.discovery.find_enr_or_cache(to).unwrap(); 1808 | node.overlay 1809 | .send_result(enr, DISSEMINATION_PROTOCOL_ID.to_vec(), promise_id, msg) 1810 | .await 1811 | .unwrap(); 1812 | } 1813 | TalkWire::Libp2p => { 1814 | let (peer_id, addr) = addr_book.read().await.get(to).unwrap().clone(); 1815 | node.libp2p 1816 | .send_message( 1817 | &peer_id, 1818 | &addr, 1819 | DISSEMINATION_PROTOCOL_ID, 1820 | ElasticPacket::Result((promise_id, ElasticResult::Data(msg))).as_ssz_bytes(), 1821 | ) 1822 | .await 1823 | .unwrap(); 1824 | } 1825 | } 1826 | } 1827 | 1828 | async fn handle_sampling_request( 1829 | _from: NodeId, 1830 | key: &NodeId, 1831 | node: &DASNode, 1832 | opts: &Options, 1833 | ) -> Option> { 1834 | let mut samples = node.samples.read().await; 1835 | 1836 | debug!("receive sampling request, have {} samples total, distance to requested key={:?}, have requested key = {}", samples.len(), Key::from(node.discovery.discv5.local_enr().node_id()).log2_distance(&Key::from(key.clone())), samples.contains_key(key)); 1837 | 1838 | samples.get(key).map(|e| b"yep".to_vec()) 1839 | } 1840 | 1841 | fn get_snapshot_file>(opts: &Options, file: S) -> PathBuf { 1842 | match &*opts.snapshot { 1843 | "new" | "last" => { 1844 | let mut paths: Vec<_> = fs::read_dir(&opts.cache_dir) 1845 | .unwrap() 1846 | .map(|r| r.unwrap()) 1847 | .collect(); 1848 | paths.sort_by_key(|dir| dir.metadata().unwrap().created().unwrap()); 1849 | paths.last().unwrap().path().join(file.as_ref()) 1850 | } 1851 | snap => PathBuf::from(&opts.cache_dir) 1852 | .join(snap) 1853 | .join(file.as_ref()), 1854 | } 1855 | } 1856 | -------------------------------------------------------------------------------- /src/messages.rs: -------------------------------------------------------------------------------- 1 | use crate::DASContentKey; 2 | use discv5::Enr; 3 | use discv5_overlay::portalnet::types::messages::SszEnr; 4 | use enr::NodeId; 5 | use itertools::Itertools; 6 | use ssz::{Decode, Encode}; 7 | use ssz_derive::{Decode, Encode}; 8 | use std::fmt; 9 | use std::fmt::{Debug, Display, Formatter, Pointer}; 10 | 11 | #[derive(Clone, Decode, Encode, PartialEq)] 12 | #[ssz(enum_behaviour = "union")] 13 | pub enum DisseminationMsg { 14 | Keys((u16, Vec<[u8; 32]>)), 15 | CloserNodes(Vec), 16 | Samples(Vec<([u8; 32], Vec)>), 17 | Received(u16), 18 | } 19 | 20 | impl Debug for DisseminationMsg { 21 | fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { 22 | match self { 23 | DisseminationMsg::Keys((request_id, keys)) => { 24 | write!( 25 | f, 26 | "Keys({request_id}, {:?})", 27 | keys.into_iter() 28 | .map(|e| NodeId::new(e).to_string()) 29 | .collect_vec() 30 | ) 31 | } 32 | DisseminationMsg::CloserNodes(enrs) => write!(f, "CloserNodes({:?})", enrs), 33 | DisseminationMsg::Samples(samples) => { 34 | write!( 35 | f, 36 | "Samples({:?})", 37 | samples 38 | .into_iter() 39 | .map(|(key, val)| (NodeId::new(key).to_string(), hex::encode(val))) 40 | .collect_vec() 41 | ) 42 | } 43 | DisseminationMsg::Received(_) => write!(f, "Received"), 44 | } 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /src/overlay.rs: -------------------------------------------------------------------------------- 1 | use async_trait::async_trait; 2 | use discv5_overlay::portalnet::types::content_key::{HistoryContentKey, OverlayContentKey}; 3 | use discv5_overlay::types::validation::Validator; 4 | use enr::k256::sha2::Sha256; 5 | use sha3::{Digest, Keccak256}; 6 | use ssz::{Decode, Encode}; 7 | use ssz_derive::{Decode, Encode}; 8 | use std::fmt; 9 | use std::fmt::{Display, Formatter}; 10 | 11 | /// A content key in the DAS overlay network. 12 | #[derive(Clone, Debug, Decode, Encode, PartialEq)] 13 | #[ssz(enum_behaviour = "union")] 14 | pub enum DASContentKey { 15 | Sample([u8; 32]), 16 | } 17 | 18 | #[allow(clippy::from_over_into)] 19 | impl Into> for DASContentKey { 20 | fn into(self) -> Vec { 21 | self.as_ssz_bytes() 22 | } 23 | } 24 | 25 | impl TryFrom> for DASContentKey { 26 | type Error = &'static str; 27 | 28 | fn try_from(value: Vec) -> Result { 29 | match DASContentKey::from_ssz_bytes(&value) { 30 | Ok(key) => Ok(key), 31 | Err(_err) => { 32 | println!("unable to decode DASContentKey"); 33 | Err("Unable to decode SSZ") 34 | } 35 | } 36 | } 37 | } 38 | 39 | impl Display for DASContentKey { 40 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 41 | let s = match self { 42 | Self::Sample(b) => format!("sample: {}", hex::encode(b)), 43 | }; 44 | 45 | write!(f, "{}", s) 46 | } 47 | } 48 | 49 | impl OverlayContentKey for DASContentKey { 50 | fn content_id(&self) -> [u8; 32] { 51 | match self { 52 | DASContentKey::Sample(b) => b.clone(), 53 | } 54 | } 55 | } 56 | 57 | pub struct DASValidator; 58 | 59 | #[async_trait] 60 | impl Validator for DASValidator { 61 | async fn validate_content( 62 | &self, 63 | content_key: &DASContentKey, 64 | content: &[u8], 65 | ) -> anyhow::Result<()> 66 | // where 67 | // DASContentKey: 'async_trait, 68 | { 69 | match content_key { 70 | DASContentKey::Sample(_) => Ok(()), 71 | } 72 | } 73 | } 74 | -------------------------------------------------------------------------------- /src/utils.rs: -------------------------------------------------------------------------------- 1 | use eyre::eyre; 2 | use std::iter; 3 | use tokio::sync::oneshot; 4 | 5 | #[macro_export] 6 | macro_rules! clone_all { 7 | ($i:ident) => { 8 | let $i = $i.clone(); 9 | }; 10 | ($i:ident, $($tt:tt)*) => { 11 | clone_all!($i); 12 | clone_all!($($tt)*); 13 | }; 14 | ($this:ident . $i:ident) => { 15 | let $i = $this.$i.clone(); 16 | }; 17 | ($this:ident . $i:ident, $($tt:tt)*) => { 18 | clone_all!($this . $i); 19 | clone_all!($($tt)*); 20 | }; 21 | } 22 | 23 | #[derive(Debug)] 24 | pub enum MsgCountCmd { 25 | Reset, 26 | Get(oneshot::Sender), 27 | Increment, 28 | } 29 | 30 | pub fn encode_result_for_discv5(r: eyre::Result>) -> Vec { 31 | match r { 32 | Ok(mut b) => iter::once(1).chain(b).collect(), 33 | Err(e) => iter::once(0).chain(e.to_string().into_bytes()).collect(), 34 | } 35 | } 36 | 37 | pub fn decode_result_from_discv5(b: Vec) -> eyre::Result> { 38 | match b.first().unwrap() { 39 | 1 => Ok(b[1..b.len()].to_vec()), 40 | 0 => Err(eyre!( 41 | "error making discv5 request: {}", 42 | std::str::from_utf8(&b[1..b.len()]).unwrap() 43 | )), 44 | _ => panic!("unexpected message encoding"), 45 | } 46 | } 47 | --------------------------------------------------------------------------------