Add Cookie extension with unit test and fuzzing

pkg/protocol/extension/cookie.go

@@ -0,0 +1,65 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package extension
+
+import (
+	"errors"
+
+	"golang.org/x/crypto/cryptobyte"
+)
+
+const maxCookieSize = 0xffff - 2
+
+// CookieExt implements the cookie extension in DTLS 1.3.
+// See RFC 8446 section 4.2.2. Cookie.
+type CookieExt struct {
+	Cookie []byte
+}
+
+// TypeValue returns the extension TypeValue.
+func (c CookieExt) TypeValue() TypeValue {
+	return CookieTypeValue
+}
+
+var errCookieExtFormat = errors.New("invalid cookie format")
+
+// Marshal encodes the extension.
+func (c *CookieExt) Marshal() ([]byte, error) {
+	cookieLength := len(c.Cookie)
+	if cookieLength == 0 || cookieLength > maxCookieSize {
+		return nil, errCookieExtFormat
+	}
+	var b cryptobyte.Builder
+	b.AddUint16(uint16(c.TypeValue()))
+	b.AddUint16LengthPrefixed(func(b *cryptobyte.Builder) {
+		b.AddUint16LengthPrefixed(func(b *cryptobyte.Builder) {
+			b.AddBytes(c.Cookie)
+		})
+	})
+
+	return b.Bytes()
+}
+
+// Unmarshal populates the extension from encoded data.
+func (c *CookieExt) Unmarshal(data []byte) error { //nolint:cyclop
+	val := cryptobyte.String(data)
+	var extension uint16
+	if !val.ReadUint16(&extension) || TypeValue(extension) != c.TypeValue() {
+		return errInvalidExtensionType
+	}
+
+	var extData cryptobyte.String
+	if !val.ReadUint16LengthPrefixed(&extData) {
+		return errBufferTooSmall
+	}
+
+	var cookie cryptobyte.String
+	if !extData.ReadUint16LengthPrefixed(&cookie) || cookie.Empty() || len(cookie) > maxCookieSize {
+		return errCookieExtFormat
+	}
+
+	c.Cookie = append([]byte(nil), cookie...)
+
+	return nil
+}

pkg/protocol/extension/cookie_test.go

@@ -0,0 +1,69 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package extension
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCookieExt(t *testing.T) {
+	extension := CookieExt{Cookie: []byte{0x1, 0x42}}
+
+	raw, err := extension.Marshal()
+	assert.NoError(t, err)
+
+	expect := []byte{
+		0x00, 0x2c, // extension type
+		0x00, 0x04, // extension length
+		0x00, 0x02, // vec length
+		0x01, 0x42, // cookie
+	}
+	assert.Equal(t, raw, expect)
+
+	newExtension := CookieExt{}
+
+	assert.NoError(t, newExtension.Unmarshal(expect))
+	assert.Equal(t, extension.Cookie, newExtension.Cookie)
+}
+
+func FuzzCookieExtUnmarshal(f *testing.F) {
+	bigCookie := make([]byte, 0xffff)
+	bigCookie[0] = 0x00
+	bigCookie[1] = 0x2c
+	bigCookie[2] = 0xff
+	bigCookie[3] = 0xff
+	bigCookie[4] = 0xff
+	bigCookie[5] = 0xfd
+
+	testcases := [][]byte{
+		{
+			0x00, 0x2c, // extension type
+			0x00, 0x04, // extension length
+			0x00, 0x02, // vec length
+			0x01, 0x42, // cookie
+		},
+		{
+			0x00, 0x2c, // extension type
+			0x00, 0x04, // extension length
+			0x00, 0x01, // vec length
+			0x01, // cookie
+		},
+		bigCookie,
+	}
+
+	for _, tc := range testcases {
+		f.Add(tc)
+	}
+	f.Fuzz(func(t *testing.T, a []byte) {
+		cookieExt := CookieExt{}
+		err := cookieExt.Unmarshal(a)
+		if err == nil {
+			length := len(cookieExt.Cookie)
+			assert.NotZero(t, length)
+			assert.LessOrEqual(t, length, 0xfffd)
+		}
+	})
+}

pkg/protocol/extension/extension.go

@@ -22,6 +22,7 @@ const (
 	ALPNTypeValue                         TypeValue = 16
 	UseExtendedMasterSecretTypeValue      TypeValue = 23
 	SupportedVersionsTypeValue            TypeValue = 43
+	CookieTypeValue                       TypeValue = 44
 	KeyShareTypeValue                     TypeValue = 51
 	ConnectionIDTypeValue                 TypeValue = 54
 	RenegotiationInfoTypeValue            TypeValue = 65281
@@ -89,6 +90,8 @@ func Unmarshal(buf []byte) ([]Extension, error) { //nolint:cyclop
 			err = unmarshalAndAppend(bufView, &SupportedVersions{})
 		case KeyShareTypeValue:
 			err = unmarshalAndAppend(bufView, &KeyShare{})
+		case CookieTypeValue:
+			err = unmarshalAndAppend(bufView, &CookieExt{})
 		default:
 		}
 

pkg/protocol/extension/testdata/fuzz/FuzzCookieExtUnmarshal/314ff2d111e92adb

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00,\x00\x04\x00\x0000")