srs/trunk/3rdparty/srs-bench/vendor/github.com/yapingcat/gomedia/mpeg2/pes-proto.go

package mpeg2

import (
    "fmt"
    "os"

    "github.com/yapingcat/gomedia/codec"
)

var H264_AUD_NALU []byte = []byte{0x00, 0x00, 0x00, 0x01, 0x09, 0xF0} //ffmpeg mpegtsenc.c mpegts_write_packet_internal
var H265_AUD_NALU []byte = []byte{0x00, 0x00, 0x00, 0x01, 0x46, 0x01, 0x50}

type PES_STREMA_ID int

const (
    PES_STREAM_END         PES_STREMA_ID = 0xB9
    PES_STREAM_START       PES_STREMA_ID = 0xBA
    PES_STREAM_SYSTEM_HEAD PES_STREMA_ID = 0xBB
    PES_STREAM_MAP         PES_STREMA_ID = 0xBC
    PES_STREAM_PRIVATE     PES_STREMA_ID = 0xBD
    PES_STREAM_AUDIO       PES_STREMA_ID = 0xC0
    PES_STREAM_VIDEO       PES_STREMA_ID = 0xE0
)

func findPESIDByStreamType(cid TS_STREAM_TYPE) PES_STREMA_ID {
    if cid == TS_STREAM_AAC {
        return PES_STREAM_AUDIO
    } else if cid == TS_STREAM_H264 || cid == TS_STREAM_H265 {
        return PES_STREAM_VIDEO
    } else {
        return PES_STREAM_PRIVATE
    }
}

type PesPacket struct {
    Stream_id                 uint8
    PES_packet_length         uint16
    PES_scrambling_control    uint8
    PES_priority              uint8
    Data_alignment_indicator  uint8
    Copyright                 uint8
    Original_or_copy          uint8
    PTS_DTS_flags             uint8
    ESCR_flag                 uint8
    ES_rate_flag              uint8
    DSM_trick_mode_flag       uint8
    Additional_copy_info_flag uint8
    PES_CRC_flag              uint8
    PES_extension_flag        uint8
    PES_header_data_length    uint8
    Pts                       uint64
    Dts                       uint64
    ESCR_base                 uint64
    ESCR_extension            uint16
    ES_rate                   uint32
    Trick_mode_control        uint8
    Trick_value               uint8
    Additional_copy_info      uint8
    Previous_PES_packet_CRC   uint16
    Pes_payload               []byte
    //TODO
    //if ( PES_extension_flag == '1')
    // PES_private_data_flag                uint8
    // pack_header_field_flag               uint8
    // program_packet_sequence_counter_flag uint8
    // P_STD_buffer_flag                    uint8
    // PES_extension_flag_2                 uint8
    // PES_private_data                     [16]byte
}

func NewPesPacket() *PesPacket {
    return new(PesPacket)
}

func (pkg *PesPacket) PrettyPrint(file *os.File) {
    file.WriteString(fmt.Sprintf("stream_id:%d\n", pkg.Stream_id))
    file.WriteString(fmt.Sprintf("PES_packet_length:%d\n", pkg.PES_packet_length))
    file.WriteString(fmt.Sprintf("PES_scrambling_control:%d\n", pkg.PES_scrambling_control))
    file.WriteString(fmt.Sprintf("PES_priority:%d\n", pkg.PES_priority))
    file.WriteString(fmt.Sprintf("data_alignment_indicator:%d\n", pkg.Data_alignment_indicator))
    file.WriteString(fmt.Sprintf("copyright:%d\n", pkg.Copyright))
    file.WriteString(fmt.Sprintf("original_or_copy:%d\n", pkg.Original_or_copy))
    file.WriteString(fmt.Sprintf("PTS_DTS_flags:%d\n", pkg.PTS_DTS_flags))
    file.WriteString(fmt.Sprintf("ESCR_flag:%d\n", pkg.ESCR_flag))
    file.WriteString(fmt.Sprintf("ES_rate_flag:%d\n", pkg.ES_rate_flag))
    file.WriteString(fmt.Sprintf("DSM_trick_mode_flag:%d\n", pkg.DSM_trick_mode_flag))
    file.WriteString(fmt.Sprintf("additional_copy_info_flag:%d\n", pkg.Additional_copy_info_flag))
    file.WriteString(fmt.Sprintf("PES_CRC_flag:%d\n", pkg.PES_CRC_flag))
    file.WriteString(fmt.Sprintf("PES_extension_flag:%d\n", pkg.PES_extension_flag))
    file.WriteString(fmt.Sprintf("PES_header_data_length:%d\n", pkg.PES_header_data_length))
    if pkg.PTS_DTS_flags&0x02 == 0x02 {
        file.WriteString(fmt.Sprintf("PTS:%d\n", pkg.Pts))
    }
    if pkg.PTS_DTS_flags&0x03 == 0x03 {
        file.WriteString(fmt.Sprintf("DTS:%d\n", pkg.Dts))
    }

    if pkg.ESCR_flag == 1 {
        file.WriteString(fmt.Sprintf("ESCR_base:%d\n", pkg.ESCR_base))
        file.WriteString(fmt.Sprintf("ESCR_extension:%d\n", pkg.ESCR_extension))
    }

    if pkg.ES_rate_flag == 1 {
        file.WriteString(fmt.Sprintf("ES_rate:%d\n", pkg.ES_rate))
    }

    if pkg.DSM_trick_mode_flag == 1 {
        file.WriteString(fmt.Sprintf("trick_mode_control:%d\n", pkg.Trick_mode_control))
    }

    if pkg.Additional_copy_info_flag == 1 {
        file.WriteString(fmt.Sprintf("additional_copy_info:%d\n", pkg.Additional_copy_info))
    }

    if pkg.PES_CRC_flag == 1 {
        file.WriteString(fmt.Sprintf("previous_PES_packet_CRC:%d\n", pkg.Previous_PES_packet_CRC))
    }
    file.WriteString("PES_packet_data_byte:\n")
    file.WriteString(fmt.Sprintf("  Size: %d\n", len(pkg.Pes_payload)))
    file.WriteString("  data:")
    for i := 0; i < 12 && i < len(pkg.Pes_payload); i++ {
        if i%4 == 0 {
            file.WriteString("\n")
            file.WriteString("      ")
        }
        file.WriteString(fmt.Sprintf("0x%02x ", pkg.Pes_payload[i]))
    }
    file.WriteString("\n")
}

