]> git.sesse.net Git - nageru/blobdiff - nageru/ffmpeg_capture.cpp
Fix a buffer overrun when receiving 4K (or 8K etc.) FFmpeg streams.
[nageru] / nageru / ffmpeg_capture.cpp
index cda801a240d7e81147465fe37050d8ea50c31e7f..dd053f89292a409c84f7064c422478d7bfd32990 100644 (file)
@@ -26,6 +26,7 @@ extern "C" {
 #include <cstdint>
 #include <utility>
 #include <vector>
+#include <unordered_set>
 
 #include <Eigen/Core>
 #include <Eigen/LU>
@@ -43,8 +44,6 @@ extern "C" {
 #include <srt/srt.h>
 #endif
 
-#define FRAME_SIZE (8 << 20)  // 8 MB.
-
 using namespace std;
 using namespace std::chrono;
 using namespace bmusb;
@@ -143,7 +142,7 @@ AVPixelFormat decide_dst_format(AVPixelFormat src_format, bmusb::PixelFormat dst
        return av_get_pix_fmt(best_format);
 }
 
-YCbCrFormat decode_ycbcr_format(const AVPixFmtDescriptor *desc, const AVFrame *frame, bool is_mjpeg)
+YCbCrFormat decode_ycbcr_format(const AVPixFmtDescriptor *desc, const AVFrame *frame, bool is_mjpeg, AVColorSpace *last_colorspace, AVChromaLocation *last_chroma_location)
 {
        YCbCrFormat format;
        AVColorSpace colorspace = frame->colorspace;
@@ -163,11 +162,14 @@ YCbCrFormat decode_ycbcr_format(const AVPixFmtDescriptor *desc, const AVFrame *f
                format.luma_coefficients = (frame->height >= 720 ? YCBCR_REC_709 : YCBCR_REC_601);
                break;
        default:
-               fprintf(stderr, "Unknown Y'CbCr coefficient enum %d from FFmpeg; choosing Rec. 709.\n",
-                       colorspace);
+               if (colorspace != *last_colorspace) {
+                       fprintf(stderr, "Unknown Y'CbCr coefficient enum %d from FFmpeg; choosing Rec. 709.\n",
+                               colorspace);
+               }
                format.luma_coefficients = YCBCR_REC_709;
                break;
        }
+       *last_colorspace = colorspace;
 
        format.full_range = is_full_range(desc);
        format.num_levels = 1 << desc->comp[0].depth;
@@ -200,12 +202,15 @@ YCbCrFormat decode_ycbcr_format(const AVPixFmtDescriptor *desc, const AVFrame *f
                format.cb_y_position = 1.0;
                break;
        default:
-               fprintf(stderr, "Unknown chroma location coefficient enum %d from FFmpeg; choosing center.\n",
-                       frame->chroma_location);
+               if (frame->chroma_location != *last_chroma_location) {
+                       fprintf(stderr, "Unknown chroma location coefficient enum %d from FFmpeg; choosing center.\n",
+                               frame->chroma_location);
+               }
                format.cb_x_position = 0.5;
                format.cb_y_position = 0.5;
                break;
        }
+       *last_chroma_location = frame->chroma_location;
 
        if (is_mjpeg && !format.full_range) {
                // Limited-range MJPEG is only detected by FFmpeg whenever a special
@@ -286,6 +291,11 @@ FFmpegCapture::~FFmpegCapture()
                dequeue_cleanup_callback();
        }
        swr_free(&resampler);
+#ifdef HAVE_SRT
+       if (srt_sock != -1) {
+               srt_close(srt_sock);
+       }
+#endif
 }
 
 void FFmpegCapture::configure_card()
@@ -427,9 +437,11 @@ void FFmpegCapture::send_disconnected_frame()
                        memset(video_frame.data + frame_width * frame_height, 128, frame_width * frame_height);  // Valid for both NV12 and planar.
                }
 
-               frame_callback(-1, AVRational{1, TIMEBASE}, -1, AVRational{1, TIMEBASE}, timecode++,
-                       video_frame, /*video_offset=*/0, video_format,
-                       FrameAllocator::Frame(), /*audio_offset=*/0, AudioFormat());
+               if (frame_callback != nullptr) {
+                       frame_callback(-1, AVRational{1, TIMEBASE}, -1, AVRational{1, TIMEBASE}, timecode++,
+                               video_frame, /*video_offset=*/0, video_format,
+                               FrameAllocator::Frame(), /*audio_offset=*/0, AudioFormat());
+               }
                last_frame_was_connected = false;
        }
 
