From a7c666f108655f442c3e64db179464b79e778cb3 Mon Sep 17 00:00:00 2001
From: Franck Royer <franck@royer.one>
Date: Fri, 12 Mar 2021 14:23:21 +1100
Subject: [PATCH] Fix WakuMessage (de)serialisation

---
 src/lib/node.spec.ts         |  4 ++--
 src/lib/waku_message.spec.ts | 29 +++++++++++++++++++++++++++--
 src/lib/waku_message.ts      | 28 ++++++++++++++++++++++------
 3 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/src/lib/node.spec.ts b/src/lib/node.spec.ts
index 1b93a913eb..9baf76f1ec 100644
--- a/src/lib/node.spec.ts
+++ b/src/lib/node.spec.ts
@@ -9,7 +9,7 @@ import { Message } from './waku_message';
 import { CODEC, TOPIC, WakuRelay } from './waku_relay';
 
 test('Publishes message', async (t) => {
-  const message = Message.fromString('Bird bird bird, bird is the word!');
+  const message = Message.fromUtf8String('Bird bird bird, bird is the word!');
 
   const [node1, node2] = await Promise.all([createNode(), createNode()]);
   const wakuRelayNode1 = new WakuRelay(node1.pubsub);
@@ -123,7 +123,7 @@ test('Nim-interop: nim node sends message', async (t) => {
   // Setup the promise before publishing to ensure the event is not missed
   const promise = waitForNextData(node.pubsub);
 
-  const message = Message.fromString('This is a message.');
+  const message = Message.fromUtf8String('This is a message.');
 
   await delay(500);
 
diff --git a/src/lib/waku_message.spec.ts b/src/lib/waku_message.spec.ts
index 6444b81627..6f6f43125f 100644
--- a/src/lib/waku_message.spec.ts
+++ b/src/lib/waku_message.spec.ts
@@ -1,15 +1,40 @@
 import { fc, testProp } from 'ava-fast-check';
 
+import { WakuMessage } from '../gen/proto/waku/v2/waku_pb';
+
 import { Message } from './waku_message';
 
+// This test is more about documenting how protobuf library works than testing it
+testProp('Protobuf round trip binary serialisation', [fc.string()], (t, s) => {
+  const wakuMsg = new WakuMessage();
+  wakuMsg.setPayload(Buffer.from(s, 'utf-8'));
+
+  const binary = wakuMsg.serializeBinary();
+  const actual = WakuMessage.deserializeBinary(binary);
+
+  const payload = actual.getPayload();
+
+  let buf;
+  if (typeof payload === 'string') {
+    buf = Buffer.from(payload, 'base64');
+  } else {
+    buf = Buffer.from(payload);
+  }
+
+  t.deepEqual(s, buf.toString('utf-8'));
+});
+
 testProp(
   'Waku message round trip binary serialisation',
   [fc.string()],
   (t, s) => {
-    const msg = Message.fromString(s);
+    const msg = Message.fromUtf8String(s);
     const binary = msg.toBinary();
     const actual = Message.fromBinary(binary);
 
-    t.true(actual.isEqualTo(msg));
+    t.true(
+      actual.isEqualTo(msg),
+      `${JSON.stringify(actual)}\n${JSON.stringify(msg)}`
+    );
   }
 );
diff --git a/src/lib/waku_message.ts b/src/lib/waku_message.ts
index cc5ee9d93d..778d5f3faf 100644
--- a/src/lib/waku_message.ts
+++ b/src/lib/waku_message.ts
@@ -3,7 +3,7 @@ import { WakuMessage } from '../gen/proto/waku/v2/waku_pb';
 // Ensure that this class matches the proto interface while
 // Protecting the user from protobuf oddities
 export class Message {
-  public payload: Uint8Array | string;
+  public payload: Uint8Array;
   public contentTopic: number;
   public version: number;
 
@@ -12,12 +12,24 @@ export class Message {
 
     const msg = protobuf.toObject();
 
-    this.payload = msg.payload;
+    // Let's make is easier to avoid mistakes and only store in Uint8Array format
+    let payload;
+    if (typeof msg.payload === 'string') {
+      payload = Buffer.from(msg.payload, 'base64');
+    } else {
+      payload = msg.payload;
+    }
+    this.payload = payload;
     this.contentTopic = msg.contentTopic;
     this.version = msg.version;
   }
 
-  static fromString(message: string): Message {
+  /**
+   * Create Message from utf-8 string
+   * @param message
+   * @returns {Message}
+   */
+  static fromUtf8String(message: string): Message {
     const wakuMsg = new WakuMessage();
 
     // Only Version 0 is implemented in Waku 2.
@@ -27,7 +39,10 @@ export class Message {
     // This is the content topic commonly used at this time
     wakuMsg.setContentTopic(1);
 
-    wakuMsg.setPayload(message);
+    const buf = Buffer.from(message, 'utf-8');
+
+    // Only accepts Uint8Array or base64 string
+    wakuMsg.setPayload(buf);
 
     return new Message(wakuMsg);
   }
@@ -42,10 +57,11 @@ export class Message {
   }
 
   // Purely for tests purposes.
-  // We do not care about protobuf field when checking equality
+  // We do consider protobuf field when checking equality
+  // As the content is held by the other fields.
   isEqualTo(other: Message) {
     return (
-      this.payload === other.payload &&
+      Buffer.compare(this.payload, other.payload) === 0 &&
       this.contentTopic === other.contentTopic &&
       this.version === other.version
     );