add marhsalInto across all content

Author: noboruma

pkg/protocol/alert/alert.go

@@ -145,9 +145,26 @@ func (a Alert) ContentType() protocol.ContentType {
 	return protocol.ContentTypeAlert
 }
 
+// Size returns the minimal buffer size required for MarshalInto.
+func (a Alert) Size() int {
+	return 2
+}
+
 // Marshal returns the encoded alert.
 func (a *Alert) Marshal() ([]byte, error) {
-	return []byte{byte(a.Level), byte(a.Description)}, nil
+	out := make([]byte, 2)
+	err := a.MarshalInto(out)
+	return out, err
+}
+
+// MarshalInto returns the encoded alert.
+func (a *Alert) MarshalInto(out []byte) error {
+	if len(out) < 2 {
+		return errBufferTooSmall
+	}
+	out[0] = byte(a.Level)
+	out[1] = byte(a.Description)
+	return nil
 }
 
 // Unmarshal populates the alert from binary data.

pkg/protocol/application_data.go

@@ -19,7 +19,20 @@ func (a ApplicationData) ContentType() ContentType {
 
 // Marshal encodes the ApplicationData to binary.
 func (a *ApplicationData) Marshal() ([]byte, error) {
-	return a.Data, nil
+	out := make([]byte, len(a.Data))
+	err := a.MarshalInto(out)
+	return out, err
+}
+
+// Marshal encodes the ApplicationData to binary.
+func (a *ApplicationData) MarshalInto(out []byte) error {
+	copy(out, a.Data)
+	return nil
+}
+
+// Size returns the size required for MarshalInto.
+func (a ApplicationData) Size() int {
+	return len(a.Data)
 }
 
 // Unmarshal populates the ApplicationData from binary.

pkg/protocol/change_cipher_spec.go

@@ -15,9 +15,25 @@ func (c ChangeCipherSpec) ContentType() ContentType {
 	return ContentTypeChangeCipherSpec
 }
 
+// Size returns the minimal buffer size required for MarshalInto.
+func (c ChangeCipherSpec) Size() int {
+	return 1
+}
+
 // Marshal encodes the ChangeCipherSpec to binary.
 func (c *ChangeCipherSpec) Marshal() ([]byte, error) {
-	return []byte{0x01}, nil
+	out := make([]byte, 1)
+	err := c.MarshalInto(out)
+	return out, err
+}
+
+// MarshalInto encodes the ChangeCipherSpec to binary into a pre-allocated buffer.
+func (c *ChangeCipherSpec) MarshalInto(out []byte) error {
+	if len(out) < 1 {
+		return errBufferTooSmall
+	}
+	out[0] = 0x01
+	return nil
 }
 
 // Unmarshal populates the ChangeCipherSpec from binary.

pkg/protocol/content.go

@@ -22,4 +22,5 @@ type Content interface {
 	ContentType() ContentType
 	Marshal() ([]byte, error)
 	Unmarshal(data []byte) error
+	//Size() int
 }

pkg/protocol/extension/extension.go

@@ -134,3 +134,16 @@ func Marshal(e []Extension) ([]byte, error) {
 
 	return append(out, extensions...), nil
 }
+
+// Size returns the length of extensions marshal.
+func Size(e []Extension) int {
+	total := 2
+	for _, e := range e {
+		raw, err := e.Marshal()
+		if err != nil {
+			return 0
+		}
+		total += len(raw)
+	}
+	return total
+}

pkg/protocol/handshake/handshake.go

@@ -62,6 +62,8 @@ func (t Type) String() string { //nolint:cyclop
 // Message is the body of a Handshake datagram.
 type Message interface {
 	Marshal() ([]byte, error)
+	MarshalInto([]byte) error
+	Size() int
 	Unmarshal(data []byte) error
 	Type() Type
 }
@@ -84,28 +86,49 @@ func (h Handshake) ContentType() protocol.ContentType {
 	return protocol.ContentTypeHandshake
 }
 
+// Size returns the minimal buffer size required for MarshalInto.
+func (h *Handshake) Size() int {
+	return HeaderLength + h.Message.Size()
+}
+
 // Marshal encodes a handshake into a binary message.
 func (h *Handshake) Marshal() ([]byte, error) {
 	if h.Message == nil {
 		return nil, errHandshakeMessageUnset
 	} else if h.Header.FragmentOffset != 0 {
 		return nil, errUnableToMarshalFragmented
 	}
+	out := make([]byte, h.Size())
+	err := h.MarshalInto(out)
+	return out, err
+}
+
+// Marshal encodes a handshake into a binary message.
+func (h *Handshake) MarshalInto(out []byte) error {
+	if h.Message == nil {
+		return errHandshakeMessageUnset
+	} else if h.Header.FragmentOffset != 0 {
+		return errUnableToMarshalFragmented
+	}
+
+	if len(out) < h.Size() {
+		return errBufferTooSmall
+	}
 
-	msg, err := h.Message.Marshal()
+	err := h.Message.MarshalInto(out[HeaderLength:])
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	h.Header.Length = uint32(len(msg)) //nolint:gosec // G115
+	h.Header.Length = uint32(h.Message.Size()) //nolint:gosec // G115
 	h.Header.FragmentLength = h.Header.Length
 	h.Header.Type = h.Message.Type()
-	header, err := h.Header.Marshal()
+	err = h.Header.MarshalInto(out[:])
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	return append(header, msg...), nil
+	return nil
 }
 
 // Unmarshal decodes a handshake from a binary message.

pkg/protocol/handshake/message_certificate.go

@@ -13,7 +13,6 @@ import (
 // https://tools.ietf.org/html/rfc5246#section-7.4.2
 type MessageCertificate struct {
 	Certificate [][]byte
-	cache       []byte
 }
 
 // Type returns the Handshake Type.
@@ -25,22 +24,29 @@ const (
 	handshakeMessageCertificateLengthFieldSize = 3
 )
 
-// Marshal encodes the Handshake.
-func (m *MessageCertificate) Marshal() ([]byte, error) {
-	if m.cache != nil {
-		return m.cache, nil
-	}
+// MarshalInto returns the minimal size required for MarshalInto.
+func (m *MessageCertificate) Size() int {
 	total := handshakeMessageCertificateLengthFieldSize
 
 	for _, cert := range m.Certificate {
 		total += handshakeMessageCertificateLengthFieldSize + len(cert)
 	}
 
-	out := make([]byte, total)
+	return total
+}
+
+// Marshal encodes the Handshake.
+func (m *MessageCertificate) Marshal() ([]byte, error) {
+	out := make([]byte, m.Size())
+	err := m.MarshalInto(out)
+	return out, err
+}
 
+// MarshalInto encodes the Handshake into a pre-allocated buffer.
+func (m *MessageCertificate) MarshalInto(out []byte) error {
 	// Total Payload Size
 	//nolint:gosec // G115
-	util.PutBigEndianUint24(out, uint32(total-handshakeMessageCertificateLengthFieldSize))
+	util.PutBigEndianUint24(out, uint32(m.Size()-handshakeMessageCertificateLengthFieldSize))
 	offset := handshakeMessageCertificateLengthFieldSize
 
 	for _, cert := range m.Certificate {
@@ -54,9 +60,7 @@ func (m *MessageCertificate) Marshal() ([]byte, error) {
 		offset += len(cert)
 	}
 
-	m.cache = out
-
-	return out, nil
+	return nil
 }
 
 // Unmarshal populates the message from encoded data.

pkg/protocol/handshake/message_certificate_13.go

@@ -71,44 +71,77 @@ func (m *MessageCertificate13) Marshal() ([]byte, error) {
 		return nil, errCertificateRequestContextTooLong
 	}
 
+	out := make([]byte, m.Size())
+	err := m.MarshalInto(out)
+	return out, err
+}
+
+func (m *MessageCertificate13) Size() int {
+	return 1 + len(m.CertificateRequestContext) + cert13CertLengthFieldSize + m.certsSize()
+}
+
+func (m *MessageCertificate13) certsSize() int {
+	certificateListSize := 0
+	for _, entry := range m.CertificateList {
+		certificateListSize += cert13CertLengthFieldSize
+		certificateListSize += len(entry.CertificateData)
+		certificateListSize += cert13ExtLengthFieldSize
+		certificateListSize += extension.Size(entry.Extensions)
+	}
+	return certificateListSize
+}
+
+// MarshalInto, same as Marshal but uses a pre-allocated buffer.
+func (m *MessageCertificate13) MarshalInto(out []byte) error {
+	// Validate certificate_request_context length
+	if len(m.CertificateRequestContext) > cert13ContextMaxLength {
+		return errCertificateRequestContextTooLong
+	}
+
+	if len(out) < m.Size() {
+		return errBufferTooSmall
+	}
+
+	// Check size of certificate_list is still within bounds
+	if m.certsSize() > maxUint24 {
+		return errCertificateListTooLong
+	}
+
 	// Start with certificate_request_context (1-byte length prefix)
 	//nolint:gosec // G115: certificate_request_context length is validated to be <= 255 above.
-	out := []byte{byte(len(m.CertificateRequestContext))}
-	out = append(out, m.CertificateRequestContext...)
+	offset := 0
+	out[0] = byte(len(m.CertificateRequestContext))
+	offset += 1
+	n := copy(out[offset:], m.CertificateRequestContext)
+	offset += n
+
+	// Add certificate_list with 3-byte length prefix
+	util.PutBigEndianUint24(out[offset:], uint32(m.certsSize())) //nolint:gosec // G115
+	offset += 3
 
 	// Build certificate_list
-	certificateList := []byte{}
 	for _, entry := range m.CertificateList {
 		// Add cert_data as a 3-byte length prefix
 		certDataLen := len(entry.CertificateData)
 		if certDataLen == 0 || certDataLen > maxUint24 {
-			return nil, errInvalidCertificateEntry
+			return errInvalidCertificateEntry
 		}
-		certDataLenBytes := make([]byte, cert13CertLengthFieldSize)
-		util.PutBigEndianUint24(certDataLenBytes, uint32(certDataLen)) //nolint:gosec // G115
-		certificateList = append(certificateList, certDataLenBytes...)
-		certificateList = append(certificateList, entry.CertificateData...)
+		util.PutBigEndianUint24(out[offset:], uint32(certDataLen)) //nolint:gosec // G115
+		offset += 3
+		n = copy(out[offset:], entry.CertificateData)
+		offset += n
 
 		// Marshal extensions (includes a 2-byte length prefix)
 		extensionsData, err := extension.Marshal(entry.Extensions)
 		if err != nil {
-			return nil, err
+			return err
 		}
-		certificateList = append(certificateList, extensionsData...)
+		n = copy(out[offset:], extensionsData)
+		offset += n
 
-		// Check size of certificate_list is still within bounds
-		if len(certificateList) > maxUint24 {
-			return nil, errCertificateListTooLong
-		}
 	}
 
-	// Add certificate_list with 3-byte length prefix
-	certificateListLenBytes := make([]byte, cert13CertLengthFieldSize)
-	util.PutBigEndianUint24(certificateListLenBytes, uint32(len(certificateList))) //nolint:gosec // G115
-	out = append(out, certificateListLenBytes...)
-	out = append(out, certificateList...)
-
-	return out, nil
+	return nil
 }
 
 // parseCertificate13Entry parses a single certificate entry from the cryptobyte string.