From 354e7436621d86731f0af4bb2a59629fdaa9259c Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Sat, 25 Jun 2016 18:23:38 +0800
Subject: [PATCH] add demuxer

---
 flv.go | 155 ++++++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 148 insertions(+), 7 deletions(-)

diff --git a/flv.go b/flv.go
index ab6cfa3..a0df656 100644
--- a/flv.go
+++ b/flv.go
@@ -1,11 +1,15 @@
 package flv
 
 import (
-	_ "fmt"
+	"time"
+	"fmt"
 	"github.com/nareix/av"
+	"github.com/nareix/codec/h264parser"
+	"github.com/nareix/codec/aacparser"
 	"github.com/nareix/pio"
 	"github.com/nareix/flv/flvio"
 	"io"
+	"bufio"
 )
 
 type Muxer struct {
@@ -19,24 +23,161 @@ func NewMuxer(w io.Writer) *Muxer {
 }
 
 func (self *Muxer) WriteHeader(streams []av.CodecData) (err error) {
-	hasVideo := false
-	hasAudio := false
+	var flags uint8
 	for _, stream := range streams {
 		if stream.Type().IsVideo() {
-			hasVideo = true
+			flags |= flvio.FILE_HAS_VIDEO
 		} else if stream.Type().IsAudio() {
-			hasAudio = true
+			flags |= flvio.FILE_HAS_AUDIO
 		}
 	}
 
-	if err = flvio.WriteFileHeader(self.pw, hasVideo, hasAudio); err != nil {
+	if err = flvio.WriteFileHeader(self.pw, flags); err != nil {
 		return
 	}
 
 	return
 }
 
-func (self *Muxer) WritePacket(i int, pkt av.Packet) (err error) {
+type flvStream struct {
+	av.CodecData
+	lastts int32
+	tm time.Duration
+}
+
+type Demuxer struct {
+	streams []*flvStream
+	videostreamidx int
+	audiostreamidx int
+	pr *pio.Reader
+}
+
+func Open(r io.Reader) (demuxer *Demuxer, err error) {
+	demuxer = NewDemuxer(r)
+	if err = demuxer.ReadHeader(); err != nil {
+		return
+	}
+	return
+}
+
+func NewDemuxer(r io.Reader) *Demuxer {
+	return &Demuxer{
+		pr: pio.NewReader(bufio.NewReaderSize(r, 128)),
+	}
+}
+
+func (self *Demuxer) ReadHeader() (err error) {
+	var flags, got uint8
+	if flags, err = flvio.ReadFileHeader(self.pr); err != nil {
+		return
+	}
+	flags &= flvio.FILE_HAS_AUDIO|flvio.FILE_HAS_VIDEO
+
+	for {
+		var _tag flvio.Tag
+		if _tag, _, err = flvio.ReadTag(self.pr); err != nil {
+			return
+		}
+
+		switch tag := _tag.(type) {
+		case *flvio.Videodata:
+			switch tag.CodecID {
+			case flvio.VIDEO_H264:
+				if tag.AVCPacketType == flvio.AVC_SEQHDR {
+					var codec h264parser.CodecData
+					if codec, err = h264parser.NewCodecDataFromAVCDecoderConfRecord(tag.Data); err != nil {
+						err = fmt.Errorf("flv: h264 seqhdr invalid")
+						return
+					}
+					self.videostreamidx = len(self.streams)
+					self.streams = append(self.streams, &flvStream{CodecData: codec})
+					got |= flvio.FILE_HAS_VIDEO
+				}
+
+			default:
+				err = fmt.Errorf("flv: unspported video CodecID=%d", tag.CodecID)
+				return
+			}
+
+		case *flvio.Audiodata:
+			switch tag.SoundFormat {
+			case flvio.SOUND_AAC:
+				if tag.AACPacketType == flvio.AAC_SEQHDR {
+					var codec aacparser.CodecData
+					if codec, err = aacparser.NewCodecDataFromMPEG4AudioConfigBytes(tag.Data); err != nil {
+						err = fmt.Errorf("flv: aac seqhdr invalid")
+						return
+					}
+					self.audiostreamidx = len(self.streams)
+					self.streams = append(self.streams, &flvStream{CodecData: codec})
+					got |= flvio.FILE_HAS_AUDIO
+				}
+
+			default:
+				err = fmt.Errorf("flv: unspported audio SoundFormat=%d", tag.SoundFormat)
+				return
+			}
+		}
+
+		if got == flags {
+			break
+		}
+	}
+
+	return
+}
+
+func (self *Demuxer) Streams() (streams []av.CodecData, err error) {
+	for _, stream := range self.streams {
+		streams = append(streams, stream.CodecData)
+	}
+	return
+}
+
+func tsToTime(ts int32) time.Duration {
+	return time.Millisecond*time.Duration(ts)
+}
+
+func (self *Demuxer) ReadPacket() (pkt av.Packet, err error) {
+	var timestamp int32
+	var stream *flvStream
+
+	loop: for {
+		var _tag flvio.Tag
+		if _tag, timestamp, err = flvio.ReadTag(self.pr); err != nil {
+			return
+		}
+
+		switch tag := _tag.(type) {
+		case *flvio.Videodata:
+			if tag.AVCPacketType == flvio.AVC_NALU {
+				stream = self.streams[self.videostreamidx]
+				pkt.Idx = int8(self.videostreamidx)
+				pkt.CompositionTime = tsToTime(tag.CompositionTime)
+				pkt.IsKeyFrame = true
+				pkt.Data = tag.Data
+				break loop
+			}
+
+		case *flvio.Audiodata:
+			if tag.AACPacketType == flvio.AAC_RAW {
+				stream = self.streams[self.audiostreamidx]
+				pkt.Idx = int8(self.audiostreamidx)
+				pkt.Data = tag.Data
+				break loop
+			}
+		}
+	}
+
+	if stream.lastts == 0 {
+		stream.tm = tsToTime(timestamp)
+	} else {
+		diff := timestamp - stream.lastts
+		stream.tm += tsToTime(diff)
+	}
+	stream.lastts = timestamp
+	pkt.Time = stream.tm
+
 	return
 }