Index: mpegtsenc.c
===================================================================
--- mpegtsenc.c	(revision 2130)
+++ mpegtsenc.c	(working copy)
@@ -151,7 +151,26 @@
    0xbd         /**< CODEC_TYPE_NB       */
 };
 
+/**
+ *  define a struct to collect information for AVC descriptors
+ *  users can add new fields in it
+ */
+typedef struct MpegTSH264Desc{
+    int profile_idc;
+    int constrain_flag;
+    int level_idc;
+    int hrd_management_valid_flag;
+    int timing_info_present_flag;
+    uint32_t num_units_in_tick;
+    uint32_t time_scale;
+    int fixed_frame_rate_flag;
+    int has_aud;
+    int sps_ready;
+    GetBitContext gb;
+}MpegTSH264Desc;
+
 typedef struct MpegTSWriteStream {
+    MpegTSH264Desc *h_des;
     PESStream pes_stream;
     int packet_size;
     int packet_number;
@@ -288,6 +307,34 @@
                 put16(&q, 1); /* ancillary page id */
             }
             break;
+        case CODEC_TYPE_VIDEO:
+            {
+                if(st->codec->codec_id == CODEC_ID_H264){
+                    //write AVC descriptor
+                    if(ts_st->h_des->sps_ready){
+                        bytestream_put_byte(&q,0x28);        /* write AVC video descriptor_tag */
+                        bytestream_put_byte(&q,4);           /* write descriptor_length */
+                        bytestream_put_byte(&q,ts_st->h_des->profile_idc);
+                        bytestream_put_byte(&q,ts_st->h_des->constrain_flag);
+                        bytestream_put_byte(&q,ts_st->h_des->level_idc);
+                        bytestream_put_byte(&q,0xc0);        /* write still/24 hours picture */
+                        // write HRD descriptor
+                        if(ts_st->h_des->hrd_management_valid_flag){
+                            bytestream_put_byte(&q,0x2a);    /* write timing and HRD descriptor_tag */
+                            bytestream_put_byte(&q,8);       /* write descriptor_length */
+                            bytestream_put_byte(&q,0x80 & ts_st->h_des->timing_info_present_flag);
+                            if(ts_st->h_des->timing_info_present_flag){
+                                bytestream_put_byte(&q,1);   /* write 90khz_flag */
+                                bytestream_put_be32(&q,ts_st->h_des->num_units_in_tick);
+                            }
+                            if(ts_st->h_des->fixed_frame_rate_flag)
+                                bytestream_put_byte(&q,0xe0);/* write last byte */
+                            else
+                                bytestream_put_byte(&q,0x20);/* write last byte */
+                        }
+                    }
+                }
+            }
         }
 
         val = 0xf000 | (q - desc_length_ptr - 2);
@@ -721,9 +768,304 @@
     return 1;
 }
 
