Ticket #9502: ffmpeg.c

#include <libavcodec/avcodec.h>
#include <libavdevice/avdevice.h>
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
#include <libavformat/avformat.h>
#include <libavutil/channel_layout.h>
#include <libavutil/opt.h>

#include <sys/time.h>
#include <unistd.h>

#define OUTPUT "icecast://source:hackme@host:5008/live"
#define RUNTIME 5
#define WAITTIME 1

typedef struct ffmpeg {
  AVFormatContext *ifmt_ctx;
  AVFormatContext *ofmt_ctx;
  AVCodecContext *dec_ctx;
  AVCodecContext *enc_ctx;
  AVFilterContext *buffersink_ctx;
  AVFilterContext *buffersrc_ctx;
  AVFilterGraph *filter_graph;
  AVPacket *packet;
  AVFrame *frame;
  AVFrame *filt_frame;
  int istream_index;
  int ostream_index;
} ffmpeg_t;

int open_input(ffmpeg_t *ffmpeg) {
  AVCodec *dec;
  int ret;
  char *format_name;
  char *in_filename;
  AVInputFormat *format;
  AVDictionary *opts = NULL;

  format_name = "avfoundation";
  in_filename = "none:default";

  format = av_find_input_format(format_name);

  if (!format) {
    return -1;
  }

  ret = avformat_open_input(&ffmpeg->ifmt_ctx, in_filename, format, NULL);
  if (ret < 0) {
    return ret;
  }

  ret = avformat_find_stream_info(ffmpeg->ifmt_ctx, NULL);
  if (ret < 0) {
    return ret;
  }

  /* select the audio stream */
  ret = av_find_best_stream(ffmpeg->ifmt_ctx, AVMEDIA_TYPE_AUDIO, -1, -1, &dec,
                            0);
  if (ret < 0) {
    return ret;
  }
  ffmpeg->istream_index = ret;

  /* create decoding context */
  ffmpeg->dec_ctx = avcodec_alloc_context3(dec);
  if (!ffmpeg->dec_ctx)
    return AVERROR(ENOMEM);

  avcodec_parameters_to_context(
      ffmpeg->dec_ctx,
      ffmpeg->ifmt_ctx->streams[ffmpeg->istream_index]->codecpar);

  /* init the audio decoder */
  ret = avcodec_open2(ffmpeg->dec_ctx, dec, NULL);
  if (ret < 0) {
    return ret;
  }

  return 0;
}