func (pkg *PesPacket) Decode(bs *codec.BitStream) error {
    if bs.RemainBytes() < 9 {
        return errNeedMore
    }
    bs.SkipBits(24)             //packet_start_code_prefix
    pkg.Stream_id = bs.Uint8(8) //stream_id
    pkg.PES_packet_length = bs.Uint16(16)
    bs.SkipBits(2) //'10'
    pkg.PES_scrambling_control = bs.Uint8(2)
    pkg.PES_priority = bs.Uint8(1)
    pkg.Data_alignment_indicator = bs.Uint8(1)
    pkg.Copyright = bs.Uint8(1)
    pkg.Original_or_copy = bs.Uint8(1)
    pkg.PTS_DTS_flags = bs.Uint8(2)
    pkg.ESCR_flag = bs.Uint8(1)
    pkg.ES_rate_flag = bs.Uint8(1)
    pkg.DSM_trick_mode_flag = bs.Uint8(1)
    pkg.Additional_copy_info_flag = bs.Uint8(1)
    pkg.PES_CRC_flag = bs.Uint8(1)
    pkg.PES_extension_flag = bs.Uint8(1)
    pkg.PES_header_data_length = bs.Uint8(8)
    if bs.RemainBytes() < int(pkg.PES_header_data_length) {
        bs.UnRead(9 * 8)
        return errNeedMore
    }
    bs.Markdot()
    if pkg.PTS_DTS_flags&0x02 == 0x02 {
        bs.SkipBits(4)
        pkg.Pts = bs.GetBits(3)
        bs.SkipBits(1)
        pkg.Pts = (pkg.Pts << 15) | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.Pts = (pkg.Pts << 15) | bs.GetBits(15)
        bs.SkipBits(1)
    }
    if pkg.PTS_DTS_flags&0x03 == 0x03 {
        bs.SkipBits(4)
        pkg.Dts = bs.GetBits(3)
        bs.SkipBits(1)
        pkg.Dts = (pkg.Dts << 15) | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.Dts = (pkg.Dts << 15) | bs.GetBits(15)
        bs.SkipBits(1)
    } else {
        pkg.Dts = pkg.Pts
    }

    if pkg.ESCR_flag == 1 {
        bs.SkipBits(2)
        pkg.ESCR_base = bs.GetBits(3)
        bs.SkipBits(1)
        pkg.ESCR_base = (pkg.Pts << 15) | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.ESCR_base = (pkg.Pts << 15) | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.ESCR_extension = bs.Uint16(9)
        bs.SkipBits(1)
    }

    if pkg.ES_rate_flag == 1 {
        bs.SkipBits(1)
        pkg.ES_rate = bs.Uint32(22)
        bs.SkipBits(1)
    }

    if pkg.DSM_trick_mode_flag == 1 {
        pkg.Trick_mode_control = bs.Uint8(3)
        pkg.Trick_value = bs.Uint8(5)
    }

    if pkg.Additional_copy_info_flag == 1 {
        pkg.Additional_copy_info = bs.Uint8(7)
    }

    if pkg.PES_CRC_flag == 1 {
        pkg.Previous_PES_packet_CRC = bs.Uint16(16)
    }

    loc := bs.DistanceFromMarkDot()
    bs.SkipBits(int(pkg.PES_header_data_length)*8 - loc) // skip remaining header

    // the -3 bytes are the combined lengths
    // of all fields between PES_packet_length and PES_header_data_length (2 bytes)
    // and the PES_header_data_length itself (1 byte)
    dataLen := int(pkg.PES_packet_length - 3 - uint16(pkg.PES_header_data_length))

    if bs.RemainBytes() < dataLen {
        pkg.Pes_payload = bs.RemainData()
        bs.UnRead((9 + int(pkg.PES_header_data_length)) * 8)
        return errNeedMore
    }

    if pkg.PES_packet_length == 0 || bs.RemainBytes() <= dataLen {
        pkg.Pes_payload = bs.RemainData()
        bs.SkipBits(bs.RemainBits())
    } else {
        pkg.Pes_payload = bs.RemainData()[:dataLen]
        bs.SkipBits(dataLen * 8)
    }

    return nil
}

