Expose Floodsub target_peer list.

src/lib.rs

@@ -268,6 +268,8 @@ enum IpfsEvent {
  Disconnect(MultiaddrWithPeerId, Channel<()>),
  /// Request background task to return the listened and external addresses
  GetAddresses(OneshotSender<Vec<Multiaddr>>),
+ PubsubAddPeer(PeerId, OneshotSender<()>),
+ PubsubRemovePeer(PeerId, OneshotSender<()>),
  PubsubSubscribe(String, OneshotSender<Option<SubscriptionStream>>),
  PubsubUnsubscribe(String, OneshotSender<bool>),
  PubsubPublish(String, Vec<u8>, OneshotSender<()>),
@@ -755,6 +757,43 @@ impl<Types: IpfsTypes> Ipfs<Types> {
  .await
  }
 
+ /// Add a peer to list of nodes to propagate messages to.
+ ///
+ /// A peer will not receive any pubsub messages from this node until it is added using this function,
+ /// unless it has added this node in the same way.
+ pub async fn pubsub_add_peer(&self, peer_id: PeerId) -> Result<(), Error> {
+ async move {
+ let (tx, rx) = oneshot_channel::<()>();
+
+ self.to_task
+ .clone()
+ .send(IpfsEvent::PubsubAddPeer(peer_id, tx))
+ .await?;
+
+ Ok(rx.await?)
+ }
+ .instrument(self.span.clone())
+ .await
+ }
+
+ /// Remove a peer from the list of nodes that messages are propagated to.
+ ///
+ /// Calling this function will not stop messages being sent to the specified peers for subscribed topics which have already been communicated.
+ pub async fn pubsub_remove_peer(&self, peer_id: PeerId) -> Result<(), Error> {
+ async move {
+ let (tx, rx) = oneshot_channel::<()>();
+
+ self.to_task
+ .clone()
+ .send(IpfsEvent::PubsubRemovePeer(peer_id, tx))
+ .await?;
+
+ Ok(rx.await?)
+ }
+ .instrument(self.span.clone())
+ .await
+ }
+
  /// Subscribes to a given topic. Can be done at most once without unsubscribing in the between.
  /// The subscription can be unsubscribed by dropping the stream or calling
  /// [`Ipfs::pubsub_unsubscribe`].
@@ -1431,6 +1470,20 @@ impl<TRepoTypes: RepoTypes> Future for IpfsFuture<TRepoTypes> {
  // ignore error, perhaps caller went away already
  let _ = ret.send(addresses);
  }
+ IpfsEvent::PubsubAddPeer(peer_id, ret) => {
+ self.swarm
+ .behaviour_mut()
+ .pubsub()
+ .add_node_to_partial_view(peer_id);
+ let _ = ret.send(());
+ }
+ IpfsEvent::PubsubRemovePeer(peer_id, ret) => {
+ self.swarm
+ .behaviour_mut()
+ .pubsub()
+ .remove_node_from_partial_view(&peer_id);
+ let _ = ret.send(());
+ }
  IpfsEvent::PubsubSubscribe(topic, ret) => {
  let _ = ret.send(self.swarm.behaviour_mut().pubsub().subscribe(topic));
  }
@@ -1780,8 +1833,11 @@ mod node {
 
 #[cfg(test)]
 mod tests {
+ use std::time::Duration;
+
  use super::*;
  use crate::make_ipld;
+ use futures::{stream::poll_immediate, StreamExt};
  use multihash::Sha2_256;
 
  #[tokio::test]
@@ -1819,4 +1875,47 @@ mod tests {
  ipfs.remove_pin(&cid, false).await.unwrap();
  assert!(!ipfs.is_pinned(&cid).await.unwrap());
  }
+
+ #[tokio::test]
+ async fn test_pubsub_send_and_receive() {
+ let alice = Node::new("alice").await;
+ let bob = Node::new("bob").await;
+ let bob_addr = bob.addrs_local().await.unwrap()[0].clone();
+
+ let topic = String::from("test_topic");
+ alice
+ .connect(bob_addr.with(Protocol::P2p(bob.id.into())))
+ .await
+ .expect("alice failed to connect to bob");
+ let _alice_messages = alice.pubsub_subscribe(topic.clone()).await.unwrap();
+ let mut bob_messages = poll_immediate(bob.pubsub_subscribe(topic.clone()).await.unwrap());
+
+ let data = vec![1, 2, 3];
+
+ alice
+ .pubsub_publish(topic.clone(), data.clone())
+ .await
+ .unwrap();
+ tokio::time::sleep(Duration::from_millis(100)).await;
+
+ assert_eq!(bob_messages.next().await, Some(Poll::Pending));
+
+ bob.pubsub_add_peer(alice.id).await.unwrap();
+ tokio::time::sleep(Duration::from_millis(100)).await;
+
+ assert_eq!(bob_messages.next().await, Some(Poll::Pending));
+
+ alice
+ .pubsub_publish(topic.clone(), data.clone())
+ .await
+ .unwrap();
+ tokio::time::sleep(Duration::from_millis(100)).await;
+
+ let received_data = bob_messages
+ .next()
+ .await
+ .expect("unexpected end of stream")
+ .map(|msg| msg.data.clone());
+ assert_eq!(received_data, Poll::Ready(data.clone()));
+ }
 }