+void pre_decode_vui(MpegTSH264Desc * h_dsc, GetBitContext *gb)
+{
+    int aspect_ratio_info_present_flag;
+    unsigned int aspect_ratio_idc;
+    int nal_hrd_parameters_present_flag, vcl_hrd_parameters_present_flag;
+    aspect_ratio_info_present_flag= get_bits1(gb);
 
+    if(aspect_ratio_info_present_flag){
+        aspect_ratio_idc= get_bits(gb, 8);
+        if(aspect_ratio_idc == EXTENDED_SAR){
+            get_bits(gb, 16);
+            get_bits(gb, 16);
+        }else if(sizeof(pixel_aspect)/sizeof(*pixel_aspect) <= aspect_ratio_idc){
+            return -1;
+        }
+    }
+
+    if(get_bits1(gb)){      /* overscan_info_present_flag */
+        get_bits1(gb);      /* overscan_appropriate_flag */
+    }
+
+    if(get_bits1(gb)){      /* video_signal_type_present_flag */
+        get_bits(gb, 3);    /* video_format */
+        get_bits1(gb);      /* video_full_range_flag */
+        if(get_bits1(gb)){  /* colour_description_present_flag */
+            get_bits(gb, 8); /* colour_primaries */
+            get_bits(gb, 8); /* transfer_characteristics */
+            get_bits(gb, 8); /* matrix_coefficients */
+        }
+    }
+
+    if(get_bits1(gb)){      /* chroma_location_info_present_flag */
+        get_ue_golomb(gb);  /* chroma_sample_location_type_top_field */
+        get_ue_golomb(gb);  /* chroma_sample_location_type_bottom_field */
+    }
+
+    h_dsc->timing_info_present_flag = get_bits1(gb);
+    if(h_dsc->timing_info_present_flag){
+        h_dsc->num_units_in_tick = get_bits_long(gb, 32);
+        h_dsc->time_scale = get_bits_long(gb, 32);
+        h_dsc->fixed_frame_rate_flag = get_bits1(gb);
+    }
+
+    nal_hrd_parameters_present_flag = get_bits1(gb);
+    if(nal_hrd_parameters_present_flag)
+        h_dsc->hrd_management_valid_flag = 1;
+    else
+        h_dsc->hrd_management_valid_flag = 0;
+    vcl_hrd_parameters_present_flag = get_bits1(gb);
+    if(vcl_hrd_parameters_present_flag)
+        h_dsc->hrd_management_valid_flag = 1;
+    else
+        h_dsc->hrd_management_valid_flag = 0;
+    return 0;
+}
+
+void pre_decode_sps(MpegTSH264Desc *h_dsc, GetBitContext *gb)
+{
+    unsigned int tmp,mb_width,mb_height;
+    int i,poc_type,poc_cycle_length,frame_mbs_only_flag,mb_aff,crop;
+    int vui_parameters_present_flag;
+    short offset_for_ref_frame[256];
+    h_dsc->profile_idc = get_bits(gb,8);
+    h_dsc->constrain_flag = get_bits(gb,8);
+    h_dsc->level_idc = get_bits(gb,8);
+    get_ue_golomb(gb);
+    if(h_dsc->profile_idc >= 100){ //high profile
+        if(get_ue_golomb(gb) == 3) //chroma_format_idc
+            get_bits1(gb);  //residual_color_transform_flag
+        get_ue_golomb(gb);  //bit_depth_luma_minus8
+        get_ue_golomb(gb);  //bit_depth_chroma_minus8
+        get_bits1(gb);
+    }
+
+    get_ue_golomb(gb);
+    poc_type= get_ue_golomb(gb);
+
+    if(poc_type == 0){ //FIXME #define
+        get_ue_golomb(gb);
+    } else if(poc_type == 1){//FIXME #define
+        get_bits1(gb);
+        get_se_golomb(gb);
+        get_se_golomb(gb);
+        tmp= get_ue_golomb(gb);
+        if(tmp >= sizeof(offset_for_ref_frame) / sizeof(offset_for_ref_frame[0])){
+            return -1;
+        }
+        poc_cycle_length= tmp;
+
+        for(i=0; i<poc_cycle_length; i++)
+            offset_for_ref_frame[i] = get_se_golomb(gb);
+    } else if(poc_type != 2){
+        return -1;
+    }
+
+    tmp= get_ue_golomb(gb);
+    if(tmp > MAX_PICTURE_COUNT-2 || tmp >= 32){
+        return -1;
+    }
+    get_bits1(gb);
+    mb_width= get_ue_golomb(gb) + 1;
+    mb_height= get_ue_golomb(gb) + 1;
+    if(mb_width >= INT_MAX/16 || mb_height >= INT_MAX/16 ||
+        avcodec_check_dimensions(NULL, 16*mb_width, 16*mb_height)){
+        return -1;
+    }
+    frame_mbs_only_flag= get_bits1(gb);
+    if(!frame_mbs_only_flag)
+        mb_aff= get_bits1(gb);
+    else
+        mb_aff= 0;
+
+    get_bits1(gb);
+    crop= get_bits1(gb);
+    if(crop){
+        get_ue_golomb(gb);
+        get_ue_golomb(gb);
+        get_ue_golomb(gb);
+        get_ue_golomb(&gb);
+    }
+
+    vui_parameters_present_flag= get_bits1(gb);
+    if( vui_parameters_present_flag )
+        decode_vui_parameters(h_dsc, gb);
+    return 0;
+}
+
+/**
+ * identifies the exact end of the bitstream
+ * @return the length of the trailing, or 0 if damaged
+ */
+static int decode_rbsp_trailing(const uint8_t *src){
+    int v= *src;
+    int r;
+    for(r=1; r<9; r++){
+        if(v&1) return r;
+        v>>=1;
+    }
+    return 0;
+}
+
+uint8_t *pre_decode_nal(uint8_t *src, int *dst_length, int length)
+{
+    int i, si, di;
+    src++; 
+    length--;
+    for(i=0; i+1<length; i+=2){
+        if(src[i]) continue;
+        if(i>0 && src[i-1]==0) i--;
+        if(i+2<length && src[i+1]==0 && src[i+2]<=3){
+            if(src[i+2]!=3){
+                /* startcode, so we must be past the end */
+                length=i;
+            }
+            break;
+        }
+    }
+
+    if(i>=length-1){ //no escaped 0
+        *dst_length= length;
+        return src;
+    }
+
+    si=di=0;
+    while(si<length){
+        //remove escapes (very rare 1:2^22)
+        if(si+2<length && src[si]==0 && src[si+1]==0 && src[si+2]<=3){
+            if(src[si+2]==3){ //escape
+                di++;
+                di++;
+                si+=3;
+                continue;
+            }
+            break;
+        }
+    }
+    *dst_length= di;
+    return src;
+}
+
+int get_slice_type(MpegTSH264Desc *h_dsc, uint8_t *buf, int length)
+{
+    int slice_type;
+    init_get_bits(&h_dsc->gb, buf, length);
+    get_ue_golomb(&h_dsc->gb);
+    slice_type = get_ue_golomb(&h_dsc->gb);
+    if(slice_type > 9)
+        return -1;
+    return slice_type;
+}
+
+/**
+ * write aud_rbsp().
+ * @param [in] buf  the buffer to write.
+ * @return  NULL.
+ */
+void write_aud(uint8_t *buf)
+{
+    int primary_pic_type = 7;
+    *buf++ = 0x00;/* startcode */
+    *buf++ = 0x00;
+    *buf++ = 0x01;
+    *buf++ = 0x29;/* nal_ref_idc and nal_uint_type */
+    *buf++ = primary_pic_type << 5 & 0x10;
+}
+
+/**
+ * Process the packet data to get descriptor information.
+ * @param [in] h_dsc  the struct to collect descriptor information.
+ * @param [in] pkt  the packet to decode.
+ * @return  when success the return value is "0" and "-1" indicates failure.
+ */
+int pre_decodeAVC(MpegTSH264Desc *h_dsc,AVPacket *pkt)
+{
+    int slice_start,slice_type,bit_length,dst_length,buf_index=0;
+    uint8_t *ptr;
+    uint8_t *buf = pkt->data;
+    int buf_size = pkt->size;
+    int insert_aud[buf_size], slice_type_arry[buf_size];
+    uint8_t *new_buf;
+    slice_start = 0;
+    h_dsc->has_aud = 0;
+    h_dsc->sps_ready = 0;
+
+    for(;;){
+       int nal_type;
+        // start code prefix search
+        for(; buf_index + 3 < buf_size; buf_index++){
+            // This should always succeed in the first iteration.
+            if(buf[buf_index] == 0 && buf[buf_index+1] == 0 && buf[buf_index+2] == 1)
+                break;
+        }
+        if(buf_index+3 >= buf_size) break;
+        buf_index+=3;
+        nal_type = buf[buf_index] & 0x1f;
+        ptr = pre_decode_nal(buf + buf_index, &dst_length, buf_size - buf_index);
+
+        if(ptr==NULL || dst_length < 0){
+            return -1;
+        }
+        while(ptr[dst_length - 1] == 0 && dst_length > 0)
+            dst_length--;
+        bit_length= !dst_length ? 0 : (8*dst_length - decode_rbsp_trailing(ptr + dst_length - 1));
+        if(nal_type == NAL_AUD){
+            h_dsc->has_aud = 1;
+        } else if(nal_type == NAL_SLICE || nal_type == NAL_DPA || nal_type == NAL_IDR_SLICE){
+            slice_type = get_slice_type(h_dsc,ptr,bit_length);
+            if(slice_type){
+                insert_aud[slice_start] = buf_index;
+                slice_type_arry[slice_start] = slice_type;
+                slice_start++;
+            }
+        }
+        if(nal_type == NAL_SPS){
+            init_get_bits(&h_dsc->gb, ptr, bit_length);
+            pre_decode_sps(h_dsc, &h_dsc->gb);
+            h_dsc->sps_ready = 1;
+        }
+    }
+    // could not find AUD nal unit, instert into the stream before nal slice
+    if(!h_dsc->has_aud){
+        uint8_t *p_new_buf;
+        uint8_t *p_buf = buf;
+        new_buf = av_mallocz(buf_size + slice_start*5);
+        if(!new_buf)
+            return -1;
+        p_new_buf = new_buf;
+        for(int i = 0; i < slice_start; i++){
+            write_aud(p_new_buf);
+            memcpy(p_new_buf,p_buf,insert_aud[i]);
+            p_buf = p_buf + insert_aud[i];
+        }
+        av_free(buf);
+        pkt->data = new_buf;
+        pkt->size = buf_size + slice_start*5;
+    }
+    return 0;
+}
+
+/**
+ * Process the packet data to get descriptor information.
+ * @param [in] ctx  the AVFormatContext which contains streams.
+ * @param [in] pkt  the packet to decode.
+ * @return  NULL
+ */
+void mpegts_set_AVCdesc(AVFormatContext *ctx, AVPacket *pkt)
+{
+    MpegTSWriteStream *ts_st = ctx->priv_data;
+    MpegTSH264Desc *h_dsc = ts_st->h_des;
+    pre_decodeAVC(h_dsc,pkt);
+}
+
 static int mpegts_write_packet(AVFormatContext *ctx, AVPacket *pkt)
 {
+    int stream_index = pkt->stream_index;
+    AVStream * st = ctx->streams[stream_index];
+    if(st->codec->codec_id == CODEC_ID_H264)
+        mpegts_set_AVCdesc(ctx,pkt); 
     ff_pes_write_packet(ctx, pkt);
     for(;;){
         int ret = output_packet(ctx, 0);
