From a2ed6ef0a6ef5f39f8408326a0d728b01f38626a Mon Sep 17 00:00:00 2001
From: Ryan <ryanford@poluglottos.com>
Date: Fri, 28 Oct 2022 14:55:30 +0200
Subject: [PATCH] feat: custom JSON marshal/unmarshal for ExtendedHeader
 wrapping PubKey w/ amino codec

---
 api/rpc_test.go      |  4 ++++
 header/header.go     | 48 +++++++++++++++++++++++++++++++++++++++++---
 header/serde_test.go | 22 +++++++++++++++++---
 3 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/api/rpc_test.go b/api/rpc_test.go
index c29b233512..b61c4705fa 100644
--- a/api/rpc_test.go
+++ b/api/rpc_test.go
@@ -73,6 +73,10 @@ func TestAllReturnValuesAreMarshalable(t *testing.T) {
 }
 
 func implementsMarshaler(t *testing.T, typ reflect.Type) { //nolint:unused
+	if typ.Implements(reflect.TypeOf(new(json.Marshaler)).Elem()) {
+		return
+	}
+
 	switch typ.Kind() {
 	case reflect.Struct:
 		for i := 0; i < typ.NumField(); i++ {
diff --git a/header/header.go b/header/header.go
index f62bc72487..2fa93b4412 100644
--- a/header/header.go
+++ b/header/header.go
@@ -3,18 +3,20 @@ package header
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 
-	"github.com/celestiaorg/celestia-node/share"
-
 	"github.com/ipfs/go-blockservice"
 	logging "github.com/ipfs/go-log/v2"
 
 	bts "github.com/tendermint/tendermint/libs/bytes"
+	amino "github.com/tendermint/tendermint/libs/json"
 	core "github.com/tendermint/tendermint/types"
 
-	"github.com/celestiaorg/celestia-app/pkg/da"
 	appshares "github.com/celestiaorg/celestia-app/pkg/shares"
+
+	"github.com/celestiaorg/celestia-app/pkg/da"
+	"github.com/celestiaorg/celestia-node/share"
 )
 
 var log = logging.Logger("header")
@@ -144,3 +146,43 @@ func (eh *ExtendedHeader) UnmarshalBinary(data []byte) error {
 	*eh = *out
 	return nil
 }
+
+// MarshalJSON marshals an ExtendedHeader to JSON. The ValidatorSet is wrapped with amino encoding, to be able to
+// unmarshal the crypto.PubKey type back from JSON.
+func (eh *ExtendedHeader) MarshalJSON() ([]byte, error) {
+	type Alias ExtendedHeader
+	validatorSet, err := amino.Marshal(eh.ValidatorSet)
+	if err != nil {
+		return nil, err
+	}
+	return json.Marshal(&struct {
+		ValidatorSet json.RawMessage `json:"validator_set"`
+		*Alias
+	}{
+		ValidatorSet: validatorSet,
+		Alias:        (*Alias)(eh),
+	})
+}
+
+// UnmarshalJSON unmarshals an ExtendedHeader from JSON. The ValidatorSet is wrapped with amino encoding, to be able to
+// unmarshal the crypto.PubKey type back from JSON.
+func (eh *ExtendedHeader) UnmarshalJSON(data []byte) error {
+	type Alias ExtendedHeader
+	aux := &struct {
+		ValidatorSet json.RawMessage `json:"validator_set"`
+		*Alias
+	}{
+		Alias: (*Alias)(eh),
+	}
+	if err := json.Unmarshal(data, &aux); err != nil {
+		return err
+	}
+
+	valSet := new(core.ValidatorSet)
+	if err := amino.Unmarshal(aux.ValidatorSet, valSet); err != nil {
+		return err
+	}
+
+	eh.ValidatorSet = valSet
+	return nil
+}
diff --git a/header/serde_test.go b/header/serde_test.go
index 36c0301eb9..2972b8c706 100644
--- a/header/serde_test.go
+++ b/header/serde_test.go
@@ -9,13 +9,29 @@ import (
 
 func TestMarshalUnmarshalExtendedHeader(t *testing.T) {
 	in := RandExtendedHeader(t)
-	data, err := in.MarshalBinary()
+	binaryData, err := in.MarshalBinary()
 	require.NoError(t, err)
 
 	out := &ExtendedHeader{}
-	err = out.UnmarshalBinary(data)
+	err = out.UnmarshalBinary(binaryData)
 	require.NoError(t, err)
-	assert.Equal(t, in.ValidatorSet, out.ValidatorSet)
+	equalExtendedHeader(t, in, out)
+
+	// A custom JSON marshal/unmarshal is necessary which wraps the ValidatorSet with amino
+	// encoding, to be able to marshal the crypto.PubKey type back from JSON.
+	jsonData, err := in.MarshalJSON()
+	require.NoError(t, err)
+
+	out = &ExtendedHeader{}
+	err = out.UnmarshalJSON(jsonData)
+	require.NoError(t, err)
+	equalExtendedHeader(t, in, out)
+}
+
+func equalExtendedHeader(t *testing.T, in, out *ExtendedHeader) {
+	// ValidatorSet.totalVotingPower is not set (is a cached value that can be recomputed client side)
+	assert.Equal(t, in.ValidatorSet.Validators, out.ValidatorSet.Validators)
+	assert.Equal(t, in.ValidatorSet.Proposer, out.ValidatorSet.Proposer)
 	assert.True(t, in.DAH.Equals(out.DAH))
 	// not the check for equality as time.Time is not serialized exactly 1:1
 	assert.NotZero(t, out.RawHeader)