extractors/bilibili: use ffmpeg to merge files once downloaded

6 files changed, 279 insertions, 40 deletions

diff --git a/src/extractors/bilibili.c b/src/extractors/bilibili.c
index 675a035..fd95884 100644
--- a/src/extractors/bilibili.c
+++ b/src/extractors/bilibili.c
@@ -13,6 +13,7 @@
 #include "../logger.h"
 #include "../process_url.h"
 #include "../utils.h"
+#include "../utils/ffmpeg.h"
 #include "bilibili.h"
 #include "extractor.h"
 
@@ -380,13 +381,13 @@ static void dash_cleanup(Dash *dash) {
   free_array(&dash->dashinfo.dash.video);
 }
 
-static int ffmpeg_merge(callback_struct_t *cb_struct) {
-  char *filename =
-      malloc(strlen(cb_struct->title) + strlen(cb_struct->ext) + 2);
-  sprintf(filename, "%s.%s", cb_struct->title, cb_struct->ext);
-  DEBUG_PRINT("Callback gets filename: %s\n", filename);
-  free_and_nullify(cb_struct->title);
-  free_and_nullify(filename);
+static int bilibili_merge(callback_struct_t *cb_struct) {
+  LOG("Bilibili", "Using ffmpeg to merge downloaded files ...");
+  merge_av(cb_struct->videofn, cb_struct->audiofn, cb_struct->filename);
+  free_and_nullify(cb_struct->videofn);
+  free_and_nullify(cb_struct->audiofn);
+  free_and_nullify(cb_struct->filename);
+  append_log("All done!");
   return 0;
 }
 
@@ -407,27 +408,26 @@ static int download(Bilibili_options *bilibili_options) {
   const char *quality_desc = id2quality_desc(video->id);
 
   static callback_struct_t callback_struct = {0};
-  callback_struct.ext = mimeType2ext(video->mimeType);
-  callback_struct.title =
-      malloc(strlen(bilibili_options->title) + strlen(quality_desc) + 3);
-  sprintf(callback_struct.title, "%s[%s]", bilibili_options->title,
-          quality_desc);
-  DEBUG_PRINT("Outside the callback, title: %s ext: %s\n",
-              callback_struct.title, callback_struct.ext);
+  const char *ext = mimeType2ext(video->mimeType);
+  callback_struct.filename = malloc(strlen(bilibili_options->title) +
+                                    strlen(quality_desc) + strlen(ext) + 1);
+  sprintf(callback_struct.filename, "%s[%s].%s", bilibili_options->title,
+          quality_desc, ext);
 
   {
-    char fn[USHRT_MAX];
-    sprintf(fn, "%s[%s]-%s.%s", bilibili_options->title, quality_desc, "video",
-            callback_struct.ext);
-    add_url(video->baseUrl, NULL, fn, "https://www.bilibili.com", NULL, NULL);
+    callback_struct.videofn = malloc(strlen(callback_struct.filename) + 6);
+    sprintf(callback_struct.videofn, "%s[%s]-%s.%s", bilibili_options->title,
+            quality_desc, "video", ext);
+    add_url(video->baseUrl, NULL, callback_struct.videofn,
+            "https://www.bilibili.com", NULL, NULL);
   }
 
   {
-    char fn[USHRT_MAX];
-    sprintf(fn, "%s[%s]-%s.%s", bilibili_options->title, quality_desc, "audio",
-            mimeType2ext(audio->mimeType));
-    add_url(audio->baseUrl, NULL, fn, "https://www.bilibili.com", &ffmpeg_merge,
-            &callback_struct);
+    callback_struct.audiofn = malloc(strlen(callback_struct.filename) + 6);
+    sprintf(callback_struct.audiofn, "%s[%s]-%s.%s", bilibili_options->title,
+            quality_desc, "audio", mimeType2ext(audio->mimeType));
+    add_url(audio->baseUrl, NULL, callback_struct.audiofn,
+            "https://www.bilibili.com", &bilibili_merge, &callback_struct);
   }
   free_and_nullify(resp);
   dash_cleanup(&dash);
diff --git a/src/process_url.c b/src/process_url.c
index 1b391f2..881375f 100644
--- a/src/process_url.c
+++ b/src/process_url.c
@@ -1,3 +1,4 @@
+#include "utils.h"
 #include <curl/curl.h>
 #include <curl/easy.h>
 #include <curl/header.h>
@@ -136,6 +137,12 @@ static size_t write2str(void *ptr, size_t size, size_t nmemb, str_data_t *s) {
   return size * nmemb;
 }
 
+static void gen_fullpathfn(char *fullpathfn, const char *outdir,
+                           const char *fn) {
+  sprintf(fullpathfn, "%s%s%s", outdir,
+          outdir[strlen(outdir) - 1] == SPLITTER_CHAR ? "" : SPLITTER_STR, fn);
+}
+
 static int parse_url(const char *URL, const char *outdir, char *fn) {
   CURLUcode ue = logerr(curl_url_set(h, CURLUPART_URL, URL, 0));
   if (ue && ue != CURLUE_NO_QUERY) {
@@ -203,9 +210,7 @@ static int parse_url(const char *URL, const char *outdir, char *fn) {
     for (unsigned short i = 0; i < results.n; i++) {
       if (results.str[i]) {
         DEBUG_PRINT("[%d] %s\n", i, results.str[i]);
-        sprintf(curl_c->outfn, "%s%s%s", outdir,
-                outdir[strlen(outdir) - 1] == SPLITTER_CHAR ? "" : SPLITTER_STR,
-                results.str[i]);
+        gen_fullpathfn(curl_c->outfn, outdir, results.str[i]);
       }
     }
     free_str_array(&results);
@@ -422,6 +427,14 @@ static void replace_illegal_char(char *str) {
   }
 }
 
+static char *callback_struct_convert_fullpath(char *filename) {
+  char *tmp = malloc(strlen(outdir_g) + strlen(filename) + 2);
+  replace_illegal_char(filename);
+  gen_fullpathfn(tmp, outdir_g, filename);
+  free_and_nullify(filename);
+  return tmp;
+}
+
 void curl_init(curl_conf_t *curl) {
   curl_global_init(CURL_GLOBAL_ALL);
   h = curl_url();
@@ -521,6 +534,16 @@ int get(const char *URL, char **pdstr) {
 void add_url(const char *URL, const char *outdir, const char *fn,
              const char *referer, callback_t callback,
              callback_struct_t *p_callback_struct) {
+
+  char *filename;
+  if (fn == NULL || fn[0] == '\0') {
+    filename = NULL;
+  } else {
+    filename = malloc(strlen(fn) + 1);
+    strcpy(filename, fn);
+    replace_illegal_char(filename);
+  }
+
   if (outdir && outdir[0] != '\0') {
     outdir_g = outdir;
   }
@@ -529,19 +552,16 @@ void add_url(const char *URL, const char *outdir, const char *fn,
     referer_g = NULL;
   }
   DEBUG_PRINT("referer_g: %s\n", referer_g);
+
   callback_g = callback;
   if (p_callback_struct) {
+    p_callback_struct->videofn =
+        callback_struct_convert_fullpath(p_callback_struct->videofn);
+    p_callback_struct->audiofn =
+        callback_struct_convert_fullpath(p_callback_struct->audiofn);
+    p_callback_struct->filename =
+        callback_struct_convert_fullpath(p_callback_struct->filename);
     p_callback_struct_g = p_callback_struct;
-    replace_illegal_char(p_callback_struct_g->title);
-  }
-
-  char *filename;
-  if (fn == NULL || fn[0] == '\0') {
-    filename = NULL;
-  } else {
-    filename = malloc(strlen(fn) + 1);
-    strcpy(filename, fn);
-    replace_illegal_char(filename);
   }
 
   // Pass our cache (outdir_g) to parse_url()
diff --git a/src/process_url.h b/src/process_url.h
index 53d8ed3..de81729 100644
--- a/src/process_url.h
+++ b/src/process_url.h
@@ -45,8 +45,9 @@ typedef struct str_data {
 } str_data_t;
 
 typedef struct callback_struct {
-  char *title;
-  const char *ext;
+  char *videofn;
+  char *audiofn;
+  char *filename;
 } callback_struct_t;
 
 typedef int (*callback_t)(callback_struct_t *);
diff --git a/src/utils/ffmpeg.c b/src/utils/ffmpeg.c
new file mode 100644
index 0000000..a5010ee
--- /dev/null
+++ b/src/utils/ffmpeg.c
@@ -0,0 +1,212 @@
+// Adapted from
+// https://github.com/leandromoreira/ffmpeg-libav-tutorial/blob/master/2_remuxing.c
+// based on https://ffmpeg.org/doxygen/trunk/remuxing_8c-example.html
+// FIXME: looks ugly, further refactor needed :(
+#include <libavformat/avformat.h>
+#include <libavutil/timestamp.h>
+#include <stdio.h>
+
+#include "../logger.h"
+#include "ffmpeg.h"
+
+int merge_av(const char *videofn, const char *audiofn, const char *outfn) {
+  AVFormatContext *input1_format_context = NULL, *input2_format_context = NULL,
+                  *output_format_context = NULL;
+  AVPacket packet;
+  int ret, i;
+  int stream_index = 0;
+  int *streams_list = NULL;
+  int number_of_streams = 0;
+
+  if ((ret = avformat_open_input(&input1_format_context, videofn, NULL, NULL)) <
+      0) {
+    append_log("Could not open input file '%s'\n", videofn);
+    goto end;
+  }
+  if ((ret = avformat_open_input(&input2_format_context, audiofn, NULL, NULL)) <
+      0) {
+    append_log("Could not open input file '%s'\n", audiofn);
+    goto end;
+  }
+  if ((ret = avformat_find_stream_info(input1_format_context, NULL)) < 0) {
+    append_log("Failed to retrieve input stream information\n");
+    goto end;
+  }
+  if ((ret = avformat_find_stream_info(input2_format_context, NULL)) < 0) {
+    append_log("Failed to retrieve input stream information\n");
+    goto end;
+  }
+
+  avformat_alloc_output_context2(&output_format_context, NULL, NULL, outfn);
+  if (!output_format_context) {
+    append_log("Could not create output context\n");
+    ret = AVERROR_UNKNOWN;
+    goto end;
+  }
+
+  number_of_streams =
+      input1_format_context->nb_streams + input2_format_context->nb_streams;
+  streams_list = av_malloc_array(number_of_streams, sizeof(*streams_list));
+
+  if (!streams_list) {
+    ret = AVERROR(ENOMEM);
+    goto end;
+  }
+
+  for (i = 0; i < input1_format_context->nb_streams; i++) {
+    AVStream *out_stream;
+    AVStream *in_stream = input1_format_context->streams[i];
+    AVCodecParameters *in_codecpar = in_stream->codecpar;
+    if (in_codecpar->codec_type != AVMEDIA_TYPE_AUDIO &&
+        in_codecpar->codec_type != AVMEDIA_TYPE_VIDEO &&
+        in_codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
+      streams_list[i] = -1;
+      continue;
+    }
+    streams_list[i] = stream_index++;
+    out_stream = avformat_new_stream(output_format_context, NULL);
+    if (!out_stream) {
+      append_log("Failed allocating output stream\n");
+      ret = AVERROR_UNKNOWN;
+      goto end;
+    }
+    ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
+    if (ret < 0) {
+      append_log("Failed to copy codec parameters\n");
+      goto end;
+    }
+  }
+  for (i = 0; i < input2_format_context->nb_streams; i++) {
+    AVStream *out_stream;
+    AVStream *in_stream = input2_format_context->streams[i];
+    AVCodecParameters *in_codecpar = in_stream->codecpar;
+    if (in_codecpar->codec_type != AVMEDIA_TYPE_AUDIO &&
+        in_codecpar->codec_type != AVMEDIA_TYPE_VIDEO &&
+        in_codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
+      streams_list[i + stream_index] = -1;
+      continue;
+    }
+    streams_list[i + stream_index] = stream_index;
+    stream_index++;
+    out_stream = avformat_new_stream(output_format_context, NULL);
+    if (!out_stream) {
+      append_log("Failed allocating output stream\n");
+      ret = AVERROR_UNKNOWN;
+      goto end;
+    }
+    ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
+    if (ret < 0) {
+      append_log("Failed to copy codec parameters\n");
+      goto end;
+    }
+  }
+#ifdef DEBUG
+  av_dump_format(input1_format_context, 0, videofn, 0);
+  av_dump_format(input2_format_context, 0, audiofn, 0);
+  av_dump_format(output_format_context, 0, outfn, 1);
+#endif
+  if (!(output_format_context->oformat->flags & AVFMT_NOFILE)) {
+    ret = avio_open(&output_format_context->pb, outfn, AVIO_FLAG_WRITE);
+    if (ret < 0) {
+      append_log("Could not open output file '%s'", outfn);
+      goto end;
+    }
+  }
+  AVDictionary *opts = NULL;
+  /* Support fragmented MP4
+   * https://developer.mozilla.org/en-US/docs/Web/API/Media_Source_Extensions_API/Transcoding_assets_for_MSE
+   */
+  av_dict_set(&opts, "movflags", "frag_keyframe+empty_moov+default_base_moof",
+              0);
+  ret = avformat_write_header(output_format_context, &opts);
+  if (ret < 0) {
+    append_log("Error occurred when opening output file\n");
+    goto end;
+  }
+  while (1) {
+    AVStream *in_stream, *out_stream;
+    ret = av_read_frame(input1_format_context, &packet);
+    if (ret < 0)
+      break;
+    in_stream = input1_format_context->streams[packet.stream_index];
+    if (packet.stream_index >= number_of_streams ||
+        streams_list[packet.stream_index] < 0) {
+      av_packet_unref(&packet);
+      continue;
+    }
+    packet.stream_index = streams_list[packet.stream_index];
+    out_stream = output_format_context->streams[packet.stream_index];
+    /* copy packet */
+    packet.pts = av_rescale_q_rnd(packet.pts, in_stream->time_base,
+                                  out_stream->time_base,
+                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
+    packet.dts = av_rescale_q_rnd(packet.dts, in_stream->time_base,
+                                  out_stream->time_base,
+                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
+    packet.duration = av_rescale_q(packet.duration, in_stream->time_base,
+                                   out_stream->time_base);
+    packet.pos = -1;
+
+    ret = av_interleaved_write_frame(output_format_context, &packet);
+    if (ret < 0) {
+      append_log("Error muxing packet\n");
+      break;
+    }
+    av_packet_unref(&packet);
+  }
+  while (1) {
+    AVStream *in_stream, *out_stream;
+    ret = av_read_frame(input2_format_context, &packet);
+    if (ret < 0)
+      break;
+    in_stream = input2_format_context->streams[packet.stream_index];
+    if (packet.stream_index >= number_of_streams ||
+        streams_list[packet.stream_index + input1_format_context->nb_streams] <
+            0) {
+      av_packet_unref(&packet);
+      continue;
+    }
+    packet.stream_index =
+        streams_list[packet.stream_index + input1_format_context->nb_streams];
+    out_stream = output_format_context->streams[packet.stream_index];
+    /* copy packet */
+    packet.pts = av_rescale_q_rnd(packet.pts, in_stream->time_base,
+                                  out_stream->time_base,
+                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
+    packet.dts = av_rescale_q_rnd(packet.dts, in_stream->time_base,
+                                  out_stream->time_base,
+                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
+    packet.duration = av_rescale_q(packet.duration, in_stream->time_base,
+                                   out_stream->time_base);
+    packet.pos = -1;
+
+    ret = av_interleaved_write_frame(output_format_context, &packet);
+    if (ret < 0) {
+      append_log("Error muxing packet\n");
+      break;
+    }
+    av_packet_unref(&packet);
+  }
+  av_write_trailer(output_format_context);
+end:
+  avformat_close_input(&input1_format_context);
+  /* close output */
+  if (output_format_context &&
+      !(output_format_context->oformat->flags & AVFMT_NOFILE))
+    avio_closep(&output_format_context->pb);
+  avformat_free_context(output_format_context);
+  av_freep(&streams_list);
+  if (ret < 0 && ret != AVERROR_EOF) {
+    append_log("Error occurred: %s\n", av_err2str(ret));
+    return 1;
+  }
+
+  // Delete seperate files
+  if (remove(videofn) != 0) {
+    append_log("Error deleting partial file %s\n", videofn);
+  }
+  if (remove(audiofn) != 0) {
+    append_log("Error deleting partial file %s\n", audiofn);
+  }
+  return 0;
+}
diff --git a/src/utils/ffmpeg.h b/src/utils/ffmpeg.h
new file mode 100644
index 0000000..91d79df
--- /dev/null
+++ b/src/utils/ffmpeg.h
@@ -0,0 +1,6 @@
+#ifndef FFMPEG_H_
+#define FFMPEG_H_
+
+int merge_av(const char *videofn, const char *audiofn, const char *outfn);
+
+#endif
diff --git a/xmake.lua b/xmake.lua
index 5d9da51..98ba483 100644
--- a/xmake.lua
+++ b/xmake.lua
@@ -2,7 +2,7 @@ add_rules("mode.debug", "mode.release", "mode.releasedbg", "mode.minsizerel")
 
 add_requires("nuklear", "nuklear_glfw_gl3", "nuklear_fonts", "glew", "glfw", "tinyfiledialogs",
     "c11threads", "pcre2",
-    "cjson")
+    "cjson", "ffmpeg")
 add_requires("libcurl", { configs = { zlib = true } })
 
 if is_plat("linux") then
@@ -24,7 +24,7 @@ if is_mode("debug") then
 end
 add_packages("nuklear", "nuklear_glfw_gl3", "nuklear_fonts", "glew", "glfw", "libcurl", "tinyfiledialogs",
     "c11threads", "pcre2",
-    "cjson")
+    "cjson", "ffmpeg")
 
 
 package("tinyfiledialogs")