From 8b5e22002c0e93e636ea370d016d25c256bdf3f1 Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Fri, 17 Feb 2023 11:00:46 -0400
Subject: [PATCH] fix: use base64url encoding

---
 waku/v2/rpc/rpc_type.go  | 36 ++++++++++++++++++++++++++++++++++++
 waku/v2/rpc/util_test.go | 24 ++++++++++++++++++++++++
 waku/v2/rpc/utils.go     | 16 ++++++++--------
 3 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/waku/v2/rpc/rpc_type.go b/waku/v2/rpc/rpc_type.go
index 720d08ec..0bd4a3d0 100644
--- a/waku/v2/rpc/rpc_type.go
+++ b/waku/v2/rpc/rpc_type.go
@@ -1,8 +1,44 @@
 package rpc
 
+import (
+	"encoding/base64"
+	"strings"
+)
+
 type SuccessReply = bool
 
 type Empty struct {
 }
 
 type MessagesReply = []*RPCWakuMessage
+
+type Base64URLByte []byte
+
+func (u Base64URLByte) MarshalJSON() ([]byte, error) {
+	base64Value := base64.URLEncoding.EncodeToString(u)
+	return []byte("\"" + base64Value + "\""), nil
+}
+
+func (h *Base64URLByte) UnmarshalText(b []byte) error {
+	inputValue := ""
+	if b != nil {
+		inputValue = string(b)
+	}
+
+	enc := base64.StdEncoding
+	if strings.ContainsAny(inputValue, "-_") {
+		enc = base64.URLEncoding
+	}
+	if len(inputValue)%4 != 0 {
+		enc = enc.WithPadding(base64.NoPadding)
+	}
+
+	decodedBytes, err := enc.DecodeString(inputValue)
+	if err != nil {
+		return err
+	}
+
+	*h = decodedBytes
+
+	return nil
+}
diff --git a/waku/v2/rpc/util_test.go b/waku/v2/rpc/util_test.go
index adf0abd4..1abc2f99 100644
--- a/waku/v2/rpc/util_test.go
+++ b/waku/v2/rpc/util_test.go
@@ -2,10 +2,13 @@ package rpc
 
 import (
 	"bytes"
+	"encoding/base64"
+	"encoding/json"
 	"net/http"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/waku-org/go-waku/waku/v2/protocol/pb"
 )
 
 func makeRequest(t *testing.T) *http.Request {
@@ -13,3 +16,24 @@ func makeRequest(t *testing.T) *http.Request {
 	require.NoError(t, err)
 	return request
 }
+
+func TestBase64Encoding(t *testing.T) {
+	input := "Hello World"
+
+	rpcMsg := ProtoToRPC(&pb.WakuMessage{
+		Payload: []byte(input),
+	})
+
+	jsonBytes, err := json.Marshal(rpcMsg)
+	require.NoError(t, err)
+
+	m := make(map[string]interface{})
+	err = json.Unmarshal(jsonBytes, &m)
+	require.NoError(t, err)
+	require.Equal(t, base64.URLEncoding.EncodeToString([]byte(input)), m["payload"])
+
+	decodedRpcMsg := new(RPCWakuMessage)
+	err = json.Unmarshal(jsonBytes, decodedRpcMsg)
+	require.NoError(t, err)
+	require.Equal(t, input, string(decodedRpcMsg.Payload))
+}
diff --git a/waku/v2/rpc/utils.go b/waku/v2/rpc/utils.go
index b836e071..895ac1b2 100644
--- a/waku/v2/rpc/utils.go
+++ b/waku/v2/rpc/utils.go
@@ -5,17 +5,17 @@ import (
 )
 
 type RateLimitProof struct {
-	Proof         []byte `json:"proof,omitempty"`
-	MerkleRoot    []byte `json:"merkle_root,omitempty"`
-	Epoch         []byte `json:"epoch,omitempty"`
-	ShareX        []byte `json:"share_x,omitempty"`
-	ShareY        []byte `json:"share_y,omitempty"`
-	Nullifier     []byte `json:"nullifier,omitempty"`
-	RlnIdentifier []byte `json:"rln_identifier,omitempty"`
+	Proof         Base64URLByte `json:"proof,omitempty"`
+	MerkleRoot    Base64URLByte `json:"merkle_root,omitempty"`
+	Epoch         Base64URLByte `json:"epoch,omitempty"`
+	ShareX        Base64URLByte `json:"share_x,omitempty"`
+	ShareY        Base64URLByte `json:"share_y,omitempty"`
+	Nullifier     Base64URLByte `json:"nullifier,omitempty"`
+	RlnIdentifier Base64URLByte `json:"rln_identifier,omitempty"`
 }
 
 type RPCWakuMessage struct {
-	Payload        []byte          `json:"payload,omitempty"`
+	Payload        Base64URLByte   `json:"payload,omitempty"`
 	ContentTopic   string          `json:"contentTopic,omitempty"`
 	Version        uint32          `json:"version"`
 	Timestamp      int64           `json:"timestamp,omitempty"`