Merge pull request #142 from xmtp/ea/all-messages

Introduce API to stream all messages for all conversations

src/Client.ts

@@ -10,7 +10,7 @@ import {
   promiseWithTimeout,
 } from './utils'
 import { sleep } from '../test/helpers'
-import Stream, { MessageFilter } from './Stream'
+import Stream, { MessageFilter, noTransformation } from './Stream'
 import { Signer } from 'ethers'
 import {
   EncryptedStore,
@@ -385,7 +385,10 @@ export default class Client {
     return Message.encode(this.keys, recipient, payload, timestamp)
   }
 
-  async decodeMessage(payload: Uint8Array): Promise<Message> {
+  async decodeMessage(
+    payload: Uint8Array,
+    contentTopic: string | undefined
+  ): Promise<Message> {
     const message = await Message.decode(this.keys, payload)
     if (message.error) {
       return message
@@ -400,6 +403,7 @@ export default class Client {
     }
     const contentType = new ContentTypeId(encoded.type)
     const codec = this.codecFor(contentType)
+    message.contentTopic = contentTopic
     if (codec) {
       message.content = codec.decode(encoded as EncodedContent, this)
       message.contentType = contentType
@@ -416,18 +420,18 @@ export default class Client {
   streamIntroductionMessages(): Promise<Stream<Message>> {
     return Stream.create<Message>(
       this,
-      buildUserIntroTopic(this.address),
+      [buildUserIntroTopic(this.address)],
       noTransformation
     )
   }
 
   streamConversationMessages(peerAddress: string): Promise<Stream<Message>> {
-    const topic = buildDirectMessageTopic(peerAddress, this.address)
+    const topics = [buildDirectMessageTopic(peerAddress, this.address)]
     return Stream.create<Message>(
       this,
-      topic,
+      topics,
       noTransformation,
-      filterForTopic(topic)
+      filterForTopics(topics)
     )
   }
 
@@ -476,11 +480,11 @@ export default class Client {
     wakuMsgs = wakuMsgs.filter((wakuMsg) => wakuMsg?.payload)
     let msgs = await Promise.all(
       wakuMsgs.map((wakuMsg) =>
-        this.decodeMessage(wakuMsg.payload as Uint8Array)
+        this.decodeMessage(wakuMsg.payload as Uint8Array, wakuMsg.contentTopic)
       )
     )
     if (opts?.checkAddresses) {
-      msgs = msgs.filter(filterForTopic(topic))
+      msgs = msgs.filter(filterForTopics([topic]))
     }
     return msgs
   }
@@ -627,18 +631,15 @@ async function getNodeList(env: keyof NodesList): Promise<string[]> {
   return Object.values(nodesList[env])
 }
 
-function noTransformation(msg: Message) {
-  return msg
-}
-
-function filterForTopic(topic: string): MessageFilter {
+// Ensure the message didn't have a spoofed address
+function filterForTopics(topics: string[]): MessageFilter {
   return (msg) => {
     const senderAddress = msg.senderAddress
     const recipientAddress = msg.recipientAddress
     return (
       senderAddress !== undefined &&
       recipientAddress !== undefined &&
-      buildDirectMessageTopic(senderAddress, recipientAddress) === topic
+      topics.includes(buildDirectMessageTopic(senderAddress, recipientAddress))
     )
   }
 }

src/Message.ts

@@ -31,6 +31,7 @@ export default class Message implements proto.V1Message {
   // the message receiving APIs need to return a Message to provide access to the header fields like sender/recipient
   contentType?: ContentTypeId
   content?: any // eslint-disable-line @typescript-eslint/no-explicit-any
+  contentTopic?: string // content topic that triggered the message
   error?: Error
   /**
    * Identifier that is deterministically derived from the bytes of the message

src/Stream.ts

@@ -10,12 +10,18 @@ export type MessageTransformer<T> = (msg: Message) => T
 
 export type MessageFilter = (msg: Message) => boolean
 
+export type ContentTopicUpdater = (msg: Message) => string[] | undefined
+
+export const noTransformation = (msg: Message): Message => {
+  return msg
+}
+
 /**
  * Stream implements an Asynchronous Iterable over messages received from a topic.
  * As such can be used with constructs like for-await-of, yield*, array destructing, etc.
  */
 export default class Stream<T> {
-  topic: string
+  topics: string[]
   client: Client
   // queue of incoming Waku messages
   messages: T[]
@@ -31,31 +37,47 @@ export default class Stream<T> {
 
   constructor(
     client: Client,
-    topic: string,
+    topics: string[],
     messageTransformer: MessageTransformer<T>,
-    messageFilter?: MessageFilter
+    messageFilter?: MessageFilter,
+    contentTopicUpdater?: ContentTopicUpdater
   ) {
     this.messages = []
     this.resolvers = []
-    this.topic = topic
+    this.topics = topics
     this.client = client
-    this.callback = this.newMessageCallback(messageTransformer, messageFilter)
+    this.callback = this.newMessageCallback(
+      messageTransformer,
+      messageFilter,
+      contentTopicUpdater
+    )
   }
 
   // returns new closure to handle incoming Waku messages
   private newMessageCallback(
     transformer: MessageTransformer<T>,
-    filter?: MessageFilter
+    filter?: MessageFilter,
+    contentTopicUpdater?: ContentTopicUpdater
   ): (wakuMsg: WakuMessage) => Promise<void> {
     return async (wakuMsg: WakuMessage) => {
       if (!wakuMsg.payload) {
         return
       }
-      const msg = await this.client.decodeMessage(wakuMsg.payload)
+      const msg = await this.client.decodeMessage(
+        wakuMsg.payload,
+        wakuMsg.contentTopic
+      )
       // If there is a filter on the stream, and the filter returns false, ignore the message
       if (filter && !filter(msg)) {
         return
       }
+      // Check to see if we should update the stream's content topic subscription
+      if (contentTopicUpdater) {
+        const topics = contentTopicUpdater(msg)
+        if (topics) {
+          this.resubscribeToTopics(topics)
+        }
+      }
       // is there a Promise already pending?
       const resolver = this.resolvers.pop()
       if (resolver) {
@@ -75,7 +97,7 @@ export default class Stream<T> {
 
     this.unsubscribeFn = await this.client.waku.filter.subscribe(
       this.callback,
-      [this.topic]
+      this.topics
     )
     await this.listenForDisconnect()
   }
@@ -93,7 +115,7 @@ export default class Stream<T> {
             }
             this.unsubscribeFn = await this.client.waku.filter.subscribe(
               this.callback,
-              [this.topic]
+              this.topics
             )
             console.log(`Connection to peer ${connection.remoteAddr} restored`)
             return
@@ -113,11 +135,18 @@ export default class Stream<T> {
 
   static async create<T>(
     client: Client,
-    topic: string,
+    topics: string[],
     messageTransformer: MessageTransformer<T>,
-    messageFilter?: MessageFilter
+    messageFilter?: MessageFilter,
+    contentTopicUpdater?: ContentTopicUpdater
   ): Promise<Stream<T>> {
-    const stream = new Stream(client, topic, messageTransformer, messageFilter)
+    const stream = new Stream(
+      client,
+      topics,
+      messageTransformer,
+      messageFilter,
+      contentTopicUpdater
+    )
     await stream.start()
     return stream
   }
@@ -167,4 +196,17 @@ export default class Stream<T> {
     // otherwise return empty Promise and queue its resolver
     return new Promise((resolve) => this.resolvers.unshift(resolve))
   }
+
+  // Unsubscribe from the existing content topics and resubscribe to the given topics.
+  private async resubscribeToTopics(topics: string[]): Promise<void> {
+    if (!this.callback || !this.unsubscribeFn) {
+      throw new Error('Missing callback for stream')
+    }
+    await this.unsubscribeFn()
+    this.topics = topics
+    this.unsubscribeFn = await this.client.waku.filter.subscribe(
+      this.callback,
+      this.topics
+    )
+  }
 }

src/conversations/Conversations.ts

@@ -1,8 +1,12 @@
 import Conversation from './Conversation'
 import Message from '../Message'
-import Stream, { MessageFilter, MessageTransformer } from '../Stream'
+import Stream, {
+  MessageFilter,
+  MessageTransformer,
+  noTransformation,
+} from '../Stream'
 import Client from '../Client'
-import { buildUserIntroTopic } from '../utils'
+import { buildDirectMessageTopic, buildUserIntroTopic } from '../utils'
 
 const messageHasHeaders: MessageFilter = (msg: Message) => {
   return Boolean(msg.recipientAddress && msg.senderAddress)
@@ -74,12 +78,76 @@ export default class Conversations {
 
     return Stream.create<Conversation>(
       this.client,
-      buildUserIntroTopic(this.client.address),
+      [buildUserIntroTopic(this.client.address)],
       messageTransformer,
       filter
     )
   }
 
+  /**
+   * Returns a stream for all new messages from existing and new conversations.
+   */
+  async streamAllMessages(): Promise<Stream<Message>> {
+    const conversations = await this.list()
+    const dmAddresses: string[] = conversations.map(
+      (conversation) => conversation.peerAddress
+    )
+    const introTopic = buildUserIntroTopic(this.client.address)
+    const topics = this.buildTopicsForAllMessages(dmAddresses, introTopic)
+
+    // If we detect a new intro topic, update the stream's direct message topics to include the new topic
+    const contentTopicUpdater = (msg: Message): string[] | undefined => {
+      if (msg.contentTopic !== introTopic || !messageHasHeaders(msg)) {
+        return undefined
+      }
+      const peerAddress = this.getPeerAddress(msg)
+      if (
+        dmAddresses.includes(peerAddress) ||
+        peerAddress === this.client.address
+      ) {
+        // No need to update if we're already subscribed
+        return undefined
+      }
+      dmAddresses.push(peerAddress)
+      return this.buildTopicsForAllMessages(dmAddresses, introTopic)
+    }
+
+    // Filter intro topics if already streaming direct messages for that address to avoid duplicates
+    const filter = (msg: Message): boolean => {
+      if (
+        msg.contentTopic === introTopic &&
+        messageHasHeaders(msg) &&
+        dmAddresses.includes(this.getPeerAddress(msg))
+      ) {
+        return false
+      }
+      return true
+    }
+
+    return Stream.create<Message>(
+      this.client,
+      topics,
+      noTransformation,
+      filter,
+      contentTopicUpdater
+    )
+  }
+
+  /**
+   * Builds a list of topics for existing conversations and new intro topics
+   */
+  private buildTopicsForAllMessages(
+    peerAddresses: string[],
+    introTopic: string
+  ): string[] {
+    const topics = peerAddresses.map((address) =>
+      buildDirectMessageTopic(address, this.client.address)
+    )
+    // Ensure we listen for new conversation topics as well
+    topics.push(introTopic)
+    return topics
+  }
+
   /**
    * Creates a new conversation for the given address. Will throw an error if the peer is not found in the XMTP network
    */