stream/mts/psi: finished writing AddDescriptor

stream/mts/encoder.go

@@ -200,8 +200,8 @@ func (e *Encoder) writePSI() error {
 		PID:  PatPid,
 		CC:   e.ccFor(PatPid),
 		AFC:  HasPayload,
-		Payload: patTable,
 	}
+	patPkt.FillPayload(patTable)
 	_, err := e.dst.Write(patPkt.Bytes(e.tsSpace[:PacketSize]))
 	if err != nil {
 		return err
@@ -215,8 +215,8 @@ func (e *Encoder) writePSI() error {
 		PID:  PmtPid,
 		CC:   e.ccFor(PmtPid),
 		AFC:  HasPayload,
-		Payload: pmtTable,
 	}
+	pmtPkt.FillPayload(pmtTable)
 	_, err = e.dst.Write(pmtPkt.Bytes(e.tsSpace[:PacketSize]))
 	if err != nil {
 		return err

stream/mts/psi/psi.go

@@ -26,6 +26,12 @@ LICENSE
 
 package psi
 
+import (
+	"errors"
+
+	"github.com/Comcast/gots/psi"
+)
+
 const (
 	PacketSize = 184 // packet size of a psi.
 )
@@ -234,7 +240,6 @@ func asByte(b bool) byte {
 }
 
 func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
-	/*
 	if psi.TableID(*p) != pmtID {
 		return errors.New("trying to add descriptor, but not pmt")
 	}
@@ -244,21 +249,36 @@ func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 		p.createDescriptor(tag, data)
 		return nil
 	}
+
+	// Old lengths
 	oldDescLen := desc.len()
-		// update the descriptor
 	oldDataLen := int(desc[1])
+
+	// New lengths
 	newDataLen := len(data)
 	newDescLen := 2 + newDataLen
 
-		// Shift data
-		//	copy(p[i+newDescLen:], data)
+	delta := newDescLen - oldDescLen
 
-		deltaLen := desc.len() - oldDescLen
-		newProgInfoLen := p.ProgramInfoLen() + deltaLen
+	if oldDataLen > newDataLen {
+		copy((*p)[i+newDescLen:], (*p)[i+oldDescLen:])
+		*p = (*p)[:len(*p)+delta]
+
+	} else if oldDataLen < newDataLen {
+		tmp := make([]byte, len(*p)+delta)
+		copy(tmp, *p)
+		*p = tmp
+		copy((*p)[i+newDescLen:], (*p)[i+oldDescLen:])
+
+	} else {
+		copy((*p)[i+2:], data)
+	}
+
+	newProgInfoLen := p.ProgramInfoLen() + delta
 	p.SetProgInfoLen(newProgInfoLen)
-		newSectionLen := int(psi.SectionLength(*p)) + deltaLen
+	newSectionLen := int(psi.SectionLength(*p)) + delta
 	p.SetSectionLen(newSectionLen)
-	*/
+
 	return nil
 }
 

stream/mts/psi/psi_test.go

@@ -282,7 +282,7 @@ var bytesTests = []struct {
 func TestBytes(t *testing.T) {
 	for _, test := range bytesTests {
 		got := test.input.Bytes()
-		if !bytes.Equal(got, addPadding(addCrc(test.want))) {
+		if !bytes.Equal(got, addCrc(test.want)) {
 			t.Errorf("unexpected error for test %v: got:%v want:%v", test.name, got,
 				test.want)
 		}