fix bug: pes packet length

yapingcat · Apr 12, 2022 · 5fbdc43 · 5fbdc43
1 parent 035c3cc
commit 5fbdc43
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/mpeg2/pes-proto.go b/mpeg2/pes-proto.go
@@ -129,16 +129,12 @@ func (pkg *PesPacket) PrettyPrint(file *os.File) {
 }
 
 func (pkg *PesPacket) Decode(bs *mpeg.BitStream) error {
-    if bs.RemainBytes() < 6 {
+    if bs.RemainBytes() < 9 {
         return errNeedMore
     }
     bs.SkipBits(24)             //packet_start_code_prefix
     pkg.Stream_id = bs.Uint8(8) //stream_id
     pkg.PES_packet_length = bs.Uint16(16)
-    if pkg.PES_packet_length != 0 && bs.RemainBytes() < int(pkg.PES_packet_length) {
-        bs.UnRead(6)
-        return errNeedMore
-    }
     bs.SkipBits(2) //'10'
     pkg.PES_scrambling_control = bs.Uint8(2)
     pkg.PES_priority = bs.Uint8(1)
@@ -153,6 +149,10 @@ func (pkg *PesPacket) Decode(bs *mpeg.BitStream) error {
     pkg.PES_CRC_flag = bs.Uint8(1)
     pkg.PES_extension_flag = bs.Uint8(1)
     pkg.PES_header_data_length = bs.Uint8(8)
+    if bs.RemainBytes() < int(pkg.PES_header_data_length) {
+        bs.UnRead(9 * 8)
+        return errNeedMore
+    }
     bs.Markdot()
     if pkg.PTS_DTS_flags&0x02 == 0x02 {
         bs.SkipBits(4)
@@ -205,6 +205,11 @@ func (pkg *PesPacket) Decode(bs *mpeg.BitStream) error {
     }
     loc := bs.DistanceFromMarkDot()
     bs.SkipBits(int(pkg.PES_header_data_length)*8 - loc)
+    if int(pkg.PES_packet_length)-int(pkg.PES_header_data_length) > bs.RemainBytes() {
+        pkg.Pes_payload = bs.RemainData()
+        bs.UnRead((9 + int(pkg.PES_header_data_length)) * 8)
+        return errNeedMore
+    }
     if pkg.PES_packet_length == 0 ||
         bs.RemainBits() <= int(pkg.PES_packet_length-3-uint16(pkg.PES_header_data_length))*8 {
         pkg.Pes_payload = bs.RemainData()
@@ -213,6 +218,7 @@ func (pkg *PesPacket) Decode(bs *mpeg.BitStream) error {
         pkg.Pes_payload = bs.RemainData()[:pkg.PES_packet_length-3-uint16(pkg.PES_header_data_length)]
         bs.SkipBits(int(pkg.PES_packet_length-3-uint16(pkg.PES_header_data_length)) * 8)
     }
+
     return nil
 }
 

diff --git a/mpeg2/ts-demuxer.go b/mpeg2/ts-demuxer.go
@@ -1,9 +1,9 @@
 package mpeg2
 
 import (
-	"errors"
+    "errors"
 
-	"github.com/yapingcat/gomedia/mpeg"
+    "github.com/yapingcat/gomedia/mpeg"
 )
 
 type pakcet_t struct {
@@ -103,13 +103,14 @@ func (demuxer *TSDemuxer) Input(data []byte) error {
                             stream.pes_pkg.Decode(bs)
                             pkg.Payload = stream.pes_pkg
                         } else {
+                            stream.pes_pkg.Pes_payload = bs.RemainData()
                             pkg.Payload = bs.RemainData()
                         }
                         stype := findPESIDByStreamType(stream.cid)
                         if stype == PES_STREAM_AUDIO {
-                            demuxer.doAudioPesPacket(bs, stream, pkg.Payload_unit_start_indicator)
+                            demuxer.doAudioPesPacket(stream, pkg.Payload_unit_start_indicator)
                         } else if stype == PES_STREAM_VIDEO {
-                            demuxer.doVideoPesPacket(bs, stream, pkg.Payload_unit_start_indicator)
+                            demuxer.doVideoPesPacket(stream, pkg.Payload_unit_start_indicator)
                         }
                     }
                 }
@@ -136,21 +137,17 @@ func (demuxer *TSDemuxer) Flush() {
     }
 }
 
-func (demuxer *TSDemuxer) doVideoPesPacket(bs *mpeg.BitStream, stream *tsstream, start uint8) {
+func (demuxer *TSDemuxer) doVideoPesPacket(stream *tsstream, start uint8) {
     if stream.cid != TS_STREAM_H264 && stream.cid != TS_STREAM_H265 {
         return
     }
-    if start == 1 {
-        stream.pkg.payload = append(stream.pkg.payload, stream.pes_pkg.Pes_payload...)
-    } else {
-        stream.pkg.payload = append(stream.pkg.payload, bs.RemainData()...)
-    }
+    stream.pkg.payload = append(stream.pkg.payload, stream.pes_pkg.Pes_payload...)
     stream.pkg.pts = stream.pes_pkg.Pts
     stream.pkg.dts = stream.pes_pkg.Dts
     demuxer.splitH26XFrame(stream)
 }
 
-func (demuxer *TSDemuxer) doAudioPesPacket(bs *mpeg.BitStream, stream *tsstream, start uint8) {
+func (demuxer *TSDemuxer) doAudioPesPacket(stream *tsstream, start uint8) {
     if stream.cid != TS_STREAM_AAC {
         return
     }
@@ -160,10 +157,8 @@ func (demuxer *TSDemuxer) doAudioPesPacket(bs *mpeg.BitStream, stream *tsstream,
             demuxer.OnFrame(stream.cid, stream.pkg.payload, stream.pkg.pts/90, stream.pkg.dts/90)
         }
         stream.pkg.payload = stream.pkg.payload[:0]
-        stream.pkg.payload = append(stream.pkg.payload, stream.pes_pkg.Pes_payload...)
-    } else {
-        stream.pkg.payload = append(stream.pkg.payload, bs.RemainData()...)
     }
+    stream.pkg.payload = append(stream.pkg.payload, stream.pes_pkg.Pes_payload...)
     stream.pkg.pts = stream.pes_pkg.Pts
     stream.pkg.dts = stream.pes_pkg.Dts
 }