Repository: pion/dtls
Pull request #462 (2022-05-10 19:38:11)

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.

codecov[bot]: # [Codecov](https://codecov.io/gh/pion/dtls/pull/462?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) Report
> Merging [#462](https://codecov.io/gh/pion/dtls/pull/462?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) (12d7cae) into [master](https://codecov.io/gh/pion/dtls/commit/e0b2ce3592e8e7d73713ac67b363a2e192a4cecf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) (e0b2ce3) will **decrease** coverage by `9.59%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #462      +/-   ##
==========================================
- Coverage   75.77%   66.18%   -9.60%     
==========================================
  Files          96       65      -31     
  Lines        5586     4687     -899     
==========================================
- Hits         4233     3102    -1131     
- Misses       1021     1317     +296     
+ Partials      332      268      -64     
```

| Flag | Coverage Δ | |
|---|---|---|
| go | `?` | |
| wasm | `66.18% <100.00%> (+0.07%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/pion/dtls/pull/462?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | Coverage Δ | |
|---|---|---|
| [errors.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-ZXJyb3JzLmdv) | `60.97% <ø> (-12.20%)` | :arrow_down: |
| [fragment\_buffer.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-ZnJhZ21lbnRfYnVmZmVyLmdv) | `87.14% <100.00%> (+2.14%)` | :arrow_up: |
| [pkg/protocol/errors.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2Vycm9ycy5nbw==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pkg/protocol/version.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL3ZlcnNpb24uZ28=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pkg/protocol/application\_data.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2FwcGxpY2F0aW9uX2RhdGEuZ28=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pkg/protocol/extension/srtp\_protection\_profile.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2V4dGVuc2lvbi9zcnRwX3Byb3RlY3Rpb25fcHJvZmlsZS5nbw==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pkg/protocol/handshake/header.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2hhbmRzaGFrZS9oZWFkZXIuZ28=) | `0.00% <0.00%> (-84.22%)` | :arrow_down: |
| [pkg/protocol/compression\_method.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2NvbXByZXNzaW9uX21ldGhvZC5nbw==) | `16.00% <0.00%> (-72.00%)` | :arrow_down: |
| [pkg/protocol/handshake/handshake.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2hhbmRzaGFrZS9oYW5kc2hha2UuZ28=) | `0.00% <0.00%> (-67.91%)` | :arrow_down: |
| [pkg/protocol/extension/use\_master\_secret.go](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-cGtnL3Byb3RvY29sL2V4dGVuc2lvbi91c2VfbWFzdGVyX3NlY3JldC5nbw==) | `0.00% <0.00%> (-60.00%)` | :arrow_down: |
| ... and [62 more](https://codecov.io/gh/pion/dtls/pull/462/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/dtls/pull/462?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pion/dtls/pull/462?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion). Last update [e0b2ce3...12d7cae](https://codecov.io/gh/pion/dtls/pull/462?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).


Patch:
diff --git a/errors.go b/errors.go
index 042996c..09c6ffe 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 8a00ee7..c86a7e8 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


Test patch:
diff --git a/fragment_buffer_test.go b/fragment_buffer_test.go
index 0f49c75..f1013cc 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)
+	}
+}