@@ -441,6 +453,66 @@ void FFmpegCapture::send_disconnected_frame()
        }
 }
 
+template<AVHWDeviceType type>
+AVPixelFormat get_hw_format(AVCodecContext *ctx, const AVPixelFormat *fmt)
+{
+       bool found_config_of_right_type = false;
+       for (int i = 0;; ++i) {  // Termination condition inside loop.
+               const AVCodecHWConfig *config = avcodec_get_hw_config(ctx->codec, i);
+               if (config == nullptr) {  // End of list.
+                       break;
+               }
+               if (!(config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) ||
+                   config->device_type != type) {
+                       // Not interesting for us.
+                       continue;
+               }
+
+               // We have a config of the right type, but does it actually support
+               // the pixel format we want? (Seemingly, FFmpeg's way of signaling errors
+               // is to just replace the pixel format with a software-decoded one,
+               // such as yuv420p.)
+               found_config_of_right_type = true;
+               for (const AVPixelFormat *fmt_ptr = fmt; *fmt_ptr != -1; ++fmt_ptr) {
+                       if (config->pix_fmt == *fmt_ptr) {
+                               fprintf(stderr, "Initialized '%s' hardware decoding for codec '%s'.\n",
+                                       av_hwdevice_get_type_name(type), ctx->codec->name);
+                               if (ctx->profile == FF_PROFILE_H264_BASELINE) {
+                                       fprintf(stderr, "WARNING: Stream claims to be H.264 Baseline, which is generally poorly supported in hardware decoders.\n");
+                                       fprintf(stderr, "         Consider encoding it as Constrained Baseline, Main or High instead.\n");
+                                       fprintf(stderr, "         Decoding might fail and fall back to software.\n");
+                               }
+                               return config->pix_fmt;
+                       }
+               }
+               fprintf(stderr, "Decoder '%s' supports only these pixel formats:", ctx->codec->name);
+               unordered_set<AVPixelFormat> seen;
+               for (const AVPixelFormat *fmt_ptr = fmt; *fmt_ptr != -1; ++fmt_ptr) {
+                       if (!seen.count(*fmt_ptr)) {
+                               fprintf(stderr, " %s", av_get_pix_fmt_name(*fmt_ptr));
+                               seen.insert(*fmt_ptr);
+                       }
+               }
+               fprintf(stderr, " (wanted %s for hardware acceleration)\n", av_get_pix_fmt_name(config->pix_fmt));
+
+       }
+
+       if (!found_config_of_right_type) {
+               fprintf(stderr, "Decoder '%s' does not support device type '%s'.\n", ctx->codec->name, av_hwdevice_get_type_name(type));
+       }
+
+       // We found no VA-API formats, so take the first software format.
+       for (const AVPixelFormat *fmt_ptr = fmt; *fmt_ptr != -1; ++fmt_ptr) {
+               if ((av_pix_fmt_desc_get(*fmt_ptr)->flags & AV_PIX_FMT_FLAG_HWACCEL) == 0) {
+                       fprintf(stderr, "Falling back to software format %s.\n", av_get_pix_fmt_name(*fmt_ptr));
+                       return *fmt_ptr;
+               }
+       }
+
+       // Fallback: Just return anything. (Should never really happen.)
+       return fmt[0];
+}
+
 bool FFmpegCapture::play_video(const string &pathname)
 {
        // Note: Call before open, not after; otherwise, there's a race.
@@ -454,6 +526,8 @@ bool FFmpegCapture::play_video(const string &pathname)
        } else {
                last_modified = buf.st_mtim;
        }
