Merge pull request #278 from nickysemenza/fix-protocol-bytes-size

fix: correctly calculate protocol.Operation Bytes() size

protocol/protocol.go

@@ -494,6 +494,12 @@ func (o *Operation) Bytes() uint16 {
 	if o.CertID != "" {
 		add(tlvLen(len([]byte(o.CertID))))
 	}
+	if o.CustomFuncName != "" {
+		add(tlvLen(len([]byte(o.CustomFuncName))))
+	}
+	if o.JaegerSpan != nil {
+		add(tlvLen(len(o.JaegerSpan)))
+	}
 	if int(length)+headerSize < paddedLength {
 		// TODO: Are we sure that's the right behavior?
 

protocol/protocol_test.go

@@ -1,6 +1,8 @@
 package protocol
 
 import (
+	"bytes"
+	"crypto/rand"
 	"crypto/sha1"
 	"crypto/sha256"
 	"net"
@@ -12,24 +14,31 @@ import (
 func TestMarshalBinary(t *testing.T) {
 	require := require.New(t)
 
+	// we want to push the payload over the paddedLength of 1024
+	// to ensure that the size is calculated correctly.
+	extra := make([]byte, 100)
+	payload := make([]byte, 1000)
+	rand.Read(extra)
+	rand.Read(payload)
 	op := Operation{
 		Opcode:         OpECDSASignSHA256,
-		Payload:        []byte("Payload"),
-		Extra:          []byte("Extra"),
+		Payload:        payload,
+		Extra:          extra,
 		Digest:         sha256.Sum256([]byte("Digest")),
 		SKI:            sha1.Sum([]byte("SKI")),
 		ClientIP:       net.ParseIP("1.1.1.1").To4(),
 		ServerIP:       net.ParseIP("2.2.2.2").To4(),
 		SNI:            "SNI",
 		CertID:         "SNI",
 		CustomFuncName: "CustomFuncName",
+		JaegerSpan:     []byte("615f730ad5fe896f:615f730ad5fe896f:1"),
 	}
 	pkt := NewPacket(42, op)
 	b, err := pkt.MarshalBinary()
 	require.NoError(err)
 
 	var pkt2 Packet
-	err = pkt2.UnmarshalBinary(b)
+	_, err = pkt2.ReadFrom(bytes.NewReader(b))
 	require.NoError(err)
 	require.Equal(pkt.ID, pkt2.ID)
 	require.Equal(op, pkt2.Operation)