int init_filters(ffmpeg_t *ffmpeg) {
  int ret = 0;
  const AVFilter *abuffersrc = avfilter_get_by_name("abuffer");
  const AVFilter *abuffersink = avfilter_get_by_name("abuffersink");
  AVFilterInOut *outputs = avfilter_inout_alloc();
  AVFilterInOut *inputs = avfilter_inout_alloc();
  AVRational time_base =
      ffmpeg->ifmt_ctx->streams[ffmpeg->istream_index]->time_base;

  char filter_descr[1024];
  snprintf(filter_descr, sizeof(filter_descr),
          "aformat=sample_fmts=%s:channel_layouts=stereo,astats=metadata=true:"
          "reset=10:measure_perchannel=RMS_level:measure_overall=none,"
          "aresample=%d,asetnsamples=nb_out_samples=%d,asetpts=PTS-STARTPTS",
          av_get_sample_fmt_name(ffmpeg->enc_ctx->sample_fmt),
          ffmpeg->enc_ctx->sample_rate,
          ffmpeg->enc_ctx->frame_size);

  ffmpeg->filter_graph = avfilter_graph_alloc();
  if (!outputs || !inputs || !ffmpeg->filter_graph) {
    ret = AVERROR(ENOMEM);
    goto end;
  }

  if (!ffmpeg->dec_ctx->channel_layout)
    ffmpeg->dec_ctx->channel_layout =
        av_get_default_channel_layout(ffmpeg->dec_ctx->channels);

  char filter_args[1024];
  snprintf(
      filter_args, sizeof(filter_args),
      "time_base=%d/%d:sample_rate=%d:sample_fmt=%s:channel_layout=%" PRId64,
      time_base.num, time_base.den, ffmpeg->dec_ctx->sample_rate,
      av_get_sample_fmt_name(ffmpeg->dec_ctx->sample_fmt),
      ffmpeg->dec_ctx->channel_layout);

  ret = avfilter_graph_create_filter(&ffmpeg->buffersrc_ctx, abuffersrc, "in",
                                     filter_args, NULL, ffmpeg->filter_graph);
  if (ret < 0) {
    goto end;
  }

  /* buffer audio sink: to terminate the filter chain. */
  ret = avfilter_graph_create_filter(&ffmpeg->buffersink_ctx, abuffersink,
                                     "out", NULL, NULL, ffmpeg->filter_graph);
  if (ret < 0) {
    goto end;
  }

  /*
   * Set the endpoints for the filter graph. The filter_graph will
   * be linked to the graph described by filter_descr.
   */

  /*
   * The buffer source output must be connected to the input pad of
   * the first filter described by filter_descr; since the first
   * filter input label is not specified, it is set to "in" by
   * default.
   */
  outputs->name = av_strdup("in");
  outputs->filter_ctx = ffmpeg->buffersrc_ctx;
  outputs->pad_idx = 0;
  outputs->next = NULL;

  /*
   * The buffer sink input must be connected to the output pad of
   * the last filter described by filter_descr; since the last
   * filter output label is not specified, it is set to "out" by
   * default.
   */
  inputs->name = av_strdup("out");
  inputs->filter_ctx = ffmpeg->buffersink_ctx;
  inputs->pad_idx = 0;
  inputs->next = NULL;

  ret = avfilter_graph_parse_ptr(ffmpeg->filter_graph, filter_descr, &inputs,
                                 &outputs, NULL);
  if (ret < 0)
    goto end;

  ret = avfilter_graph_config(ffmpeg->filter_graph, NULL);

end:
  avfilter_inout_free(&inputs);
  avfilter_inout_free(&outputs);

  return ret;
}

int encode_audio_frame(ffmpeg_t *ffmpeg) {
  int ret;

  ret = avcodec_send_frame(ffmpeg->enc_ctx, ffmpeg->filt_frame);
  if (ret == AVERROR_EOF) {
    return ret;
  }

  if (ret < 0) {
    return ret;
  }

  while (ret >= 0) {
    ret = avcodec_receive_packet(ffmpeg->enc_ctx, ffmpeg->packet);
    if (ret == AVERROR(EAGAIN)) {
      return 0;
    }

    if (ret == AVERROR_EOF) {
      return AVERROR_EOF;
    }

    if (ret < 0) {
      return ret;
    }

    ffmpeg->packet->stream_index = ffmpeg->ostream_index;
    av_packet_rescale_ts(ffmpeg->packet, ffmpeg->enc_ctx->time_base,
                         ffmpeg->ofmt_ctx->streams[0]->time_base);

    ret = av_interleaved_write_frame(ffmpeg->ofmt_ctx, ffmpeg->packet);
    if (ret < 0) {
      return ret;
    }
  }

  return 0;
}

