Add limit to fragmentBuffer

Before we imposed no limit on the amount of data we would buffer during
the handshake. This changes adds a 2 megabyte. When the limit is
exceeded the Conn returns an error.

errors.go

@@ -62,6 +62,7 @@ var (
 	errSequenceNumberOverflow            = &InternalError{Err: errors.New("sequence number overflow")}                        //nolint:goerr113
 	errInvalidFSMTransition              = &InternalError{Err: errors.New("invalid state machine transition")}                //nolint:goerr113
 	errFailedToAccessPoolReadBuffer      = &InternalError{Err: errors.New("failed to access pool read buffer")}               //nolint:goerr113
+	errFragmentBufferOverflow            = &InternalError{Err: errors.New("fragment buffer overflow")}                        //nolint:goerr113
 )
 
 // FatalError indicates that the DTLS connection is no longer available.

fragment_buffer.go

@@ -6,6 +6,9 @@ import (
 	"github.com/pion/dtls/v2/pkg/protocol/recordlayer"
 )
 
+// 2 megabytes
+const fragmentBufferMaxSize = 2000000
+
 type fragment struct {
 	recordLayerHeader recordlayer.Header
 	handshakeHeader   handshake.Header
@@ -23,10 +26,25 @@ func newFragmentBuffer() *fragmentBuffer {
 	return &fragmentBuffer{cache: map[uint16][]*fragment{}}
 }
 
+// current total size of buffer
+func (f *fragmentBuffer) size() int {
+	size := 0
+	for i := range f.cache {
+		for j := range f.cache[i] {
+			size += len(f.cache[i][j].data)
+		}
+	}
+	return size
+}
+
 // Attempts to push a DTLS packet to the fragmentBuffer
 // when it returns true it means the fragmentBuffer has inserted and the buffer shouldn't be handled
 // when an error returns it is fatal, and the DTLS connection should be stopped
 func (f *fragmentBuffer) push(buf []byte) (bool, error) {
+	if f.size()+len(buf) >= fragmentBufferMaxSize {
+		return false, errFragmentBufferOverflow
+	}
+
 	frag := new(fragment)
 	if err := frag.recordLayerHeader.Unmarshal(buf); err != nil {
 		return false, err

fragment_buffer_test.go

@@ -1,6 +1,7 @@
 package dtls
 
 import (
+	"errors"
 	"reflect"
 	"testing"
 )
@@ -57,7 +58,7 @@ func TestFragmentBuffer(t *testing.T) {
 			Epoch: 0,
 		},
 		{
-			Name: "Multiple Handshakes in Signle Fragment",
+			Name: "Multiple Handshakes in Single Fragment",
 			In: [][]byte{
 				{
 					0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x30, /* record header */
@@ -113,3 +114,18 @@ func TestFragmentBuffer(t *testing.T) {
 		}
 	}
 }
+
+func TestFragmentBuffer_Overflow(t *testing.T) {
+	fragmentBuffer := newFragmentBuffer()
+
+	// Push a buffer that doesn't exceed size limits
+	if _, err := fragmentBuffer.push([]byte{0x16, 0xfe, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff, 0x00}); err != nil {
+		t.Fatal(err)
+	}
+
+	// Allocate a buffer that exceeds cache size
+	largeBuffer := make([]byte, fragmentBufferMaxSize)
+	if _, err := fragmentBuffer.push(largeBuffer); !errors.Is(err, errFragmentBufferOverflow) {
+		t.Fatalf("Pushing a large buffer returned (%s) expected(%s)", err, errFragmentBufferOverflow)
+	}
+}