FFmpeg stream forwarding in C++

example.cpp

#include <cstdint>
#include <string>
#include <memory>

extern "C" {
    #include <libavcodec/avcodec.h>
    #include <libavformat/avformat.h>
    #include <libavutil/imgutils.h>
    #include <libavutil/samplefmt.h>
    #include <libavutil/timestamp.h>
    #include <libavutil/opt.h>
    #include <libswscale/swscale.h>
    #include <libswresample/swresample.h>
    #include <sys/time.h>
}

struct timebase_t {
    int numerator;
    int denominator;
};

enum class media_type_t {
    VIDEO,
    AUDIO
};

struct stream_desc_t {
    int64_t bit_rate;
    int width;
    int height;
    int channels;
    int sample_rate;
    timebase_t timebase;
    media_type_t media_type;
    std::string codec_name;
};

enum class packet_type_t { UNKNOWN, VIDEO, AUDIO };



class ffmpeg_input {
public:
    ffmpeg_input(std::string __url);
    ~ffmpeg_input();

    bool read(AVPacket* packet, packet_type_t* type);

    std::shared_ptr<stream_desc_t> video_desc() const;
    std::shared_ptr<stream_desc_t> audio_desc() const;

private:
    std::string m_url;

    AVFormatContext*   m_format_context;
    AVFrame*           m_frame;
    AVPacket           m_packet;

    int                m_video_stream_id;
    int                m_audio_stream_id;
    AVStream*          m_video_stream;
    AVStream*          m_audio_stream;

    bool _M_open(const std::string& __source);
    bool _M_open_input_stream(const std::string& __source);
    inline bool _M_open_media_context();
    int _M_open_codec_context(
        int*             __stream_id,
        AVFormatContext* __format_context,
        AVStream**       __stream,
        AVMediaType      __type);
    std::shared_ptr<stream_desc_t> _M_stream2desc(int __index, bool __audio) const;
};



class ffmpeg_output {
public:
    ffmpeg_output(
        std::string video_url,
        std::shared_ptr<stream_desc_t> video_desc,
        std::string audio_url,
        std::shared_ptr<stream_desc_t> audio_desc);
    ~ffmpeg_output();

    bool write_video(AVPacket* packet);
    bool write_audio(AVPacket* packet);

private:
    std::string m_video_url;
    std::string m_audio_url;

    AVOutputFormat* m_format;
    AVFormatContext* m_video_context;
    AVFormatContext* m_audio_context;
    AVStream* m_video_stream;
    AVStream* m_audio_stream;

    bool _M_open_video(
        std::string                    __url,
        std::shared_ptr<stream_desc_t> __desc,
        AVFormatContext**              __context,
        AVStream**                     __stream);
    bool _M_open_audio(
        std::string                    __url,
        std::shared_ptr<stream_desc_t> __desc,
        AVFormatContext**              __context,
        AVStream**                     __stream);
};



int main(int argc, char** argv) {
    try {
    AVPacket packet;
    packet_type_t packet_type;

    ffmpeg_input input(argv[1]);
    ffmpeg_output output(
        "rtp://localhost:5008/", input.video_desc(),
        "rtp://localhost:5006/", input.audio_desc());

    while (input.read(&packet, &packet_type)) {
        AVPacket out_packet;
        av_init_packet(&out_packet);
        out_packet.pts          = packet.pts;
        out_packet.flags        = packet.flags;
        out_packet.data         = packet.data;
        out_packet.size         = packet.size;

        switch (packet_type) {
        case packet_type_t::VIDEO: output.write_video(&out_packet); break;
        case packet_type_t::AUDIO: output.write_audio(&out_packet); break;
        default: av_packet_unref(&out_packet); break;
        }
    }
    } catch (std::runtime_error& e) {
        std::cerr << e.what() << std::endl;
    }
}


// ==================================
// ---------- ffmpeg_input ----------
// ==================================


ffmpeg_input::ffmpeg_input(std::string __url)
: m_url(__url),
  m_format_context(nullptr),
  m_frame(nullptr),
  m_video_stream_id(-1),
  m_video_stream(nullptr),
  m_audio_stream_id(-1),
  m_audio_stream(nullptr)
{
    // WARNING: Memory leak on error - dtor not called
    std::cout << "URL: " << m_url << std::endl;
    if (!_M_open(m_url))
        throw std::runtime_error("failed to open input stream");
    // status = av_read_play(m_format_context);
}


ffmpeg_input::~ffmpeg_input()
{
    if (m_format_context)
        avformat_close_input(&m_format_context);    
    if (m_frame)
        av_frame_free(&m_frame);
}


bool
ffmpeg_input::read(
    AVPacket*      __packet,
    packet_type_t* __type)
{
    packet_type_t type = packet_type_t::UNKNOWN;
    auto success = av_read_frame(m_format_context, __packet) >= 0;
    if (__packet->stream_index == m_video_stream_id)
        type = packet_type_t::VIDEO;
    else if (__packet->stream_index == m_audio_stream_id)
        type = packet_type_t::AUDIO;
    if (__type) *__type = type;
    return success;
}


