From 1786ed26614df36521eb4e08671add6790273f9a Mon Sep 17 00:00:00 2001
From: saxon <saxon.milton@gmail.com>
Date: Tue, 29 Jan 2019 16:14:00 +1030
Subject: [PATCH] stream/mts/psi: finished writing AddDescriptor

---
 stream/mts/encoder.go      | 20 ++++++------
 stream/mts/psi/psi.go      | 64 +++++++++++++++++++++++++-------------
 stream/mts/psi/psi_test.go |  2 +-
 3 files changed, 53 insertions(+), 33 deletions(-)

diff --git a/stream/mts/encoder.go b/stream/mts/encoder.go
index 8f46a15a..33194ab8 100644
--- a/stream/mts/encoder.go
+++ b/stream/mts/encoder.go
@@ -196,12 +196,12 @@ func (e *Encoder) Encode(nalu []byte) error {
 func (e *Encoder) writePSI() error {
 	// Write PAT.
 	patPkt := Packet{
-		PUSI:    true,
-		PID:     PatPid,
-		CC:      e.ccFor(PatPid),
-		AFC:     HasPayload,
-		Payload: patTable,
+		PUSI: true,
+		PID:  PatPid,
+		CC:   e.ccFor(PatPid),
+		AFC:  HasPayload,
 	}
+	patPkt.FillPayload(patTable)
 	_, err := e.dst.Write(patPkt.Bytes(e.tsSpace[:PacketSize]))
 	if err != nil {
 		return err
@@ -211,12 +211,12 @@ func (e *Encoder) writePSI() error {
 
 	// Create mts packet from pmt table.
 	pmtPkt := Packet{
-		PUSI:    true,
-		PID:     PmtPid,
-		CC:      e.ccFor(PmtPid),
-		AFC:     HasPayload,
-		Payload: pmtTable,
+		PUSI: true,
+		PID:  PmtPid,
+		CC:   e.ccFor(PmtPid),
+		AFC:  HasPayload,
 	}
+	pmtPkt.FillPayload(pmtTable)
 	_, err = e.dst.Write(pmtPkt.Bytes(e.tsSpace[:PacketSize]))
 	if err != nil {
 		return err
diff --git a/stream/mts/psi/psi.go b/stream/mts/psi/psi.go
index c41e1d98..9637594d 100644
--- a/stream/mts/psi/psi.go
+++ b/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,31 +240,45 @@ 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")
-		}
+	if psi.TableID(*p) != pmtID {
+		return errors.New("trying to add descriptor, but not pmt")
+	}
 
-		i, desc := p.HasDescriptor(tag)
-		if desc == nil {
-			p.createDescriptor(tag, data)
-			return nil
-		}
-		oldDescLen := desc.len()
-		// update the descriptor
-		oldDataLen := int(desc[1])
-		newDataLen := len(data)
-		newDescLen := 2 + newDataLen
+	i, desc := p.HasDescriptor(tag)
+	if desc == nil {
+		p.createDescriptor(tag, data)
+		return nil
+	}
 
-		// Shift data
-		//	copy(p[i+newDescLen:], data)
+	// Old lengths
+	oldDescLen := desc.len()
+	oldDataLen := int(desc[1])
+
+	// New lengths
+	newDataLen := len(data)
+	newDescLen := 2 + newDataLen
+
+	delta := newDescLen - oldDescLen
+
+	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)) + delta
+	p.SetSectionLen(newSectionLen)
 
-		deltaLen := desc.len() - oldDescLen
-		newProgInfoLen := p.ProgramInfoLen() + deltaLen
-		p.SetProgInfoLen(newProgInfoLen)
-		newSectionLen := int(psi.SectionLength(*p)) + deltaLen
-		p.SetSectionLen(newSectionLen)
-	*/
 	return nil
 }
 
diff --git a/stream/mts/psi/psi_test.go b/stream/mts/psi/psi_test.go
index e437066e..61a740a4 100644
--- a/stream/mts/psi/psi_test.go
+++ b/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)
 		}