1/*
2 *  Copyright (c) 2014 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 <string.h>
12
13#include "webrtc/base/logging.h"
14#include "webrtc/modules/include/module_common_types.h"
15#include "webrtc/modules/rtp_rtcp/source/byte_io.h"
16#include "webrtc/modules/rtp_rtcp/source/h264_sps_parser.h"
17#include "webrtc/modules/rtp_rtcp/source/rtp_format_h264.h"
18
19namespace webrtc {
20namespace {
21
22enum Nalu {
23  kSlice = 1,
24  kIdr = 5,
25  kSei = 6,
26  kSps = 7,
27  kPps = 8,
28  kStapA = 24,
29  kFuA = 28
30};
31
32static const size_t kNalHeaderSize = 1;
33static const size_t kFuAHeaderSize = 2;
34static const size_t kLengthFieldSize = 2;
35static const size_t kStapAHeaderSize = kNalHeaderSize + kLengthFieldSize;
36
37// Bit masks for FU (A and B) indicators.
38enum NalDefs { kFBit = 0x80, kNriMask = 0x60, kTypeMask = 0x1F };
39
40// Bit masks for FU (A and B) headers.
41enum FuDefs { kSBit = 0x80, kEBit = 0x40, kRBit = 0x20 };
42
43// TODO(pbos): Avoid parsing this here as well as inside the jitter buffer.
44bool VerifyStapANaluLengths(const uint8_t* nalu_ptr, size_t length_remaining) {
45  while (length_remaining > 0) {
46    // Buffer doesn't contain room for additional nalu length.
47    if (length_remaining < sizeof(uint16_t))
48      return false;
49    uint16_t nalu_size = nalu_ptr[0] << 8 | nalu_ptr[1];
50    nalu_ptr += sizeof(uint16_t);
51    length_remaining -= sizeof(uint16_t);
52    if (nalu_size > length_remaining)
53      return false;
54    nalu_ptr += nalu_size;
55    length_remaining -= nalu_size;
56  }
57  return true;
58}
59
60bool ParseSingleNalu(RtpDepacketizer::ParsedPayload* parsed_payload,
61                     const uint8_t* payload_data,
62                     size_t payload_data_length) {
63  parsed_payload->type.Video.width = 0;
64  parsed_payload->type.Video.height = 0;
65  parsed_payload->type.Video.codec = kRtpVideoH264;
66  parsed_payload->type.Video.isFirstPacket = true;
67  RTPVideoHeaderH264* h264_header =
68      &parsed_payload->type.Video.codecHeader.H264;
69
70  const uint8_t* nalu_start = payload_data + kNalHeaderSize;
71  size_t nalu_length = payload_data_length - kNalHeaderSize;
72  uint8_t nal_type = payload_data[0] & kTypeMask;
73  if (nal_type == kStapA) {
74    // Skip the StapA header (StapA nal type + length).
75    if (payload_data_length <= kStapAHeaderSize) {
76      LOG(LS_ERROR) << "StapA header truncated.";
77      return false;
78    }
79    if (!VerifyStapANaluLengths(nalu_start, nalu_length)) {
80      LOG(LS_ERROR) << "StapA packet with incorrect NALU packet lengths.";
81      return false;
82    }
83
84    nal_type = payload_data[kStapAHeaderSize] & kTypeMask;
85    nalu_start += kStapAHeaderSize;
86    nalu_length -= kStapAHeaderSize;
87    h264_header->packetization_type = kH264StapA;
88  } else {
89    h264_header->packetization_type = kH264SingleNalu;
90  }
91  h264_header->nalu_type = nal_type;
92
93  // We can read resolution out of sps packets.
94  if (nal_type == kSps) {
95    H264SpsParser parser(nalu_start, nalu_length);
96    if (parser.Parse()) {
97      parsed_payload->type.Video.width = parser.width();
98      parsed_payload->type.Video.height = parser.height();
99    }
100  }
101  switch (nal_type) {
102    case kSps:
103    case kPps:
104    case kIdr:
105      parsed_payload->frame_type = kVideoFrameKey;
106      break;
107    default:
108      parsed_payload->frame_type = kVideoFrameDelta;
109      break;
110  }
111  return true;
112}
113
114bool ParseFuaNalu(RtpDepacketizer::ParsedPayload* parsed_payload,
115                  const uint8_t* payload_data,
116                  size_t payload_data_length,
117                  size_t* offset) {
118  if (payload_data_length < kFuAHeaderSize) {
119    LOG(LS_ERROR) << "FU-A NAL units truncated.";
120    return false;
121  }
122  uint8_t fnri = payload_data[0] & (kFBit | kNriMask);
123  uint8_t original_nal_type = payload_data[1] & kTypeMask;
124  bool first_fragment = (payload_data[1] & kSBit) > 0;
125
126  uint8_t original_nal_header = fnri | original_nal_type;
127  if (first_fragment) {
128    *offset = kNalHeaderSize;
129    uint8_t* payload = const_cast<uint8_t*>(payload_data + *offset);
130    payload[0] = original_nal_header;
131  } else {
132    *offset = kFuAHeaderSize;
133  }
134
135  if (original_nal_type == kIdr) {
136    parsed_payload->frame_type = kVideoFrameKey;
137  } else {
138    parsed_payload->frame_type = kVideoFrameDelta;
139  }
140  parsed_payload->type.Video.width = 0;
141  parsed_payload->type.Video.height = 0;
142  parsed_payload->type.Video.codec = kRtpVideoH264;
143  parsed_payload->type.Video.isFirstPacket = first_fragment;
144  RTPVideoHeaderH264* h264_header =
145      &parsed_payload->type.Video.codecHeader.H264;
146  h264_header->packetization_type = kH264FuA;
147  h264_header->nalu_type = original_nal_type;
148  return true;
149}
150}  // namespace
151
152RtpPacketizerH264::RtpPacketizerH264(FrameType frame_type,
153                                     size_t max_payload_len)
154    : payload_data_(NULL),
155      payload_size_(0),
156      max_payload_len_(max_payload_len) {
157}
158
159RtpPacketizerH264::~RtpPacketizerH264() {
160}
161
162void RtpPacketizerH264::SetPayloadData(
163    const uint8_t* payload_data,
164    size_t payload_size,
165    const RTPFragmentationHeader* fragmentation) {
166  assert(packets_.empty());
167  assert(fragmentation);
168  payload_data_ = payload_data;
169  payload_size_ = payload_size;
170  fragmentation_.CopyFrom(*fragmentation);
171  GeneratePackets();
172}
173
174void RtpPacketizerH264::GeneratePackets() {
175  for (size_t i = 0; i < fragmentation_.fragmentationVectorSize;) {
176    size_t fragment_offset = fragmentation_.fragmentationOffset[i];
177    size_t fragment_length = fragmentation_.fragmentationLength[i];
178    if (fragment_length > max_payload_len_) {
179      PacketizeFuA(fragment_offset, fragment_length);
180      ++i;
181    } else {
182      i = PacketizeStapA(i, fragment_offset, fragment_length);
183    }
184  }
185}
186
187void RtpPacketizerH264::PacketizeFuA(size_t fragment_offset,
188                                     size_t fragment_length) {
189  // Fragment payload into packets (FU-A).
190  // Strip out the original header and leave room for the FU-A header.
191  fragment_length -= kNalHeaderSize;
192  size_t offset = fragment_offset + kNalHeaderSize;
193  size_t bytes_available = max_payload_len_ - kFuAHeaderSize;
194  size_t fragments =
195      (fragment_length + (bytes_available - 1)) / bytes_available;
196  size_t avg_size = (fragment_length + fragments - 1) / fragments;
197  while (fragment_length > 0) {
198    size_t packet_length = avg_size;
199    if (fragment_length < avg_size)
200      packet_length = fragment_length;
201    uint8_t header = payload_data_[fragment_offset];
202    packets_.push(Packet(offset,
203                         packet_length,
204                         offset - kNalHeaderSize == fragment_offset,
205                         fragment_length == packet_length,
206                         false,
207                         header));
208    offset += packet_length;
209    fragment_length -= packet_length;
210  }
211}
212
213int RtpPacketizerH264::PacketizeStapA(size_t fragment_index,
214                                      size_t fragment_offset,
215                                      size_t fragment_length) {
216  // Aggregate fragments into one packet (STAP-A).
217  size_t payload_size_left = max_payload_len_;
218  int aggregated_fragments = 0;
219  size_t fragment_headers_length = 0;
220  assert(payload_size_left >= fragment_length);
221  while (payload_size_left >= fragment_length + fragment_headers_length) {
222    assert(fragment_length > 0);
223    uint8_t header = payload_data_[fragment_offset];
224    packets_.push(Packet(fragment_offset,
225                         fragment_length,
226                         aggregated_fragments == 0,
227                         false,
228                         true,
229                         header));
230    payload_size_left -= fragment_length;
231    payload_size_left -= fragment_headers_length;
232
233    // Next fragment.
234    ++fragment_index;
235    if (fragment_index == fragmentation_.fragmentationVectorSize)
236      break;
237    fragment_offset = fragmentation_.fragmentationOffset[fragment_index];
238    fragment_length = fragmentation_.fragmentationLength[fragment_index];
239
240    fragment_headers_length = kLengthFieldSize;
241    // If we are going to try to aggregate more fragments into this packet
242    // we need to add the STAP-A NALU header and a length field for the first
243    // NALU of this packet.
244    if (aggregated_fragments == 0)
245      fragment_headers_length += kNalHeaderSize + kLengthFieldSize;
246    ++aggregated_fragments;
247  }
248  packets_.back().last_fragment = true;
249  return fragment_index;
250}
251
252bool RtpPacketizerH264::NextPacket(uint8_t* buffer,
253                                   size_t* bytes_to_send,
254                                   bool* last_packet) {
255  *bytes_to_send = 0;
256  if (packets_.empty()) {
257    *bytes_to_send = 0;
258    *last_packet = true;
259    return false;
260  }
261
262  Packet packet = packets_.front();
263
264  if (packet.first_fragment && packet.last_fragment) {
265    // Single NAL unit packet.
266    *bytes_to_send = packet.size;
267    memcpy(buffer, &payload_data_[packet.offset], packet.size);
268    packets_.pop();
269    assert(*bytes_to_send <= max_payload_len_);
270  } else if (packet.aggregated) {
271    NextAggregatePacket(buffer, bytes_to_send);
272    assert(*bytes_to_send <= max_payload_len_);
273  } else {
274    NextFragmentPacket(buffer, bytes_to_send);
275    assert(*bytes_to_send <= max_payload_len_);
276  }
277  *last_packet = packets_.empty();
278  return true;
279}
280
281void RtpPacketizerH264::NextAggregatePacket(uint8_t* buffer,
282                                            size_t* bytes_to_send) {
283  Packet packet = packets_.front();
284  assert(packet.first_fragment);
285  // STAP-A NALU header.
286  buffer[0] = (packet.header & (kFBit | kNriMask)) | kStapA;
287  int index = kNalHeaderSize;
288  *bytes_to_send += kNalHeaderSize;
289  while (packet.aggregated) {
290    // Add NAL unit length field.
291    ByteWriter<uint16_t>::WriteBigEndian(&buffer[index], packet.size);
292    index += kLengthFieldSize;
293    *bytes_to_send += kLengthFieldSize;
294    // Add NAL unit.
295    memcpy(&buffer[index], &payload_data_[packet.offset], packet.size);
296    index += packet.size;
297    *bytes_to_send += packet.size;
298    packets_.pop();
299    if (packet.last_fragment)
300      break;
301    packet = packets_.front();
302  }
303  assert(packet.last_fragment);
304}
305
306void RtpPacketizerH264::NextFragmentPacket(uint8_t* buffer,
307                                           size_t* bytes_to_send) {
308  Packet packet = packets_.front();
309  // NAL unit fragmented over multiple packets (FU-A).
310  // We do not send original NALU header, so it will be replaced by the
311  // FU indicator header of the first packet.
312  uint8_t fu_indicator = (packet.header & (kFBit | kNriMask)) | kFuA;
313  uint8_t fu_header = 0;
314
315  // S | E | R | 5 bit type.
316  fu_header |= (packet.first_fragment ? kSBit : 0);
317  fu_header |= (packet.last_fragment ? kEBit : 0);
318  uint8_t type = packet.header & kTypeMask;
319  fu_header |= type;
320  buffer[0] = fu_indicator;
321  buffer[1] = fu_header;
322
323  if (packet.last_fragment) {
324    *bytes_to_send = packet.size + kFuAHeaderSize;
325    memcpy(buffer + kFuAHeaderSize, &payload_data_[packet.offset], packet.size);
326  } else {
327    *bytes_to_send = packet.size + kFuAHeaderSize;
328    memcpy(buffer + kFuAHeaderSize, &payload_data_[packet.offset], packet.size);
329  }
330  packets_.pop();
331}
332
333ProtectionType RtpPacketizerH264::GetProtectionType() {
334  return kProtectedPacket;
335}
336
337StorageType RtpPacketizerH264::GetStorageType(
338    uint32_t retransmission_settings) {
339  return kAllowRetransmission;
340}
341
342std::string RtpPacketizerH264::ToString() {
343  return "RtpPacketizerH264";
344}
345
346bool RtpDepacketizerH264::Parse(ParsedPayload* parsed_payload,
347                                const uint8_t* payload_data,
348                                size_t payload_data_length) {
349  assert(parsed_payload != NULL);
350  if (payload_data_length == 0) {
351    LOG(LS_ERROR) << "Empty payload.";
352    return false;
353  }
354
355  uint8_t nal_type = payload_data[0] & kTypeMask;
356  size_t offset = 0;
357  if (nal_type == kFuA) {
358    // Fragmented NAL units (FU-A).
359    if (!ParseFuaNalu(
360            parsed_payload, payload_data, payload_data_length, &offset)) {
361      return false;
362    }
363  } else {
364    // We handle STAP-A and single NALU's the same way here. The jitter buffer
365    // will depacketize the STAP-A into NAL units later.
366    if (!ParseSingleNalu(parsed_payload, payload_data, payload_data_length))
367      return false;
368  }
369
370  parsed_payload->payload = payload_data + offset;
371  parsed_payload->payload_length = payload_data_length - offset;
372  return true;
373}
374}  // namespace webrtc
375