Override STR.SerializeWithAd of the STR struct for signing

client/encoding_test.go

@@ -43,7 +43,7 @@ func TestUnmarshalSampleMessage(t *testing.T) {
 	msg, _ := keyserver.MarshalResponse(res)
 	response := UnmarshalResponse(protocol.RegistrationType, []byte(msg))
 	str := response.DirectoryResponse.(*protocol.DirectoryProof).STR
-	if str.Ad == nil || !bytes.Equal(str.Ad.Serialize(), str.Policies.Serialize()) {
+	if !bytes.Equal(d.LatestSTR().Serialize(), str.Serialize()) {
 		t.Error("Cannot unmarshal Associate Data properly")
 	}
 }

merkletree/str.go

@@ -51,17 +51,26 @@ func NewSTR(key sign.PrivateKey, ad AssocData, m *MerkleTree, epoch uint64, prev
 	return str
 }
 
-// Serialize serializes the signed tree root into
-// a specified format for signing.
+// Serialize serializes the signed tree root
+// and its associated data into a specified format for signing.
+// One should use this function for signing as well as
+// verifying the signature.
+// Any composition struct of SignedTreeRoot with
+// a specific AssocData should override this method.
 func (str *SignedTreeRoot) Serialize() []byte {
+	return append(str.SerializeInternal(), str.Ad.Serialize()...)
+}
+
+// SerializeInternal serializes the signed tree root into
+// a specified format.
+func (str *SignedTreeRoot) SerializeInternal() []byte {
 	var strBytes []byte
 	strBytes = append(strBytes, utils.ULongToBytes(str.Epoch)...) // t - epoch number
 	if str.Epoch > 0 {
 		strBytes = append(strBytes, utils.ULongToBytes(str.PreviousEpoch)...) // t_prev - previous epoch number
 	}
 	strBytes = append(strBytes, str.TreeHash...)        // root
 	strBytes = append(strBytes, str.PreviousSTRHash...) // previous STR hash
-	strBytes = append(strBytes, str.Ad.Serialize()...)
 	return strBytes
 }
 

merkletree/str_test.go

@@ -7,7 +7,7 @@ func TestVerifyHashChain(t *testing.T) {
 
 	keyPrefix := "key"
 	valuePrefix := []byte("value")
-	pad, err := NewPAD(TestAd{""}, signKey, vrfPrivKey1, 10)
+	pad, err := NewPAD(TestAd{"abc"}, signKey, vrfPrivKey1, 10)
 	if err != nil {
 		t.Fatal(err)
 	}

protocol/str.go

@@ -1,10 +1,6 @@
 package protocol
 
-import (
-	"encoding/json"
-
-	"github.com/coniks-sys/coniks-go/merkletree"
-)
+import "github.com/coniks-sys/coniks-go/merkletree"
 
 // DirSTR disambiguates merkletree.SignedTreeRoot's AssocData interface,
 // for the purpose of exporting and unmarshalling.
@@ -21,22 +17,12 @@ func NewDirSTR(str *merkletree.SignedTreeRoot) *DirSTR {
 	}
 }
 
+// Serialize overrides merkletree.SignedTreeRoot.Serialize
+func (str *DirSTR) Serialize() []byte {
+	return append(str.SerializeInternal(), str.Policies.Serialize()...)
+}
+
 // VerifyHashChain wraps merkletree.SignedTreeRoot.VerifyHashChain
 func (str *DirSTR) VerifyHashChain(savedSTR *DirSTR) bool {
 	return str.SignedTreeRoot.VerifyHashChain(savedSTR.SignedTreeRoot)
 }
-
-// UnmarshalJSON fills in the unexported Ad interface from the underlying
-// merkletree.SignedTreeRoot.  This is necessary since, for now, Serialize
-// and VerifyHashChain dispatch to methods which dereference it.
-func (str *DirSTR) UnmarshalJSON(m []byte) error {
-	// Use an alias to avoid an infinite loop
-	type DirSTR2 DirSTR
-	str2 := &DirSTR2{}
-	if err := json.Unmarshal(m, str2); err != nil {
-		return err
-	}
-	str2.Ad = str2.Policies
-	*str = DirSTR(*str2)
-	return nil
-}