Repository: pion/rtp
Pull request #162 (2022-01-03 00:01:01)

Adds zero-copy header extensions
This change adds an interface and impl for zero-copy extensions. The code is forked from the existing `packet.go` but those files will be migrated later.

codecov[bot]: # [Codecov](https://codecov.io/gh/pion/rtp/pull/162?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) Report
> Merging [#162](https://codecov.io/gh/pion/rtp/pull/162?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) (9f25f03) into [master](https://codecov.io/gh/pion/rtp/commit/2c6c078f9741ddb70f6bbcaccec946135d7b1d26?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) (2c6c078) will **decrease** coverage by `7.19%`.
> The diff coverage is `37.32%`.

[![Impacted file tree graph](https://codecov.io/gh/pion/rtp/pull/162/graphs/tree.svg?width=650&height=150&src=pr&token=4KrbiBsspj&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion)](https://codecov.io/gh/pion/rtp/pull/162?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion)

```diff
@@            Coverage Diff             @@
##           master     #162      +/-   ##
==========================================
- Coverage   93.45%   86.25%   -7.20%     
==========================================
  Files          14       15       +1     
  Lines        1421     1630     +209     
==========================================
+ Hits         1328     1406      +78     
- Misses         75      195     +120     
- Partials       18       29      +11     
```

| Flag | Coverage Δ | |
|---|---|---|
| go | `86.25% <37.32%> (-7.20%)` | :arrow_down: |
| wasm | `85.76% <37.32%> (-7.13%)` | :arrow_down: |

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/rtp/pull/162?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | Coverage Δ | |
|---|---|---|
| [header\_extension.go](https://codecov.io/gh/pion/rtp/pull/162/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-aGVhZGVyX2V4dGVuc2lvbi5nbw==) | `37.32% <37.32%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/rtp/pull/162?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/rtp/pull/162?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion). Last update [2c6c078...9f25f03](https://codecov.io/gh/pion/rtp/pull/162?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/header_extension.go b/header_extension.go
new file mode 100644
index 0000000..cb38324
--- /dev/null
+++ b/header_extension.go
@@ -0,0 +1,350 @@
+package rtp
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+)
+
+const (
+	headerExtensionProfileOneByte = 0xBEDE
+	headerExtensionProfileTwoByte = 0x1000
+	headerExtensionIDReserved     = 0xF
+)
+
+// HeaderExtension represents an RTP extension header.
+type HeaderExtension interface {
+	Set(id uint8, payload []byte) error
+	GetIDs() []uint8
+	Get(id uint8) []byte
+	Del(id uint8) error
+
+	Unmarshal(buf []byte) (int, error)
+	Marshal() ([]byte, error)
+	MarshalTo(buf []byte) (int, error)
+	MarshalSize() int
+}
+
+// OneByteHeaderExtension is an RFC8285 one-byte header extension.
+type OneByteHeaderExtension struct {
+	payload []byte
+}
+
+// Set sets the extension payload for the specified ID.
+func (e *OneByteHeaderExtension) Set(id uint8, buf []byte) error {
+	if id < 1 || id > 14 {
+		return fmt.Errorf("%w actual(%d)", errRFC8285OneByteHeaderIDRange, id)
+	}
+	if len(buf) > 16 {
+		return fmt.Errorf("%w actual(%d)", errRFC8285OneByteHeaderSize, len(buf))
+	}
+
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n] >> 4
+		len := int(e.payload[n]&^0xF0 + 1)
+		n++
+
+		if extid == id {
+			e.payload = append(e.payload[:n+1], append(buf, e.payload[n+1+len:]...)...)
+			return nil
+		}
+		n += len
+	}
+	e.payload = append(e.payload, (id<<4 | uint8(len(buf)-1)))
+	e.payload = append(e.payload, buf...)
+	binary.BigEndian.PutUint16(e.payload[2:4], binary.BigEndian.Uint16(e.payload[2:4])+1)
+	return nil
+}
+
+// GetIDs returns the available IDs.
+func (e *OneByteHeaderExtension) GetIDs() []uint8 {
+	ids := make([]uint8, 0, binary.BigEndian.Uint16(e.payload[2:4]))
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n] >> 4
+		len := int(e.payload[n]&^0xF0 + 1)
+		n++
+
+		if extid == headerExtensionIDReserved {
+			break
+		}
+
+		ids = append(ids, extid)
+		n += len
+	}
+	return ids
+}
+
+// Get returns the payload of the extension with the given ID.
+func (e *OneByteHeaderExtension) Get(id uint8) []byte {
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n] >> 4
+		len := int(e.payload[n]&^0xF0 + 1)
+		n++
+
+		if extid == id {
+			return e.payload[n : n+len]
+		}
+		n += len
+	}
+	return nil
+}
+
+// Del deletes the extension with the specified ID.
+func (e *OneByteHeaderExtension) Del(id uint8) error {
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n] >> 4
+		len := int(e.payload[n]&^0xF0 + 1)
+
+		if extid == id {
+			e.payload = append(e.payload[:n], e.payload[n+1+len:]...)
+			return nil
+		}
+		n += len + 1
+	}
+	return errHeaderExtensionNotFound
+}
+
+// Unmarshal parses the extension payload.
+func (e *OneByteHeaderExtension) Unmarshal(buf []byte) (int, error) {
+	profile := binary.BigEndian.Uint16(buf[0:2])
+	if profile != headerExtensionProfileOneByte {
+		return 0, fmt.Errorf("%w actual(%x)", errHeaderExtensionNotFound, buf[0:2])
+	}
+	e.payload = buf
+	return len(buf), nil
+}
+
+// Marshal returns the extension payload.
+func (e *OneByteHeaderExtension) Marshal() ([]byte, error) {
+	return e.payload, nil
+}
+
+// MarshalTo writes the extension payload to the given buffer.
+func (e *OneByteHeaderExtension) MarshalTo(buf []byte) (int, error) {
+	size := e.MarshalSize()
+	if size > len(buf) {
+		return 0, io.ErrShortBuffer
+	}
+	return copy(buf, e.payload), nil
+}
+
+// MarshalSize returns the size of the extension payload.
+func (e *OneByteHeaderExtension) MarshalSize() int {
+	return len(e.payload)
+}
+
+// TwoByteHeaderExtension is an RFC8285 two-byte header extension.
+type TwoByteHeaderExtension struct {
+	payload []byte
+}
+
+// Set sets the extension payload for the specified ID.
+func (e *TwoByteHeaderExtension) Set(id uint8, buf []byte) error {
+	if id < 1 || id > 255 {
+		return fmt.Errorf("%w actual(%d)", errRFC8285TwoByteHeaderIDRange, id)
+	}
+	if len(buf) > 255 {
+		return fmt.Errorf("%w actual(%d)", errRFC8285TwoByteHeaderSize, len(buf))
+	}
+
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n]
+		n++
+
+		len := int(e.payload[n])
+		n++
+
+		if extid == id {
+			e.payload = append(e.payload[:n+2], append(buf, e.payload[n+2+len:]...)...)
+			return nil
+		}
+		n += len
+	}
+	e.payload = append(e.payload, id, uint8(len(buf)))
+	e.payload = append(e.payload, buf...)
+	binary.BigEndian.PutUint16(e.payload[2:4], binary.BigEndian.Uint16(e.payload[2:4])+1)
+	return nil
+}
+
+// GetIDs returns the available IDs.
+func (e *TwoByteHeaderExtension) GetIDs() []uint8 {
+	ids := make([]uint8, 0, binary.BigEndian.Uint16(e.payload[2:4]))
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n]
+		n++
+
+		len := int(e.payload[n])
+		n++
+
+		ids = append(ids, extid)
+		n += len
+	}
+	return ids
+}
+
+// Get returns the payload of the extension with the given ID.
+func (e *TwoByteHeaderExtension) Get(id uint8) []byte {
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n]
+		n++
+
+		len := int(e.payload[n])
+		n++
+
+		if extid == id {
+			return e.payload[n : n+len]
+		}
+		n += len
+	}
+	return nil
+}
+
+// Del deletes the extension with the specified ID.
+func (e *TwoByteHeaderExtension) Del(id uint8) error {
+	for n := 4; n < len(e.payload); {
+		if e.payload[n] == 0x00 { // padding
+			n++
+			continue
+		}
+
+		extid := e.payload[n]
+
+		len := int(e.payload[n+1])
+
+		if extid == id {
+			e.payload = append(e.payload[:n], e.payload[n+2+len:]...)
+			return nil
+		}
+		n += len + 2
+	}
+	return errHeaderExtensionNotFound
+}
+
+// Unmarshal parses the extension payload.
+func (e *TwoByteHeaderExtension) Unmarshal(buf []byte) (int, error) {
+	profile := binary.BigEndian.Uint16(buf[0:2])
+	if profile != headerExtensionProfileTwoByte {
+		return 0, fmt.Errorf("%w actual(%x)", errHeaderExtensionNotFound, buf[0:2])
+	}
+	e.payload = buf
+	return len(buf), nil
+}
+
+// Marshal returns the extension payload.
+func (e *TwoByteHeaderExtension) Marshal() ([]byte, error) {
+	return e.payload, nil
+}
+
+// MarshalTo marshals the extension to the given buffer.
+func (e *TwoByteHeaderExtension) MarshalTo(buf []byte) (int, error) {
+	size := e.MarshalSize()
+	if size > len(buf) {
+		return 0, io.ErrShortBuffer
+	}
+	return copy(buf, e.payload), nil
+}
+
+// MarshalSize returns the size of the extension payload.
+func (e *TwoByteHeaderExtension) MarshalSize() int {
+	return len(e.payload)
+}
+
+// RawExtension represents an RFC3550 header extension.
+type RawExtension struct {
+	payload []byte
+}
+
+// Set sets the extension payload for the specified ID.
+func (e *RawExtension) Set(id uint8, payload []byte) error {
+	if id != 0 {
+		return fmt.Errorf("%w actual(%d)", errRFC3550HeaderIDRange, id)
+	}
+	e.payload = payload
+	return nil
+}
+
+// GetIDs returns the available IDs.
+func (e *RawExtension) GetIDs() []uint8 {
+	return []uint8{0}
+}
+
+// Get returns the payload of the extension with the given ID.
+func (e *RawExtension) Get(id uint8) []byte {
+	if id == 0 {
+		return e.payload
+	}
+	return nil
+}
+
+// Del deletes the extension with the specified ID.
+func (e *RawExtension) Del(id uint8) error {
+	if id == 0 {
+		e.payload = nil
+		return nil
+	}
+	return fmt.Errorf("%w actual(%d)", errRFC3550HeaderIDRange, id)
+}
+
+// Unmarshal parses the extension from the given buffer.
+func (e *RawExtension) Unmarshal(buf []byte) (int, error) {
+	profile := binary.BigEndian.Uint16(buf[0:2])
+	if profile == headerExtensionProfileOneByte || profile == headerExtensionProfileTwoByte {
+		return 0, fmt.Errorf("%w actual(%x)", errHeaderExtensionNotFound, buf[0:2])
+	}
+	e.payload = buf
+	return len(buf), nil
+}
+
+// Marshal returns the raw extension payload.
+func (e *RawExtension) Marshal() ([]byte, error) {
+	return e.payload, nil
+}
+
+// MarshalTo marshals the extension to the given buffer.
+func (e *RawExtension) MarshalTo(buf []byte) (int, error) {
+	size := e.MarshalSize()
+	if size > len(buf) {
+		return 0, io.ErrShortBuffer
+	}
+	return copy(buf, e.payload), nil
+}
+
+// MarshalSize returns the size of the extension when marshaled.
+func (e *RawExtension) MarshalSize() int {
+	return len(e.payload)
+}


