demo live stream

server/ingest/fmp4_assembler.go

@@ -0,0 +1,139 @@
+package ingest
+
+import (
+	"encoding/binary"
+	"fmt"
+)
+
+type liveChunk struct {
+	FrameType byte
+	Payload   []byte
+}
+
+type fmp4Assembler struct {
+	pending []byte
+
+	initReady bool
+	initBuf   []byte
+
+	mediaPrefix []byte
+	inFragment  bool
+	fragmentBuf []byte
+}
+
+func newFMP4Assembler() *fmp4Assembler {
+	return &fmp4Assembler{}
+}
+
+func (a *fmp4Assembler) Push(data []byte) ([]liveChunk, error) {
+	if len(data) == 0 {
+		return nil, nil
+	}
+	a.pending = append(a.pending, data...)
+
+	result := make([]liveChunk, 0, 2)
+	for {
+		box, boxType, ok, err := a.nextBox()
+		if err != nil {
+			return nil, err
+		}
+		if !ok {
+			break
+		}
+		chunks, err := a.consumeBox(boxType, box)
+		if err != nil {
+			return nil, err
+		}
+		result = append(result, chunks...)
+	}
+	return result, nil
+}
+
+func (a *fmp4Assembler) nextBox() ([]byte, string, bool, error) {
+	if len(a.pending) < 8 {
+		return nil, "", false, nil
+	}
+
+	headerLen := 8
+	size := uint64(binary.BigEndian.Uint32(a.pending[:4]))
+	if size == 1 {
+		if len(a.pending) < 16 {
+			return nil, "", false, nil
+		}
+		headerLen = 16
+		size = binary.BigEndian.Uint64(a.pending[8:16])
+	}
+	if size == 0 {
+		// 0 means box runs to stream end; wait for more data.
+		return nil, "", false, nil
+	}
+	if size < uint64(headerLen) {
+		return nil, "", false, fmt.Errorf("invalid MP4 box size %d", size)
+	}
+	if size > uint64(len(a.pending)) {
+		return nil, "", false, nil
+	}
+
+	boxLen := int(size)
+	boxType := string(a.pending[4:8])
+	box := make([]byte, boxLen)
+	copy(box, a.pending[:boxLen])
+	a.pending = a.pending[boxLen:]
+	return box, boxType, true, nil
+}
+
+func (a *fmp4Assembler) consumeBox(boxType string, box []byte) ([]liveChunk, error) {
+	if !a.initReady {
+		a.initBuf = append(a.initBuf, box...)
+		if boxType != "moov" {
+			return nil, nil
+		}
+
+		a.initReady = true
+		initPayload := cloneBytes(a.initBuf)
+		a.initBuf = nil
+		return []liveChunk{{FrameType: liveFrameTypeInit, Payload: initPayload}}, nil
+	}
+
+	switch boxType {
+	case "styp", "sidx", "prft":
+		if a.inFragment {
+			a.fragmentBuf = append(a.fragmentBuf, box...)
+		} else {
+			a.mediaPrefix = append(a.mediaPrefix, box...)
+		}
+		return nil, nil
+	case "moof":
+		a.fragmentBuf = a.fragmentBuf[:0]
+		if len(a.mediaPrefix) > 0 {
+			a.fragmentBuf = append(a.fragmentBuf, a.mediaPrefix...)
+			a.mediaPrefix = a.mediaPrefix[:0]
+		}
+		a.fragmentBuf = append(a.fragmentBuf, box...)
+		a.inFragment = true
+		return nil, nil
+	case "mdat":
+		if !a.inFragment {
+			return nil, nil
+		}
+		a.fragmentBuf = append(a.fragmentBuf, box...)
+		mediaPayload := cloneBytes(a.fragmentBuf)
+		a.fragmentBuf = a.fragmentBuf[:0]
+		a.inFragment = false
+		return []liveChunk{{FrameType: liveFrameTypeMedia, Payload: mediaPayload}}, nil
+	default:
+		if a.inFragment {
+			a.fragmentBuf = append(a.fragmentBuf, box...)
+		}
+		return nil, nil
+	}
+}
+
+func cloneBytes(data []byte) []byte {
+	if len(data) == 0 {
+		return nil
+	}
+	out := make([]byte, len(data))
+	copy(out, data)
+	return out
+}