CodecMTS.h

Go to the documentation of this file.

#pragma once
 
#define TS_PACKET_SIZE 188
 
#ifndef MTS_WRITE_BUFFER_SIZE
#define MTS_WRITE_BUFFER_SIZE 2000
#endif
 
#include "AudioTools/AudioCodecs/AudioCodecsBase.h"
#include "AudioTools/CoreAudio/AudioTypes.h"
#include "AudioToolsConfig.h"
#include "stdlib.h"
 
namespace audio_tools {
 
enum class MTSStreamType {
  NONE = 0x00,
  VIDEO = 0x01,
  VIDEO_H262 = 0x02,
  AUDIO_MP3 = 0x03,
  AUDIO_MP3_LOW_BITRATE = 0x04,
  PRV_SECTIONS = 0x05,
  PES_PRV = 0x06,
  MHEG = 0x07,
  H222_0_DSM_CC = 0x08,
  H222_1 = 0x09,
  A = 0x0A,
  B = 0x0B,
  C = 0x0C,
  D = 0x0D,
  H222_0_AUX = 0x0E,
  AUDIO_AAC = 0x0F,
  VISUAL = 0x10,
  AUDIO_AAC_LATM = 0x11,
  SL_PES = 0x12,
  SL_SECTIONS = 0x13,
  SYNC_DOWNLOAD = 0x14,
  PES_METADATA = 0x15,
  METDATA_SECTIONS = 0x16,
  METADATA_DATA_CAROUSEL = 0x17,
  METADATA_OBJ_CAROUSEL = 0x18,
  METADATA_SYNC_DOWNLOAD = 0x19,
  IPMP = 0x1A,
  VIDEO_AVC = 0X1B,
  VIDEO_H222_0 = 0x1C,
  DCII_VIDEO = 0x80,
  AUDIO_A53 = 0x81,
  SCTE_STD_SUBTITLE = 0x82,
  SCTE_ISOCH_DATA = 0x83,
  ATSC_PROG_ID = 0x85,
  SCTE_25 = 0x86,
  AUDIO_EAC3 = 0x87,
  AUDIO_DTS_HD = 0x88,
  DVB_MPE_FEC = 0x90,
  ULE = 0x91,
  VEI = 0x92,
  ATSC_DATA_SERVICE_TABLE = 0x95,
  SCTE_IP_DATA = 0xA0,
  DCII_TEXT = 0xC0,
  ATSC_SYNC_DATA = 0xC2,
  SCTE_AYSNC_DATA = 0xC3,
  ATSC_USER_PRIV_PROG_ELEMENTS = 0xC4,
  VC1 = 0xEA,
  ATSC_USER_PRIV = 0xEB,
};
 
// enum class AACProfile : uint8_t {
//   MAIN = 0,  // AAC Main (High complexity, rarely used)
//   LC = 1,    // AAC Low Complexity (Most common)
//   SSR = 2,   // AAC Scalable Sample Rate (Rare)
//   LTP = 3    // AAC Long Term Prediction (Not widely supported)
// };
 
class MTSDecoder : public AudioDecoder {
 public:
  MTSDecoder() = default;
  MTSDecoder(AudioDecoder &dec) { p_dec = &dec; };
  bool begin() override {
    TRACED();
    pmt_pid = 0xFFFF; // undefined
    pes_count = 0;
    is_adts_missing = false;
    open_pes_data_size = 0;
    frame_length = 0;
    selected_stream_type = MTSStreamType::NONE;
 
    // default supported stream types
    if (stream_types.empty()) {
      addStreamType(MTSStreamType::AUDIO_AAC);
      addStreamType(MTSStreamType::AUDIO_AAC_LATM);
    }
 
    // automatically close when called multiple times
    if (is_active) {
      end();
    }
 
    if (p_dec) p_dec->begin();
    is_active = true;
    return true;
  }
 
  void end() override {
    TRACED();
    if (p_dec) p_dec->end();
    is_active = false;
  }
 
  virtual operator bool() override { return is_active; }
 
  const char *mime() { return "video/MP2T"; }
 
