diff --git a/data.go b/data.go index e330170..8d158e1 100644 --- a/data.go +++ b/data.go @@ -1,6 +1,8 @@ package astits -import "github.com/pkg/errors" +import ( + "github.com/pkg/errors" +) // PIDs const ( @@ -61,7 +63,10 @@ func parseData(ps []*Packet, prs PacketsParser, pm programMap) (ds []*Data, err } ds = psiData.toData(pid) } else if isPESPayload(payload) { - ds = append(ds, &Data{PES: parsePESData(payload), PID: pid}) + d, err := parsePESData(payload) + if err == nil { + ds = append(ds, &Data{PES: d, PID: pid}) + } } return } diff --git a/data_pes.go b/data_pes.go index 06511cd..0591e5a 100644 --- a/data_pes.go +++ b/data_pes.go @@ -1,6 +1,10 @@ package astits -import "github.com/asticode/go-astilog" +import ( + "fmt" + + "github.com/pkg/errors" +) // P-STD buffer scales const ( @@ -99,13 +103,16 @@ type DSMTrickMode struct { } // parsePESData parses a PES data -func parsePESData(i []byte) (d *PESData) { +func parsePESData(i []byte) (d *PESData, err error) { // Init d = &PESData{} // Parse header var offset, dataStart, dataEnd = 3, 0, 0 - d.Header, dataStart, dataEnd = parsePESHeader(i, &offset) + if d.Header, dataStart, dataEnd, err = parsePESHeader(i, &offset); err != nil { + err = errors.Wrap(err, "astits: parsing PES header failed") + return + } // Parse data d.Data = i[dataStart:dataEnd] @@ -118,7 +125,7 @@ func hasPESOptionalHeader(streamID uint8) bool { } // parsePESData parses a PES header -func parsePESHeader(i []byte, offset *int) (h *PESHeader, dataStart, dataEnd int) { +func parsePESHeader(i []byte, offset *int) (h *PESHeader, dataStart, dataEnd int, err error) { // Init h = &PESHeader{} @@ -138,10 +145,9 @@ func parsePESHeader(i []byte, offset *int) (h *PESHeader, dataStart, dataEnd int } // Check for incomplete data - // TODO Throw away the data? if dataEnd > len(i) { - astilog.Debug("PES dataEnd > len(i), needs fixing") - dataEnd = len(i) + err = fmt.Errorf("astits: pes dataEnd (%d) > len(i) (%d)", dataEnd, len(i)) + return } // Optional header diff --git a/data_pes_test.go b/data_pes_test.go index 0a2fd66..15ab4d8 100644 --- a/data_pes_test.go +++ b/data_pes_test.go @@ -215,10 +215,12 @@ func pesWithHeaderBytes() []byte { func TestParsePESSection(t *testing.T) { // No optional header and specific packet length - d := parsePESData(pesWithoutHeaderBytes()) + d, err := parsePESData(pesWithoutHeaderBytes()) + assert.NoError(t, err) assert.Equal(t, d, pesWithoutHeader) // Optional header and no specific header length - d = parsePESData(pesWithHeaderBytes()) + d, err = parsePESData(pesWithHeaderBytes()) + assert.NoError(t, err) assert.Equal(t, d, pesWithHeader) }