From 82b17612dd676e21d204d6fe03a47ff13bcc70ea Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Mon, 18 Apr 2016 00:20:01 +0800
Subject: [PATCH] add frag atom support

---
 atom/genStruct.js   |  66 +++++++++
 atom/otherStruct.go | 318 ++++++++++++++++++++++++++++++++++++++++----
 atom/struct.go      | 223 +++++++++++++++++++++++++++++++
 3 files changed, 581 insertions(+), 26 deletions(-)

diff --git a/atom/genStruct.js b/atom/genStruct.js
index 4c1c2c7..6a2122f 100644
--- a/atom/genStruct.js
+++ b/atom/genStruct.js
@@ -340,6 +340,72 @@ var atoms = {
 		],
 	},
 
+	movieFrag: {
+		cc4: 'moof',
+		fields: [
+			['$atoms', [
+				['header', '*movieFragHeader'],
+				['tracks', '[]*trackFrag'],
+			]],
+		],
+	},
+
+	trackFragDecodeTime: {
+		cc4: 'tfdt',
+	},
+
+	movieFragHeader: {
+		cc4: 'mfhd',
+		fields: [
+			['version', 'int8'],
+			['flags', 'int24'],
+			['seqNum', 'int32'],
+		],
+	},
+
+	trackFrag: {
+		cc4: 'traf',
+		fields: [
+			['$atoms', [
+				['header', '*trackFragHeader'],
+				['decodeTime', '*trackFragDecodeTime'],
+				['run', '*trackFragRun'],
+			]],
+		],
+	},
+
+	trackFragRun: {
+		cc4: 'trun',
+	},
+
+	trackFragHeader: {
+		cc4: 'tfhd',
+	},
+
+	/*
+	// need hand write
+	trackFragRun: {
+		cc4: 'trun',
+		fields: [
+			['version', 'int8'],
+			['flags', 'int24'],
+			['sampleCount', 'int32'],
+			['dataOffset', 'int32'],
+			['entries', '[]int32'],
+		],
+	},
+
+	trackFragHeader: {
+		cc4: 'tfhd',
+		fields: [
+			['version', 'int8'],
+			['flags', 'int24'],
+			['id', 'int32'],
+			['sampleDescriptionIndex', 'int32'],
+			['_', '[12]byte'],
+		],
+	},
+	*/
 };
 
 var DeclReadFunc = (opts) => {
diff --git a/atom/otherStruct.go b/atom/otherStruct.go
index ec73293..c5b298e 100644
--- a/atom/otherStruct.go
+++ b/atom/otherStruct.go
@@ -1,15 +1,15 @@
-
 package atom
 
 import (
-	"io"
-	"fmt"
 	"bytes"
+	"fmt"
+	"io"
+	"github.com/nareix/bits"
 )
 
 type GolombBitReader struct {
-	R io.Reader
-	buf [1]byte
+	R    io.Reader
+	buf  [1]byte
 	left byte
 }
 
@@ -21,7 +21,7 @@ func (self *GolombBitReader) ReadBit() (res uint, err error) {
 		self.left = 8
 	}
 	self.left--
-	res = uint(self.buf[0]>>self.left)&1
+	res = uint(self.buf[0]>>self.left) & 1
 	return
 }
 
@@ -31,7 +31,7 @@ func (self *GolombBitReader) ReadBits(n int) (res uint, err error) {
 		if bit, err = self.ReadBit(); err != nil {
 			return
 		}
-		res |= bit<<uint(n-i-1)
+		res |= bit << uint(n-i-1)
 	}
 	return
 }
@@ -51,7 +51,7 @@ func (self *GolombBitReader) ReadExponentialGolombCode() (res uint, err error) {
 	if res, err = self.ReadBits(i); err != nil {
 		return
 	}
-	res += (1<<uint(i)) - 1
+	res += (1 << uint(i)) - 1
 	return
 }
 
@@ -59,27 +59,27 @@ func (self *GolombBitReader) ReadSE() (res uint, err error) {
 	if res, err = self.ReadExponentialGolombCode(); err != nil {
 		return
 	}
-	if res & 0x01 != 0 {
-		res = (res+1)/2
+	if res&0x01 != 0 {
+		res = (res + 1) / 2
 	} else {
-		res = -res/2
+		res = -res / 2
 	}
 	return
 }
 
 type H264SPSInfo struct {
 	ProfileIdc uint
-	LevelIdc uint
+	LevelIdc   uint
 
-	MbWidth uint
+	MbWidth  uint
 	MbHeight uint
 
-	CropLeft uint
-	CropRight uint
-	CropTop uint
+	CropLeft   uint
+	CropRight  uint
+	CropTop    uint
 	CropBottom uint
 
-	Width uint
+	Width  uint
 	Height uint
 }
 
@@ -176,7 +176,7 @@ func ParseH264SPS(data []byte) (res *H264SPSInfo, err error) {
 		}
 	}
 
