From 7dc2af56736b663e76f1400ba403532ba590bceb Mon Sep 17 00:00:00 2001
From: Sean DuBois <sean@siobud.com>
Date: Thu, 9 Nov 2023 10:53:12 -0500
Subject: [PATCH] Combine OneByte and TwoByte extension parsing

Reduces duplicated logic and increase coverage. Missing tests for bounds
check on Padding and Payload Sizes

Co-Authored-By: Juho Nurminen <juhonurm@gmail.com>
---
 AUTHORS.txt    |  3 +++
 packet.go      | 69 ++++++++++++++++++++++----------------------------
 packet_test.go | 28 ++++++++++++++++++++
 3 files changed, 61 insertions(+), 39 deletions(-)

diff --git a/AUTHORS.txt b/AUTHORS.txt
index 727aaaa..f8dbb32 100644
--- a/AUTHORS.txt
+++ b/AUTHORS.txt
@@ -18,6 +18,7 @@ Guilherme <gqgs@protonmail.com>
 Haiyang Wang <ocean2811@outlook.com>
 Hugo Arregui <hugo@decentraland.org>
 John Bradley <jrb@turrettech.com>
+Juho Nurminen <juhonurm@gmail.com>
 Juliusz Chroboczek <jch@irif.fr>
 kawaway <kawawa.y@gmail.com>
 Kazuyuki Honda <hakobera@gmail.com>
@@ -31,6 +32,7 @@ Raphael Derosso Pereira <raphaelpereira@gmail.com>
 Rob Lofthouse <ri.lofthouse@gmail.com>
 Robin Raymond <robin@goheadroom.com>
 Sean <sean@pion.ly>
+Sean <sean@siobud.com>
 Sean DuBois <duboisea@twitch.tv>
 Sean DuBois <seaduboi@amazon.com>
 Sean DuBois <sean@pion.ly>
@@ -40,6 +42,7 @@ Steffen Vogel <post@steffenvogel.de>
 Tarrence van As <tarrence13@gmail.com>
 wangzixiang <wangzixiang@onething.net>
 Woodrow Douglass <wdouglass@carnegierobotics.com>
+訾明华 <565209960@qq.com>
 
 # List of contributors not appearing in Git history
 
diff --git a/packet.go b/packet.go
index 7aebb14..af88af3 100644
--- a/packet.go
+++ b/packet.go
@@ -149,64 +149,55 @@ func (h *Header) Unmarshal(buf []byte) (n int, err error) { //nolint:gocognit
 		n += 2
 		extensionLength := int(binary.BigEndian.Uint16(buf[n:])) * 4
 		n += 2
+		extensionEnd := n + extensionLength
 
-		if expected := n + extensionLength; len(buf) < expected {
-			return n, fmt.Errorf("size %d < %d: %w",
-				len(buf), expected,
-				errHeaderSizeInsufficientForExtension,
-			)
+		if len(buf) < extensionEnd {
+			return n, fmt.Errorf("size %d < %d: %w", len(buf), extensionEnd, errHeaderSizeInsufficientForExtension)
 		}
 
-		switch h.ExtensionProfile {
-		// RFC 8285 RTP One Byte Header Extension
-		case extensionProfileOneByte:
-			end := n + extensionLength
-			for n < end {
+		if h.ExtensionProfile == extensionProfileOneByte || h.ExtensionProfile == extensionProfileTwoByte {
+			var (
+				extid      uint8
+				payloadLen int
+			)
+
+			for n < extensionEnd {
 				if buf[n] == 0x00 { // padding
 					n++
 					continue
 				}
 
-				extid := buf[n] >> 4
-				payloadLen := int(buf[n]&^0xF0 + 1)
-				n++
+				if h.ExtensionProfile == extensionProfileOneByte {
+					extid = buf[n] >> 4
+					payloadLen = int(buf[n]&^0xF0 + 1)
+					n++
 
-				if extid == extensionIDReserved {
-					break
-				}
+					if extid == extensionIDReserved {
+						break
+					}
+				} else {
+					extid = buf[n]
+					n++
 
-				extension := Extension{id: extid, payload: buf[n : n+payloadLen]}
-				h.Extensions = append(h.Extensions, extension)
-				n += payloadLen
-			}
+					if len(buf) <= n {
+						return n, fmt.Errorf("size %d < %d: %w", len(buf), n, errHeaderSizeInsufficientForExtension)
+					}
 
-		// RFC 8285 RTP Two Byte Header Extension
-		case extensionProfileTwoByte:
-			end := n + extensionLength
-			for n < end {
-				if buf[n] == 0x00 { // padding
+					payloadLen = int(buf[n])
 					n++
-					continue
 				}
 
-				extid := buf[n]
-				n++
-
-				payloadLen := int(buf[n])
-				n++
+				if extensionPayloadEnd := n + payloadLen; len(buf) <= extensionPayloadEnd {
+					return n, fmt.Errorf("size %d < %d: %w", len(buf), extensionPayloadEnd, errHeaderSizeInsufficientForExtension)
+				}
 
 				extension := Extension{id: extid, payload: buf[n : n+payloadLen]}
 				h.Extensions = append(h.Extensions, extension)
 				n += payloadLen
 			}
-
-		default: // RFC3550 Extension
-			if len(buf) < n+extensionLength {
-				return n, fmt.Errorf("%w: %d < %d",
-					errHeaderSizeInsufficientForExtension, len(buf), n+extensionLength)
-			}
-
-			extension := Extension{id: 0, payload: buf[n : n+extensionLength]}
+		} else {
+			// RFC3550 Extension
+			extension := Extension{id: 0, payload: buf[n:extensionEnd]}
 			h.Extensions = append(h.Extensions, extension)
 			n += len(h.Extensions[0].payload)
 		}
diff --git a/packet_test.go b/packet_test.go
index 3044c0d..98dfb40 100644
--- a/packet_test.go
+++ b/packet_test.go
@@ -1192,6 +1192,34 @@ func TestRFC8285TwoByteSetExtensionShouldErrorWhenPayloadTooLarge(t *testing.T)
 	}
 }
 
+func TestRFC8285Padding(t *testing.T) {
+	header := &Header{}
+
+	for _, payload := range [][]byte{
+		{
+			0b00010000,                      // header.Extension = true
+			0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // SequenceNumber, Timestamp, SSRC
+			0xBE, 0xDE, // header.ExtensionProfile = extensionProfileOneByte
+			0, 1, // extensionLength
+			0, 0, 0, // padding
+			1, // extid
+		},
+		{
+			0b00010000,                      // header.Extension = true
+			0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // SequenceNumber, Timestamp, SSRC
+			0x10, 0x00, // header.ExtensionProfile = extensionProfileOneByte
+			0, 1, // extensionLength
+			0, 0, 0, // padding
+			1, // extid
+		},
+	} {
+		_, err := header.Unmarshal(payload)
+		if !errors.Is(err, errHeaderSizeInsufficientForExtension) {
+			t.Fatal("Expected errHeaderSizeInsufficientForExtension")
+		}
+	}
+}
+
 func TestRFC3550SetExtensionShouldErrorWhenNonZero(t *testing.T) {
 	payload := []byte{
 		// Payload