int open_output(ffmpeg_t *ffmpeg) {
  AVStream *stream = NULL;
  const AVCodec *output_codec = NULL;
  AVDictionary *codec_opts = NULL;
  AVDictionary *out_opts = NULL;
  int error = -1;

  av_dict_set(&out_opts, "content_type", "audio/mpeg", 0);
  // av_dict_set(&out_opts, "tls", "1", 0);

  /* Create a new format context for the output container format. */
  avformat_alloc_output_context2(&ffmpeg->ofmt_ctx, NULL, "mp3", OUTPUT);
  if (!ffmpeg->ofmt_ctx) {
    return AVERROR(ENOMEM);
  }

  error = avio_open2(&ffmpeg->ofmt_ctx->pb, OUTPUT, AVIO_FLAG_WRITE, NULL,
                     &out_opts);
  if (error < 0) {
    goto cleanup;
  }

  av_dict_copy(&ffmpeg->ofmt_ctx->metadata, ffmpeg->ifmt_ctx->metadata, 0);

  /* Find the encoder to be used by its name. */
  if (!(output_codec = avcodec_find_encoder_by_name("libmp3lame"))) {
    error = AVERROR_ENCODER_NOT_FOUND;
    goto cleanup;
  }

  if (!(stream = avformat_new_stream(ffmpeg->ofmt_ctx, output_codec))) {
    error = AVERROR(ENOMEM);
    goto cleanup;
  }

  stream->id = ffmpeg->ofmt_ctx->nb_streams - 1;
  ffmpeg->ostream_index = stream->index;

  ffmpeg->enc_ctx = avcodec_alloc_context3(output_codec);
  if (!ffmpeg->enc_ctx) {
    error = AVERROR(ENOMEM);
    goto cleanup;
  }

  /* Set the basic encoder parameters.
   * The input file's sample rate is used to avoid a sample rate conversion. */
  ffmpeg->enc_ctx->channel_layout = AV_CH_LAYOUT_STEREO;
  ffmpeg->enc_ctx->channels = 2;
  ffmpeg->enc_ctx->sample_rate = 44100;
  ;
  ffmpeg->enc_ctx->sample_fmt = output_codec->sample_fmts[0];
  ffmpeg->enc_ctx->time_base = (AVRational){1, ffmpeg->enc_ctx->sample_rate};
  av_dict_set(&codec_opts, "b", "128k", 0);

  /* Some container formats (like MP4) require global headers to be present.
   * Mark the encoder so that it behaves accordingly. */
  if ((ffmpeg->ofmt_ctx)->oformat->flags & AVFMT_GLOBALHEADER)
    ffmpeg->enc_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

  /* Open the encoder for the audio stream to use it later. */
  if ((error = avcodec_open2(ffmpeg->enc_ctx, output_codec, &codec_opts)) < 0) {
    goto cleanup;
  }

  stream->time_base = ffmpeg->enc_ctx->time_base;

  error = avcodec_parameters_from_context(stream->codecpar, ffmpeg->enc_ctx);
  if (error < 0) {
    goto cleanup;
  }

  error = avformat_write_header(ffmpeg->ofmt_ctx, NULL);
  if (error < 0) {
    goto cleanup;
  }

  return 0;

cleanup:
  avcodec_free_context(&ffmpeg->enc_ctx);
  avio_closep(&(ffmpeg->ofmt_ctx)->pb);
  avformat_free_context(ffmpeg->ofmt_ctx);
  ffmpeg->ofmt_ctx = NULL;
  return error;
}

