Add MarhsalInto & Size across all content

Author: noboruma

conn.go

@@ -537,7 +537,7 @@ func (c *Conn) writePackets(ctx context.Context, pkts []*packet) error {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
-	var rawPackets [][]byte
+	rawPackets := make([][]byte, 0, len(pkts))
 
 	for _, pkt := range pkts {
 		if dtlsHandshake, ok := pkt.record.Content.(*handshake.Handshake); ok {
@@ -591,8 +591,8 @@ func (c *Conn) compactRawPackets(rawPackets [][]byte) [][]byte {
 		return rawPackets
 	}
 
-	combinedRawPackets := make([][]byte, 0)
-	currentCombinedRawPacket := make([]byte, 0)
+	combinedRawPackets := make([][]byte, 0, len(rawPackets))
+	var currentCombinedRawPacket []byte
 
 	for _, rawPacket := range rawPackets {
 		if len(currentCombinedRawPacket) == 0 && len(rawPacket) >= c.maximumTransmissionUnit {
@@ -679,8 +679,6 @@ func (c *Conn) processPacket(pkt *packet) ([]byte, error) { //nolint:cyclop
 
 //nolint:cyclop
 func (c *Conn) processHandshakePacket(pkt *packet, dtlsHandshake *handshake.Handshake) ([][]byte, error) {
-	var rawPackets [][]byte
-
 	handshakeFragments, err := c.fragmentHandshake(dtlsHandshake)
 	if err != nil {
 		return nil, err
@@ -690,6 +688,7 @@ func (c *Conn) processHandshakePacket(pkt *packet, dtlsHandshake *handshake.Hand
 		c.state.localSequenceNumber = append(c.state.localSequenceNumber, uint64(0))
 	}
 
+	rawPackets := make([][]byte, 0, len(handshakeFragments))
 	for _, handshakeFragment := range handshakeFragments {
 		seq := atomic.AddUint64(&c.state.localSequenceNumber[epoch], 1) - 1
 		if seq > recordlayer.MaxSequenceNumber {
@@ -764,11 +763,10 @@ func (c *Conn) fragmentHandshake(dtlsHandshake *handshake.Handshake) ([][]byte,
 		return nil, err
 	}
 
-	var fragmentedHandshakes [][]byte
-
 	contentFragments := splitBytes(content, c.maximumTransmissionUnit)
 
 	offset := 0
+	fragmentedHandshakes := make([][]byte, 0, len(contentFragments))
 	for _, contentFragment := range contentFragments {
 		contentFragmentLen := len(contentFragment)
 

handshaker.go

@@ -187,7 +187,8 @@ func (s *handshakeFSM) Run(ctx context.Context, conn flightConn, initialState ha
 		close(s.closed)
 	}()
 	for {
-		s.cfg.log.Tracef("[handshake:%s] %s: %s", srvCliStr(s.state.isClient), s.currentFlight.String(), state.String())
+		s.cfg.log.Tracef("[handshake:%s] %s: %s",
+			srvCliStr(s.state.isClient), s.currentFlight.String(), state.String())
 		if s.cfg.onFlightState != nil {
 			s.cfg.onFlightState(s.currentFlight, state)
 		}

pkg/protocol/alert/alert.go

@@ -145,9 +145,28 @@ 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, a.Size())
+	err := a.MarshalInto(out)
+
+	return out, err
+}
+
+// MarshalInto returns the encoded alert.
+func (a *Alert) MarshalInto(out []byte) error {
+	if len(out) < a.Size() {
+		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,22 @@ 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
+}
+
+// MarshalInto encodes the ApplicationData to binary into a pre-allocated buffer.
+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,27 @@ 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) < c.Size() {
+		return errBufferTooSmall
+	}
+	out[0] = 0x01
+
+	return nil
 }
 
 // Unmarshal populates the ChangeCipherSpec from binary.

pkg/protocol/content.go

@@ -21,5 +21,7 @@ const (
 type Content interface {
 	ContentType() ContentType
 	Marshal() ([]byte, error)
+	MarshalInto([]byte) error
 	Unmarshal(data []byte) error
+	Size() int
 }

pkg/protocol/extension/extension.go

@@ -134,3 +134,17 @@ 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,50 @@ 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)
 
-	msg, err := h.Message.Marshal()
-	if err != nil {
-		return nil, err
+	return out, err
+}
+
+// MarshalInto encodes a handshake into a binary message into a pre-allocated buffer.
+func (h *Handshake) MarshalInto(out []byte) error {
+	if h.Message == nil {
+		return errHandshakeMessageUnset
+	} else if h.Header.FragmentOffset != 0 {
+		return errUnableToMarshalFragmented
 	}
 
-	h.Header.Length = uint32(len(msg)) //nolint:gosec // G115
+	if len(out) < h.Size() {
+		return errBufferTooSmall
+	}
+
+	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
+	}
+
+	err = h.Message.MarshalInto(out[HeaderLength:])
+	if err != nil {
+		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,30 @@ const (
 	handshakeMessageCertificateLengthFieldSize = 3
 )
 
-// Marshal encodes the Handshake.
-func (m *MessageCertificate) Marshal() ([]byte, error) {
-	if m.cache != nil {
-		return m.cache, nil
-	}
+// Size 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 +61,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 += extension.Size(entry.Extensions)
+	}
+
+	return certificateListSize
+}
+
+// MarshalInto is 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)) //nolint:gosec // G115
+	offset += 1
+	n := copy(out[offset:], m.CertificateRequestContext) //nolint:gosec // G115
+	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
-		}
-		certificateList = append(certificateList, extensionsData...)
-
-		// Check size of certificate_list is still within bounds
-		if len(certificateList) > maxUint24 {
-			return nil, errCertificateListTooLong
+			return err
 		}
+		n = copy(out[offset:], extensionsData)
+		offset += n
 	}
 
-	// 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.