std::shared_ptr<stream_desc_t>
ffmpeg_input::video_desc() const
{ return _M_stream2desc(m_video_stream_id, false); }


std::shared_ptr<stream_desc_t>
ffmpeg_input::audio_desc() const
{ return _M_stream2desc(m_audio_stream_id, true); }


bool
ffmpeg_input::_M_open(const std::string& __source)
{
    if (!_M_open_input_stream(__source)) return false;

    m_frame = av_frame_alloc();
    av_init_packet(&m_packet);
    m_packet.data = nullptr;
    m_packet.size = 0;

    return _M_open_media_context();
}


bool
ffmpeg_input::_M_open_input_stream(const std::string& __source)
{
    AVDictionary* format_options = NULL;
    av_dict_set(&format_options, "rtsp_transport", "tcp",      0);
    av_dict_set(&format_options, "stimeout",       "10000000", 0);

    bool success = true;
    auto result = 0;
    result = avformat_open_input(&m_format_context,
        __source.c_str(), NULL, &format_options);

    if (result < 0) {
        char buffer[100];
        av_strerror(result, buffer, 100);
        std::cerr << "ERROR: Could not open source " << __source
            << " [ " << result << ": " << buffer << " ]" << std::endl;
        success = false;
    }
    else {
        result = avformat_find_stream_info(m_format_context, NULL);
        if (result < 0) {
            std::cerr << "ERROR: Could not find stream information" << std::endl;
            success = false;
        }
    }

    av_dict_free(&format_options);
    return success;
}


inline bool
ffmpeg_input::_M_open_media_context()
{
    bool video_found = false;
    auto status = 0;
    
    status = _M_open_codec_context(
        &m_video_stream_id,
        m_format_context,
        &m_video_stream,
        AVMEDIA_TYPE_VIDEO);
    if (status >= 0) video_found = true;

    status = _M_open_codec_context(
        &m_audio_stream_id,
        m_format_context,
        &m_audio_stream,
        AVMEDIA_TYPE_AUDIO);
    if (status < 0) {
        std::cerr << "WARNING: failed to open audio stream" << std::endl;
        m_audio_stream_id = -1;
    }

    return video_found;
}


int
ffmpeg_input::_M_open_codec_context(
    int*             __stream_id,
    AVFormatContext* __format_context,
    AVStream**       __stream,
    AVMediaType      __type)
{
    int status;

    status = av_find_best_stream(__format_context, __type, -1, -1, NULL, 0);
    if (status < 0) {
        std::cerr << "ERROR: Could not find " <<
            av_get_media_type_string(__type) << " stream" << std::endl;
        return status;
    }

    *__stream_id = status;
    *__stream = __format_context->streams[*__stream_id];

    return 0;
}


std::shared_ptr<stream_desc_t>
ffmpeg_input::_M_stream2desc(
    int  __index,
    bool __isaudio) const
{
    if (__index < 0) nullptr;

    auto desc = std::make_shared<stream_desc_t>();
    memset(desc.get(), 0, sizeof(stream_desc_t));

    const auto& stream         = m_format_context->streams[__index];
    desc->bit_rate             = stream->codecpar->bit_rate;
    if (__isaudio) {
        desc->channels         = stream->codecpar->channels;
        desc->media_type       = media_type_t::VIDEO;
    }
    else {
        desc->width            = stream->codecpar->width;
        desc->height           = stream->codecpar->height;
        desc->media_type       = media_type_t::AUDIO;
    }
    desc->sample_rate          = stream->codecpar->sample_rate;
    desc->timebase.numerator   = stream->time_base.num;
    desc->timebase.denominator = stream->time_base.den;
    desc->codec_name           = avcodec_get_name(stream->codecpar->codec_id);

    return desc;
}


// ==================================
// ---------- ffmpeg_output ----------
// ==================================


ffmpeg_output::ffmpeg_output(
    std::string                    __video_url,
    std::shared_ptr<stream_desc_t> __video_desc,
    std::string                    __audio_url,
    std::shared_ptr<stream_desc_t> __audio_desc)
: m_video_url(__video_url),
  m_audio_url(__audio_url),
  m_format(nullptr),
  m_video_context(nullptr),
  m_audio_context(nullptr),
  m_video_stream(nullptr),
  m_audio_stream(nullptr)
{
    // WARNING: Memory leak on error - dtor not called

    int status;
    m_format = av_guess_format("rtp", NULL, NULL);
    if (m_format == NULL)
        throw std::runtime_error("format not available");

    if (__video_desc) {
        auto success = _M_open_video(
            __video_url,
            __video_desc,
            &m_video_context,
            &m_video_stream);
        if (!success)
            throw std::runtime_error("Failed to open video stream");
        status = avformat_write_header(m_video_context, /*options=*/NULL);
        // if (status == AVERROR)
    }

    if (__audio_desc) {
        auto success = _M_open_audio(
            __audio_url,
            __audio_desc,
            &m_audio_context,
            &m_audio_stream);
        if (!success)
            throw std::runtime_error("Failed to open audio stream");
        status = avformat_write_header(m_audio_context, /*options=*/NULL);
        // if (status == AVERROR)
    }
}


