stream-server/rtmp/flv/writer.go

package flv

import (
	"io"
	"encoding/binary"
	"stream_server/transcoder"
)

type FLVWriter struct {
	w io.WriteCloser // abstraction as a wrapper to the actual transcoder command
}

func NewFLVWriter(stream_key string) (*FLVWriter, error) {
	writer := new(FLVWriter)
	flvpipe, err := transcoder.NewTranscoder(stream_key) // give control over the ffmpeg input as a stdin pipe (will push in flv packets as they come)
	if err != nil {
		return nil, err
	}
	writer.w = flvpipe
	if err = writer.write_flv_header(); err != nil { // init its header now, might as well
		return nil, err
	}
	return writer, nil
}

func (writer *FLVWriter) Close() (error) {
	return writer.w.Close()
}

// first data tag, AMF0 message but FLV spec requires the encoded data
// just requires skipping the RTMP specific bits, then just writing as is
func (writer *FLVWriter) WriteMetadataTag(data *[]byte) (err error) {
	uint24_buf := make([]byte, 4)
	tag_header := make([]byte, 11)
	tag_header[0] = 18 // RTMP msg_type 18
	binary.BigEndian.PutUint32(uint24_buf, uint32(len((*data)[16:]))) // "@setDataFrame", null, actual data
	copy(tag_header[1:4], uint24_buf[1:])

	if _, err = writer.w.Write(tag_header); err != nil {
		return
	}
	if _, err = writer.w.Write((*data)[16:]); err != nil {
		return
	}


	tag_len_buf := make([]byte, 4)
	binary.BigEndian.PutUint32(tag_len_buf, uint32(len((*data)[16:]) + 11))
	_, err = writer.w.Write(tag_len_buf)
	return	
}

// write the actual audio + video data, same as metadata except no skipping
func (writer *FLVWriter) WriteMediaTag(data *[]byte, timestamp uint32, media_type uint8) (err error) {
	uint24_buf := make([]byte, 4)
	tag_header := make([]byte, 11)
	tag_header[0] = media_type
	data_len := uint32(len(*data))
	binary.BigEndian.PutUint32(uint24_buf, data_len)
	copy(tag_header[1:4], uint24_buf[1:])

	binary.BigEndian.PutUint32(uint24_buf, timestamp)
	copy(tag_header[4:7], uint24_buf[1:])
	tag_header[7] = uint24_buf[0]

	if _, err = writer.w.Write(tag_header); err != nil {
		return
	}
	if _, err = writer.w.Write(*data); err != nil {
		return
	}

	tag_len_buf := make([]byte, 4)
	binary.BigEndian.PutUint32(tag_len_buf, data_len + 11)
	_, err = writer.w.Write(tag_len_buf)
	return 
}

// FLV file header always 9 bytes + 4 bytes for first ghost data tag
func (writer *FLVWriter) write_flv_header() (err error) {
	header := make([]byte, 13)
	copy(header[:3], "FLV")
	header[3] = 1 // FLV version, only 1 is ever used
	header[4] = 5 // 1 and 4 -> audio OR video | 5 -> audio + video
	header[8] = 9 // len of header
	
	// flv spec requires each tag + tag_len. defines a ghost tag at start with len 0
	// likely for player seeking?

	_, err = writer.w.Write(header)
	return
}