From 9dd998318013abe4dd12137f75724bae5e9ee9eb Mon Sep 17 00:00:00 2001
From: chunningham <c.a.cunningham6@gmail.com>
Date: Wed, 5 Apr 2023 23:52:57 +1000
Subject: [PATCH] fix orbit behaviour builder and use specific errors

---
 src/p2p/behaviour/builder.rs | 74 +++++++++++++++++-------------------
 1 file changed, 35 insertions(+), 39 deletions(-)

diff --git a/src/p2p/behaviour/builder.rs b/src/p2p/behaviour/builder.rs
index c05d2838..4cfa437f 100644
--- a/src/p2p/behaviour/builder.rs
+++ b/src/p2p/behaviour/builder.rs
@@ -1,18 +1,12 @@
 use crate::p2p::{
-    behaviour::{BaseBehaviour, Behaviour},
-    transport::IntoTransport,
+    behaviour::{poll_swarm, BaseBehaviour, Behaviour},
+    transport::{build_transport, IntoTransport},
     IdentifyConfig,
 };
-use futures::{
-    channel::{mpsc, oneshot},
-    future::{select, Either},
-    io::{AsyncRead, AsyncWrite},
-    sink::SinkExt,
-    stream::StreamExt,
-};
+use futures::io::{AsyncRead, AsyncWrite};
 use libp2p::{
     autonat::{Behaviour as AutoNat, Config as AutoNatConfig},
-    core::{upgrade, Transport},
+    core::Transport,
     dcutr::Behaviour as Dcutr,
     gossipsub::{
         Behaviour as Gossipsub, Config as GossipsubConfig, ConfigBuilder as GossipsubConfigBuilder,
@@ -24,12 +18,12 @@ use libp2p::{
         record::store::{MemoryStore, MemoryStoreConfig, RecordStore},
         Kademlia, KademliaConfig,
     },
-    mplex, noise,
+    noise::NoiseError,
     ping::{Behaviour as Ping, Config as PingConfig},
     relay::client::{new, Behaviour as Client},
-    swarm::{Swarm, SwarmBuilder},
-    yamux,
+    swarm::SwarmBuilder,
 };
+use std::time::Duration;
 use thiserror::Error;
 
 #[derive(Debug, Clone, Default)]
@@ -41,6 +35,7 @@ pub struct BehaviourConfig<KSC = MemoryStoreConfig> {
     kademlia_store: KSC,
     autonat: AutoNatConfig,
     relay: bool,
+    transport_timeout: Duration,
 }
 
 impl<KSC> BehaviourConfig<KSC> {
@@ -53,6 +48,7 @@ impl<KSC> BehaviourConfig<KSC> {
             kademlia_store: ksc.into(),
             autonat: Default::default(),
             relay: Default::default(),
+            transport_timeout: Duration::from_secs(20),
         }
     }
     pub fn identify(&mut self, i: impl Into<IdentifyConfig>) -> &mut Self {
@@ -94,7 +90,6 @@ impl<KSC> BehaviourConfig<KSC> {
     {
         let peer_id = keypair.public().to_peer_id();
         Ok(Behaviour {
-            exchange: todo!(),
             base: BaseBehaviour {
                 identify: Identify::new(self.identify.to_config(keypair.public())),
                 ping: Ping::new(self.ping),
@@ -118,11 +113,8 @@ impl<KSC> BehaviourConfig<KSC> {
             },
         })
     }
-    pub fn launch<T, KS>(
-        self,
-        keypair: Keypair,
-        transport: T,
-    ) -> Result<(), OrbitLaunchError<T::Error>>
+
+    pub fn launch<T, KS>(self, keypair: Keypair, transport: T) -> Result<(), OrbitLaunchError<T>>
     where
         T: IntoTransport,
         T::T: 'static + Send + Unpin,
@@ -132,32 +124,29 @@ impl<KSC> BehaviourConfig<KSC> {
         <T::T as Transport>::Dial: Send,
         <T::T as Transport>::ListenerUpgrade: Send,
         KS: RecordStore + Send + 'static,
+        KSC: RecordStoreConfig<KS>,
     {
         let local_public_key = keypair.public();
         let id = local_public_key.to_peer_id();
-        let transport = transport.into_transport()?;
+        let transport = transport
+            .into_transport()
+            .map_err(OrbitLaunchError::TransportConfig)?;
         let (transport, behaviour) = if self.relay {
             let (t, b) = new(id);
-            (transport.or_transport(t), self.build(keypair, Some(b))?)
+            (
+                build_transport(transport.or_transport(t), self.transport_timeout, &keypair)?,
+                self.build(keypair, Some(b))?,
+            )
         } else {
-            (transport, self.build(keypair, None)?)
+            (
+                build_transport(transport, self.transport_timeout, &keypair)?,
+                self.build(keypair, None)?,
+            )
         };
 
-        let mut swarm = SwarmBuilder::with_tokio_executor(
-            transport
-                .upgrade(upgrade::Version::V1)
-                // TODO replace with AWAKE protcol (or similar)
-                .authenticate(noise::NoiseAuthenticated::xx(&keypair).unwrap())
-                .multiplex(upgrade::SelectUpgrade::new(
-                    yamux::YamuxConfig::default(),
-                    mplex::MplexConfig::default(),
-                ))
-                .timeout(std::time::Duration::from_secs(20))
-                .boxed(),
-            behaviour,
-            id,
-        )
-        .build();
+        let swarm = SwarmBuilder::with_tokio_executor(transport, behaviour, id).build();
+        tokio::spawn(poll_swarm(swarm));
+
         Ok(())
     }
 }
@@ -169,11 +158,18 @@ pub enum OrbitBehaviourBuildError {
 }
 
 #[derive(Error, Debug)]
-pub enum OrbitLaunchError<T> {
+pub enum OrbitLaunchError<T>
+where
+    T: IntoTransport,
+{
     #[error(transparent)]
     Config(#[from] OrbitBehaviourBuildError),
     #[error(transparent)]
-    Transport(T),
+    Transport(<T::T as Transport>::Error),
+    #[error(transparent)]
+    TransportConfig(T::Error),
+    #[error(transparent)]
+    Noise(#[from] NoiseError),
 }
 
 pub trait RecordStoreConfig<S>