1 files changed, 26 insertions, 10 deletions
diff --git a/client-native-lib/src/state.rs b/client-native-lib/src/state.rs
index 57c4b29..c5e9365 100644
--- a/client-native-lib/src/state.rs
+++ b/client-native-lib/src/state.rs
@@ -1,5 +1,6 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, pin::Pin, sync::Arc};
 
+use futures_util::Future;
 use log::warn;
 use tokio::sync::{mpsc::UnboundedSender, RwLock};
 use webrtc::api::API;
@@ -7,18 +8,31 @@ use webrtc::api::API;
 use crate::{
     crypto::Key,
     peer::Peer,
-    protocol::{self, ClientboundPacket, RelayMessage, RelayMessageWrapper, ServerboundPacket}, Config,
+    protocol::{self, ClientboundPacket, RelayMessage, RelayMessageWrapper, ServerboundPacket},
+    Config,
 };
 
-pub struct State {
+pub trait HasPeer {
+    fn peer(&self) -> &Arc<Peer>;
+}
+pub trait PeerInit<P> {
+    fn add_peer(
+        &self,
+        p: Arc<Peer>,
+    ) -> Pin<Box<dyn Future<Output = Arc<P>> + 'static + Send + Sync>>;
+}
+
+pub struct State<P: HasPeer, I: PeerInit<P>> {
+    pub sup: Arc<I>,
     pub config: Config,
     pub api: API,
     pub key: Key,
     pub my_id: RwLock<Option<usize>>,
     pub sender: UnboundedSender<ServerboundPacket>,
-    pub peers: RwLock<HashMap<usize, Arc<Peer>>>,
+    pub peers: RwLock<HashMap<usize, Arc<P>>>,
+    pub relay_tx: UnboundedSender<(usize, RelayMessage)>,
 }
-impl State {
+impl<P: HasPeer, I: PeerInit<P>> State<P, I> {
     pub async fn my_id(&self) -> usize {
         self.my_id.read().await.expect("not initialized yet")
     }
@@ -35,10 +49,12 @@ impl State {
                 if id == self.my_id().await {
                     // we joined - YAY!
                 } else {
-                    self.peers
-                        .write()
-                        .await
-                        .insert(id, Peer::create(self.clone(), id).await);
+                    self.peers.write().await.insert(
+                        id,
+                        self.sup
+                            .add_peer(Peer::create(self.clone(), self.relay_tx.clone(), id).await)
+                            .await,
+                    );
                 }
             }
             protocol::ClientboundPacket::ClientLeave { id: _ } => {}
@@ -52,7 +68,7 @@ impl State {
 
     pub async fn on_relay(&self, sender: usize, p: RelayMessage) {
         if let Some(peer) = self.peers.read().await.get(&sender).cloned() {
-            peer.on_relay(p).await
+            peer.peer().on_relay(p).await
         } else {
             warn!("got a packet from a non-existent peer")
         }