diff --git a/data_psi.go b/data_psi.go index face09b..7b306e6 100644 --- a/data_psi.go +++ b/data_psi.go @@ -119,8 +119,7 @@ func parsePSIData(i *astikit.BytesIterator) (d *PSIData, err error) { if s, stop, err = parsePSISection(i); err != nil { err = fmt.Errorf("astits: parsing PSI table failed: %w", err) return - } - if stop { + } else if stop { break } d.Sections = append(d.Sections, s) @@ -135,14 +134,10 @@ func parsePSISection(i *astikit.BytesIterator) (s *PSISection, stop bool, err er // Parse header var offsetStart, offsetSectionsEnd, offsetEnd int - if s.Header, offsetStart, _, offsetSectionsEnd, offsetEnd, err = parsePSISectionHeader(i); err != nil { + if s.Header, offsetStart, _, offsetSectionsEnd, offsetEnd, stop, err = parsePSISectionHeader(i); err != nil { err = fmt.Errorf("astits: parsing PSI section header failed: %w", err) return - } - - // Check whether we need to stop the parsing - if shouldStopPSIParsing(s.Header.TableID) { - stop = true + } else if stop { return } @@ -206,7 +201,7 @@ func shouldStopPSIParsing(tableID PSITableID) bool { } // parsePSISectionHeader parses a PSI section header -func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd int, err error) { +func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd int, stop bool, err error) { // Init h = &PSISectionHeader{} offsetStart = i.Offset() @@ -225,7 +220,8 @@ func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offse h.TableType = h.TableID.Type() // Check whether we need to stop the parsing - if shouldStopPSIParsing(h.TableID) { + if h.TableID == PSITableIDNull { + stop = true return } diff --git a/data_psi_test.go b/data_psi_test.go index 1a1112f..2e7ea07 100644 --- a/data_psi_test.go +++ b/data_psi_test.go @@ -212,7 +212,7 @@ func TestParsePSISectionHeader(t *testing.T) { w.Write(uint8(254)) // Table ID w.Write("1") // Syntax section indicator w.Write("0000000") // Finish the byte - d, _, _, _, _, err := parsePSISectionHeader(astikit.NewBytesIterator(buf.Bytes())) + d, _, _, _, _, _, err := parsePSISectionHeader(astikit.NewBytesIterator(buf.Bytes())) assert.Equal(t, d, &PSISectionHeader{ TableID: 254, TableType: PSITableTypeUnknown, @@ -220,12 +220,13 @@ func TestParsePSISectionHeader(t *testing.T) { assert.NoError(t, err) // Valid table type - d, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd, err := parsePSISectionHeader(astikit.NewBytesIterator(psiSectionHeaderBytes())) + d, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd, stop, err := parsePSISectionHeader(astikit.NewBytesIterator(psiSectionHeaderBytes())) assert.Equal(t, d, psiSectionHeader) assert.Equal(t, 0, offsetStart) assert.Equal(t, 3, offsetSectionsStart) assert.Equal(t, 2729, offsetSectionsEnd) assert.Equal(t, 2733, offsetEnd) + assert.Equal(t, false, stop) assert.NoError(t, err) }