diff --git a/elkrem/elkrem.go b/elkrem/elkrem.go index c9ee3309..b9ae2b55 100644 --- a/elkrem/elkrem.go +++ b/elkrem/elkrem.go @@ -104,9 +104,13 @@ func NewElkremReceiver(th uint8) ElkremReceiver { // Next() increments the index to the next hash and outputs it func (e *ElkremSender) Next() (*wire.ShaHash, error) { - // increment index + sha, err := e.AtIndex(e.current) + if err != nil { + return nil, err + } + // increment index for next time e.current++ - return e.AtIndex(e.current) + return sha, nil } // w is the wanted index, i is the root index diff --git a/elkrem/elkrem_test.go b/elkrem/elkrem_test.go index 459a6dc6..e049c2cd 100644 --- a/elkrem/elkrem_test.go +++ b/elkrem/elkrem_test.go @@ -8,10 +8,11 @@ import ( // TestElkremBig makes a height 63 (max size possible) tree and tries 10K hashes func TestElkremBig(t *testing.T) { - ex := NewElkremSender(63, wire.DoubleSha256SH([]byte("elktest"))) + sndr := NewElkremSender(63, wire.DoubleSha256SH([]byte("elktest"))) rcv := NewElkremReceiver(63) + SenderSerdesTest(t, sndr) for n := uint64(0); n < 10000; n++ { - sha, err := ex.AtIndex(n) + sha, err := sndr.Next() if err != nil { t.Fatal(err) } @@ -27,22 +28,23 @@ func TestElkremBig(t *testing.T) { } } } + SenderSerdesTest(t, sndr) + ReceiverSerdesTest(t, rcv) for n := uint64(0); n < 10000; n += 500 { sha, err := rcv.AtIndex(n) if err != nil { t.Fatal(err) } - t.Logf("Retreived index %d %s\n", - n, sha.String()) + t.Logf("Retreived index %d %s\n", n, sha.String()) } } // TestElkremSmall makes a height 15 (65534 size) tree and tries 10K hashes func TestElkremSmall(t *testing.T) { - ex := NewElkremSender(15, wire.DoubleSha256SH([]byte("elktest"))) + sndr := NewElkremSender(15, wire.DoubleSha256SH([]byte("elktest"))) rcv := NewElkremReceiver(15) for n := uint64(0); n < 5000; n++ { - sha, err := ex.AtIndex(n) + sha, err := sndr.Next() if err != nil { t.Fatal(err) } diff --git a/elkrem/serdes.go b/elkrem/serdes.go index 4be0f691..84054d6b 100644 --- a/elkrem/serdes.go +++ b/elkrem/serdes.go @@ -38,7 +38,7 @@ func (e *ElkremSender) ToBytes() ([]byte, error) { return nil, fmt.Errorf("%d byte hash, expect 32", n) } - return nil, nil + return buf.Bytes(), nil } // ElkremSenderFromBytes turns a 41 byte slice into a sender, picking up at diff --git a/elkrem/serdes_test.go b/elkrem/serdes_test.go new file mode 100644 index 00000000..cb3e4173 --- /dev/null +++ b/elkrem/serdes_test.go @@ -0,0 +1,19 @@ +package elkrem + +import "testing" + +func ReceiverSerdesTest(t *testing.T, er ElkremReceiver) { + b, err := er.ToBytes() + if err != nil { + t.Fatal(err) + } + t.Logf("Serialized receiver; %d bytes, hex:\n%x\n", len(b), b) +} + +func SenderSerdesTest(t *testing.T, es ElkremSender) { + b, err := es.ToBytes() + if err != nil { + t.Fatal(err) + } + t.Logf("Serialized sender; %d bytes, hex:\n%x\n", len(b), b) +}