  size_t write(const uint8_t *data, size_t len) override {
    // only process when open
    if (!is_active) {
      TRACEE();
      return 0;
    }
 
    // free up space if the buffer can't currently hold everything
    if ((size_t)buffer.availableForWrite() < len) {
      demux();
    }
 
    // write as much as currently fits: len may exceed the total buffer
    // capacity, so a full write can never be guaranteed. Returning a
    // partial count (standard Print::write contract) lets the caller
    // resend the remainder instead of stalling forever.
    size_t toWrite = (size_t)buffer.availableForWrite() < len
                          ? (size_t)buffer.availableForWrite()
                          : len;
    if (toWrite == 0) {
      LOGI("MTSDecoder::write: Buffer full");
      return 0;
    }
    LOGI("MTSDecoder::write: %d", (int)toWrite);
    size_t result = buffer.writeArray((uint8_t *)data, toWrite);
    demux();
    return result;
  }
 
  void resizeBuffer(int size) { buffer.resize(size); }
 
  void clearStreamTypes() {
    TRACED();
    stream_types.clear();
  }
 
  void addStreamType(MTSStreamType type) {
    TRACED();
    stream_types.push_back(type);
  }
 
  bool isStreamTypeActive(MTSStreamType type) {
    for (int j = 0; j < stream_types.size(); j++) {
      if (stream_types[j] == type) return true;
    }
    return false;
  }
 
  void setOutput(AudioStream &out_stream) override {
    if (p_dec) {
      p_dec->setOutput(out_stream);
    } else {
      AudioDecoder::setOutput(out_stream);
    }
  }
 
  void setOutput(AudioOutput &out_stream) override {
    if (p_dec) {
      p_dec->setOutput(out_stream);
    } else {
      AudioDecoder::setOutput(out_stream);
    }
  }
 
  void setOutput(Print &out_stream) override {
    if (p_dec) {
      p_dec->setOutput(out_stream);
    } else {
      AudioDecoder::setOutput(out_stream);
    }
  }
 
 protected:
  bool is_active = false;
  SingleBuffer<uint8_t> buffer{MTS_WRITE_BUFFER_SIZE};
  Vector<MTSStreamType> stream_types;
  Vector<int> pids{0};
  AudioDecoder *p_dec = nullptr;
  uint16_t pmt_pid = 0xFFFF;
  // AACProfile aac_profile = AACProfile::LC;
  MTSStreamType selected_stream_type = MTSStreamType::NONE;
  int open_pes_data_size = 0;
  int frame_length = 0;
  bool is_adts_missing = false;
  size_t pes_count = 0;
 
  void addPID(uint16_t pid) {
    if (pid == 0) return;
    for (int j = 0; j < pids.size(); j++) {
      if (pids[j] == pid) return;
    }
    LOGI("-> PMT PID: 0x%04X(%d)", pid, pid);
    pids.push_back(pid);
  }
 
  void demux() {
    TRACED();
    int count = 0;
    while (parse()) {
      LOGI("demux: step #%d with PES #%d", ++count, (int)pes_count);
    }
    LOGI("Number of demux calls: %d", count);
  }
 
  int syncPos() {
    int len = buffer.available();
    if (len < TS_PACKET_SIZE) return -1;
    for (int j = 0; j < len; j++) {
      if (buffer.data()[j] == 0x47) {
        return j;
      }
    }
    return -1;
  }
 
  bool parse() {
    int pos = syncPos();
    if (pos < 0) return false;
    if (pos != 0) {
      LOGW("Sync byte not found at position 0. Skipping %d bytes", pos);
      buffer.clearArray(pos);
    }
    // not enough data left for a full packet after skipping to the sync byte
    if (buffer.available() < TS_PACKET_SIZE) return false;
 
    // parse data
    uint8_t *packet = buffer.data();
    int pid = ((packet[1] & 0x1F) << 8) | (packet[2] & 0xFF);
    LOGI("PID: 0x%04X(%d)", pid, pid);
 
    // PES contains the audio data
    if (!is_adts_missing && pids.contains(pid)) {
      parsePES(packet, pid);
    } else {
      parsePacket(packet, pid);
    }
 
    // remove processed data
    buffer.clearArray(TS_PACKET_SIZE);
    return true;
  }
 
