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