data_psi: add test for parsing multiple sections

asticode · Apr 7, 2021 · d0a1f65 · d0a1f65
1 parent 9d2253c
commit d0a1f65
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/data_psi.go b/data_psi.go
@@ -120,6 +120,9 @@ func parsePSIData(i *astikit.BytesIterator) (d *PSIData, err error) {
 			err = fmt.Errorf("astits: parsing PSI table failed: %w", err)
 			return
 		}
+		if stop {
+			break
+		}
 		d.Sections = append(d.Sections, s)
 	}
 	return
@@ -199,8 +202,7 @@ func parseCRC32(i *astikit.BytesIterator) (c uint32, err error) {
 
 // shouldStopPSIParsing checks whether the PSI parsing should be stopped
 func shouldStopPSIParsing(tableID PSITableID) bool {
-	return tableID == PSITableIDNull ||
-		tableID.isUnknown()
+	return tableID == PSITableIDNull
 }
 
 // parsePSISectionHeader parses a PSI section header
@@ -241,7 +243,7 @@ func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offse
 	h.PrivateBit = bs[0]&0x40 > 0
 
 	// Section length
-	h.SectionLength = uint16(bs[0]&0xf)<<8 | uint16(bs[1])
+	h.SectionLength = uint16(bs[0]&3)<<8 | uint16(bs[1])
 
 	// Offsets
 	offsetSectionsStart = i.Offset()

diff --git a/data_psi_test.go b/data_psi_test.go
@@ -366,6 +366,27 @@ var psiDataTestCases = []psiDataTestCase{
 	},
 }
 
+func TestParsePSIDataPMTMultipleSections(t *testing.T) {
+	pmt := hexToBytes(`00C0001500000100FF000000
+000000010000000000038D646B02B07B
+0001C90000EF9BF02109044749E10B05
+04474139348713C1010100F30D01656E
+670100000554562D504702EF9BF00E11
+01FF1006C0BD62C0080006010281EF9C
+F018050441432D33810A083805FF0F01
+BF656E670A04656E670081EF9DF01805
+0441432D33810A082885FF0001BF7370
+610A0473706100082F08E3FFFFFFFFFF
+FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF
+FFFFFFFFFFFFFFFFFFFFFFFF`)
+	d, err := parsePSIData(astikit.NewBytesIterator(pmt))
+	assert.NoError(t, err)
+	assert.NotNil(t, d)
+	assert.Len(t, d.Sections, 2)
+	assert.Equal(t, PSITableID(0xc0), d.Sections[0].Header.TableID)
+	assert.Equal(t, PSITableID(0x02), d.Sections[1].Header.TableID)
+}
+
 func TestWritePSIData(t *testing.T) {
 	for _, tc := range psiDataTestCases {
 		t.Run(tc.name, func(t *testing.T) {