func (pkg *PesPacket) DecodeMpeg1(bs *codec.BitStream) error {
    if bs.RemainBytes() < 6 {
        return errNeedMore
    }
    bs.SkipBits(24)             //packet_start_code_prefix
    pkg.Stream_id = bs.Uint8(8) //stream_id
    pkg.PES_packet_length = bs.Uint16(16)
    if pkg.PES_packet_length != 0 && bs.RemainBytes() < int(pkg.PES_packet_length) {
        bs.UnRead(6 * 8)
        return errNeedMore
    }
    bs.Markdot()
    for bs.NextBits(8) == 0xFF {
        bs.SkipBits(8)
    }
    if bs.NextBits(2) == 0x01 {
        bs.SkipBits(16)
    }
    if bs.NextBits(4) == 0x02 {
        bs.SkipBits(4)
        pkg.Pts = bs.GetBits(3)
        bs.SkipBits(1)
        pkg.Pts = pkg.Pts<<15 | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.Pts = pkg.Pts<<15 | bs.GetBits(15)
        bs.SkipBits(1)
    } else if bs.NextBits(4) == 0x03 {
        bs.SkipBits(4)
        pkg.Pts = bs.GetBits(3)
        bs.SkipBits(1)
        pkg.Pts = pkg.Pts<<15 | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.Pts = pkg.Pts<<15 | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.Dts = bs.GetBits(3)
        bs.SkipBits(1)
        pkg.Dts = pkg.Pts<<15 | bs.GetBits(15)
        bs.SkipBits(1)
        pkg.Dts = pkg.Pts<<15 | bs.GetBits(15)
        bs.SkipBits(1)
    } else if bs.NextBits(8) == 0x0F {
        bs.SkipBits(8)
    } else {
        return errParser
    }
    loc := bs.DistanceFromMarkDot() / 8
    if pkg.PES_packet_length < uint16(loc) {
        return errParser
    }
    if pkg.PES_packet_length == 0 ||
        bs.RemainBits() <= int(pkg.PES_packet_length-uint16(loc))*8 {
        pkg.Pes_payload = bs.RemainData()
        bs.SkipBits(bs.RemainBits())
    } else {
        pkg.Pes_payload = bs.RemainData()[:pkg.PES_packet_length-uint16(loc)]
        bs.SkipBits(int(pkg.PES_packet_length-uint16(loc)) * 8)
    }
    return nil
}

func (pkg *PesPacket) Encode(bsw *codec.BitStreamWriter) {
    bsw.PutBytes([]byte{0x00, 0x00, 0x01})
    bsw.PutByte(pkg.Stream_id)
    bsw.PutUint16(pkg.PES_packet_length, 16)
    bsw.PutUint8(0x02, 2)
    bsw.PutUint8(pkg.PES_scrambling_control, 2)
    bsw.PutUint8(pkg.PES_priority, 1)
    bsw.PutUint8(pkg.Data_alignment_indicator, 1)
    bsw.PutUint8(pkg.Copyright, 1)
    bsw.PutUint8(pkg.Original_or_copy, 1)
    bsw.PutUint8(pkg.PTS_DTS_flags, 2)
    bsw.PutUint8(pkg.ESCR_flag, 1)
    bsw.PutUint8(pkg.ES_rate_flag, 1)
    bsw.PutUint8(pkg.DSM_trick_mode_flag, 1)
    bsw.PutUint8(pkg.Additional_copy_info_flag, 1)
    bsw.PutUint8(pkg.PES_CRC_flag, 1)
    bsw.PutUint8(pkg.PES_extension_flag, 1)
    bsw.PutByte(pkg.PES_header_data_length)
    if pkg.PTS_DTS_flags == 0x02 {
        bsw.PutUint8(0x02, 4)
        bsw.PutUint64(pkg.Pts>>30, 3)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.Pts>>15, 15)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.Pts, 15)
        bsw.PutUint8(0x01, 1)
    }

    if pkg.PTS_DTS_flags == 0x03 {
        bsw.PutUint8(0x03, 4)
        bsw.PutUint64(pkg.Pts>>30, 3)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.Pts>>15, 15)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.Pts, 15)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint8(0x01, 4)
        bsw.PutUint64(pkg.Dts>>30, 3)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.Dts>>15, 15)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.Dts, 15)
        bsw.PutUint8(0x01, 1)
    }

    if pkg.ESCR_flag == 1 {
        bsw.PutUint8(0x03, 2)
        bsw.PutUint64(pkg.ESCR_base>>30, 3)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.ESCR_base>>15, 15)
        bsw.PutUint8(0x01, 1)
        bsw.PutUint64(pkg.ESCR_base, 15)
        bsw.PutUint8(0x01, 1)
    }
    bsw.PutBytes(pkg.Pes_payload)
}