From d1f8a96c02d311684fb93ed419f54b624476c2f6 Mon Sep 17 00:00:00 2001
From: Aman Karmani <aman@tmm1.net>
Date: Wed, 14 Apr 2021 14:16:40 -0700
Subject: [PATCH] cleanups

---
 data_psi.go      | 16 ++++++----------
 data_psi_test.go |  5 +++--
 2 files changed, 9 insertions(+), 12 deletions(-)

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)
 }