libdatachannel/include/rtc/nalunit.hpp

/*
 * libdatachannel streamer example
 * Copyright (c) 2020 Filip Klembara (in2core)
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef NAL_UNIT_H
#define NAL_UNIT_H

#if RTC_ENABLE_MEDIA

#include "common.hpp"

namespace rtc {

#pragma pack(push, 1)

/// Nalu header
struct RTC_CPP_EXPORT NalUnitHeader {
	bool forbiddenBit() { return _first >> 7; }
	uint8_t nri() { return _first >> 5 & 0x03; }
	uint8_t unitType() { return _first & 0x1F; }

	void setForbiddenBit(bool isSet) { _first = (_first & 0x7F) | (isSet << 7); }
	void setNRI(uint8_t nri) { _first = (_first & 0x9F) | ((nri & 0x03) << 5); }
	void setUnitType(uint8_t type) { _first = (_first & 0xE0) | (type & 0x1F); }

private:
	uint8_t _first = 0;
};

/// Nalu fragment header
struct RTC_CPP_EXPORT NalUnitFragmentHeader {
	bool isStart() { return _first >> 7; }
	bool reservedBit6() { return (_first >> 6) & 0x01; }
	bool isEnd() { return (_first >> 5) & 0x01; }
	uint8_t unitType() { return _first & 0x1F; }

	void setStart(bool isSet) { _first = (_first & 0x7F) | (isSet << 7); }
	void setEnd(bool isSet) { _first = (_first & 0xDF) | (isSet << 6); }
	void setReservedBit6(bool isSet) { _first = (_first & 0xBF) | (isSet << 5); }
	void setUnitType(uint8_t type) { _first = (_first & 0xE0) | (type & 0x1F); }

private:
	uint8_t _first = 0;
};

#pragma pack(pop)

/// Nal unit
struct RTC_CPP_EXPORT NalUnit : binary {
	NalUnit(const NalUnit &unit) = default;
	NalUnit(size_t size, bool includingHeader = true)
		: binary(size + (includingHeader ? 0 : 1)) {}

	template <typename Iterator>
	NalUnit(Iterator begin_, Iterator end_) : binary(begin_, end_) {}

	NalUnit(binary &&data) : binary(std::move(data)) {}

	NalUnit() : binary(1) {}

	bool forbiddenBit() { return header()->forbiddenBit(); }
	uint8_t nri() { return header()->nri(); }
	uint8_t unitType() { return header()->unitType(); }
	binary payload() {
		assert(size() >= 1);
		return {begin() + 1, end()};
	}

	void setForbiddenBit(bool isSet) { header()->setForbiddenBit(isSet); }
	void setNRI(uint8_t nri) { header()->setNRI(nri); }
	void setUnitType(uint8_t type) { header()->setUnitType(type); }
	void setPayload(binary payload) {
		assert(size() >= 1);
		erase(begin() + 1, end());
		insert(end(), payload.begin(), payload.end());
	}

protected:
	NalUnitHeader *header() {
		assert(size() >= 1);
		return (NalUnitHeader *)data();
	}
};

/// Nal unit fragment A
struct RTC_CPP_EXPORT NalUnitFragmentA : NalUnit {
	enum class FragmentType { Start, Middle, End };

	NalUnitFragmentA(FragmentType type, bool forbiddenBit, uint8_t nri, uint8_t unitType,
					 binary data);

	static std::vector<std::shared_ptr<NalUnitFragmentA>> fragmentsFrom(std::shared_ptr<NalUnit> nalu, uint16_t maximumFragmentSize);

	uint8_t unitType() { return fragmentHeader()->unitType(); }

	binary payload() {
		assert(size() >= 2);
		return {begin() + 2, end()};
	}

	FragmentType type() {
		if (fragmentHeader()->isStart()) {
			return FragmentType::Start;
		} else if (fragmentHeader()->isEnd()) {
			return FragmentType::End;
		} else {
			return FragmentType::Middle;
		}
	}

	void setUnitType(uint8_t type) { fragmentHeader()->setUnitType(type); }

	void setPayload(binary payload) {
		assert(size() >= 2);
		erase(begin() + 2, end());
		insert(end(), payload.begin(), payload.end());
	}

	void setFragmentType(FragmentType type);

protected:
	NalUnitHeader *fragmentIndicator() { return (NalUnitHeader *)data(); }

	NalUnitFragmentHeader *fragmentHeader() {
		return (NalUnitFragmentHeader *)fragmentIndicator() + 1;
	}

	const uint8_t nal_type_fu_A = 28;
};

class RTC_CPP_EXPORT NalUnits : public std::vector<std::shared_ptr<NalUnit>> {
public:
	static const uint16_t defaultMaximumFragmentSize = 1400;
	std::vector<std::shared_ptr<binary>> generateFragments(uint16_t maximumFragmentSize);
};

} // namespace rtc

#endif /* RTC_ENABLE_MEDIA */

#endif /* NAL_UNIT_H */