Added first packet to data

asticode · Feb 18, 2018 · 6e7bd26 · 6e7bd26
1 parent dd1bd94
commit 6e7bd26
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 28 deletions.
diff --git a/data.go b/data.go
@@ -14,14 +14,15 @@ const (
 
 // Data represents a data
 type Data struct {
-	EIT *EITData
-	NIT *NITData
-	PAT *PATData
-	PES *PESData
-	PID uint16
-	PMT *PMTData
-	SDT *SDTData
-	TOT *TOTData
+	EIT         *EITData
+	FirstPacket *Packet
+	NIT         *NITData
+	PAT         *PATData
+	PES         *PESData
+	PID         uint16
+	PMT         *PMTData
+	SDT         *SDTData
+	TOT         *TOTData
 }
 
 // parseData parses a payload spanning over multiple packets and returns a set of data
@@ -61,11 +62,15 @@ func parseData(ps []*Packet, prs PacketsParser, pm programMap) (ds []*Data, err
 			err = errors.Wrap(err, "astits: parsing PSI data failed")
 			return
 		}
-		ds = psiData.toData(pid)
+		ds = psiData.toData(ps[0], pid)
 	} else if isPESPayload(payload) {
 		d, err := parsePESData(payload)
 		if err == nil {
-			ds = append(ds, &Data{PES: d, PID: pid})
+			ds = append(ds, &Data{
+				FirstPacket: ps[0],
+				PES:         d,
+				PID:         pid,
+			})
 		}
 	}
 	return

diff --git a/data_psi.go b/data_psi.go
@@ -333,23 +333,23 @@ func parsePSISectionSyntaxData(i []byte, offset *int, h *PSISectionHeader, sh *P
 }
 
 // toData parses the PSI tables and returns a set of Data
-func (d *PSIData) toData(pid uint16) (ds []*Data) {
+func (d *PSIData) toData(firstPacket *Packet, pid uint16) (ds []*Data) {
 	// Loop through sections
 	for _, s := range d.Sections {
 		// Switch on table type
 		switch s.Header.TableType {
 		case PSITableTypeEIT:
-			ds = append(ds, &Data{EIT: s.Syntax.Data.EIT, PID: pid})
+			ds = append(ds, &Data{EIT: s.Syntax.Data.EIT, FirstPacket: firstPacket, PID: pid})
 		case PSITableTypeNIT:
-			ds = append(ds, &Data{NIT: s.Syntax.Data.NIT, PID: pid})
+			ds = append(ds, &Data{FirstPacket: firstPacket, NIT: s.Syntax.Data.NIT, PID: pid})
 		case PSITableTypePAT:
-			ds = append(ds, &Data{PAT: s.Syntax.Data.PAT, PID: pid})
+			ds = append(ds, &Data{FirstPacket: firstPacket, PAT: s.Syntax.Data.PAT, PID: pid})
 		case PSITableTypePMT:
-			ds = append(ds, &Data{PID: pid, PMT: s.Syntax.Data.PMT})
+			ds = append(ds, &Data{FirstPacket: firstPacket, PID: pid, PMT: s.Syntax.Data.PMT})
 		case PSITableTypeSDT:
-			ds = append(ds, &Data{PID: pid, SDT: s.Syntax.Data.SDT})
+			ds = append(ds, &Data{FirstPacket: firstPacket, PID: pid, SDT: s.Syntax.Data.SDT})
 		case PSITableTypeTOT:
-			ds = append(ds, &Data{PID: pid, TOT: s.Syntax.Data.TOT})
+			ds = append(ds, &Data{FirstPacket: firstPacket, PID: pid, TOT: s.Syntax.Data.TOT})
 		}
 	}
 	return

diff --git a/data_psi_test.go b/data_psi_test.go
@@ -260,12 +260,13 @@ func TestParsePSISectionSyntaxHeader(t *testing.T) {
 }
 
 func TestPSIToData(t *testing.T) {
+	p := &Packet{}
 	assert.Equal(t, []*Data{
-		{EIT: eit, PID: 2},
-		{NIT: nit, PID: 2},
-		{PAT: pat, PID: 2},
-		{PMT: pmt, PID: 2},
-		{SDT: sdt, PID: 2},
-		{TOT: tot, PID: 2},
-	}, psi.toData(uint16(2)))
+		{EIT: eit, FirstPacket: p, PID: 2},
+		{FirstPacket: p, NIT: nit, PID: 2},
+		{FirstPacket: p, PAT: pat, PID: 2},
+		{FirstPacket: p, PMT: pmt, PID: 2},
+		{FirstPacket: p, SDT: sdt, PID: 2},
+		{FirstPacket: p, TOT: tot, PID: 2},
+	}, psi.toData(p, uint16(2)))
 }
diff --git a/data_test.go b/data_test.go
@@ -43,7 +43,7 @@ func TestParseData(t *testing.T) {
 	}
 	ds, err = parseData(ps, nil, pm)
 	assert.NoError(t, err)
-	assert.Equal(t, []*Data{{PES: pesWithHeader, PID: uint16(256)}}, ds)
+	assert.Equal(t, []*Data{{FirstPacket: ps[0], PES: pesWithHeader, PID: uint16(256)}}, ds)
 
 	// PSI
 	pm.set(uint16(256), uint16(1))
@@ -60,7 +60,7 @@ func TestParseData(t *testing.T) {
 	}
 	ds, err = parseData(ps, nil, pm)
 	assert.NoError(t, err)
-	assert.Equal(t, psi.toData(uint16(256)), ds)
+	assert.Equal(t, psi.toData(ps[0], uint16(256)), ds)
 }
 
 func TestIsPSIPayload(t *testing.T) {

diff --git a/demuxer_test.go b/demuxer_test.go
@@ -62,6 +62,10 @@ func TestDemuxerNextData(t *testing.T) {
 	b3, _ := packet(PacketHeader{ContinuityCounter: uint8(2), PayloadUnitStartIndicator: true, PID: PIDPAT}, PacketAdaptationField{}, []byte{})
 	w.Write(b3)
 	dmx := New(context.Background(), bytes.NewReader(w.Bytes()))
+	p, err := dmx.NextPacket()
+	assert.NoError(t, err)
+	_, err = dmx.Rewind()
+	assert.NoError(t, err)
 
 	// Next data
 	var ds []*Data
@@ -72,11 +76,11 @@ func TestDemuxerNextData(t *testing.T) {
 			ds = append(ds, d)
 		}
 	}
-	assert.Equal(t, psi.toData(PIDPAT), ds)
+	assert.Equal(t, psi.toData(p, PIDPAT), ds)
 	assert.Equal(t, map[uint16]uint16{0x3: 0x2, 0x5: 0x4}, dmx.programMap.p)
 
 	// No more packets
-	_, err := dmx.NextData()
+	_, err = dmx.NextData()
 	assert.EqualError(t, err, ErrNoMorePackets.Error())
 }