Fix SPS/PPS insertion logic in MP4StreamParser.

media/formats/mp4/avc_unittest.cc

Index: media/formats/mp4/avc_unittest.cc
diff --git a/media/formats/mp4/avc_unittest.cc b/media/formats/mp4/avc_unittest.cc
index f6a1d569c11b97a830af382c4218d9558f008ce0..d0381cd3bccb7d617548b0c50c16f755e686e857 100644
--- a/media/formats/mp4/avc_unittest.cc
+++ b/media/formats/mp4/avc_unittest.cc
@@ -5,7 +5,9 @@
 #include <string.h>
 
 #include "base/basictypes.h"
+#include "base/strings/string_util.h"
 #include "media/base/stream_parser_buffer.h"
+#include "media/filters/h264_parser.h"
 #include "media/formats/mp4/avc.h"
 #include "media/formats/mp4/box_definitions.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -24,28 +26,182 @@ static const uint8 kExpectedParamSets[] = {
   0x00, 0x00, 0x00, 0x01, 0x67, 0x34,
   0x00, 0x00, 0x00, 0x01, 0x68, 0x56, 0x78};
 
+
+static H264NALU::Type StringToNALUType(const std::string& name) {
+  if (name == "P")
+    return H264NALU::kNonIDRSlice;
+
+  if (name == "I")
+    return H264NALU::kIDRSlice;
+
+  if (name == "SEI")
+    return H264NALU::kSEIMessage;
+
+  if (name == "SPS")
+    return H264NALU::kSPS;
+
+  if (name == "SPSExt")
+    return H264NALU::kSPSExt;
+
+  if (name == "PPS")
+    return H264NALU::kPPS;
+
+  if (name == "AUD")
+    return H264NALU::kAUD;
+
+  if (name == "EOSeq")
+    return H264NALU::kEOSeq;
+
+  if (name == "EOStr")
+    return H264NALU::kEOStream;
+
+  if (name == "FILL")
+    return H264NALU::kFiller;
+
+  if (name == "R14")
+    return H264NALU::kReserved14;
+
+  NOTREACHED() << "Unexpected name: " << name;
+  return H264NALU::kUnspecified;
+}
+
+static std::string NALUTypeToString(int type) {
+  switch (type) {
+    case H264NALU::kNonIDRSlice:
+      return "P";
+    case H264NALU::kSliceDataA:
+      return "SDA";
+    case H264NALU::kSliceDataB:
+      return "SDB";
+    case H264NALU::kSliceDataC:
+      return "SDC";
+    case H264NALU::kIDRSlice:
+      return "I";
+    case H264NALU::kSEIMessage:
+      return "SEI";
+    case H264NALU::kSPS:
+      return "SPS";
+    case H264NALU::kSPSExt:
+      return "SPSExt";
+    case H264NALU::kPPS:
+      return "PPS";
+    case H264NALU::kAUD:
+      return "AUD";
+    case H264NALU::kEOSeq:
+      return "EOSeq";
+    case H264NALU::kEOStream:
+      return "EOStr";
+    case H264NALU::kFiller:
+      return "FILL";
+    case H264NALU::kReserved14:
+      return "R14";
+    default:
+      NOTREACHED() << "Unexpected type: " << type;
+      return "UnsupportedType";
+  };
+
+  NOTREACHED() << "Unexpected type: " << type;
+  return "UnsupportedType";
+}
+
+void StringToAnnexB(const std::string& str, std::vector<uint8>* buffer) {
+  DCHECK(!str.empty());
+
+  std::vector<std::string> tokens;
+  EXPECT_GT(Tokenize(str, " ", &tokens), 0u);
+
+  buffer->clear();
+  for (size_t i = 0; i < tokens.size(); ++i) {
+    // Write the start code.
+    buffer->push_back(0x00);
+    buffer->push_back(0x00);
+    buffer->push_back(0x00);
+    buffer->push_back(0x01);
+
+    // Write NALU type.
+    buffer->push_back(StringToNALUType(tokens[i]));
+
+    // Write junk for the payload since the current code doesn't
+    // actually look at it.
+    buffer->push_back(0x32);
+    buffer->push_back(0x12);
+    buffer->push_back(0x67);
+  }
+}
+
+std::string AnnexBToString(const std::vector<uint8>& buffer) {
+  std::stringstream ss;
+
+  H264Parser parser;
+  parser.SetStream(&buffer[0], buffer.size());
+
+  H264NALU nalu;
+  bool first = true;
+  while (parser.AdvanceToNextNALU(&nalu) == H264Parser::kOk) {
+    if (!first)
+      ss << " ";
+    else
+      first = false;
+
+    ss << NALUTypeToString(nalu.nal_unit_type);
+  }
+  return ss.str();
+}
+
 class AVCConversionTest : public testing::TestWithParam<int> {
  protected:
-  void MakeInputForLength(int length_size, std::vector<uint8>* buf) {
-    buf->clear();
+  void WriteLength(int length_size, int length, std::vector<uint8>* buf) {
+    DCHECK_GE(length, 0);
+    DCHECK_LE(length, 255);
+
     for (int i = 1; i < length_size; i++)
       buf->push_back(0);
-    buf->push_back(sizeof(kNALU1));
+    buf->push_back(length);
+  }
+
+  void MakeInputForLength(int length_size, std::vector<uint8>* buf) {
+    buf->clear();
+
+    WriteLength(length_size, sizeof(kNALU1), buf);
     buf->insert(buf->end(), kNALU1, kNALU1 + sizeof(kNALU1));
 
-    for (int i = 1; i < length_size; i++)
-      buf->push_back(0);
-    buf->push_back(sizeof(kNALU2));
+    WriteLength(length_size, sizeof(kNALU2), buf);
     buf->insert(buf->end(), kNALU2, kNALU2 + sizeof(kNALU2));
   }
+
 };
 
 TEST_P(AVCConversionTest, ParseCorrectly) {
   std::vector<uint8> buf;
   MakeInputForLength(GetParam(), &buf);
   EXPECT_TRUE(AVC::ConvertFrameToAnnexB(GetParam(), &buf));
+  EXPECT_TRUE(AVC::IsValidAnnexB(buf));
   EXPECT_EQ(buf.size(), sizeof(kExpected));
   EXPECT_EQ(0, memcmp(kExpected, &buf[0], sizeof(kExpected)));
+  EXPECT_EQ("P SDC", AnnexBToString(buf));
+}
+
+// Intentionally write NALU sizes that are larger than the buffer.
+TEST_P(AVCConversionTest, NALUSizeTooLarge) {
+  std::vector<uint8> buf;
+  WriteLength(GetParam(), 10 * sizeof(kNALU1), &buf);
+  buf.insert(buf.end(), kNALU1, kNALU1 + sizeof(kNALU1));
+  EXPECT_FALSE(AVC::ConvertFrameToAnnexB(GetParam(), &buf));
+}
+
+TEST_P(AVCConversionTest, NALUSizeIsZero) {
+  std::vector<uint8> buf;
+  WriteLength(GetParam(), 0, &buf);
+
+  WriteLength(GetParam(), sizeof(kNALU1), &buf);
+  buf.insert(buf.end(), kNALU1, kNALU1 + sizeof(kNALU1));
+
+  WriteLength(GetParam(), 0, &buf);
+
+  WriteLength(GetParam(), sizeof(kNALU2), &buf);
+  buf.insert(buf.end(), kNALU2, kNALU2 + sizeof(kNALU2));
+
+  EXPECT_FALSE(AVC::ConvertFrameToAnnexB(GetParam(), &buf));
 }
 
 TEST_P(AVCConversionTest, ParsePartial) {
@@ -88,6 +244,104 @@ TEST_F(AVCConversionTest, ConvertConfigToAnnexB) {
   EXPECT_TRUE(AVC::ConvertConfigToAnnexB(avc_config, &buf));
   EXPECT_EQ(0, memcmp(kExpectedParamSets, &buf[0],
                       sizeof(kExpectedParamSets)));
+  EXPECT_EQ("SPS SPS PPS", AnnexBToString(buf));
+}
+
+// Verify that we can round trip string -> Annex B -> string.
+TEST_F(AVCConversionTest, StringConversionFunctions) {
+  std::string str =
+      "AUD SPS SPSExt SPS PPS SEI SEI R14 I P FILL EOSeq EOStr";
+  std::vector<uint8> buf;
+
+  StringToAnnexB(str, &buf);
+
+  EXPECT_TRUE(AVC::IsValidAnnexB(buf));
+
+  EXPECT_EQ(str, AnnexBToString(buf));
+}
+
+TEST_F(AVCConversionTest, ValidAnnexBConstructs) {
+  const char* test_cases[] = {
+    "I",
+    "I I I I",
+    "AUD I",
+    "AUD SPS PPS I",
+    "I EOSeq",
+    "I EOSeq EOStr",
+    "I EOStr",
+    "P",
+    "P P P P",
+    "AUD SPS PPS P",
+    "SEI SEI I",
+    "SEI SEI R14 I",
+    "SPS SPSExt SPS PPS I P",
+  };
+
+  for (size_t i = 0; i < arraysize(test_cases); ++i) {
+    std::vector<uint8> buf;
+    StringToAnnexB(test_cases[i], &buf);
+    EXPECT_TRUE(AVC::IsValidAnnexB(buf)) << "'" <<test_cases[i] << "' failed";
+  }
+}
+
+TEST_F(AVCConversionTest, InvalidAnnexBConstructs) {
+  const char* test_cases[] = {
+    "SPS PPS AUD I",  // Parame sets must come after AUD.
+    "PPS SPS",  // SPS must come before PPS.
+    "SPSExt SPS",  // SPS must come before SPSExt.
+    "SPS PPS SPSExt",  // SPSExt must follow an SPS.
+    "EOSeq",  // EOSeq must come after a VCL.
+    "EOStr",  // EOStr must come after a VCL.
+    "I EOStr EOSeq",  // EOSeq must come before EOStr.
+    "P P SPS",  // SPS must come before a VCL.
+    "R14 SEI",  // Reserved14-18 must come after SEI.
+    "SEI I R14",  // Reserved14-18 must come before VCL.
+  };
+
+  for (size_t i = 0; i < arraysize(test_cases); ++i) {
+    std::vector<uint8> buf;
+    StringToAnnexB(test_cases[i], &buf);
+    EXPECT_FALSE(AVC::IsValidAnnexB(buf)) << "'" <<test_cases[i] << "' failed";
+  }
+}
+
+TEST_F(AVCConversionTest, InsertParamSetsAnnexB) {
+  typedef struct {
+    const char* input;
+    const char* expected;
+  } InsertTestCases;
+
+  InsertTestCases test_cases[] = {
+    { "I", "SPS SPS PPS I" },
+    { "AUD I", "AUD SPS SPS PPS I" },
+    { "SPS PPS I", "SPS PPS I" },  // Skip insert since params already exist.
+    { "AUD SPS PPS I", "AUD SPS PPS I" },  // Skip insert since params
+                                           // already exist.
+  };
+
+  AVCDecoderConfigurationRecord avc_config;
+  avc_config.sps_list.resize(2);
+  avc_config.sps_list[0].push_back(0x67);
+  avc_config.sps_list[0].push_back(0x12);
+  avc_config.sps_list[1].push_back(0x67);
+  avc_config.sps_list[1].push_back(0x34);
+  avc_config.pps_list.resize(1);
+  avc_config.pps_list[0].push_back(0x68);
+  avc_config.pps_list[0].push_back(0x56);
+  avc_config.pps_list[0].push_back(0x78);
+
+  for (size_t i = 0; i < arraysize(test_cases); ++i) {
+    std::vector<uint8> buf;
+    StringToAnnexB(test_cases[i].input, &buf);
+
+    EXPECT_TRUE(AVC::InsertParamSetsAnnexB(avc_config, &buf))
+        << "'" << test_cases[i].input << "' insert failed.";
+    EXPECT_TRUE(AVC::IsValidAnnexB(buf))
+        << "'" << test_cases[i].input << "' created invalid AnnexB.";
+    EXPECT_EQ(test_cases[i].expected, AnnexBToString(buf))
+        << "'" << test_cases[i].input << "' generated unexpected output.";
+  }
+
 }
 
 }  // namespace mp4