  void parsePacket(uint8_t *packet, int pid) {
    TRACEI();
    bool payloadUnitStartIndicator = false;
 
    int payloadStart =
        getPayloadStart(packet, false, payloadUnitStartIndicator);
    int len = TS_PACKET_SIZE - payloadStart;
 
    // if we are at the beginning we start with a pat
    if (pid == 0 && payloadUnitStartIndicator) {
      pids.clear();
    }
 
    // PID 0 is for PAT
    if (pid == 0) {
      parsePAT(&packet[payloadStart], len);
    } else if (pid == pmt_pid && packet[payloadStart] == 0x02) {
      parsePMT(&packet[payloadStart], len);
    } else {
      LOGE("-> Packet ignored for PID 0x%x", pid);
    }
  }
 
  int getPayloadStart(uint8_t *packet, bool isPES,
                      bool &payloadUnitStartIndicator) {
    uint8_t adaptionField = (packet[3] & 0x30) >> 4;
    int adaptationSize = 0;
    int offset = 4;  // Start after TS header (4 bytes)
 
    // Check for adaptation field
    // 00 (0) → Invalid (should never happen).
    // 01 (1) → Payload only (no adaptation field).
    // 10 (2) → Adaptation field only (no payload).
    // 11 (3) → Adaptation field + payload.
    if (adaptionField == 0b11) {  // Adaptation field exists
      adaptationSize = packet[4] + 1;
      offset += adaptationSize;
    }
 
    // If PUSI is set, there's a pointer field (skip it)
    if (packet[1] & 0x40) {
      if (!isPES) offset += packet[offset] + 1;
      payloadUnitStartIndicator = true;
    }
 
    LOGI("Payload Unit Start Indicator (PUSI): %d", payloadUnitStartIndicator);
    LOGI("Adaption Field Control: 0x%x / size: %d", adaptionField,
         adaptationSize);
 
    return offset;
  }
 
  void parsePAT(uint8_t *pat, int len) {
    TRACEI();
    assert(pat[0] == 0);  // Program Association section
    int startOfProgramNums = 8;
    int lengthOfPATValue = 4;
    int sectionLength = ((pat[1] & 0x0F) << 8) | (pat[2] & 0xFF);
    LOGI("PAT Section Length: %d", sectionLength);
    if (sectionLength >= len) {
      LOGE("Unexpected PAT Section Length: %d", sectionLength);
      sectionLength = len;
    }
    int indexOfPids = 0;
    for (int i = startOfProgramNums; i <= sectionLength;
         i += lengthOfPATValue) {
      int program_number = ((pat[i] & 0xFF) << 8) | (pat[i + 1] & 0xFF);
      int pid = ((pat[i + 2] & 0x1F) << 8) | (pat[i + 3] & 0xFF);
      LOGI("Program Num: 0x%04X(%d) / PID: 0x%04X(%d) ", program_number,
           program_number, pid, pid);
 
      if (pmt_pid == 0xFFFF && pid >= 0x0020 && pid <= 0x1FFE) {
        pmt_pid = pid;
      }
    }
    LOGI("Using PMT PID: 0x%04X(%d)", pmt_pid, pmt_pid);
  }
 
  void parsePMT(uint8_t *pmt, int len) {
    TRACEI();
    assert(pmt[0] == 0x02);  // Program Association section
    int staticLengthOfPMT = 12;
    int sectionLength = ((pmt[1] & 0x0F) << 8) | (pmt[2] & 0xFF);
    LOGI("- PMT Section Length: %d", sectionLength);
    if (sectionLength >= len) {
      LOGE("Unexpected PMT Section Length: %d", sectionLength);
      sectionLength = len;
    }
    int programInfoLength = ((pmt[10] & 0x0F) << 8) | (pmt[11] & 0xFF);
    LOGI("- PMT Program Info Length: %d", programInfoLength);
 
    int cursor = staticLengthOfPMT + programInfoLength;
    while (cursor < sectionLength - 1) {
      if (cursor + 4 >= len) break;
      MTSStreamType streamType = static_cast<MTSStreamType>(pmt[cursor] & 0xFF);
      int elementaryPID =
          ((pmt[cursor + 1] & 0x1F) << 8) | (pmt[cursor + 2] & 0xFF);
      LOGI("-- Stream Type: 0x%02X(%d) [%s] for Elementary PID: 0x%04X(%d)",
           (int)streamType, (int)streamType, toStr(streamType), elementaryPID,
           elementaryPID);
 
      if (isStreamTypeActive(streamType)) {
        selected_stream_type = streamType;
        addPID(elementaryPID);
      }
 
      int esInfoLength =
          ((pmt[cursor + 3] & 0x0F) << 8) | (pmt[cursor + 4] & 0xFF);
      LOGI("-- ES Info Length: 0x%04X(%d)", esInfoLength, esInfoLength);
      cursor += 5 + esInfoLength;
    }
  }
 
