cmd/ts-repair: updated comments and made funcs more robust

cmd/ts-repair/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"fmt"
 	"io"
@@ -10,21 +11,26 @@ import (
 	"github.com/Comcast/gots/packet"
 )
 
+// Various errors that we can encounter.
 const (
-	errBadInPath     = "No file path provided, or file does not exist"
-	errCantCreateOut = "Can't create output file"
-	errCantGetPid    = "Can't get pid from packet"
-	errReadFail      = "Read failed"
-	errWriteFail     = "Write to file failed"
-	errBadMode       = "Bad fix mode"
+	errBadInPath         = "No file path provided, or file does not exist"
+	errCantCreateOut     = "Can't create output file"
+	errCantGetPid        = "Can't get pid from packet"
+	errReadFail          = "Read failed"
+	errWriteFail         = "Write to file failed"
+	errBadMode           = "Bad fix mode"
+	errAdaptationPresent = "Adaptation field is already present in packet"
+	errNoAdaptationField = "No adaptation field in this packet"
 )
 
+// Consts describing flag usage.
 const (
 	inUsage   = "The path to the file to be repaired"
 	outUsage  = "Output file path"
 	modeUsage = "Fix mode: 0 = cc-shift, 1 = di-update"
 )
 
+// Repair modes.
 const (
 	ccShift = iota
 	diUpdate
@@ -36,20 +42,27 @@ var ccMap = map[int]byte{
 	mts.VideoPid: 16,
 }
 
+// packetNo will keep track of the ts packet number for reference.
 var packetNo int
 
+// Option defines a func that performs an action on p in order to change a ts option.
 type Option func(p *Packet)
 
+// Packet is a byte array of size mts.PacketSize i.e. 188 bytes. We define this
+// to allow us to write receiver funcs for the [mts.PacketSize]byte type.
 type Packet [mts.PacketSize]byte
 
+// CC returns the CC of p.
 func (p *Packet) CC() byte {
 	return (*p)[3] & 0x0f
 }
 
+// setCC sets the CC of p.
 func (p *Packet) setCC(cc byte) {
 	(*p)[3] |= cc & 0xf
 }
 
+// setDI sets the discontinuity counter of p.
 func (p *Packet) setDI(di bool) {
 	if di {
 		p[5] |= 0x80
@@ -58,27 +71,41 @@ func (p *Packet) setDI(di bool) {
 	}
 }
 
-func (p *Packet) addAdaptationField(options ...Option) {
-	// Create space for adaptation field
+// addAdaptationField adds an adaptation field to p, and applys the passed options to this field.
+// TODO: this will probably break if we already have adaptation field.
+func (p *Packet) addAdaptationField(options ...Option) error {
+	if p.hasAdaptation() {
+		return errors.New(errAdaptationPresent)
+	}
+	// Create space for adaptation field.
 	copy(p[mts.HeadSize+mts.DefaultAdaptationSize:], p[mts.HeadSize:len(p)-mts.DefaultAdaptationSize])
 
 	// TODO: seperate into own function
-	// Update adaptation field control
+	// Update adaptation field control.
 	p[mts.AdaptationControlIdx] &= 0xff ^ mts.AdaptationControlMask
 	p[mts.AdaptationControlIdx] |= mts.AdaptationControlMask
-	// Default the adaptationfield
+	// Default the adaptationfield.
 	p.resetAdaptation()
 
+	// Apply and options that have bee passed.
 	for _, option := range options {
 		option(p)
 	}
+	return nil
 }
 
-func (p *Packet) resetAdaptation() {
+// resetAdaptation sets fields in ps adaptation field to 0 if the adaptation field
+// exists, otherwise an error is returned.
+func (p *Packet) resetAdaptation() error {
+	if !p.hasAdaptation() {
+		return errors.New(errNoAdaptationField)
+	}
 	p[mts.AdaptationIdx] = mts.DefaultAdaptationBodySize
 	p[mts.AdaptationBodyIdx] = 0x00
+	return nil
 }
 
+// hasAdaptation returns true if p has an adaptation field and false otherwise.
 func (p *Packet) hasAdaptation() bool {
 	afc := p[mts.AdaptationControlIdx] & mts.AdaptationControlMask
 	if afc == 0x20 || afc == 0x30 {
@@ -88,6 +115,8 @@ func (p *Packet) hasAdaptation() bool {
 	}
 }
 
+// DiscontinuityIndicator returns and Option that will set p's discontinuity
+// indicator according to f.
 func DiscontinuityIndicator(f bool) Option {
 	return func(p *Packet) {
 		set := byte(mts.DiscontinuityIndicatorMask)
@@ -130,23 +159,26 @@ func main() {
 			panic(errReadFail + ": " + err.Error())
 		}
 		packetNo++
-		// Get the pid from the packet and set the cc based on this pid using our map
+
+		// Get the pid from the packet
 		pid, err := packet.Pid((*packet.Packet)(&p))
 		if err != nil {
 			panic(errCantGetPid)
 		}
 
+		// Get the cc from the packet and also the expected cc (if exists)
 		cc := p.CC()
 		expect, exists := expectedCC(int(pid))
 		if !exists {
 			updateCCMap(int(pid), cc)
 		} else {
 			switch *modePtr {
+			// ccShift mode shifts all CC regardless of presence of Discontinuities or not
 			case ccShift:
 				p.setCC(expect)
-
+				// diUpdate mode finds discontinuities and sets the discontinuity indicator to true.
+				// If we have a pat or pmt then we need to add an adaptation field and then set the DI.
 			case diUpdate:
-
 				if cc != expect {
 					fmt.Printf("***** Discontinuity found (packetNo: %v pid: %v, cc: %v, expect: %v)\n", packetNo, pid, cc, expect)
 					if p.hasAdaptation() {
@@ -156,19 +188,20 @@ func main() {
 					}
 					updateCCMap(int(pid), p.CC())
 				}
-
 			default:
 				panic(errBadMode)
 			}
 		}
-		// Write this packet to the output file
+
+		// Write this packet to the output file.
 		if _, err := outFile.Write(p[:]); err != nil {
 			panic(errWriteFail + ": " + err.Error())
 		}
 	}
 }
 
-// ccFor gets the next cc for the given pid
+// expectedCC returns the expected cc for the given pid. If the cc hasn't been
+// used yet, then 16 and false is returned.
 func expectedCC(pid int) (byte, bool) {
 	cc := ccMap[pid]
 	if cc == 16 {
@@ -178,6 +211,7 @@ func expectedCC(pid int) (byte, bool) {
 	return cc, true
 }
 
+// updateCCMap updates the cc for the passed pid.
 func updateCCMap(pid int, cc byte) {
 	ccMap[pid] = (cc + 1) & 0xf
 }