-	// log2_max_frame_num_minus4 
+	// log2_max_frame_num_minus4
 	if _, err = r.ReadExponentialGolombCode(); err != nil {
 		return
 	}
@@ -269,8 +269,8 @@ func ParseH264SPS(data []byte) (res *H264SPSInfo, err error) {
 		}
 	}
 
-	self.Width = (self.MbWidth*16) - self.CropLeft*2 - self.CropRight*2
-	self.Height = ((2-frame_mbs_only_flag)*self.MbHeight*16) - self.CropTop*2 - self.CropBottom*2
+	self.Width = (self.MbWidth * 16) - self.CropLeft*2 - self.CropRight*2
+	self.Height = ((2 - frame_mbs_only_flag) * self.MbHeight * 16) - self.CropTop*2 - self.CropBottom*2
 
 	res = self
 	return
@@ -279,10 +279,10 @@ func ParseH264SPS(data []byte) (res *H264SPSInfo, err error) {
 type AVCDecoderConfRecord struct {
 	AVCProfileIndication int
 	ProfileCompatibility int
-	AVCLevelIndication int
-	LengthSizeMinusOne int
-	SPS [][]byte
-	PPS [][]byte
+	AVCLevelIndication   int
+	LengthSizeMinusOne   int
+	SPS                  [][]byte
+	PPS                  [][]byte
 }
 
 func WriteSampleByNALU(w io.Writer, nalu []byte) (size int, err error) {
@@ -327,11 +327,11 @@ func WriteAVCDecoderConfRecord(w io.Writer, self AVCDecoderConfRecord) (err erro
 	if err = WriteInt(w, self.AVCLevelIndication, 1); err != nil {
 		return
 	}
-	if err = WriteInt(w, self.LengthSizeMinusOne | 0xfc, 1); err != nil {
+	if err = WriteInt(w, self.LengthSizeMinusOne|0xfc, 1); err != nil {
 		return
 	}
 
-	if err = WriteInt(w, len(self.SPS) | 0xe0, 1); err != nil {
+	if err = WriteInt(w, len(self.SPS)|0xe0, 1); err != nil {
 		return
 	}
 	for _, data := range self.SPS {
@@ -412,3 +412,269 @@ func ReadAVCDecoderConfRecord(r *io.LimitedReader) (self AVCDecoderConfRecord, e
 	return
 }
 
+const (
+	TFHD_BASE_DATA_OFFSET     = 0x01
+	TFHD_STSD_ID              = 0x02
+	TFHD_DEFAULT_DURATION     = 0x08
+	TFHD_DEFAULT_SIZE         = 0x10
+	TFHD_DEFAULT_FLAGS        = 0x20
+	TFHD_DURATION_IS_EMPTY    = 0x010000
+	TFHD_DEFAULT_BASE_IS_MOOF = 0x020000
+)
+
+type TrackFragHeader struct {
+	Version   int
+	Flags     int
+	Id        int
+	DefaultSize      int
+	DefaultDuration  int
+	DefaultFlags int
+	BaseDataOffset int64
+	StsdId int
+}
+
+func WalkTrackFragHeader(w Walker, self *TrackFragHeader) {
+	w.StartStruct("TrackFragHeader")
+	w.Name("Flags")
+	w.HexInt(self.Flags)
+	w.Name("Id")
+	w.Int(self.Id)
+	w.Name("DefaultDuration")
+	w.Int(self.DefaultDuration)
+	w.Name("DefaultSize")
+	w.Int(self.DefaultSize)
+	w.Name("DefaultFlags")
+	w.HexInt(self.DefaultFlags)
+	w.EndStruct()
+}
+
+func WriteTrackFragHeader(w io.WriteSeeker, self *TrackFragHeader) (err error) {
+	panic("unimplmented")
+	return
+}
+
+func ReadTrackFragHeader(r *io.LimitedReader) (res *TrackFragHeader, err error) {
+	self := &TrackFragHeader{}
+
+	if self.Version, err = ReadInt(r, 1); err != nil {
+		return
+	}
+	if self.Flags, err = ReadInt(r, 3); err != nil {
+		return
+	}
+	if self.Id, err = ReadInt(r, 4); err != nil {
+		return
+	}
+
+	if self.Flags&TFHD_BASE_DATA_OFFSET != 0 {
+		if self.BaseDataOffset, err = bits.ReadInt64BE(r, 64); err != nil {
+			return
+		}
+	}
+
+	if self.Flags&TFHD_STSD_ID != 0 {
+		if self.StsdId, err = ReadInt(r, 4); err != nil {
+			return
+		}
+	}
+
+	if self.Flags&TFHD_DEFAULT_DURATION != 0 {
+		if self.DefaultDuration, err = ReadInt(r, 4); err != nil {
+			return
+		}
+	}
+
+	if self.Flags&TFHD_DEFAULT_SIZE != 0 {
+		if self.DefaultSize, err = ReadInt(r, 4); err != nil {
+			return
+		}
+	}
+
+	if self.Flags&TFHD_DEFAULT_FLAGS != 0 {
+		if self.DefaultFlags,err = ReadInt(r, 4); err != nil {
+			return
+		}
+	}
+
+	res = self
+	return
+}
+
+const (
+	TRUN_DATA_OFFSET        = 0x01
+	TRUN_FIRST_SAMPLE_FLAGS = 0x04
+	TRUN_SAMPLE_DURATION    = 0x100
+	TRUN_SAMPLE_SIZE        = 0x200
+	TRUN_SAMPLE_FLAGS       = 0x400
+	TRUN_SAMPLE_CTS         = 0x800
+)
+
+type TrackFragRunEntry struct {
+	Duration int
+	Size     int
+	Flags    int
+	Cts      int
+}
+
+type TrackFragRun struct {
+	Version          int
+	Flags            int
+	FirstSampleFlags int
+	DataOffset       int
+	Entries          []TrackFragRunEntry
+}
+
+func WalkTrackFragRun(w Walker, self *TrackFragRun) {
+	w.StartStruct("TrackFragRun")
+	w.Name("Flags")
+	w.HexInt(self.Flags)
+	w.Name("FirstSampleFlags")
+	w.HexInt(self.FirstSampleFlags)
+	w.Name("DataOffset")
+	w.Int(self.DataOffset)
+	w.Name("EntriesCount")
+	w.Int(len(self.Entries))
+	for i := 0; i < 10 && i < len(self.Entries); i++ {
+		entry := self.Entries[i]
+		w.Println(fmt.Sprintf("Entry[%d] Flags=%x Duration=%d Size=%d Cts=%d",
+			i, entry.Flags, entry.Duration, entry.Size, entry.Cts))
+	}
+	w.EndStruct()
+}
+
+func WriteTrackFragRun(w io.WriteSeeker, self *TrackFragRun) (err error) {
+	panic("unimplmented")
+	return
+}
+
+func ReadTrackFragRun(r *io.LimitedReader) (res *TrackFragRun, err error) {
+	self := &TrackFragRun{}
+
+	if self.Version, err = ReadInt(r, 1); err != nil {
+		return
+	}
+	if self.Flags, err = ReadInt(r, 3); err != nil {
+		return
+	}
+
+	var count int
+	if count, err = ReadInt(r, 4); err != nil {
+		return
+	}
+
+	if self.Flags&TRUN_DATA_OFFSET != 0 {
+		if self.DataOffset, err = ReadInt(r, 4); err != nil {
+			return
+		}
+	}
+	if self.Flags&TRUN_FIRST_SAMPLE_FLAGS != 0 {
+		if self.FirstSampleFlags, err = ReadInt(r, 4); err != nil {
+			return
+		}
+	}
+
+	for i := 0; i < count; i++ {
+		var flags int
+
+		if i > 0 {
+			flags = self.Flags
+		} else {
+			flags = self.FirstSampleFlags
+		}
+
+		entry := TrackFragRunEntry{}
+		if flags&TRUN_SAMPLE_DURATION != 0 {
+			if entry.Duration, err = ReadInt(r, 4); err != nil {
+				return
+			}
+		}
+		if flags&TRUN_SAMPLE_SIZE != 0 {
+			if entry.Size, err = ReadInt(r, 4); err != nil {
+				return
+			}
+		}
+		if flags&TRUN_SAMPLE_FLAGS != 0 {
+			if entry.Flags, err = ReadInt(r, 4); err != nil {
+				return
+			}
+		}
+		if flags&TRUN_SAMPLE_CTS != 0 {
+			if entry.Cts, err = ReadInt(r, 4); err != nil {
+				return
+			}
+		}
+
+		self.Entries = append(self.Entries, entry)
+	}
+
+	res = self
+	return
+}
+
+type TrackFragDecodeTime struct {
+	Version int
+	Flags   int
+	Time    int64
+}
+
+func ReadTrackFragDecodeTime(r *io.LimitedReader) (res *TrackFragDecodeTime, err error) {
+
+	self := &TrackFragDecodeTime{}
+	if self.Version, err = ReadInt(r, 1); err != nil {
+		return
+	}
+	if self.Flags, err = ReadInt(r, 3); err != nil {
+		return
+	}
+	if self.Version != 0 {
+		if self.Time, err = bits.ReadInt64BE(r, 64); err != nil {
+			return
+		}
+	} else {
+		if self.Time, err = bits.ReadInt64BE(r, 32); err != nil {
+			return
+		}
+	}
+	res = self
+	return
+}
+
+func WriteTrackFragDecodeTime(w io.WriteSeeker, self *TrackFragDecodeTime) (err error) {
+	var aw *Writer
+	if aw, err = WriteAtomHeader(w, "tfdt"); err != nil {
+		return
+	}
+	w = aw
+	if err = WriteInt(w, self.Version, 1); err != nil {
+		return
+	}
+	if err = WriteInt(w, self.Flags, 3); err != nil {
+		return
+	}
+	if self.Version != 0 {
+		if err = bits.WriteInt64BE(w, self.Time, 64); err != nil {
+			return
+		}
+	} else {
+		if err = bits.WriteInt64BE(w, self.Time, 32); err != nil {
+			return
+		}
+	}
+	if err = aw.Close(); err != nil {
+		return
+	}
+	return
+}
+
+func WalkTrackFragDecodeTime(w Walker, self *TrackFragDecodeTime) {
+	w.StartStruct("TrackFragDecodeTime")
+	w.Name("Version")
+	w.Int(self.Version)
+	w.Name("Flags")
+	w.Int(self.Flags)
+	w.Name("Time")
+	w.Int64(self.Time)
+	w.EndStruct()
+	return
+}
+
diff --git a/atom/struct.go b/atom/struct.go
index 17128b5..3e75a92 100644
--- a/atom/struct.go
+++ b/atom/struct.go
@@ -2475,3 +2475,226 @@ func WalkChunkOffset(w Walker, self *ChunkOffset) {
 	w.EndStruct()
 	return
 }
+
+type MovieFrag struct {
+	Header *MovieFragHeader
+	Tracks []*TrackFrag
+}
+
+func ReadMovieFrag(r *io.LimitedReader) (res *MovieFrag, err error) {
+
+	self := &MovieFrag{}
+	for r.N > 0 {
+		var cc4 string
+		var ar *io.LimitedReader
+		if ar, cc4, err = ReadAtomHeader(r, ""); err != nil {
+			return
+		}
+		switch cc4 {
+		case "mfhd":
+			{
+				if self.Header, err = ReadMovieFragHeader(ar); err != nil {
+					return
+				}
+			}
+		case "traf":
+			{
+				var item *TrackFrag
+				if item, err = ReadTrackFrag(ar); err != nil {
+					return
+				}
+				self.Tracks = append(self.Tracks, item)
+			}
+
+		}
+		if _, err = ReadDummy(ar, int(ar.N)); err != nil {
+			return
+		}
+	}
+	res = self
+	return
+}
+func WriteMovieFrag(w io.WriteSeeker, self *MovieFrag) (err error) {
+
+	var aw *Writer
+	if aw, err = WriteAtomHeader(w, "moof"); err != nil {
+		return
+	}
+	w = aw
+	if self.Header != nil {
+		if err = WriteMovieFragHeader(w, self.Header); err != nil {
+			return
+		}
+	}
+	if self.Tracks != nil {
+		for _, elem := range self.Tracks {
+			if err = WriteTrackFrag(w, elem); err != nil {
+				return
+			}
+		}
+	}
+	if err = aw.Close(); err != nil {
+		return
+	}
+	return
+}
+func WalkMovieFrag(w Walker, self *MovieFrag) {
+
+	w.StartStruct("MovieFrag")
+	if self.Header != nil {
+		WalkMovieFragHeader(w, self.Header)
+	}
+	for i, item := range self.Tracks {
+		if w.FilterArrayItem("MovieFrag", "Tracks", i, len(self.Tracks)) {
+			if item != nil {
+				WalkTrackFrag(w, item)
+			}
+		} else {
+			w.ArrayLeft(i, len(self.Tracks))
+			break
+		}
+	}
+	w.EndStruct()
+	return
+}
+
+type MovieFragHeader struct {
+	Version int
+	Flags   int
+	SeqNum  int
+}
+
+func ReadMovieFragHeader(r *io.LimitedReader) (res *MovieFragHeader, err error) {
+
+	self := &MovieFragHeader{}
+	if self.Version, err = ReadInt(r, 1); err != nil {
+		return
+	}
+	if self.Flags, err = ReadInt(r, 3); err != nil {
+		return
+	}
+	if self.SeqNum, err = ReadInt(r, 4); err != nil {
+		return
+	}
+	res = self
+	return
+}
+func WriteMovieFragHeader(w io.WriteSeeker, self *MovieFragHeader) (err error) {
+
+	var aw *Writer
+	if aw, err = WriteAtomHeader(w, "mfhd"); err != nil {
+		return
+	}
+	w = aw
+	if err = WriteInt(w, self.Version, 1); err != nil {
+		return
+	}
+	if err = WriteInt(w, self.Flags, 3); err != nil {
+		return
+	}
+	if err = WriteInt(w, self.SeqNum, 4); err != nil {
+		return
+	}
+	if err = aw.Close(); err != nil {
+		return
+	}
+	return
+}
+func WalkMovieFragHeader(w Walker, self *MovieFragHeader) {
+
+	w.StartStruct("MovieFragHeader")
+	w.Name("Version")
+	w.Int(self.Version)
+	w.Name("Flags")
+	w.Int(self.Flags)
+	w.Name("SeqNum")
+	w.Int(self.SeqNum)
+	w.EndStruct()
+	return
+}
+
+type TrackFrag struct {
+	Header     *TrackFragHeader
+	DecodeTime *TrackFragDecodeTime
+	Run        *TrackFragRun
+}
+
+func ReadTrackFrag(r *io.LimitedReader) (res *TrackFrag, err error) {
+
+	self := &TrackFrag{}
+	for r.N > 0 {
+		var cc4 string
+		var ar *io.LimitedReader
+		if ar, cc4, err = ReadAtomHeader(r, ""); err != nil {
+			return
+		}
+		switch cc4 {
+		case "tfhd":
+			{
+				if self.Header, err = ReadTrackFragHeader(ar); err != nil {
+					return
+				}
+			}
+		case "tfdt":
+			{
+				if self.DecodeTime, err = ReadTrackFragDecodeTime(ar); err != nil {
+					return
+				}
+			}
+		case "trun":
+			{
+				if self.Run, err = ReadTrackFragRun(ar); err != nil {
+					return
+				}
+			}
+
+		}
+		if _, err = ReadDummy(ar, int(ar.N)); err != nil {
+			return
+		}
+	}
+	res = self
+	return
+}
+func WriteTrackFrag(w io.WriteSeeker, self *TrackFrag) (err error) {
+
+	var aw *Writer
+	if aw, err = WriteAtomHeader(w, "traf"); err != nil {
+		return
+	}
+	w = aw
+	if self.Header != nil {
+		if err = WriteTrackFragHeader(w, self.Header); err != nil {
+			return
+		}
+	}
+	if self.DecodeTime != nil {
+		if err = WriteTrackFragDecodeTime(w, self.DecodeTime); err != nil {
+			return
+		}
+	}
+	if self.Run != nil {
+		if err = WriteTrackFragRun(w, self.Run); err != nil {
+			return
+		}
+	}
+	if err = aw.Close(); err != nil {
+		return
+	}
+	return
+}
+func WalkTrackFrag(w Walker, self *TrackFrag) {
+
+	w.StartStruct("TrackFrag")
+	if self.Header != nil {
+		WalkTrackFragHeader(w, self.Header)
+	}
+	if self.DecodeTime != nil {
+		WalkTrackFragDecodeTime(w, self.DecodeTime)
+	}
+	if self.Run != nil {
+		WalkTrackFragRun(w, self.Run)
+	}
+	w.EndStruct()
+	return
+}