Support multiple NAL fragments per .mkv payload

Apparently the payload for an h.264 track in an .mkv file can contain multiple
NAL fragments, we used to discard everything after the first one before, now
we preserve them all.

Change-Id: Ic7187365309f3880a3256982879a45df50db697d
related-to-bug: 5337218

1 files changed, 61 insertions, 47 deletions

diff --git a/media/libstagefright/matroska/MatroskaExtractor.cpp b/media/libstagefright/matroska/MatroskaExtractor.cpp
index ffa3356..20a25d7d 100644
--- a/media/libstagefright/matroska/MatroskaExtractor.cpp
+++ b/media/libstagefright/matroska/MatroskaExtractor.cpp
@@ -423,74 +423,88 @@ again:
     MediaBuffer *frame = *mPendingFrames.begin();
     mPendingFrames.erase(mPendingFrames.begin());
 
-    size_t size = frame->range_length();
-
     if (mType != AVC) {
         *out = frame;
 
         return OK;
     }
 
-    if (size < mNALSizeLen) {
-        frame->release();
-        frame = NULL;
+    // Each input frame contains one or more NAL fragments, each fragment
+    // is prefixed by mNALSizeLen bytes giving the fragment length,
+    // followed by a corresponding number of bytes containing the fragment.
+    // We output all these fragments into a single large buffer separated
+    // by startcodes (0x00 0x00 0x00 0x01).
+
+    const uint8_t *srcPtr =
+        (const uint8_t *)frame->data() + frame->range_offset();
+
+    size_t srcSize = frame->range_length();
+
+    size_t dstSize = 0;
+    MediaBuffer *buffer = NULL;
+    uint8_t *dstPtr = NULL;
+
+    for (int32_t pass = 0; pass < 2; ++pass) {
+        size_t srcOffset = 0;
+        size_t dstOffset = 0;
+        while (srcOffset + mNALSizeLen <= srcSize) {
+            size_t NALsize;
+            switch (mNALSizeLen) {
+                case 1: NALsize = srcPtr[srcOffset]; break;
+                case 2: NALsize = U16_AT(srcPtr + srcOffset); break;
+                case 3: NALsize = U24_AT(srcPtr + srcOffset); break;
+                case 4: NALsize = U32_AT(srcPtr + srcOffset); break;
+                default:
+                    TRESPASS();
+            }
 
-        return ERROR_MALFORMED;
-    }
+            if (srcOffset + mNALSizeLen + NALsize > srcSize) {
+                break;
+            }
 
-    // In the case of AVC content, each NAL unit is prefixed by
-    // mNALSizeLen bytes of length. We want to prefix the data with
-    // a four-byte 0x00000001 startcode instead of the length prefix.
-    // mNALSizeLen ranges from 1 through 4 bytes, so add an extra
-    // 3 bytes of padding to the buffer start.
-    static const size_t kPadding = 3;
+            if (pass == 1) {
+                memcpy(&dstPtr[dstOffset], "\x00\x00\x00\x01", 4);
 
-    MediaBuffer *buffer = new MediaBuffer(size + kPadding);
+                memcpy(&dstPtr[dstOffset + 4],
+                       &srcPtr[srcOffset + mNALSizeLen],
+                       NALsize);
+            }
 
-    int64_t timeUs;
-    CHECK(frame->meta_data()->findInt64(kKeyTime, &timeUs));
-    int32_t isSync;
-    CHECK(frame->meta_data()->findInt32(kKeyIsSyncFrame, &isSync));
+            dstOffset += 4;  // 0x00 00 00 01
+            dstOffset += NALsize;
 
-    buffer->meta_data()->setInt64(kKeyTime, timeUs);
-    buffer->meta_data()->setInt32(kKeyIsSyncFrame, isSync);
+            srcOffset += mNALSizeLen + NALsize;
+        }
 
-    memcpy((uint8_t *)buffer->data() + kPadding,
-           (const uint8_t *)frame->data() + frame->range_offset(),
-           size);
+        if (srcOffset < srcSize) {
+            // There were trailing bytes or not enough data to complete
+            // a fragment.
 
-    buffer->set_range(kPadding, size);
+            frame->release();
+            frame = NULL;
 
-    frame->release();
-    frame = NULL;
+            return ERROR_MALFORMED;
+        }
 
-    uint8_t *data = (uint8_t *)buffer->data();
+        if (pass == 0) {
+            dstSize = dstOffset;
 
-    size_t NALsize;
-    switch (mNALSizeLen) {
-        case 1: NALsize = data[kPadding]; break;
-        case 2: NALsize = U16_AT(&data[kPadding]); break;
-        case 3: NALsize = U24_AT(&data[kPadding]); break;
-        case 4: NALsize = U32_AT(&data[kPadding]); break;
-        default:
-            TRESPASS();
-    }
+            buffer = new MediaBuffer(dstSize);
 
-    if (size < NALsize + mNALSizeLen) {
-        buffer->release();
-        buffer = NULL;
+            int64_t timeUs;
+            CHECK(frame->meta_data()->findInt64(kKeyTime, &timeUs));
+            int32_t isSync;
+            CHECK(frame->meta_data()->findInt32(kKeyIsSyncFrame, &isSync));
 
-        return ERROR_MALFORMED;
-    }
+            buffer->meta_data()->setInt64(kKeyTime, timeUs);
+            buffer->meta_data()->setInt32(kKeyIsSyncFrame, isSync);
 
-    if (size > NALsize + mNALSizeLen) {
-        LOGW("discarding %d bytes of data.", size - NALsize - mNALSizeLen);
+            dstPtr = (uint8_t *)buffer->data();
+        }
     }
 
-    // actual data starts at &data[kPadding + mNALSizeLen]
-
-    memcpy(&data[mNALSizeLen - 1], "\x00\x00\x00\x01", 4);
-    buffer->set_range(mNALSizeLen - 1, NALsize + 4);
+    frame->release();
+    frame = NULL;
 
     *out = buffer;