int runFfmpeg() {
  int ret;
  ffmpeg_t ffmpeg;
  AVDictionaryEntry *left_volume = NULL;
  AVDictionaryEntry *right_volume = NULL;
  double volume;
  char *end_ptr;
  int first;
  double encoding_time;
  AVRational in_time_base;
  struct timeval start, current;

  gettimeofday(&start, NULL);

  memset(&ffmpeg, 0, sizeof(ffmpeg));

  ffmpeg.packet = av_packet_alloc();
  ffmpeg.frame = av_frame_alloc();
  ffmpeg.filt_frame = av_frame_alloc();

  if (!ffmpeg.packet || !ffmpeg.frame || !ffmpeg.filt_frame) {
    return AVERROR(ENOMEM);
  }

  ret = open_input(&ffmpeg);
  if (ret < 0)
    goto end;

  ret = open_output(&ffmpeg);
  if (ret < 0)
    goto end;

  ret = init_filters(&ffmpeg);
  if (ret < 0)
    goto end;

  av_dump_format(ffmpeg.ifmt_ctx, 0, "default", 0);
  av_dump_format(ffmpeg.ofmt_ctx, 0, "icecast", 1);

  in_time_base = ffmpeg.ifmt_ctx->streams[ffmpeg.istream_index]->time_base;
  /* read all packets */
  while (1) {
    gettimeofday(&current, NULL);
    if (RUNTIME < current.tv_sec - start.tv_sec) {
      printf("Reached max time!\n");
      ret = 0;
      goto end;
    }

    do {
      ret = av_read_frame(ffmpeg.ifmt_ctx, ffmpeg.packet);
    } while (ret > 0 || ret == AVERROR(EAGAIN));

    if (ret < 0)
      break;

    if (ffmpeg.packet->stream_index == ffmpeg.istream_index) {
      ret = avcodec_send_packet(ffmpeg.dec_ctx, ffmpeg.packet);
      av_packet_unref(ffmpeg.packet);

      if (ret < 0)
        goto end;

      while (ret >= 0) {
        ret = avcodec_receive_frame(ffmpeg.dec_ctx, ffmpeg.frame);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
          break;
        if (ret < 0)
          goto end;

        /* push the audio data from decoded frame into the filtergraph */
        ret = av_buffersrc_add_frame_flags(ffmpeg.buffersrc_ctx, ffmpeg.frame,
                                           AV_BUFFERSRC_FLAG_KEEP_REF);
        av_frame_unref(ffmpeg.frame);

        if (ret < 0)
          goto end;

        /* pull filtered audio from the filtergraph */
        while (1) {
          ret =
              av_buffersink_get_frame(ffmpeg.buffersink_ctx, ffmpeg.filt_frame);
          if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            break;
          if (ret < 0)
            goto end;

          left_volume = av_dict_get(ffmpeg.filt_frame->metadata,
                                    "lavfi.astats.1.RMS_level", NULL,
                                    AV_DICT_IGNORE_SUFFIX);

          right_volume = av_dict_get(ffmpeg.filt_frame->metadata,
                                     "lavfi.astats.2.RMS_level", NULL,
                                     AV_DICT_IGNORE_SUFFIX);

          ret = encode_audio_frame(&ffmpeg);
          av_frame_unref(ffmpeg.filt_frame);
        }
        av_frame_unref(ffmpeg.frame);

        if (ret == AVERROR_EOF)
          goto end;
      }
    }
  }

end:
  if (ret < 0 && ret != AVERROR_EOF) {
    printf("errror: %s\n", av_err2str(ret));
  }

  if (ffmpeg.packet)
    av_packet_free(&ffmpeg.packet);

  if (ffmpeg.filter_graph)
    avfilter_graph_free(&ffmpeg.filter_graph);

  if (ffmpeg.dec_ctx) {
    avcodec_free_context(&ffmpeg.dec_ctx);
  }

  if (ffmpeg.ifmt_ctx)
    avformat_close_input(&ffmpeg.ifmt_ctx);

  if (ffmpeg.enc_ctx) {
    avcodec_free_context(&ffmpeg.enc_ctx);
  }

  if (ffmpeg.ofmt_ctx) {
    if (ffmpeg.ofmt_ctx->pb)
      avio_close(ffmpeg.ofmt_ctx->pb);
    avformat_free_context(ffmpeg.ofmt_ctx);
  }

  if (ffmpeg.frame)
    av_frame_free(&ffmpeg.frame);

  if (ffmpeg.filt_frame)
    av_frame_free(&ffmpeg.filt_frame);

  return ret;
}

int main() {
  int err = 0;

  printf("Starting ffmpeg connection test..\n");

  av_log_set_level(AV_LOG_INFO);
  avdevice_register_all();

  while (err == 0) {
    err = runFfmpeg();
    printf("One iteration done..\n");
    sleep(WAITTIME);
  }

  return err;
}