From 397ebe4144b4a7a11c424e3c954ae94529dfb74d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20L=C3=B3pez=20Guimaraes?= <danielectra@outlook.com>
Date: Wed, 26 Feb 2025 23:51:21 +0000
Subject: [PATCH] fix(prudp): Add missing size checks on PRUDP packets

PRUDPLite and PRUDPv1 were missing some size checks while parsing their
data. We also need one bounds check at the beginning to choose the
proper packet type by the magic.
---
 prudp_packet_lite.go | 21 +++++++++++++++++++--
 prudp_packet_v1.go   | 18 +++++++++++++++++-
 prudp_server.go      |  5 +++++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/prudp_packet_lite.go b/prudp_packet_lite.go
index 320145c..0668dbe 100644
--- a/prudp_packet_lite.go
+++ b/prudp_packet_lite.go
@@ -4,6 +4,7 @@ import (
 	"crypto/hmac"
 	"crypto/md5"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"net"
 
@@ -177,6 +178,10 @@ func (p *PRUDPPacketLite) decode() error {
 		return fmt.Errorf("Failed to decode PRUDPLite options. %s", err.Error())
 	}
 
+	if p.readStream.Remaining() < uint64(payloadLength) {
+		return errors.New("Failed to read PRUDPLite payload. Not have enough data")
+	}
+
 	p.payload = p.readStream.ReadBytesNext(int64(payloadLength))
 
 	return nil
@@ -208,6 +213,10 @@ func (p *PRUDPPacketLite) Bytes() []byte {
 }
 
 func (p *PRUDPPacketLite) decodeOptions() error {
+	if p.readStream.Remaining() < uint64(p.optionsLength) {
+		return errors.New("Not have enough data")
+	}
+
 	data := p.readStream.ReadBytesNext(int64(p.optionsLength))
 	optionsStream := NewByteStreamIn(data, p.server.LibraryVersions, p.server.ByteStreamSettings)
 
@@ -231,7 +240,11 @@ func (p *PRUDPPacketLite) decodeOptions() error {
 			}
 
 			if optionID == 1 {
-				p.connectionSignature = optionsStream.ReadBytesNext(int64(optionSize))
+				if optionsStream.Remaining() < uint64(optionSize) {
+					err = errors.New("Failed to read connection signature. Not have enough data")
+				} else {
+					p.connectionSignature = optionsStream.ReadBytesNext(int64(optionSize))
+				}
 			}
 
 			if optionID == 4 {
@@ -253,7 +266,11 @@ func (p *PRUDPPacketLite) decodeOptions() error {
 
 		if p.packetType == constants.ConnectPacket && !p.HasFlag(constants.PacketFlagAck) {
 			if optionID == 0x80 {
-				p.liteSignature = optionsStream.ReadBytesNext(int64(optionSize))
+				if optionsStream.Remaining() < uint64(optionSize) {
+					err = errors.New("Failed to read lite signature. Not have enough data")
+				} else {
+					p.liteSignature = optionsStream.ReadBytesNext(int64(optionSize))
+				}
 			}
 		}
 
diff --git a/prudp_packet_v1.go b/prudp_packet_v1.go
index d05592c..8197a22 100644
--- a/prudp_packet_v1.go
+++ b/prudp_packet_v1.go
@@ -89,6 +89,10 @@ func (p *PRUDPPacketV1) decode() error {
 		return fmt.Errorf("Failed to decode PRUDPv1 header. %s", err.Error())
 	}
 
+	if p.readStream.Remaining() < 16 {
+		return errors.New("Failed to read PRUDPv1 signature. Not have enough data")
+	}
+
 	p.signature = p.readStream.ReadBytesNext(16)
 
 	err = p.decodeOptions()
@@ -96,6 +100,10 @@ func (p *PRUDPPacketV1) decode() error {
 		return fmt.Errorf("Failed to decode PRUDPv1 options. %s", err.Error())
 	}
 
+	if p.readStream.Remaining() < uint64(p.payloadLength) {
+		return errors.New("Failed to read PRUDPv1 payload. Not have enough data")
+	}
+
 	p.payload = p.readStream.ReadBytesNext(int64(p.payloadLength))
 
 	return nil
@@ -211,6 +219,10 @@ func (p *PRUDPPacketV1) encodeHeader() []byte {
 }
 
 func (p *PRUDPPacketV1) decodeOptions() error {
+	if p.readStream.Remaining() < uint64(p.optionsLength) {
+		return errors.New("Not have enough data")
+	}
+
 	data := p.readStream.ReadBytesNext(int64(p.optionsLength))
 	optionsStream := NewByteStreamIn(data, p.server.LibraryVersions, p.server.ByteStreamSettings)
 
@@ -234,7 +246,11 @@ func (p *PRUDPPacketV1) decodeOptions() error {
 			}
 
 			if optionID == 1 {
-				p.connectionSignature = optionsStream.ReadBytesNext(16)
+				if optionsStream.Remaining() < 16 {
+					err = errors.New("Not have enough data")
+				} else {
+					p.connectionSignature = optionsStream.ReadBytesNext(16)
+				}
 			}
 
 			if optionID == 4 {
diff --git a/prudp_server.go b/prudp_server.go
index 52c0254..35a2909 100644
--- a/prudp_server.go
+++ b/prudp_server.go
@@ -127,6 +127,11 @@ func (ps *PRUDPServer) initPRUDPv1ConnectionSignatureKey() {
 }
 
 func (ps *PRUDPServer) handleSocketMessage(packetData []byte, address net.Addr, webSocketConnection *gws.Conn) error {
+	// * Check that the message is long enough for initial parsing
+	if len(packetData) < 2 {
+		return nil
+	}
+
 	readStream := NewByteStreamIn(packetData, ps.LibraryVersions, ps.ByteStreamSettings)
 
 	var packets []PRUDPPacketInterface