ffmpeg_output::~ffmpeg_output()
{
    if (m_audio_context) {
        avio_close(m_audio_context->pb);
        avformat_free_context(m_audio_context);        
    }

    if (m_video_context) {
        avio_close(m_video_context->pb);
        avformat_free_context(m_video_context);
    }
}


bool
ffmpeg_output::write_video(AVPacket* __packet)
{
    if (!m_video_context) return false;
    __packet->stream_index = m_video_stream->index;
    auto status = av_interleaved_write_frame(m_video_context, __packet);
    av_packet_unref(__packet);
    return status == 0;
}


bool
ffmpeg_output::write_audio(AVPacket* __packet)
{
    if (!m_audio_context) return false;
    __packet->stream_index = m_audio_stream->index;
    auto status = av_interleaved_write_frame(m_audio_context, __packet);
    av_packet_unref(__packet);
    return status == 0;
}


bool
ffmpeg_output::_M_open_video(
    std::string                    __url,
    std::shared_ptr<stream_desc_t> __desc,
    AVFormatContext**              __context,
    AVStream**                     __stream)
{
    int status;

    // create context
    status = avformat_alloc_output_context2(
        __context,
        m_format,
        m_format->name,
        __url.c_str());
    if (status < 0) {
        std::cerr << "error: avformat_alloc_output_context2; line: " << __LINE__ << std::endl;
        return false;
    }

    // open IO interface
    status = avio_open(
        &(*__context)->pb,
        __url.c_str(),
        AVIO_FLAG_WRITE);
    if (status < 0) {
        std::cerr << "error: avio_open; line: " << __LINE__ << std::endl;
        return false;
    }

    // set stream description
    AVCodec* codec = avcodec_find_encoder_by_name(__desc->codec_name.c_str());
    if (!codec) {
        auto name = __desc->codec_name;
        std::transform(name.begin(), name.end(), name.begin(),
            [](unsigned char c){ return std::tolower(c); });
        if (name == "h264") name = "libopenh264";
        codec = avcodec_find_encoder_by_name(name.c_str());
        if (!codec) {
            std::cerr << "error: avcodec_find_encoder_by_name; line: " << __LINE__ << std::endl;
            std::cerr << "Input: " << __desc->codec_name << std::endl;
            return false;
        }
    }

    *__stream = avformat_new_stream(*__context, codec);
    (*__stream)->codecpar->codec_id    = codec->id;
    (*__stream)->codecpar->bit_rate    = __desc->bit_rate;
    (*__stream)->codecpar->width       = __desc->width;
    (*__stream)->codecpar->height      = __desc->height;
    (*__stream)->codecpar->sample_rate = __desc->sample_rate;
    (*__stream)->time_base.num         = __desc->timebase.numerator;
    (*__stream)->time_base.den         = __desc->timebase.denominator;
    (*__stream)->codecpar->codec_type  = AVMEDIA_TYPE_VIDEO;

    return true;
}


bool
ffmpeg_output::_M_open_audio(
    std::string                    __url,
    std::shared_ptr<stream_desc_t> __desc,
    AVFormatContext**              __context,
    AVStream**                     __stream)
{
    int status;

    // create context
    status = avformat_alloc_output_context2(
        __context,
        m_format,
        m_format->name,
        __url.c_str());
    if (status < 0) {
        std::cerr << "error: avformat_alloc_output_context2; line: " << __LINE__ << std::endl;
        return false;
    }

    // open IO interface
    status = avio_open(
        &(*__context)->pb,
        __url.c_str(),
        AVIO_FLAG_WRITE);
    if (status < 0) {
        std::cerr << "error: avio_open; line: " << __LINE__ << std::endl;
        return false;
    }

    // set stream description
    AVCodec* codec = avcodec_find_encoder_by_name(__desc->codec_name.c_str());
    if (!codec) {
        std::cerr << "error: avcodec_find_encoder_by_name; line: " << __LINE__ << std::endl;
        std::cerr << "Input: " << __desc->codec_name << std::endl;
        return false;
    }

    *__stream = avformat_new_stream(*__context, codec);
    (*__stream)->codecpar->codec_id    = codec->id;
    (*__stream)->codecpar->bit_rate    = __desc->bit_rate;
    (*__stream)->codecpar->sample_rate = __desc->sample_rate;
    (*__stream)->codecpar->channels    = __desc->channels;
    (*__stream)->time_base.num         = __desc->timebase.numerator;
    (*__stream)->time_base.den         = __desc->timebase.denominator;
    (*__stream)->codecpar->codec_type  = AVMEDIA_TYPE_AUDIO;

    return true;
}