stream/mts/psi: added some more tests for AddDescriptor, and in the process fixed some bugs with AddDescriptor

stream/mts/psi/descriptor_test.go

@@ -127,6 +127,24 @@ var (
 	}
 )
 
+var (
+	pmtTimeBytesResizedBigger = []byte{
+		0x00, 0x02, 0xb0, 0x1e, 0x00, 0x01, 0xc1, 0x00, 0x00, 0xe1, 0x00, 0xf0, 0x0c,
+		TimeDescTag,                                                // Descriptor tag
+		0x0a,                                                       // Length of bytes to follow
+		0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, // timestamp
+		0x1b, 0xe1, 0x00, 0xf0, 0x00,
+	}
+
+	pmtTimeBytesResizedSmaller = []byte{
+		0x00, 0x02, 0xb0, 0x1a, 0x00, 0x01, 0xc1, 0x00, 0x00, 0xe1, 0x00, 0xf0, 0x08,
+		TimeDescTag,                        // Descriptor tag
+		0x06,                               // Length of bytes to follow
+		0x01, 0x02, 0x03, 0x04, 0x05, 0x06, // timestamp
+		0x1b, 0xe1, 0x00, 0xf0, 0x00,
+	}
+)
+
 // TestHasDescriptorExists checks that PSIBytes.HasDescriptor performs as expected
 // when the PSI we're interested in has the descriptor of interest. HasDescriptor
 // should return the descriptor bytes.
@@ -259,4 +277,39 @@ func TestAddDescriptorNonEmpty(t *testing.T) {
 	}
 }
 
-// TODO: check that we can update descriptor with AddDescriptor
+func TestAddDescriptorUpdateSame(t *testing.T) {
+	newData := [8]byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
+	want := PSIBytes(tstPsi2.Bytes())
+	want.createDescriptor(TimeDescTag, newData[:])
+
+	got := PSIBytes(tstPsi1.Bytes())
+	if err := got.AddDescriptor(TimeDescTag, newData[:]); err != nil {
+		t.Errorf(errUnexpectedErr, err.Error())
+	}
+
+	if !bytes.Equal(got, want) {
+		t.Errorf(errNotExpectedOut, got, want)
+	}
+}
+
+func TestAddDescriptorUpdateBigger(t *testing.T) {
+	got := PSIBytes(tstPsi1.Bytes())
+	if err := got.AddDescriptor(TimeDescTag, []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a}); err != nil {
+		t.Errorf(errUnexpectedErr, err.Error())
+	}
+	want := addCrc(pmtTimeBytesResizedBigger)
+	if !bytes.Equal(got, want) {
+		t.Errorf(errNotExpectedOut, got, want)
+	}
+}
+
+func TestAddDescriptorUpdateSmaller(t *testing.T) {
+	got := PSIBytes(tstPsi1.Bytes())
+	if err := got.AddDescriptor(TimeDescTag, []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06}); err != nil {
+		t.Errorf(errUnexpectedErr, err.Error())
+	}
+	want := addCrc(pmtTimeBytesResizedSmaller)
+	if !bytes.Equal(got, want) {
+		t.Errorf(errNotExpectedOut, got, want)
+	}
+}

stream/mts/psi/psi.go

@@ -263,8 +263,7 @@ func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 
 	// If the old data length is more than the new data length, we need shift data
 	// after descriptor up, and then trim the psi. If the oldDataLen is less than
-	// new data then we need reseize psi and shift data down. If data is same size
-	// just copy new data in.
+	// new data then we need reseize psi and shift data down. If same do nothing.
 	switch {
 	case oldDataLen > newDataLen:
 		copy((*p)[i+newDescLen:], (*p)[i+oldDescLen:])
@@ -274,10 +273,12 @@ func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 		copy(tmp, *p)
 		*p = tmp
 		copy((*p)[i+newDescLen:], (*p)[i+oldDescLen:])
-	default:
-		copy((*p)[i+2:], data)
 	}
 
+	// Copy in new data
+	(*p)[i+1] = byte(newDataLen)
+	copy((*p)[i+2:], data)
+
 	newProgInfoLen := p.ProgramInfoLen() + delta
 	p.setProgInfoLen(newProgInfoLen)
 	newSectionLen := int(psi.SectionLength(*p)) + delta
@@ -286,18 +287,6 @@ func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 	return nil
 }
 
-// len returns the length of a descriptor in bytes.
-func (d *Descriptor) len() int {
-	return int(2 + (*d)[1])
-}
-
-// ProgramInfoLen returns the program info length of a PSI.
-//
-// TODO: check if pmt - if not return 0 ? or -1 ?
-func (p *PSIBytes) ProgramInfoLen() int {
-	return int((((*p)[ProgramInfoLenIdx1] & ProgramInfoLenMask1) << 8) | (*p)[ProgramInfoLenIdx2])
-}
-
 // HasDescriptor checks if a descriptor of the given tag exists in a PSI. If the descriptor
 // of the given tag exists, an index of this descriptor, as well as the Descriptor is returned.
 // If the descriptor of the given tag cannot be found, -1 and a nil slice is returned.
@@ -310,19 +299,15 @@ func (p *PSIBytes) HasDescriptor(tag int) (int, Descriptor) {
 	}
 	for i := 0; i < len(descs); i += 2 + int(descs[i+1]) {
 		if int(descs[i]) == tag {
-			return i, descs[i : i+2+int(descs[i+1])]
+			return i + DescriptorsIdx, descs[i : i+2+int(descs[i+1])]
 		}
 	}
 	return -1, nil
 }
 
-// descriptors returns the descriptors in a psi if they exist, otherwise
-// a nil slice is returned.
-func (p *PSIBytes) descriptors() []byte {
-	return (*p)[DescriptorsIdx : DescriptorsIdx+p.ProgramInfoLen()]
-}
-
-// createDescriptor creates a descriptor in a psi given a tag and data.
+// createDescriptor creates a descriptor in a psi given a tag and data. It does so
+// by resizing the psi, shifting existing data down and copying in new descriptor
+// in new space.
 func (p *PSIBytes) createDescriptor(tag int, data []byte) {
 	curProgLen := p.ProgramInfoLen()
 	oldSyntaxSectionLen := SyntaxSecLenFrom(*p)
@@ -363,3 +348,21 @@ func (p *PSIBytes) setSectionLen(l int) {
 	(*p)[SyntaxSecLenIdx1] |= byte(l>>8) & SyntaxSecLenMask1
 	(*p)[SyntaxSecLenIdx2] = byte(l)
 }
+
+// descriptors returns the descriptors in a psi if they exist, otherwise
+// a nil slice is returned.
+func (p *PSIBytes) descriptors() []byte {
+	return (*p)[DescriptorsIdx : DescriptorsIdx+p.ProgramInfoLen()]
+}
+
+// len returns the length of a descriptor in bytes.
+func (d *Descriptor) len() int {
+	return int(2 + (*d)[1])
+}
+
+// ProgramInfoLen returns the program info length of a PSI.
+//
+// TODO: check if pmt - if not return 0 ? or -1 ?
+func (p *PSIBytes) ProgramInfoLen() int {
+	return int((((*p)[ProgramInfoLenIdx1] & ProgramInfoLenMask1) << 8) | (*p)[ProgramInfoLenIdx2])
+}