  void parsePES(uint8_t *packet, int pid) {
    LOGI("parsePES: %d", pid);
    ++pes_count;
 
    // calculate payload start
    bool payloadUnitStartIndicator = false;
    int payloadStart = getPayloadStart(packet, true, payloadUnitStartIndicator);
 
    // PES
    uint8_t *pes = packet + payloadStart;
    int len = TS_PACKET_SIZE - payloadStart;
    // PES (AAC) data
    uint8_t *pesData = nullptr;
    int pesDataSize = 0;
 
    if (payloadUnitStartIndicator) {
      if (len < 6) {
        LOGE("PES packet too short: %d", len);
        return;
      }
      // PES header is not alligned correctly
      if (!isPESStartCodeValid(pes)) {
        LOGE("PES header not aligned correctly");
        return;
      }
 
      int pesPacketLength =
          (static_cast<int>(pes[4]) << 8) | static_cast<int>(pes[5]);
 
      // PES Header size is at least 6 bytes, but can be larger with optional
      // fields. Only inspect the optional fields when they are actually
      // within the available data (pes[6..8]).
      int pesHeaderSize = 6;
      if (len >= 9 && (pes[6] & 0xC0) != 0) {  // Check for PTS/DTS flags
        pesHeaderSize += 3 + ((pes[7] & 0xC0) == 0xC0 ? 5 : 0);
        pesHeaderSize += pes[8];  // PES header stuffing size
      }
      LOGI("- PES Header Size: %d", pesHeaderSize);
 
      // pesHeaderSize is derived from stream-controlled bytes and can
      // legally encode a value that exceeds the available payload
      if (pesHeaderSize >= len) {
        LOGE("Unexpected PES Header Size: %d (len: %d)", pesHeaderSize, len);
        return;
      }
      pesData = pes + pesHeaderSize;
      pesDataSize = len - pesHeaderSize;
 
      if (pes_count == 1 && selected_stream_type == MTSStreamType::AUDIO_AAC) {
        is_adts_missing = findSyncWord(pesData, pesDataSize) == -1;
      }
 
      open_pes_data_size = pesPacketLength;
 
    } else {
      pesData = pes;
      pesDataSize = len;
    }
 
    // Recalculate the open data
    open_pes_data_size -= pesDataSize;
    if (open_pes_data_size < 0) {
      return;
    }
 
    LOGI("- writing %d bytes (open: %d)", pesDataSize, open_pes_data_size);
    if (p_print) {
      size_t result = writeData<uint8_t>(p_print, pesData, pesDataSize);
      assert(result == pesDataSize);
    }
    if (p_dec) {
      size_t result =
          writeDataT<uint8_t, AudioDecoder>(p_dec, pesData, pesDataSize);
      assert(result == pesDataSize);
    }
  }
 
  bool isPESStartCodeValid(uint8_t *pes) {
    if (pes[0] != 0) return false;
    if (pes[1] != 0) return false;
    if (pes[2] != 0x1) return false;
    return true;
  }
 
  const char *toStr(MTSStreamType type) {
    switch (type) {
      case MTSStreamType::AUDIO_MP3:
        return "AUDIO_MP3";
      case MTSStreamType::AUDIO_MP3_LOW_BITRATE:
        return "AUDIO_MP3_LOW_BITRATE";
      case MTSStreamType::AUDIO_AAC:
        return "AUDIO_AAC";
      case MTSStreamType::AUDIO_AAC_LATM:
        return "AUDIO_AAC_LATM";
      default:
        return "UNKNOWN";
    }
  }
 
  int findSyncWord(const uint8_t *buf, size_t nBytes, uint8_t synch = 0xFF,
                   uint8_t syncl = 0xF0) {
    // nBytes is unsigned: bail out up front so "nBytes - 1" can never
    // underflow into a huge loop bound below
    if (nBytes < 2) return -1;
    for (size_t i = 0; i < nBytes - 1; i++) {
      if ((buf[i + 0] & synch) == synch && (buf[i + 1] & syncl) == syncl)
        return (int)i;
    }
    return -1;
  }
};
 
using MPEG_TSDecoder = MTSDecoder;
 
}  // namespace audio_tools