1/*
2 *  Copyright 2018 The WebRTC project authors. All Rights Reserved.
3 *
4 *  Use of this source code is governed by a BSD-style license
5 *  that can be found in the LICENSE file in the root of the source
6 *  tree. An additional intellectual property rights grant can be found
7 *  in the file PATENTS.  All contributing project authors may
8 *  be found in the AUTHORS file in the root of the source tree.
9 */
10
11#include "common_video/h264/h264_common.h"
12#include "components/video_codec/nalu_rewriter.h"
13#include "rtc_base/arraysize.h"
14#include "rtc_base/gunit.h"
15
16#import <XCTest/XCTest.h>
17
18#if TARGET_OS_IPHONE
19#import <AVFoundation/AVFoundation.h>
20#import <UIKit/UIKit.h>
21#endif
22
23@interface NaluRewriterTests : XCTestCase
24
25@end
26
27static const uint8_t NALU_TEST_DATA_0[] = {0xAA, 0xBB, 0xCC};
28static const uint8_t NALU_TEST_DATA_1[] = {0xDE, 0xAD, 0xBE, 0xEF};
29
30// clang-format off
31static const uint8_t SPS_PPS_BUFFER[] = {
32  // SPS nalu.
33  0x00, 0x00, 0x00, 0x01, 0x27, 0x42, 0x00, 0x1E, 0xAB, 0x40, 0xF0, 0x28,
34  0xD3, 0x70, 0x20, 0x20, 0x20, 0x20,
35  // PPS nalu.
36  0x00, 0x00, 0x00, 0x01, 0x28, 0xCE, 0x3C, 0x30};
37// clang-format on
38
39@implementation NaluRewriterTests
40
41- (void)testCreateVideoFormatDescription {
42  CMVideoFormatDescriptionRef description =
43      webrtc::CreateVideoFormatDescription(SPS_PPS_BUFFER, arraysize(SPS_PPS_BUFFER));
44  XCTAssertTrue(description);
45  if (description) {
46    CFRelease(description);
47    description = nullptr;
48  }
49
50  // clang-format off
51  const uint8_t sps_pps_not_at_start_buffer[] = {
52    // Add some non-SPS/PPS NALUs at the beginning
53    0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x01, 0xFF, 0x00, 0x00, 0x00, 0x01,
54    0xAB, 0x33, 0x21,
55    // SPS nalu.
56    0x00, 0x00, 0x01, 0x27, 0x42, 0x00, 0x1E, 0xAB, 0x40, 0xF0, 0x28, 0xD3,
57    0x70, 0x20, 0x20, 0x20, 0x20,
58    // PPS nalu.
59    0x00, 0x00, 0x01, 0x28, 0xCE, 0x3C, 0x30};
60  // clang-format on
61  description = webrtc::CreateVideoFormatDescription(sps_pps_not_at_start_buffer,
62                                                     arraysize(sps_pps_not_at_start_buffer));
63
64  XCTAssertTrue(description);
65
66  if (description) {
67    CFRelease(description);
68    description = nullptr;
69  }
70
71  const uint8_t other_buffer[] = {0x00, 0x00, 0x00, 0x01, 0x28};
72  XCTAssertFalse(webrtc::CreateVideoFormatDescription(other_buffer, arraysize(other_buffer)));
73}
74
75- (void)testReadEmptyInput {
76  const uint8_t annex_b_test_data[] = {0x00};
77  webrtc::AnnexBBufferReader reader(annex_b_test_data, 0);
78  const uint8_t* nalu = nullptr;
79  size_t nalu_length = 0;
80  XCTAssertEqual(0u, reader.BytesRemaining());
81  XCTAssertFalse(reader.ReadNalu(&nalu, &nalu_length));
82  XCTAssertEqual(nullptr, nalu);
83  XCTAssertEqual(0u, nalu_length);
84}
85
86- (void)testReadSingleNalu {
87  const uint8_t annex_b_test_data[] = {0x00, 0x00, 0x00, 0x01, 0xAA};
88  webrtc::AnnexBBufferReader reader(annex_b_test_data, arraysize(annex_b_test_data));
89  const uint8_t* nalu = nullptr;
90  size_t nalu_length = 0;
91  XCTAssertEqual(arraysize(annex_b_test_data), reader.BytesRemaining());
92  XCTAssertTrue(reader.ReadNalu(&nalu, &nalu_length));
93  XCTAssertEqual(annex_b_test_data + 4, nalu);
94  XCTAssertEqual(1u, nalu_length);
95  XCTAssertEqual(0u, reader.BytesRemaining());
96  XCTAssertFalse(reader.ReadNalu(&nalu, &nalu_length));
97  XCTAssertEqual(nullptr, nalu);
98  XCTAssertEqual(0u, nalu_length);
99}
100
101- (void)testReadSingleNalu3ByteHeader {
102  const uint8_t annex_b_test_data[] = {0x00, 0x00, 0x01, 0xAA};
103  webrtc::AnnexBBufferReader reader(annex_b_test_data, arraysize(annex_b_test_data));
104  const uint8_t* nalu = nullptr;
105  size_t nalu_length = 0;
106  XCTAssertEqual(arraysize(annex_b_test_data), reader.BytesRemaining());
107  XCTAssertTrue(reader.ReadNalu(&nalu, &nalu_length));
108  XCTAssertEqual(annex_b_test_data + 3, nalu);
109  XCTAssertEqual(1u, nalu_length);
110  XCTAssertEqual(0u, reader.BytesRemaining());
111  XCTAssertFalse(reader.ReadNalu(&nalu, &nalu_length));
112  XCTAssertEqual(nullptr, nalu);
113  XCTAssertEqual(0u, nalu_length);
114}
115
116- (void)testReadMissingNalu {
117  // clang-format off
118  const uint8_t annex_b_test_data[] = {0x01,
119                                       0x00, 0x01,
120                                       0x00, 0x00, 0x00, 0xFF};
121  // clang-format on
122  webrtc::AnnexBBufferReader reader(annex_b_test_data, arraysize(annex_b_test_data));
123  const uint8_t* nalu = nullptr;
124  size_t nalu_length = 0;
125  XCTAssertEqual(0u, reader.BytesRemaining());
126  XCTAssertFalse(reader.ReadNalu(&nalu, &nalu_length));
127  XCTAssertEqual(nullptr, nalu);
128  XCTAssertEqual(0u, nalu_length);
129}
130
131- (void)testReadMultipleNalus {
132  // clang-format off
133  const uint8_t annex_b_test_data[] = {0x00, 0x00, 0x00, 0x01, 0xFF,
134                                       0x01,
135                                       0x00, 0x01,
136                                       0x00, 0x00, 0x00, 0xFF,
137                                       0x00, 0x00, 0x01, 0xAA, 0xBB};
138  // clang-format on
139  webrtc::AnnexBBufferReader reader(annex_b_test_data, arraysize(annex_b_test_data));
140  const uint8_t* nalu = nullptr;
141  size_t nalu_length = 0;
142  XCTAssertEqual(arraysize(annex_b_test_data), reader.BytesRemaining());
143  XCTAssertTrue(reader.ReadNalu(&nalu, &nalu_length));
144  XCTAssertEqual(annex_b_test_data + 4, nalu);
145  XCTAssertEqual(8u, nalu_length);
146  XCTAssertEqual(5u, reader.BytesRemaining());
147  XCTAssertTrue(reader.ReadNalu(&nalu, &nalu_length));
148  XCTAssertEqual(annex_b_test_data + 15, nalu);
149  XCTAssertEqual(2u, nalu_length);
150  XCTAssertEqual(0u, reader.BytesRemaining());
151  XCTAssertFalse(reader.ReadNalu(&nalu, &nalu_length));
152  XCTAssertEqual(nullptr, nalu);
153  XCTAssertEqual(0u, nalu_length);
154}
155
156- (void)testEmptyOutputBuffer {
157  const uint8_t expected_buffer[] = {0x00};
158  const size_t buffer_size = 1;
159  std::unique_ptr<uint8_t[]> buffer(new uint8_t[buffer_size]);
160  memset(buffer.get(), 0, buffer_size);
161  webrtc::AvccBufferWriter writer(buffer.get(), 0);
162  XCTAssertEqual(0u, writer.BytesRemaining());
163  XCTAssertFalse(writer.WriteNalu(NALU_TEST_DATA_0, arraysize(NALU_TEST_DATA_0)));
164  XCTAssertEqual(0, memcmp(expected_buffer, buffer.get(), arraysize(expected_buffer)));
165}
166
167- (void)testWriteSingleNalu {
168  const uint8_t expected_buffer[] = {
169      0x00, 0x00, 0x00, 0x03, 0xAA, 0xBB, 0xCC,
170  };
171  const size_t buffer_size = arraysize(NALU_TEST_DATA_0) + 4;
172  std::unique_ptr<uint8_t[]> buffer(new uint8_t[buffer_size]);
173  webrtc::AvccBufferWriter writer(buffer.get(), buffer_size);
174  XCTAssertEqual(buffer_size, writer.BytesRemaining());
175  XCTAssertTrue(writer.WriteNalu(NALU_TEST_DATA_0, arraysize(NALU_TEST_DATA_0)));
176  XCTAssertEqual(0u, writer.BytesRemaining());
177  XCTAssertFalse(writer.WriteNalu(NALU_TEST_DATA_1, arraysize(NALU_TEST_DATA_1)));
178  XCTAssertEqual(0, memcmp(expected_buffer, buffer.get(), arraysize(expected_buffer)));
179}
180
181- (void)testWriteMultipleNalus {
182  // clang-format off
183  const uint8_t expected_buffer[] = {
184    0x00, 0x00, 0x00, 0x03, 0xAA, 0xBB, 0xCC,
185    0x00, 0x00, 0x00, 0x04, 0xDE, 0xAD, 0xBE, 0xEF
186  };
187  // clang-format on
188  const size_t buffer_size = arraysize(NALU_TEST_DATA_0) + arraysize(NALU_TEST_DATA_1) + 8;
189  std::unique_ptr<uint8_t[]> buffer(new uint8_t[buffer_size]);
190  webrtc::AvccBufferWriter writer(buffer.get(), buffer_size);
191  XCTAssertEqual(buffer_size, writer.BytesRemaining());
192  XCTAssertTrue(writer.WriteNalu(NALU_TEST_DATA_0, arraysize(NALU_TEST_DATA_0)));
193  XCTAssertEqual(buffer_size - (arraysize(NALU_TEST_DATA_0) + 4), writer.BytesRemaining());
194  XCTAssertTrue(writer.WriteNalu(NALU_TEST_DATA_1, arraysize(NALU_TEST_DATA_1)));
195  XCTAssertEqual(0u, writer.BytesRemaining());
196  XCTAssertEqual(0, memcmp(expected_buffer, buffer.get(), arraysize(expected_buffer)));
197}
198
199- (void)testOverflow {
200  const uint8_t expected_buffer[] = {0x00, 0x00, 0x00};
201  const size_t buffer_size = arraysize(NALU_TEST_DATA_0);
202  std::unique_ptr<uint8_t[]> buffer(new uint8_t[buffer_size]);
203  memset(buffer.get(), 0, buffer_size);
204  webrtc::AvccBufferWriter writer(buffer.get(), buffer_size);
205  XCTAssertEqual(buffer_size, writer.BytesRemaining());
206  XCTAssertFalse(writer.WriteNalu(NALU_TEST_DATA_0, arraysize(NALU_TEST_DATA_0)));
207  XCTAssertEqual(buffer_size, writer.BytesRemaining());
208  XCTAssertEqual(0, memcmp(expected_buffer, buffer.get(), arraysize(expected_buffer)));
209}
210
211- (void)testH264AnnexBBufferToCMSampleBuffer {
212  // clang-format off
213  const uint8_t annex_b_test_data[] = {
214    0x00,
215    0x00, 0x00, 0x01,
216    0x01, 0x00, 0x00, 0xFF, // first chunk, 4 bytes
217    0x00, 0x00, 0x01,
218    0xAA, 0xFF, // second chunk, 2 bytes
219    0x00, 0x00, 0x01,
220    0xBB};  // third chunk, 1 byte, will not fit into output array
221
222  const uint8_t expected_cmsample_data[] = {
223    0x00, 0x00, 0x00, 0x04,
224    0x01, 0x00, 0x00, 0xFF, // first chunk, 4 bytes
225    0x00, 0x00, 0x00, 0x02,
226    0xAA, 0xFF}; // second chunk, 2 bytes
227  // clang-format on
228
229  CMMemoryPoolRef memory_pool = CMMemoryPoolCreate(nil);
230  CMSampleBufferRef out_sample_buffer = nil;
231  CMVideoFormatDescriptionRef description = [self createDescription];
232
233  Boolean result = webrtc::H264AnnexBBufferToCMSampleBuffer(annex_b_test_data,
234                                                            arraysize(annex_b_test_data),
235                                                            description,
236                                                            &out_sample_buffer,
237                                                            memory_pool);
238
239  XCTAssertTrue(result);
240
241  XCTAssertEqual(description, CMSampleBufferGetFormatDescription(out_sample_buffer));
242
243  char* data_ptr = nullptr;
244  CMBlockBufferRef block_buffer = CMSampleBufferGetDataBuffer(out_sample_buffer);
245  size_t block_buffer_size = CMBlockBufferGetDataLength(block_buffer);
246  CMBlockBufferGetDataPointer(block_buffer, 0, nullptr, nullptr, &data_ptr);
247  XCTAssertEqual(block_buffer_size, arraysize(annex_b_test_data));
248
249  int data_comparison_result =
250      memcmp(expected_cmsample_data, data_ptr, arraysize(expected_cmsample_data));
251
252  XCTAssertEqual(0, data_comparison_result);
253
254  if (description) {
255    CFRelease(description);
256    description = nullptr;
257  }
258
259  CMMemoryPoolInvalidate(memory_pool);
260  CFRelease(memory_pool);
261}
262
263- (void)testH264CMSampleBufferToAnnexBBuffer {
264  // clang-format off
265  const uint8_t cmsample_data[] = {
266    0x00, 0x00, 0x00, 0x04,
267    0x01, 0x00, 0x00, 0xFF, // first chunk, 4 bytes
268    0x00, 0x00, 0x00, 0x02,
269    0xAA, 0xFF}; // second chunk, 2 bytes
270
271  const uint8_t expected_annex_b_data[] = {
272    0x00, 0x00, 0x00, 0x01,
273    0x01, 0x00, 0x00, 0xFF, // first chunk, 4 bytes
274    0x00, 0x00, 0x00, 0x01,
275    0xAA, 0xFF}; // second chunk, 2 bytes
276  // clang-format on
277
278  rtc::Buffer annexb_buffer(arraysize(cmsample_data));
279  std::unique_ptr<webrtc::RTPFragmentationHeader> out_header_ptr;
280  CMSampleBufferRef sample_buffer =
281      [self createCMSampleBufferRef:(void*)cmsample_data cmsampleSize:arraysize(cmsample_data)];
282
283  Boolean result = webrtc::H264CMSampleBufferToAnnexBBuffer(sample_buffer,
284                                                            /* is_keyframe */ false,
285                                                            &annexb_buffer,
286                                                            &out_header_ptr);
287
288  XCTAssertTrue(result);
289
290  XCTAssertEqual(arraysize(expected_annex_b_data), annexb_buffer.size());
291
292  int data_comparison_result =
293      memcmp(expected_annex_b_data, annexb_buffer.data(), arraysize(expected_annex_b_data));
294
295  XCTAssertEqual(0, data_comparison_result);
296
297  webrtc::RTPFragmentationHeader* out_header = out_header_ptr.get();
298
299  XCTAssertEqual(2, (int)out_header->Size());
300
301  XCTAssertEqual(4, (int)out_header->Offset(0));
302  XCTAssertEqual(4, (int)out_header->Length(0));
303
304  XCTAssertEqual(12, (int)out_header->Offset(1));
305  XCTAssertEqual(2, (int)out_header->Length(1));
306}
307
308- (void)testH264CMSampleBufferToAnnexBBufferWithKeyframe {
309  // clang-format off
310  const uint8_t cmsample_data[] = {
311    0x00, 0x00, 0x00, 0x04,
312    0x01, 0x00, 0x00, 0xFF, // first chunk, 4 bytes
313    0x00, 0x00, 0x00, 0x02,
314    0xAA, 0xFF}; // second chunk, 2 bytes
315
316  const uint8_t expected_annex_b_data[] = {
317    0x00, 0x00, 0x00, 0x01,
318    0x01, 0x00, 0x00, 0xFF, // first chunk, 4 bytes
319    0x00, 0x00, 0x00, 0x01,
320    0xAA, 0xFF}; // second chunk, 2 bytes
321  // clang-format on
322
323  rtc::Buffer annexb_buffer(arraysize(cmsample_data));
324  std::unique_ptr<webrtc::RTPFragmentationHeader> out_header_ptr;
325  CMSampleBufferRef sample_buffer =
326      [self createCMSampleBufferRef:(void*)cmsample_data cmsampleSize:arraysize(cmsample_data)];
327
328  Boolean result = webrtc::H264CMSampleBufferToAnnexBBuffer(sample_buffer,
329                                                            /* is_keyframe */ true,
330                                                            &annexb_buffer,
331                                                            &out_header_ptr);
332
333  XCTAssertTrue(result);
334
335  XCTAssertEqual(arraysize(SPS_PPS_BUFFER) + arraysize(expected_annex_b_data),
336                 annexb_buffer.size());
337
338  XCTAssertEqual(0, memcmp(SPS_PPS_BUFFER, annexb_buffer.data(), arraysize(SPS_PPS_BUFFER)));
339
340  XCTAssertEqual(0,
341                 memcmp(expected_annex_b_data,
342                        annexb_buffer.data() + arraysize(SPS_PPS_BUFFER),
343                        arraysize(expected_annex_b_data)));
344
345  webrtc::RTPFragmentationHeader* out_header = out_header_ptr.get();
346
347  XCTAssertEqual(4, (int)out_header->Size());
348
349  XCTAssertEqual(4, (int)out_header->Offset(0));
350  XCTAssertEqual(14, (int)out_header->Length(0));
351
352  XCTAssertEqual(22, (int)out_header->Offset(1));
353  XCTAssertEqual(4, (int)out_header->Length(1));
354
355  XCTAssertEqual(30, (int)out_header->Offset(2));
356  XCTAssertEqual(4, (int)out_header->Length(2));
357
358  XCTAssertEqual(38, (int)out_header->Offset(3));
359  XCTAssertEqual(2, (int)out_header->Length(3));
360}
361
362- (CMVideoFormatDescriptionRef)createDescription {
363  CMVideoFormatDescriptionRef description =
364      webrtc::CreateVideoFormatDescription(SPS_PPS_BUFFER, arraysize(SPS_PPS_BUFFER));
365  XCTAssertTrue(description);
366  return description;
367}
368
369- (CMSampleBufferRef)createCMSampleBufferRef:(void*)cmsampleData cmsampleSize:(size_t)cmsampleSize {
370  CMSampleBufferRef sample_buffer = nil;
371  OSStatus status;
372
373  CMVideoFormatDescriptionRef description = [self createDescription];
374  CMBlockBufferRef block_buffer = nullptr;
375
376  status = CMBlockBufferCreateWithMemoryBlock(nullptr,
377                                              cmsampleData,
378                                              cmsampleSize,
379                                              nullptr,
380                                              nullptr,
381                                              0,
382                                              cmsampleSize,
383                                              kCMBlockBufferAssureMemoryNowFlag,
384                                              &block_buffer);
385
386  status = CMSampleBufferCreate(nullptr,
387                                block_buffer,
388                                true,
389                                nullptr,
390                                nullptr,
391                                description,
392                                1,
393                                0,
394                                nullptr,
395                                0,
396                                nullptr,
397                                &sample_buffer);
398
399  return sample_buffer;
400}
401
402@end
403