Support MP4 over WebSocket

pkg/mp4/consumer.go

@@ -18,11 +18,17 @@ type Consumer struct {
 
 	muxer  *Muxer
 	codecs []*streamer.Codec
-	start  bool
+	wait   byte
 
 	send int
 }
 
+const (
+	waitNone byte = iota
+	waitKeyframe
+	waitInit
+)
+
 func (c *Consumer) GetMedias() []*streamer.Media {
 	if c.Medias != nil {
 		return c.Medias
@@ -55,13 +61,18 @@ func (c *Consumer) AddTrack(media *streamer.Media, track *streamer.Track) *strea
 	codec := track.Codec
 	switch codec.Name {
 	case streamer.CodecH264:
+		c.wait = waitInit
+
 		push := func(packet *rtp.Packet) error {
 			if packet.Version != h264.RTPPacketVersionAVC {
 				return nil
 			}
 
-			if !c.start {
-				return nil
+			if c.wait != waitNone {
+				if c.wait == waitInit || !h264.IsKeyframe(packet.Payload) {
+					return nil
+				}
+				c.wait = waitNone
 			}
 
 			buf := c.muxer.Marshal(trackID, packet)
@@ -82,13 +93,18 @@ func (c *Consumer) AddTrack(media *streamer.Media, track *streamer.Track) *strea
 		return track.Bind(push)
 
 	case streamer.CodecH265:
+		c.wait = waitInit
+
 		push := func(packet *rtp.Packet) error {
 			if packet.Version != h264.RTPPacketVersionAVC {
 				return nil
 			}
 
-			if !c.start {
-				return nil
+			if c.wait != waitNone {
+				if c.wait == waitInit || !h265.IsKeyframe(packet.Payload) {
+					return nil
+				}
+				c.wait = waitNone
 			}
 
 			buf := c.muxer.Marshal(trackID, packet)
@@ -107,7 +123,7 @@ func (c *Consumer) AddTrack(media *streamer.Media, track *streamer.Track) *strea
 
 	case streamer.CodecAAC:
 		push := func(packet *rtp.Packet) error {
-			if !c.start {
+			if c.wait != waitNone {
 				return nil
 			}
 
@@ -139,7 +155,9 @@ func (c *Consumer) Init() ([]byte, error) {
 }
 
 func (c *Consumer) Start() {
-	c.start = true
+	if c.wait == waitInit {
+		c.wait = waitKeyframe
+	}
 }
 
 //

pkg/mp4/muxer.go

@@ -19,8 +19,6 @@ type Muxer struct {
 	fragIndex uint32
 	dts       []uint64
 	pts       []uint32
-	//data      []byte
-	//total     int
 }
 
 func (m *Muxer) MimeType(codecs []*streamer.Codec) string {
@@ -185,10 +183,13 @@ func (m *Muxer) GetInit(codecs []*streamer.Codec) ([]byte, error) {
 	return append(FTYP(), data...), nil
 }
 
-//func (m *Muxer) Rewind() {
-//	m.dts = 0
-//	m.pts = 0
-//}
+func (m *Muxer) Reset() {
+	m.fragIndex = 0
+	for i := range m.dts {
+		m.dts[i] = 0
+		m.pts[i] = 0
+	}
+}
 
 func (m *Muxer) Marshal(trackID byte, packet *rtp.Packet) []byte {
 	run := &mp4fio.TrackFragRun{
@@ -218,15 +219,16 @@ func (m *Muxer) Marshal(trackID byte, packet *rtp.Packet) []byte {
 	}
 
 	entry := mp4io.TrackFragRunEntry{
-		//Duration: 90000,
 		Size: uint32(len(packet.Payload)),
 	}
 
 	newTime := packet.Timestamp
 	if m.pts[trackID] > 0 {
-		//m.dts += uint64(newTime - m.pts)
 		entry.Duration = newTime - m.pts[trackID]
 		m.dts[trackID] += uint64(entry.Duration)
+	} else {
+		// important, or Safari will fail with first frame
+		entry.Duration = 1
 	}
 	m.pts[trackID] = newTime
 

pkg/mp4/segment.go

@@ -7,13 +7,20 @@ import (
 	"github.com/pion/rtp"
 )
 
-type Keyframe struct {
+type Segment struct {
 	streamer.Element
 
-	MimeType string
+	Medias       []*streamer.Media
+	MimeType     string
+	OnlyKeyframe bool
 }
 
-func (c *Keyframe) GetMedias() []*streamer.Media {
+func (c *Segment) GetMedias() []*streamer.Media {
+	if c.Medias != nil {
+		return c.Medias
+	}
+
+	// default medias
 	return []*streamer.Media{
 		{
 			Kind:      streamer.KindVideo,
@@ -26,7 +33,7 @@ func (c *Keyframe) GetMedias() []*streamer.Media {
 	}
 }
 
-func (c *Keyframe) AddTrack(media *streamer.Media, track *streamer.Track) *streamer.Track {
+func (c *Segment) AddTrack(media *streamer.Media, track *streamer.Track) *streamer.Track {
 	muxer := &Muxer{}
 
 	codecs := []*streamer.Codec{track.Codec}
@@ -40,15 +47,46 @@ func (c *Keyframe) AddTrack(media *streamer.Media, track *streamer.Track) *strea
 
 	switch track.Codec.Name {
 	case streamer.CodecH264:
-		push := func(packet *rtp.Packet) error {
-			if !h264.IsKeyframe(packet.Payload) {
+		var push streamer.WriterFunc
+
+		if c.OnlyKeyframe {
+			push = func(packet *rtp.Packet) error {
+				if !h264.IsKeyframe(packet.Payload) {
+					return nil
+				}
+
+				buf := muxer.Marshal(0, packet)
+				c.Fire(append(init, buf...))
+
 				return nil
 			}
+		} else {
+			var buf []byte
 
-			buf := muxer.Marshal(0, packet)
-			c.Fire(append(init, buf...))
+			push = func(packet *rtp.Packet) error {
+				if h264.IsKeyframe(packet.Payload) {
+					// fist frame - send only IFrame
+					// other frames - send IFrame and all PFrames
+					if buf == nil {
+						buf = append(buf, init...)
+						b := muxer.Marshal(0, packet)
+						buf = append(buf, b...)
+					}
 
-			return nil
+					c.Fire(buf)
+
+					buf = buf[:0]
+					buf = append(buf, init...)
+					muxer.Reset()
+				}
+
+				if buf != nil {
+					b := muxer.Marshal(0, packet)
+					buf = append(buf, b...)
+				}
+
+				return nil
+			}
 		}
 
 		var wrapper streamer.WrapperFunc