Add optimizations for connection & handshake

Author: noboruma

conn.go

@@ -595,7 +595,12 @@ func (c *Conn) compactRawPackets(rawPackets [][]byte) [][]byte {
 	currentCombinedRawPacket := make([]byte, 0)
 
 	for _, rawPacket := range rawPackets {
-		if len(currentCombinedRawPacket) > 0 && len(currentCombinedRawPacket)+len(rawPacket) >= c.maximumTransmissionUnit {
+		if len(currentCombinedRawPacket) == 0 && len(rawPacket) >= c.maximumTransmissionUnit {
+			combinedRawPackets = append(combinedRawPackets, rawPacket)
+
+			continue
+		} else if len(currentCombinedRawPacket) > 0 &&
+			len(currentCombinedRawPacket)+len(rawPacket) >= c.maximumTransmissionUnit {
 			combinedRawPackets = append(combinedRawPackets, currentCombinedRawPacket)
 			currentCombinedRawPacket = []byte{}
 		}
@@ -674,7 +679,7 @@ func (c *Conn) processPacket(pkt *packet) ([]byte, error) { //nolint:cyclop
 
 //nolint:cyclop
 func (c *Conn) processHandshakePacket(pkt *packet, dtlsHandshake *handshake.Handshake) ([][]byte, error) {
-	rawPackets := make([][]byte, 0)
+	var rawPackets [][]byte
 
 	handshakeFragments, err := c.fragmentHandshake(dtlsHandshake)
 	if err != nil {
@@ -710,12 +715,15 @@ func (c *Conn) processHandshakePacket(pkt *packet, dtlsHandshake *handshake.Hand
 				ConnectionID:   c.state.remoteConnectionID,
 				SequenceNumber: pkt.record.Header.SequenceNumber,
 			}
-			rawPacket, err = cidHeader.Marshal()
+
+			hs := recordlayer.FixedHeaderSize + len(cidHeader.ConnectionID)
+			rawPacket = make([]byte, hs+len(rawInner))
+			err = cidHeader.MarshalInto(rawPacket)
 			if err != nil {
 				return nil, err
 			}
 			pkt.record.Header = *cidHeader
-			rawPacket = append(rawPacket, rawInner...)
+			copy(rawPacket[hs:], rawInner)
 		} else {
 			recordlayerHeader := &recordlayer.Header{
 				Version:        pkt.record.Header.Version,
@@ -725,13 +733,15 @@ func (c *Conn) processHandshakePacket(pkt *packet, dtlsHandshake *handshake.Hand
 				SequenceNumber: seq,
 			}
 
-			rawPacket, err = recordlayerHeader.Marshal()
+			hs := recordlayer.FixedHeaderSize + len(recordlayerHeader.ConnectionID)
+			rawPacket = make([]byte, hs+len(handshakeFragment))
+			err = recordlayerHeader.MarshalInto(rawPacket)
 			if err != nil {
 				return nil, err
 			}
 
 			pkt.record.Header = *recordlayerHeader
-			rawPacket = append(rawPacket, handshakeFragment...)
+			copy(rawPacket[hs:], handshakeFragment)
 		}
 
 		if pkt.shouldEncrypt {
@@ -754,14 +764,9 @@ func (c *Conn) fragmentHandshake(dtlsHandshake *handshake.Handshake) ([][]byte,
 		return nil, err
 	}
 
-	fragmentedHandshakes := make([][]byte, 0)
+	var fragmentedHandshakes [][]byte
 
 	contentFragments := splitBytes(content, c.maximumTransmissionUnit)
-	if len(contentFragments) == 0 {
-		contentFragments = [][]byte{
-			{},
-		}
-	}
 
 	offset := 0
 	for _, contentFragment := range contentFragments {
@@ -777,12 +782,13 @@ func (c *Conn) fragmentHandshake(dtlsHandshake *handshake.Handshake) ([][]byte,
 
 		offset += contentFragmentLen
 
-		fragmentedHandshake, err := headerFragment.Marshal()
+		fragmentedHandshake := make([]byte, handshake.HeaderLength+len(contentFragment))
+		err := headerFragment.MarshalInto(fragmentedHandshake)
 		if err != nil {
 			return nil, err
 		}
 
-		fragmentedHandshake = append(fragmentedHandshake, contentFragment...)
+		copy(fragmentedHandshake[handshake.HeaderLength:], contentFragment)
 		fragmentedHandshakes = append(fragmentedHandshakes, fragmentedHandshake)
 	}
 

errors.go

@@ -143,6 +143,8 @@ var (
 	//nolint:err113
 	errFailedToAccessPoolReadBuffer = &InternalError{Err: errors.New("failed to access pool read buffer")}
 	//nolint:err113
+	errFailedToAccessPoolTimer = &InternalError{Err: errors.New("failed to access pool timer")}
+	//nolint:err113
 	errFragmentBufferOverflow = &InternalError{Err: errors.New("fragment buffer overflow")}
 
 	//nolint:err113

handshaker.go

@@ -275,6 +275,15 @@ func (s *handshakeFSM) send(ctx context.Context, c flightConn) (handshakeState,
 	return handshakeWaiting, nil
 }
 
+var timerPool = sync.Pool{ //nolint:gochecknoglobals
+	New: func() any {
+		t := time.NewTimer(time.Millisecond)
+		t.Stop()
+
+		return t
+	},
+}
+
 func (s *handshakeFSM) wait(ctx context.Context, conn flightConn) (handshakeState, error) { //nolint:gocognit,cyclop
 	parse, errFlight := s.currentFlight.getFlightParser()
 	if errFlight != nil {
@@ -285,7 +294,15 @@ func (s *handshakeFSM) wait(ctx context.Context, conn flightConn) (handshakeStat
 		return handshakeErrored, errFlight
 	}
 
-	retransmitTimer := time.NewTimer(s.retransmitInterval)
+	retransmitTimer, ok := timerPool.Get().(*time.Timer)
+	if !ok {
+		return handshakeErrored, errFailedToAccessPoolTimer
+	}
+	defer func() {
+		retransmitTimer.Stop()
+		timerPool.Put(retransmitTimer)
+	}()
+	retransmitTimer.Reset(s.retransmitInterval)
 	for {
 		select {
 		case state := <-conn.recvHandshake():

pkg/protocol/application_data.go

@@ -19,7 +19,7 @@ func (a ApplicationData) ContentType() ContentType {
 
 // Marshal encodes the ApplicationData to binary.
 func (a *ApplicationData) Marshal() ([]byte, error) {
-	return append([]byte{}, a.Data...), nil
+	return a.Data, nil
 }
 
 // Unmarshal populates the ApplicationData from binary.

pkg/protocol/handshake/header.go

@@ -29,14 +29,24 @@ type Header struct {
 // Marshal encodes the Header.
 func (h *Header) Marshal() ([]byte, error) {
 	out := make([]byte, HeaderLength)
+	err := h.MarshalInto(out)
+
+	return out, err
+}
+
+// MarshalInto encodes the Header using a pre-allocated buffer.
+func (h *Header) MarshalInto(out []byte) error {
+	if len(out) < HeaderLength {
+		return errBufferTooSmall
+	}
 
 	out[0] = byte(h.Type)
 	util.PutBigEndianUint24(out[1:], h.Length)
 	binary.BigEndian.PutUint16(out[4:], h.MessageSequence)
 	util.PutBigEndianUint24(out[6:], h.FragmentOffset)
 	util.PutBigEndianUint24(out[9:], h.FragmentLength)
 
-	return out, nil
+	return nil
 }
 
 // Unmarshal populates the header from encoded data.

pkg/protocol/handshake/message_certificate.go

@@ -13,6 +13,7 @@ import (
 // https://tools.ietf.org/html/rfc5246#section-7.4.2
 type MessageCertificate struct {
 	Certificate [][]byte
+	cache       []byte
 }
 
 // Type returns the Handshake Type.
@@ -26,6 +27,9 @@ const (
 
 // Marshal encodes the Handshake.
 func (m *MessageCertificate) Marshal() ([]byte, error) {
+	if m.cache != nil {
+		return m.cache, nil
+	}
 	total := handshakeMessageCertificateLengthFieldSize
 
 	for _, cert := range m.Certificate {
@@ -50,6 +54,8 @@ func (m *MessageCertificate) Marshal() ([]byte, error) {
 		offset += len(cert)
 	}
 
+	m.cache = out
+
 	return out, nil
 }
 

pkg/protocol/recordlayer/header.go

@@ -37,8 +37,24 @@ func (h *Header) Marshal() ([]byte, error) {
 	}
 
 	hs := FixedHeaderSize + len(h.ConnectionID)
-
 	out := make([]byte, hs)
+	err := h.MarshalInto(out)
+
+	return out, err
+}
+
+// MarshalInto encodes a TLS RecordLayer Header to binary using pre-allocated buffer.
+func (h *Header) MarshalInto(out []byte) error {
+	if h.SequenceNumber > MaxSequenceNumber {
+		return errSequenceNumberOverflow
+	}
+
+	hs := FixedHeaderSize + len(h.ConnectionID)
+
+	if len(out) < hs {
+		return errBufferTooSmall
+	}
+
 	out[0] = byte(h.ContentType)
 	out[1] = h.Version.Major
 	out[2] = h.Version.Minor
@@ -47,7 +63,7 @@ func (h *Header) Marshal() ([]byte, error) {
 	copy(out[11:11+len(h.ConnectionID)], h.ConnectionID)
 	binary.BigEndian.PutUint16(out[hs-2:], h.ContentLen)
 
-	return out, nil
+	return nil
 }
 
 // Unmarshal populates a TLS RecordLayer Header from binary.

pkg/protocol/recordlayer/recordlayer.go

@@ -58,12 +58,16 @@ func (r *RecordLayer) Marshal() ([]byte, error) {
 	r.Header.ContentLen = uint16(len(contentRaw)) //nolint:gosec // G115
 	r.Header.ContentType = r.Content.ContentType()
 
-	headerRaw, err := r.Header.Marshal()
+	out := make([]byte, len(contentRaw)+r.Header.Size())
+
+	err = r.Header.MarshalInto(out)
 	if err != nil {
 		return nil, err
 	}
 
-	return append(headerRaw, contentRaw...), nil
+	copy(out[r.Header.Size():], contentRaw)
+
+	return out, nil
 }
 
 // Unmarshal populates the RecordLayer from binary.