From a6397ff7282bc56dc37a68ea9211702edb4de1de Mon Sep 17 00:00:00 2001
From: Sean DuBois <sean@siobud.com>
Date: Tue, 10 May 2022 15:36:31 -0400
Subject: [PATCH] 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               |  1 +
 fragment_buffer.go      | 18 ++++++++++++++++++
 fragment_buffer_test.go | 18 +++++++++++++++++-
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/errors.go b/errors.go
index 042996c62..09c6ffe9f 100644
--- a/errors.go
+++ b/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.
diff --git a/fragment_buffer.go b/fragment_buffer.go
index 8a00ee7e4..c86a7e8f9 100644
--- a/fragment_buffer.go
+++ b/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
diff --git a/fragment_buffer_test.go b/fragment_buffer_test.go
index 0f49c7588..f1013cc70 100644
--- a/fragment_buffer_test.go
+++ b/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)
+	}
+}