From f0e7100cfd755abbe007afcb787f371c1456e0a4 Mon Sep 17 00:00:00 2001
From: Matthias Wright <matthias.s.wright@gmail.com>
Date: Thu, 17 Aug 2023 17:07:47 +0200
Subject: [PATCH] feat(e2e): add bootstrapper to spawn_swarm binary

---
 core/e2e/src/bin/spawn_swarm.rs | 66 +++++++++++++++++++++++++++++++--
 core/e2e/src/swarm.rs           |  2 -
 core/e2e/tests/dht.rs           |  1 +
 3 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/core/e2e/src/bin/spawn_swarm.rs b/core/e2e/src/bin/spawn_swarm.rs
index 22dbbe458..d0d829dfc 100644
--- a/core/e2e/src/bin/spawn_swarm.rs
+++ b/core/e2e/src/bin/spawn_swarm.rs
@@ -1,10 +1,24 @@
-use std::time::SystemTime;
+use std::{sync::Arc, thread, time::SystemTime};
 
 use anyhow::Result;
 use clap::Parser;
-use fleek_crypto::PublicKey;
-use lightning_e2e::{swarm::Swarm, utils::shutdown};
+use fleek_crypto::{NodeSecretKey, PublicKey, SecretKey};
+use lightning_application::query_runner::QueryRunner;
+use lightning_dht::{
+    config::{Bootstrapper, Config as DhtConfig},
+    dht::Builder as DhtBuilder,
+};
+use lightning_e2e::{
+    swarm::Swarm,
+    utils::{
+        networking::{PortAssigner, Transport},
+        shutdown,
+    },
+};
+use lightning_interfaces::WithStartAndShutdown;
+use lightning_topology::Topology;
 use resolved_pathbuf::ResolvedPathBuf;
+use tokio::sync::Notify;
 
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
@@ -22,6 +36,45 @@ struct Cli {
 async fn main() -> Result<()> {
     let args = Cli::parse();
 
+    // Start bootstrapper
+    let mut port_assigner = PortAssigner::default();
+    let bootstrapper_port = port_assigner
+        .get_port(12001, 13000, Transport::Udp)
+        .expect("Failed to assign port");
+
+    let bootstrapper_address = format!("0.0.0.0:{bootstrapper_port}").parse().unwrap();
+    let bootstrapper_config = DhtConfig {
+        address: bootstrapper_address,
+        bootstrappers: vec![],
+    };
+    let bootstrap_secret_key = NodeSecretKey::generate();
+    let bootstrap_shutdown_notify = Arc::new(Notify::new());
+    let bootstrap_ready = Arc::new(Notify::new());
+    let bootstrap_ready_rx = bootstrap_ready.clone();
+    let bootstrap_shutdown_notify_rx = bootstrap_shutdown_notify.clone();
+
+    let key_cloned = bootstrap_secret_key.clone();
+    let _bootstrap_handle = thread::spawn(move || {
+        let mut builder = tokio::runtime::Builder::new_multi_thread();
+        let runtime = builder
+            .enable_all()
+            .build()
+            .expect("Failed to build tokio runtime for node container.");
+
+        runtime.block_on(async move {
+            let builder = DhtBuilder::new(key_cloned, bootstrapper_config);
+            let dht = builder.build::<Topology<QueryRunner>>().unwrap();
+            dht.start().await;
+            bootstrap_ready_rx.notify_one();
+
+            bootstrap_shutdown_notify_rx.notified().await;
+            dht.shutdown().await;
+        });
+    });
+
+    // Wait for bootstrapper to start
+    bootstrap_ready.notified().await;
+
     let epoch_start = SystemTime::now()
         .duration_since(SystemTime::UNIX_EPOCH)
         .unwrap()
@@ -30,9 +83,16 @@ async fn main() -> Result<()> {
     let path = ResolvedPathBuf::try_from("~/.lightning-test/e2e/spawn-swarm").unwrap();
     let swarm = Swarm::builder()
         .with_directory(path)
+        .with_min_port(12001)
+        .with_max_port(13000)
         .with_num_nodes(args.num_nodes)
         .with_epoch_time(args.epoch_time)
         .with_epoch_start(epoch_start)
+        .with_bootstrappers(vec![Bootstrapper {
+            address: bootstrapper_address,
+            network_public_key: bootstrap_secret_key.to_pk(),
+        }])
+        .with_port_assigner(port_assigner)
         .build();
     swarm.launch().await.unwrap();
 
diff --git a/core/e2e/src/swarm.rs b/core/e2e/src/swarm.rs
index b9a4d60d5..98ee428d5 100644
--- a/core/e2e/src/swarm.rs
+++ b/core/e2e/src/swarm.rs
@@ -306,8 +306,6 @@ impl SwarmBuilder {
 
             nodes.insert(node_secret_key.to_pk(), (config, owner_secret_key, i));
         }
-
-        println!("GENESIS LEN: {}", genesis.committee.len());
         let nodes = nodes
             .into_iter()
             .map(|(node_pub_key, (config, owner_secret_key, index))| {
diff --git a/core/e2e/tests/dht.rs b/core/e2e/tests/dht.rs
index fabec3e27..b408dd44e 100644
--- a/core/e2e/tests/dht.rs
+++ b/core/e2e/tests/dht.rs
@@ -84,6 +84,7 @@ async fn e2e_dht() -> Result<()> {
             address: bootstrapper_address,
             network_public_key: bootstrap_secret_key.to_pk(),
         }])
+        .with_port_assigner(port_assigner)
         .build();
     swarm.launch().await.unwrap();