+       last_colorspace = static_cast<AVColorSpace>(-1);
+       last_chroma_location = static_cast<AVChromaLocation>(-1);
 
        AVFormatContextWithCloser format_ctx;
        if (srt_sock == -1) {
@@ -464,7 +538,7 @@ bool FFmpegCapture::play_video(const string &pathname)
        } else {
 #ifdef HAVE_SRT
                // SRT socket, already opened.
-               AVInputFormat *mpegts_fmt = av_find_input_format("mpegts");
+               const AVInputFormat *mpegts_fmt = av_find_input_format("mpegts");
                format_ctx = avformat_open_input_unique(&FFmpegCapture::read_srt_thunk, this,
                        mpegts_fmt, /*options=*/nullptr,
                        AVIOInterruptCB{ &FFmpegCapture::interrupt_cb_thunk, this });
@@ -494,7 +568,8 @@ bool FFmpegCapture::play_video(const string &pathname)
 
        // Open video decoder.
        const AVCodecParameters *video_codecpar = format_ctx->streams[video_stream_index]->codecpar;
-       AVCodec *video_codec = avcodec_find_decoder(video_codecpar->codec_id);
+       const AVCodec *video_codec = avcodec_find_decoder(video_codecpar->codec_id);
+
        video_timebase = format_ctx->streams[video_stream_index]->time_base;
        AVCodecContextWithDeleter video_codec_ctx = avcodec_alloc_context3_unique(nullptr);
        if (avcodec_parameters_to_context(video_codec_ctx.get(), video_codecpar) < 0) {
@@ -505,6 +580,28 @@ bool FFmpegCapture::play_video(const string &pathname)
                fprintf(stderr, "%s: Cannot find video decoder\n", pathname.c_str());
                return false;
        }
+
+       // Seemingly, it's not too easy to make something that just initializes
+       // “whatever goes”, so we don't get CUDA or VULKAN or whatever here
+       // without enumerating through several different types.
+       // VA-API and VDPAU will do for now. We prioritize VDPAU for the
+       // simple reason that there's a VA-API-via-VDPAU emulation for NVidia
+       // cards that seems to work, but just hangs when trying to transfer the frame.
+       //
+       // Note that we don't actually check codec support beforehand,
+       // so if you have a low-end VDPAU device but a high-end VA-API device,
+       // you lose out on the extra codec support from the latter.
+       AVBufferRef *hw_device_ctx = nullptr;
+       if (av_hwdevice_ctx_create(&hw_device_ctx, AV_HWDEVICE_TYPE_VDPAU, nullptr, nullptr, 0) >= 0) {
+               video_codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
+               video_codec_ctx->get_format = get_hw_format<AV_HWDEVICE_TYPE_VDPAU>;
+       } else if (av_hwdevice_ctx_create(&hw_device_ctx, AV_HWDEVICE_TYPE_VAAPI, nullptr, nullptr, 0) >= 0) {
+               video_codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
+               video_codec_ctx->get_format = get_hw_format<AV_HWDEVICE_TYPE_VAAPI>;
+       } else {
+               fprintf(stderr, "Failed to initialize VA-API or VDPAU for FFmpeg acceleration. Decoding video in software.\n");
+       }
+
        if (avcodec_open2(video_codec_ctx.get(), video_codec, nullptr) < 0) {
                fprintf(stderr, "%s: Cannot open video decoder\n", pathname.c_str());
                return false;
@@ -525,7 +622,7 @@ bool FFmpegCapture::play_video(const string &pathname)
                        fprintf(stderr, "%s: Cannot fill audio codec parameters\n", pathname.c_str());
                        return false;
                }
-               AVCodec *audio_codec = avcodec_find_decoder(audio_codecpar->codec_id);
+               const AVCodec *audio_codec = avcodec_find_decoder(audio_codecpar->codec_id);
                if (audio_codec == nullptr) {
                        fprintf(stderr, "%s: Cannot find audio decoder\n", pathname.c_str());
                        return false;
@@ -542,6 +639,7 @@ bool FFmpegCapture::play_video(const string &pathname)
 
        // Main loop.
        bool first_frame = true;
+       int consecutive_errors = 0;
        while (!producer_thread_should_quit.should_quit()) {
                if (process_queued_commands(format_ctx.get(), pathname, last_modified, /*rewound=*/nullptr)) {
                        return true;
@@ -558,7 +656,14 @@ bool FFmpegCapture::play_video(const string &pathname)
                AVFrameWithDeleter frame = decode_frame(format_ctx.get(), video_codec_ctx.get(), audio_codec_ctx.get(),
                        pathname, video_stream_index, audio_stream_index, subtitle_stream_index, audio_frame.get(), &audio_format, &audio_pts, &error);
                if (error) {
-                       return false;
+                       if (++consecutive_errors >= 100) {
+                               fprintf(stderr, "More than 100 consecutive video frames, aborting playback.\n");
+                               return false;
+                       } else {
+                               continue;
+                       }
+               } else {
+                       consecutive_errors = 0;
                }
                if (frame == nullptr) {
                        // EOF. Loop back to the start if we can.
@@ -666,7 +771,7 @@ bool FFmpegCapture::play_video(const string &pathname)
                                                1e3 * duration<double>(now - next_frame_start).count());
                                        pts_origin = frame->pts;
                                        start = next_frame_start = now;
-                                       timecode += MAX_FPS * 2 + 1;
+                                       timecode += TYPICAL_FPS * 2 + 1;
                                }
                        }
                        bool finished_wakeup;
@@ -684,12 +789,14 @@ bool FFmpegCapture::play_video(const string &pathname)
                                        // Make sure to get the audio resampler reset. (This is a hack;
                                        // ideally, the frame callback should just accept a way to signal
                                        // audio discontinuity.)
-                                       timecode += MAX_FPS * 2 + 1;
+                                       timecode += TYPICAL_FPS * 2 + 1;
                                }
                                last_neutral_color = get_neutral_color(frame->metadata);
-                               frame_callback(frame->pts, video_timebase, audio_pts, audio_timebase, timecode++,
-                                       video_frame.get_and_release(), 0, video_format,
-                                       audio_frame.get_and_release(), 0, audio_format);
+                               if (frame_callback != nullptr) {
+                                       frame_callback(frame->pts, video_timebase, audio_pts, audio_timebase, timecode++,
+                                               video_frame.get_and_release(), 0, video_format,
+                                               audio_frame.get_and_release(), 0, audio_format);
+                               }
                                first_frame = false;
                                last_frame = steady_clock::now();
                                last_frame_was_connected = true;
@@ -794,13 +901,16 @@ AVFrameWithDeleter FFmpegCapture::decode_frame(AVFormatContext *format_ctx, AVCo
                        if (pkt.stream_index == audio_stream_index && audio_callback != nullptr) {
                                audio_callback(&pkt, format_ctx->streams[audio_stream_index]->time_base);
                        }
-                       if (pkt.stream_index == video_stream_index) {
+                       if (pkt.stream_index == video_stream_index && video_callback != nullptr) {
+                               video_callback(&pkt, format_ctx->streams[video_stream_index]->time_base);
+                       }
+                       if (pkt.stream_index == video_stream_index && global_flags.transcode_video) {
                                if (avcodec_send_packet(video_codec_ctx, &pkt) < 0) {
                                        fprintf(stderr, "%s: Cannot send packet to video codec.\n", pathname.c_str());
                                        *error = true;
                                        return AVFrameWithDeleter(nullptr);
                                }
-                       } else if (pkt.stream_index == audio_stream_index) {
+                       } else if (pkt.stream_index == audio_stream_index && global_flags.transcode_audio) {
                                has_audio = true;
                                if (avcodec_send_packet(audio_codec_ctx, &pkt) < 0) {
                                        fprintf(stderr, "%s: Cannot send packet to audio codec.\n", pathname.c_str());
@@ -837,6 +947,22 @@ AVFrameWithDeleter FFmpegCapture::decode_frame(AVFormatContext *format_ctx, AVCo
                // Decode video, if we have a frame.
                int err = avcodec_receive_frame(video_codec_ctx, video_avframe.get());
                if (err == 0) {
+                       if (video_avframe->format == AV_PIX_FMT_VAAPI ||
+                           video_avframe->format == AV_PIX_FMT_VDPAU) {
+                               // Get the frame down to the CPU. (TODO: See if we can keep it
+                               // on the GPU all the way, since it will be going up again later.
+                               // However, this only works if the OpenGL GPU is the same one.)
+                               AVFrameWithDeleter sw_frame = av_frame_alloc_unique();
+                               int err = av_hwframe_transfer_data(sw_frame.get(), video_avframe.get(), 0);
+                               if (err != 0) {
+                                       fprintf(stderr, "%s: Cannot transfer hardware video frame to software.\n", pathname.c_str());
+                                       *error = true;
+                                       return AVFrameWithDeleter(nullptr);
+                               }
+                               sw_frame->pts = video_avframe->pts;
+                               sw_frame->pkt_duration = video_avframe->pkt_duration;
+                               video_avframe = move(sw_frame);
+                       }
                        frame_finished = true;
                        break;
                } else if (err != AVERROR(EAGAIN)) {
@@ -997,7 +1123,7 @@ UniqueFrame FFmpegCapture::make_video_frame(const AVFrame *frame, const string &
                video_frame->len = (frame_width(frame) * 2) * frame_height(frame);
 
                const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(sws_dst_format);
-               current_frame_ycbcr_format = decode_ycbcr_format(desc, frame, is_mjpeg);
+               current_frame_ycbcr_format = decode_ycbcr_format(desc, frame, is_mjpeg, &last_colorspace, &last_chroma_location);
        } else {
                assert(pixel_format == bmusb::PixelFormat_8BitYCbCrPlanar);
                const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(sws_dst_format);
@@ -1016,8 +1142,18 @@ UniqueFrame FFmpegCapture::make_video_frame(const AVFrame *frame, const string &
 
                video_frame->len = frame_width(frame) * frame_height(frame) + 2 * chroma_width * chroma_height;
 
-               current_frame_ycbcr_format = decode_ycbcr_format(desc, frame, is_mjpeg);
+               current_frame_ycbcr_format = decode_ycbcr_format(desc, frame, is_mjpeg, &last_colorspace, &last_chroma_location);
        }
+
+       // FIXME: Currently, if the video is too high-res for one of the allocated
+       // frames, we simply refuse to scale it here to avoid crashes. It would be better
+       // if we could somehow signal getting larger frames, especially as 4K is a thing now.
+       if (video_frame->len > FRAME_SIZE) {
+               fprintf(stderr, "%s: Decoded frame would be larger than supported FRAME_SIZE (%zu > %u), not decoding.\n", pathname.c_str(), video_frame->len, FRAME_SIZE);
+               *error = true;
+               return video_frame;
+       }
+
        sws_scale(sws_ctx.get(), frame->data, frame->linesize, 0, frame->height, pic_data, linesizes);
 
        return video_frame;
@@ -1047,7 +1183,7 @@ unsigned FFmpegCapture::frame_height(const AVFrame *frame) const
        if (height == 0) {
                return frame->height;
        } else {
-               return width;
+               return height;
        }
 }