Test patch:
diff --git a/header_extension_test.go b/header_extension_test.go
new file mode 100644
index 0000000..8aba946
--- /dev/null
+++ b/header_extension_test.go
@@ -0,0 +1,302 @@
+package rtp
+
+import (
+	"bytes"
+	"encoding/hex"
+	"testing"
+)
+
+func TestHeaderExtension_RFC8285OneByteExtension(t *testing.T) {
+	p := &OneByteHeaderExtension{}
+
+	rawPkt := []byte{
+		0xBE, 0xDE, 0x00, 0x01, 0x50, 0xAA, 0x00, 0x00,
+		0x98, 0x36, 0xbe, 0x88, 0x9e,
+	}
+	if _, err := p.Unmarshal(rawPkt); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	dstData, _ := p.Marshal()
+	if !bytes.Equal(dstData, rawPkt) {
+		t.Errorf("Marshal failed raw \nMarshaled:\n%s\nrawPkt:\n%s", hex.Dump(dstData), hex.Dump(rawPkt))
+	}
+}
+
+func TestHeaderExtension_RFC8285OneByteTwoExtensionOfTwoBytes(t *testing.T) {
+	p := &OneByteHeaderExtension{}
+
+	//  0                   1                   2                   3
+	//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |       0xBE    |    0xDE       |           length=1            |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |  ID   | L=0   |     data      |  ID   |  L=0  |   data...
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	rawPkt := []byte{
+		0xBE, 0xDE, 0x00, 0x01, 0x10, 0xAA, 0x20, 0xBB,
+	}
+	if _, err := p.Unmarshal(rawPkt); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	ext1 := p.Get(1)
+	ext1Expect := []byte{0xAA}
+	if !bytes.Equal(ext1, ext1Expect) {
+		t.Errorf("Extension has incorrect data. Got: %+v, Expected: %+v", ext1, ext1Expect)
+	}
+
+	ext2 := p.Get(2)
+	ext2Expect := []byte{0xBB}
+	if !bytes.Equal(ext2, ext2Expect) {
+		t.Errorf("Extension has incorrect data. Got: %+v, Expected: %+v", ext2, ext2Expect)
+	}
+
+	dstData, _ := p.Marshal()
+	if !bytes.Equal(dstData, rawPkt) {
+		t.Errorf("Marshal failed raw \nMarshaled:\n%s\nrawPkt:\n%s", hex.Dump(dstData), hex.Dump(rawPkt))
+	}
+}
+
+func TestHeaderExtension_RFC8285OneByteMultipleExtensionsWithPadding(t *testing.T) {
+	p := &OneByteHeaderExtension{}
+
+	//  0                   1                   2                   3
+	//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |       0xBE    |    0xDE       |           length=3            |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |  ID   | L=0   |     data      |  ID   |  L=1  |   data...
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	//       ...data   |    0 (pad)    |    0 (pad)    |  ID   | L=3   |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |                          data                                 |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	rawPkt := []byte{
+		0xBE, 0xDE, 0x00, 0x03, 0x10, 0xAA, 0x21, 0xBB,
+		0xBB, 0x00, 0x00, 0x33, 0xCC, 0xCC, 0xCC, 0xCC,
+	}
+	if _, err := p.Unmarshal(rawPkt); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	ext1 := p.Get(1)
+	ext1Expect := []byte{0xAA}
+	if !bytes.Equal(ext1, ext1Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext1, ext1Expect)
+	}
+
+	ext2 := p.Get(2)
+	ext2Expect := []byte{0xBB, 0xBB}
+	if !bytes.Equal(ext2, ext2Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext2, ext2Expect)
+	}
+
+	ext3 := p.Get(3)
+	ext3Expect := []byte{0xCC, 0xCC, 0xCC, 0xCC}
+	if !bytes.Equal(ext3, ext3Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext3, ext3Expect)
+	}
+
+	dstBuf := map[string][]byte{
+		"CleanBuffer": make([]byte, 1000),
+		"DirtyBuffer": make([]byte, 1000),
+	}
+	for i := range dstBuf["DirtyBuffer"] {
+		dstBuf["DirtyBuffer"][i] = 0xFF
+	}
+	for name, buf := range dstBuf {
+		buf := buf
+		t.Run(name, func(t *testing.T) {
+			n, err := p.MarshalTo(buf)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !bytes.Equal(buf[:n], rawPkt) {
+				t.Errorf("Marshal failed raw \nMarshaled:\n%s\nrawPkt:\n%s", hex.Dump(buf[:n]), hex.Dump(rawPkt))
+			}
+		})
+	}
+}
+
+func TestHeaderExtension_RFC8285TwoByteExtension(t *testing.T) {
+	p := &TwoByteHeaderExtension{}
+
+	rawPkt := []byte{
+		0x10, 0x00, 0x00, 0x07, 0x05, 0x18, 0xAA, 0xAA,
+		0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA,
+		0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA,
+		0xAA, 0xAA, 0x00, 0x00,
+	}
+	if _, err := p.Unmarshal(rawPkt); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	dstData, _ := p.Marshal()
+	if !bytes.Equal(dstData, rawPkt) {
+		t.Errorf("Marshal failed raw \nMarshaled:\n%s\nrawPkt:\n%s", hex.Dump(dstData), hex.Dump(rawPkt))
+	}
+}
+
+func TestHeaderExtension_RFC8285TwoByteMultipleExtensionsWithPadding(t *testing.T) {
+	p := &TwoByteHeaderExtension{}
+
+	// 0                   1                   2                   3
+	// 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |       0x10    |    0x00       |           length=3            |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |      ID=1     |     L=0       |     ID=2      |     L=1       |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |       data    |    0 (pad)    |       ID=3    |      L=4      |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |                          data                                 |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	rawPkt := []byte{
+		0x10, 0x00, 0x00, 0x03, 0x01, 0x00, 0x02, 0x01,
+		0xBB, 0x00, 0x03, 0x04, 0xCC, 0xCC, 0xCC, 0xCC,
+	}
+
+	if _, err := p.Unmarshal(rawPkt); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	ext1 := p.Get(1)
+	ext1Expect := []byte{}
+	if !bytes.Equal(ext1, ext1Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext1, ext1Expect)
+	}
+
+	ext2 := p.Get(2)
+	ext2Expect := []byte{0xBB}
+	if !bytes.Equal(ext2, ext2Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext2, ext2Expect)
+	}
+
+	ext3 := p.Get(3)
+	ext3Expect := []byte{0xCC, 0xCC, 0xCC, 0xCC}
+	if !bytes.Equal(ext3, ext3Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext3, ext3Expect)
+	}
+}
+
+func TestHeaderExtension_RFC8285TwoByteMultipleExtensionsWithLargeExtension(t *testing.T) {
+	p := &TwoByteHeaderExtension{}
+
+	// 0                   1                   2                   3
+	// 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |       0x10    |    0x00       |           length=3            |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |      ID=1     |     L=0       |     ID=2      |     L=1       |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	// |       data    |       ID=3    |      L=17      |    data...
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	//                            ...data...
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	//                            ...data...
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	//                            ...data...
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	//                            ...data...                           |
+	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+	rawPkt := []byte{
+		0x10, 0x00, 0x00, 0x06, 0x01, 0x00, 0x02, 0x01,
+		0xBB, 0x03, 0x11, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC,
+		0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC,
+	}
+
+	if _, err := p.Unmarshal(rawPkt); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	ext1 := p.Get(1)
+	ext1Expect := []byte{}
+	if !bytes.Equal(ext1, ext1Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext1, ext1Expect)
+	}
+
+	ext2 := p.Get(2)
+	ext2Expect := []byte{0xBB}
+	if !bytes.Equal(ext2, ext2Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext2, ext2Expect)
+	}
+
+	ext3 := p.Get(3)
+	ext3Expect := []byte{
+		0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC,
+		0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC,
+	}
+	if !bytes.Equal(ext3, ext3Expect) {
+		t.Errorf("Extension has incorrect data. Got: %v+, Expected: %v+", ext3, ext3Expect)
+	}
+
+	dstData, _ := p.Marshal()
+	if !bytes.Equal(dstData, rawPkt) {
+		t.Errorf("Marshal failed raw \nMarshaled: %+v,\nrawPkt:    %+v", dstData, rawPkt)
+	}
+}
+
+func TestHeaderExtension_RFC8285OneByteDelExtension(t *testing.T) {
+	p := &OneByteHeaderExtension{}
+
+	if _, err := p.Unmarshal([]byte{0xBE, 0xDE, 0x00, 0x00}); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	if err := p.Set(1, []byte{0xBB}); err != nil {
+		t.Fatal("Set err for valid extension")
+	}
+
+	ext := p.Get(1)
+	if ext == nil {
+		t.Error("Extension should exist")
+	}
+
+	err := p.Del(1)
+	if err != nil {
+		t.Error("Should successfully delete extension")
+	}
+
+	ext = p.Get(1)
+	if ext != nil {
+		t.Error("Extension should not exist")
+	}
+
+	err = p.Del(1)
+	if err == nil {
+		t.Error("Should return error when deleting extension that doesnt exist")
+	}
+}
+
+func TestHeaderExtension_RFC8285TwoByteDelExtension(t *testing.T) {
+	p := &TwoByteHeaderExtension{}
+
+	if _, err := p.Unmarshal([]byte{0x10, 0x00, 0x00, 0x00}); err != nil {
+		t.Fatal("Unmarshal err for valid extension")
+	}
+
+	if err := p.Set(1, []byte{0xBB}); err != nil {
+		t.Fatal("Set err for valid extension")
+	}
+
+	ext := p.Get(1)
+	if ext == nil {
+		t.Error("Extension should exist")
+	}
+
+	err := p.Del(1)
+	if err != nil {
+		t.Error("Should successfully delete extension")
+	}
+
+	ext = p.Get(1)
+	if ext != nil {
+		t.Error("Extension should not exist")
+	}
+
+	err = p.Del(1)
+	if err == nil {
+		t.Error("Should return error when deleting extension that doesnt exist")
+	}
+}