From 5b9cb7dc3f2f1cc36a6aadb12f140d18a14ee5fd Mon Sep 17 00:00:00 2001
From: winlin <winlin@vip.126.com>
Date: Sat, 11 Apr 2020 15:11:42 +0800
Subject: [PATCH] For #307, refactor code, parse B-frame by sample.

---
 trunk/src/app/srs_app_rtc.cpp         | 26 +++--------------
 trunk/src/kernel/srs_kernel_codec.cpp | 42 +++++++++++++++++++++++++++
 trunk/src/kernel/srs_kernel_codec.hpp |  5 ++++
 3 files changed, 51 insertions(+), 22 deletions(-)

diff --git a/trunk/src/app/srs_app_rtc.cpp b/trunk/src/app/srs_app_rtc.cpp
index c44d203df..ff3633b20 100644
--- a/trunk/src/app/srs_app_rtc.cpp
+++ b/trunk/src/app/srs_app_rtc.cpp
@@ -125,31 +125,13 @@ srs_error_t SrsRtpH264Muxer::frame_to_packet(SrsSharedPtrMessage* shared_frame,
     for (int i = 0; i < format->video->nb_samples; ++i) {
         SrsSample* sample = &format->video->samples[i];
 
-        uint8_t header = sample->bytes[0];
-        uint8_t nal_type = header & kNalTypeMask;
-
         // Because RTC does not support B-frame, so we will drop them.
         // TODO: Drop B-frame in better way, which not cause picture corruption.
-        if (discard_bframe && (nal_type == SrsAvcNaluTypeNonIDR || nal_type == SrsAvcNaluTypeDataPartitionA || nal_type == SrsAvcNaluTypeIDR)) {
-            SrsBuffer* stream = new SrsBuffer(sample->bytes, sample->size);
-            SrsAutoFree(SrsBuffer, stream);
-
-            // Skip nalu header.
-            stream->skip(1);
-
-            SrsBitBuffer bitstream(stream);
-            int32_t first_mb_in_slice = 0;
-            if ((err = srs_avc_nalu_read_uev(&bitstream, first_mb_in_slice)) != srs_success) {
-                return srs_error_wrap(err, "nalu read uev");
+        if (discard_bframe) {
+            if ((err = sample->parse_bframe()) != srs_success) {
+                return srs_error_wrap(err, "parse bframe");
             }
-
-            int32_t slice_type = 0;
-            if ((err = srs_avc_nalu_read_uev(&bitstream, slice_type)) != srs_success) {
-                return srs_error_wrap(err, "nalu read uev");
-            }
-
-            srs_verbose("nal_type=%d, slice type=%d", nal_type, slice_type);
-            if (slice_type == SrsAvcSliceTypeB || slice_type == SrsAvcSliceTypeB1) {
+            if (sample->bframe) {
                 continue;
             }
         }
diff --git a/trunk/src/kernel/srs_kernel_codec.cpp b/trunk/src/kernel/srs_kernel_codec.cpp
index 8b4b8c41e..41034e6f4 100644
--- a/trunk/src/kernel/srs_kernel_codec.cpp
+++ b/trunk/src/kernel/srs_kernel_codec.cpp
@@ -364,12 +364,53 @@ SrsSample::SrsSample()
 {
     size = 0;
     bytes = NULL;
+    bframe = false;
 }
 
 SrsSample::~SrsSample()
 {
 }
 
+srs_error_t SrsSample::parse_bframe()
+{
+    srs_error_t err = srs_success;
+
+    // H.264 nalu header type mask.
+    static uint8_t kNalTypeMask      = 0x1F;
+
+    uint8_t header = bytes[0];
+    SrsAvcNaluType nal_type = (SrsAvcNaluType)(header & kNalTypeMask);
+
+    if (nal_type != SrsAvcNaluTypeNonIDR && nal_type != SrsAvcNaluTypeDataPartitionA && nal_type != SrsAvcNaluTypeIDR) {
+        return err;
+    }
+
+    SrsBuffer* stream = new SrsBuffer(bytes, size);
+    SrsAutoFree(SrsBuffer, stream);
+
+    // Skip nalu header.
+    stream->skip(1);
+
+    SrsBitBuffer bitstream(stream);
+    int32_t first_mb_in_slice = 0;
+    if ((err = srs_avc_nalu_read_uev(&bitstream, first_mb_in_slice)) != srs_success) {
+        return srs_error_wrap(err, "nalu read uev");
+    }
+
+    int32_t slice_type_v = 0;
+    if ((err = srs_avc_nalu_read_uev(&bitstream, slice_type_v)) != srs_success) {
+        return srs_error_wrap(err, "nalu read uev");
+    }
+    SrsAvcSliceType slice_type = (SrsAvcSliceType)slice_type_v;
+
+    if (slice_type == SrsAvcSliceTypeB || slice_type == SrsAvcSliceTypeB1) {
+        bframe = true;
+        srs_verbose("nal_type=%d, slice type=%d", nal_type, slice_type);
+    }
+
+    return err;
+}
+
 SrsCodecConfig::SrsCodecConfig()
 {
 }
@@ -458,6 +499,7 @@ srs_error_t SrsFrame::add_sample(char* bytes, int size)
     SrsSample* sample = &samples[nb_samples++];
     sample->bytes = bytes;
     sample->size = size;
+    sample->bframe = false;
     
     return err;
 }
diff --git a/trunk/src/kernel/srs_kernel_codec.hpp b/trunk/src/kernel/srs_kernel_codec.hpp
index bffeee505..b265ccb2d 100644
--- a/trunk/src/kernel/srs_kernel_codec.hpp
+++ b/trunk/src/kernel/srs_kernel_codec.hpp
@@ -534,9 +534,14 @@ public:
     int size;
     // The ptr of unit, user must manage it.
     char* bytes;
+    // Whether is B frame.
+    bool bframe;
 public:
     SrsSample();
     virtual ~SrsSample();
+public:
+    // If we need to know whether sample is bframe, we have to parse the NALU payload.
+    virtual srs_error_t parse_bframe();
 };
 
 /**