implement recording the video feed

app/CMakeLists.txt

@@ -1,6 +1,6 @@
 find_package(Qt6 REQUIRED COMPONENTS Core Widgets Multimedia OpenGLWidgets)
 find_package(SDL2 REQUIRED)
-find_package(FFmpeg REQUIRED COMPONENTS avcodec avutil avfilter)
+find_package(FFmpeg REQUIRED COMPONENTS avformat avcodec avutil avfilter)
 
 set(CMAKE_AUTOMOC ON)
 set(CMAKE_AUTOUIC ON)
@@ -28,6 +28,7 @@ target_link_libraries(vanilla-gui PRIVATE
     Qt6::OpenGLWidgets
     SDL2
     vanilla
+    FFmpeg::avformat
     FFmpeg::avcodec
     FFmpeg::avutil
     FFmpeg::avfilter

app/backend.cpp

@@ -93,6 +93,17 @@ void writeNullTermString(int pipe, const QString &s)
     writeByte(pipe, 0);
 }
 
+void Backend::requestIDR()
+{
+    if (m_pipe) {
+        m_pipeMutex.lock();
+        writeByte(m_pipeOut, VANILLA_PIPE_IN_REQ_IDR);
+        m_pipeMutex.unlock();
+    } else {
+        vanilla_request_idr();
+    }
+}
+
 void Backend::connectToConsole(const QString &wirelessInterface)
 {
     if (m_pipe) {

app/backend.h

@@ -54,6 +54,7 @@ public slots:
     void connectToConsole(const QString &wirelessInterface);
     void updateTouch(int x, int y);
     void setButton(int button, int32_t value);
+    void requestIDR();
 
 private:
     BackendPipe *m_pipe;

app/mainwindow.cpp

@@ -2,6 +2,7 @@
 
 #include <QAudioDevice>
 #include <QComboBox>
+#include <QFileDialog>
 #include <QGroupBox>
 #include <QHBoxLayout>
 #include <QLabel>
@@ -92,6 +93,16 @@ MainWindow::MainWindow(QWidget *parent) : QWidget(parent)
         QPushButton *fullScreenBtn = new QPushButton(tr("Full Screen"), configSection);
         connect(fullScreenBtn, &QPushButton::clicked, this, &MainWindow::setFullScreen);
         configLayout->addWidget(fullScreenBtn, row, 0, 1, 2);
+        
+        row++;
+
+        m_recordBtn = new QPushButton(tr("Record"), configSection);
+        m_recordBtn->setCheckable(true);
+        configLayout->addWidget(m_recordBtn, row, 0);
+
+        m_screenshotBtn = new QPushButton(tr("Screenshot"), configSection);
+        connect(m_screenshotBtn, &QPushButton::clicked, this, &MainWindow::takeScreenshot);
+        configLayout->addWidget(m_screenshotBtn, row, 1);
     }
 
     {
@@ -147,6 +158,7 @@ MainWindow::MainWindow(QWidget *parent) : QWidget(parent)
     startObjectOnThread(m_backend);
 
     m_videoDecoder = new VideoDecoder();
+    connect(m_recordBtn, &QPushButton::clicked, m_videoDecoder, &VideoDecoder::enableRecording);
     startObjectOnThread(m_videoDecoder);
 
     m_gamepadHandler = new GamepadHandler();
@@ -158,8 +170,12 @@ MainWindow::MainWindow(QWidget *parent) : QWidget(parent)
     QMetaObject::invokeMethod(m_audioHandler, &AudioHandler::run, Qt::QueuedConnection);
 
     connect(m_backend, &Backend::videoAvailable, m_videoDecoder, &VideoDecoder::sendPacket);
+    connect(m_backend, &Backend::audioAvailable, m_videoDecoder, &VideoDecoder::sendAudio);
     connect(m_backend, &Backend::syncCompleted, this, [this](bool e){if (e) m_connectBtn->setEnabled(true);});
     connect(m_videoDecoder, &VideoDecoder::frameReady, m_viewer, &Viewer::setImage);
+    connect(m_videoDecoder, &VideoDecoder::recordingError, this, &MainWindow::recordingError);
+    connect(m_videoDecoder, &VideoDecoder::recordingFinished, this, &MainWindow::recordingFinished);
+    connect(m_videoDecoder, &VideoDecoder::requestIDR, m_backend, &Backend::requestIDR, Qt::DirectConnection);
     connect(m_backend, &Backend::audioAvailable, m_audioHandler, &AudioHandler::write);
     connect(m_backend, &Backend::vibrate, m_gamepadHandler, &GamepadHandler::vibrate, Qt::DirectConnection);
     connect(m_viewer, &Viewer::touch, m_backend, &Backend::updateTouch, Qt::DirectConnection);
@@ -333,3 +349,36 @@ void MainWindow::startObjectOnThread(QObject *object)
     thread->start();
     m_threadMap.insert(object, thread);
 }
+
+void MainWindow::recordingError(int err)
+{
+    QMessageBox::critical(this, tr("Recording Error"), tr("Recording failed with the following error: %0 (%1)").arg(av_err2str(err), QString::number(err)));
+}
+
+void MainWindow::recordingFinished(const QString &filename)
+{
+    QString s = QFileDialog::getSaveFileName(this, tr("Save Screenshot"), QString(), tr("MPEG-4 Video (*.mp4)"));
+    if (!s.isEmpty()) {
+        QString ext = QStringLiteral(".mp4");
+        if (!s.endsWith(ext, Qt::CaseInsensitive)) {
+            s = s.append(ext);
+        }
+
+        QFile::copy(filename, s);
+    }
+}
+
+void MainWindow::takeScreenshot()
+{
+    // Make a copy of the current image
+    QImage ss = m_viewer->image();
+
+    QString s = QFileDialog::getSaveFileName(this, tr("Save Screenshot"), QString(), tr("PNG (*.png)"));
+    if (!s.isEmpty()) {
+        QString ext = QStringLiteral(".png");
+        if (!s.endsWith(ext, Qt::CaseInsensitive)) {
+            s = s.append(ext);
+        }
+        ss.save(s);
+    }
+}

app/mainwindow.h

@@ -39,6 +39,8 @@ private:
 
     QPushButton *m_syncBtn;
     QPushButton *m_connectBtn;
+    QPushButton *m_recordBtn;
+    QPushButton *m_screenshotBtn;
 
     QSplitter *m_splitter;
 
@@ -65,6 +67,11 @@ private slots:
 
     void showInputConfigDialog();
 
+    void recordingError(int err);
+    void recordingFinished(const QString &filename);
+
+    void takeScreenshot();
+
 };
 
 #endif // MAINWINDOW_H

app/videodecoder.cpp

@@ -1,14 +1,26 @@
 #include "videodecoder.h"
 
+#include <QDateTime>
+#include <QDir>
 #include <QImage>
+#include <QStandardPaths>
+
+#include <vanilla.h>
 
 extern "C" {
 #include <libavfilter/buffersink.h>
 #include <libavfilter/buffersrc.h>
 }
 
+enum RecordingStream {
+    VIDEO_STREAM_INDEX,
+    AUDIO_STREAM_INDEX
+};
+
 VideoDecoder::VideoDecoder(QObject *parent) : QObject(parent)
 {
+    m_recordingCtx = nullptr;
+    
     m_packet = av_packet_alloc();
     m_frame = av_frame_alloc();
     
@@ -45,6 +57,13 @@ void cleanupFrame(void *v)
     av_frame_free(&f);
 }
 
+int64_t VideoDecoder::getCurrentTimestamp(AVRational timebase)
+{
+    int64_t millis = QDateTime::currentMSecsSinceEpoch() - m_recordingStartTime;
+    int64_t ts = av_rescale_q(millis, {1, 1000}, timebase);
+    return ts;
+}
+
 void VideoDecoder::sendPacket(const QByteArray &data)
 {
     int ret;
@@ -61,6 +80,20 @@ void VideoDecoder::sendPacket(const QByteArray &data)
         return;
     }
 
+    // If recording, send packet to file
+    if (m_recordingCtx) {
+        AVPacket *encPkt = av_packet_clone(m_packet);
+        encPkt->stream_index = VIDEO_STREAM_INDEX;
+
+        int64_t ts = getCurrentTimestamp(m_videoStream->time_base);
+
+        encPkt->dts = ts;
+        encPkt->pts = ts;
+
+        av_interleaved_write_frame(m_recordingCtx, encPkt);
+        av_packet_free(&encPkt);
+    }
+
     // Send packet to decoder
     ret = avcodec_send_packet(m_codecCtx, m_packet);
     av_packet_unref(m_packet);
@@ -94,4 +127,123 @@ void VideoDecoder::sendPacket(const QByteArray &data)
         QImage image(filtered->data[0], filtered->width, filtered->height, filtered->linesize[0], QImage::Format_RGB888, cleanupFrame, filtered);
         emit frameReady(image);
     }
+}
+
+void VideoDecoder::sendAudio(const QByteArray &data)
+{
+    if (m_recordingCtx) {
+        int ret;
+
+        // Copy data into buffer that FFmpeg will take ownership of
+        uint8_t *buffer = (uint8_t *) av_malloc(data.size());
+        memcpy(buffer, data.data(), data.size());
+
+        // Create AVPacket from this data
+        AVPacket *audPkt = av_packet_alloc();
+        int64_t ts;
+        ret = av_packet_from_data(audPkt, buffer, data.size());
+        if (ret < 0) {
+            fprintf(stderr, "Failed to initialize packet from data: %i\n", ret);
+            av_free(buffer);
+            goto free;
+        }
+
+        ts = getCurrentTimestamp(m_audioStream->time_base);
+
+        audPkt->stream_index = AUDIO_STREAM_INDEX;
+        audPkt->dts = ts;
+        audPkt->pts = ts;
+
+        av_interleaved_write_frame(m_recordingCtx, audPkt);
+
+free:
+        av_packet_free(&audPkt);
+    }
+}
+
+void VideoDecoder::enableRecording(bool e)
+{
+    if (e) startRecording(); else stopRecording();
+}
+
+void VideoDecoder::startRecording()
+{
+    m_recordingFilename = QDir(QStandardPaths::writableLocation(QStandardPaths::TempLocation)).filePath("vanilla-recording-%0.mp4").arg(QDateTime::currentSecsSinceEpoch());
+    
+    QByteArray filenameUtf8 = m_recordingFilename.toUtf8();
+
+    int r = avformat_alloc_output_context2(&m_recordingCtx, nullptr, nullptr, filenameUtf8.constData());
+    if (r < 0) {
+        emit recordingError(r);
+        return;
+    }
+
+    m_videoStream = avformat_new_stream(m_recordingCtx, nullptr);
+    if (!m_videoStream) {
+        emit recordingError(AVERROR(ENOMEM));
+        goto freeContext;
+    }
+
+    m_videoStream->id = VIDEO_STREAM_INDEX;
+    m_videoStream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
+    m_videoStream->codecpar->width = 854;
+    m_videoStream->codecpar->height = 480;
+    m_videoStream->codecpar->format = AV_PIX_FMT_YUV420P;
+    m_videoStream->time_base = {1, 60};
+    m_videoStream->codecpar->codec_id = AV_CODEC_ID_H264;
+
+    size_t sps_pps_size;
+    vanilla_retrieve_sps_pps_data(nullptr, &sps_pps_size);
+    m_videoStream->codecpar->extradata_size = sps_pps_size;
+    m_videoStream->codecpar->extradata = (uint8_t *) av_malloc(sps_pps_size);
+    vanilla_retrieve_sps_pps_data(m_videoStream->codecpar->extradata, &sps_pps_size);
+
+    m_audioStream = avformat_new_stream(m_recordingCtx, nullptr);
+    if (!m_audioStream) {
+        emit recordingError(AVERROR(ENOMEM));
+        goto freeContext;
+    }
+
+    m_audioStream->id = AUDIO_STREAM_INDEX;
+    m_audioStream->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
+    m_audioStream->codecpar->sample_rate = 48000;
+    m_audioStream->codecpar->ch_layout = AV_CHANNEL_LAYOUT_STEREO;
+    m_audioStream->codecpar->format = AV_SAMPLE_FMT_S16;
+    m_audioStream->time_base = {1, 48000};
+    m_audioStream->codecpar->codec_id = AV_CODEC_ID_PCM_S16LE;
+
+    r = avio_open2(&m_recordingCtx->pb, filenameUtf8.constData(), AVIO_FLAG_WRITE, nullptr, nullptr);
+    if (r < 0) {
+        emit recordingError(r);
+        goto freeContext;
+    }
+
+    r = avformat_write_header(m_recordingCtx, nullptr);
+    if (r < 0) {
+        printf("err 5\n");
+        emit recordingError(r);
+        goto freeContext;
+    }
+
+    emit requestIDR();
+
+    printf("nb streams: %i\n", m_recordingCtx->nb_streams);
+
+    m_recordingStartTime = QDateTime::currentMSecsSinceEpoch();
+    return;
+
+freeContext:
+    avformat_free_context(m_recordingCtx);
+    m_recordingCtx = nullptr;
+}
+
+void VideoDecoder::stopRecording()
+{
+    if (m_recordingCtx) {
+        av_write_trailer(m_recordingCtx);
+        avio_closep(&m_recordingCtx->pb);
+        avformat_free_context(m_recordingCtx);
+        m_recordingCtx = nullptr;
+        emit recordingFinished(m_recordingFilename);
+    }
 }

app/videodecoder.h

@@ -4,6 +4,7 @@
 #include <QObject>
 
 extern "C" {
+#include <libavformat/avformat.h>
 #include <libavcodec/avcodec.h>
 #include <libavfilter/avfilter.h>
 #include <libswscale/swscale.h>
@@ -19,11 +20,20 @@ public:
 
 signals:
     void frameReady(const QImage &image);
+    void recordingError(int err);
+    void recordingFinished(const QString &filename);
+    void requestIDR();
 
 public slots:
     void sendPacket(const QByteArray &data);
+    void sendAudio(const QByteArray &data);
+    void enableRecording(bool e);
+    void startRecording();
+    void stopRecording();
 
 private:
+    int64_t getCurrentTimestamp(AVRational timebase);
+    
     AVCodecContext *m_codecCtx;
     AVPacket *m_packet;
     AVFrame *m_frame;
@@ -32,7 +42,11 @@ private:
     AVFilterContext *m_buffersrcCtx;
     AVFilterContext *m_buffersinkCtx;
 
-    QByteArray m_currentPacket;
+    AVFormatContext *m_recordingCtx;
+    AVStream *m_videoStream;
+    AVStream *m_audioStream;
+    QString m_recordingFilename;
+    int64_t m_recordingStartTime;
 
 };
 

app/viewer.h

@@ -9,6 +9,8 @@ class Viewer : public QOpenGLWidget
 public:
     Viewer(QWidget *parent = nullptr);
 
+    const QImage &image() const { return m_image; }
+
 public slots:
     void setImage(const QImage &image);
 

lib/gamepad/gamepad.c

@@ -12,6 +12,7 @@
 #include <wpa_ctrl.h>
 
 #include "audio.h"
+#include "command.h"
 #include "input.h"
 #include "video.h"
 
@@ -64,6 +65,16 @@ int create_socket(int *socket_out, uint16_t port)
     return 1;
 }
 
+void send_stop_code(int from_socket, in_port_t port)
+{
+    struct sockaddr_in address;
+    address.sin_family = AF_INET;
+    address.sin_addr.s_addr = inet_addr("127.0.0.1");
+
+    address.sin_port = htons(port);
+    sendto(from_socket, &STOP_CODE, sizeof(STOP_CODE), 0, (struct sockaddr *)&address, sizeof(address));
+}
+
 int main_loop(vanilla_event_handler_t event_handler, void *context)
 {
     struct gamepad_thread_context info;
@@ -84,20 +95,15 @@ int main_loop(vanilla_event_handler_t event_handler, void *context)
     pthread_create(&video_thread, NULL, listen_video, &info);
     pthread_create(&audio_thread, NULL, listen_audio, &info);
     pthread_create(&input_thread, NULL, listen_input, &info);
+    pthread_create(&cmd_thread, NULL, listen_command, &info);
 
     while (1) {
         usleep(250 * 1000);
         if (is_interrupted()) {
             // Wake up any threads that might be blocked on `recv`
-            struct sockaddr_in address;
+            send_stop_code(info.socket_msg, PORT_VID);
-            address.sin_family = AF_INET;
+            send_stop_code(info.socket_msg, PORT_AUD);
-            address.sin_addr.s_addr = inet_addr("127.0.0.1");
+            send_stop_code(info.socket_msg, PORT_CMD);
-
-            address.sin_port = htons(PORT_VID);
-            sendto(info.socket_msg, &STOP_CODE, sizeof(STOP_CODE), 0, (struct sockaddr *) &address, sizeof(address));
-
-            address.sin_port = htons(PORT_AUD);
-            sendto(info.socket_msg, &STOP_CODE, sizeof(STOP_CODE), 0, (struct sockaddr *) &address, sizeof(address));
             break;
         }
     }
@@ -105,6 +111,7 @@ int main_loop(vanilla_event_handler_t event_handler, void *context)
     pthread_join(video_thread, NULL);
     pthread_join(audio_thread, NULL);
     pthread_join(input_thread, NULL);
+    pthread_join(cmd_thread, NULL);
 
     ret = VANILLA_SUCCESS;
 

lib/gamepad/video.c

@@ -10,6 +10,7 @@
 
 #include "gamepad.h"
 #include "vanilla.h"
+#include "status.h"
 #include "util.h"
 
 typedef struct
@@ -28,6 +29,23 @@ typedef struct
     uint8_t payload[2048];
 } VideoPacket;
 
+pthread_mutex_t video_mutex;
+int idr_is_queued = 0;
+
+void request_idr()
+{
+    pthread_mutex_lock(&video_mutex);
+    idr_is_queued = 1;
+    pthread_mutex_unlock(&video_mutex);
+}
+
+void send_idr_request_to_console(int socket_msg)
+{
+    // Make an IDR request to the Wii U?
+    unsigned char idr_request[] = {1, 0, 0, 0}; // Undocumented
+    send_to_console(socket_msg, idr_request, sizeof(idr_request), PORT_MSG);
+}
+
 void handle_video_packet(vanilla_event_handler_t event_handler, void *context, unsigned char *data, size_t size, int socket_msg)
 {
     // TODO: This is all really weird. Copied from drc-sim-c but I feel like there's probably a better way.
@@ -81,14 +99,19 @@ void handle_video_packet(vanilla_event_handler_t event_handler, void *context, u
             if (is_idr) {
                 is_streaming = 1;
             } else {
-                // Make an IDR request to the Wii U?
+                send_idr_request_to_console(socket_msg);
-                unsigned char idr_request[] = {1, 0, 0, 0}; // Undocumented
-                send_to_console(socket_msg, idr_request, sizeof(idr_request), PORT_MSG);
                 return;
             }
         }
     }
 
+    pthread_mutex_lock(&video_mutex);
+    if (idr_is_queued) {
+        send_idr_request_to_console(socket_msg);
+        idr_is_queued = 0;
+    }
+    pthread_mutex_unlock(&video_mutex);
+
     memcpy(video_packet + video_packet_size, vp->payload, vp->payload_size);
     video_packet_size += vp->payload_size;
 
@@ -102,16 +125,9 @@ void handle_video_packet(vanilla_event_handler_t event_handler, void *context, u
             int slice_header = is_idr ? 0x25b804ff : (0x21e003ff | ((frame_decode_num & 0xff) << 13));
             frame_decode_num++;
 
-            uint8_t params[] = {
-                    // sps
-                    0x00, 0x00, 0x00, 0x01, 0x67, 0x64, 0x00, 0x20, 0xac, 0x2b, 0x40, 0x6c, 0x1e, 0xf3, 0x68,
-                    // pps
-                    0x00, 0x00, 0x00, 0x01, 0x68, 0xee, 0x06, 0x0c, 0xe8
-            };
-
             if (is_idr) {
-                memcpy(nals_current, params, sizeof(params));
+                memcpy(nals_current, sps_pps_params, sizeof(sps_pps_params));
-                nals_current += sizeof(params);
+                nals_current += sizeof(sps_pps_params);
             }
 
             // begin slice nalu
@@ -154,6 +170,7 @@ void *listen_video(void *x)
     unsigned char data[2048];
     ssize_t size;
 
+    pthread_mutex_init(&video_mutex, NULL);
 
     do {
         size = recv(info->socket_vid, data, sizeof(data), 0);
@@ -163,6 +180,8 @@ void *listen_video(void *x)
         }
     } while (!is_interrupted());
 
+    pthread_mutex_destroy(&video_mutex);
+
     pthread_exit(NULL);
 
     return NULL;

lib/gamepad/video.h

@@ -1,6 +1,16 @@
 #ifndef GAMEPAD_VIDEO_H
 #define GAMEPAD_VIDEO_H
 
+#include <stdint.h>
+
 void *listen_video(void *x);
+void request_idr();
+
+static const uint8_t sps_pps_params[] = {
+    // sps
+    0x00, 0x00, 0x00, 0x01, 0x67, 0x64, 0x00, 0x20, 0xac, 0x2b, 0x40, 0x6c, 0x1e, 0xf3, 0x68,
+    // pps
+    0x00, 0x00, 0x00, 0x01, 0x68, 0xee, 0x06, 0x0c, 0xe8
+};
 
 #endif // GAMEPAD_VIDEO_H

lib/vanilla.c

@@ -1,11 +1,13 @@
 #include "vanilla.h"
 
 #include <signal.h>
+#include <string.h>
 #include <unistd.h>
 #include <wpa_ctrl.h>
 
 #include "gamepad/gamepad.h"
 #include "gamepad/input.h"
+#include "gamepad/video.h"
 #include "status.h"
 #include "sync.h"
 #include "util.h"
@@ -124,4 +126,17 @@ void vanilla_log_no_newline_va(const char *format, va_list args)
 void vanilla_install_logger(void (*logger)(const char *, va_list))
 {
     custom_logger = logger;
+}
+
+void vanilla_request_idr()
+{
+    request_idr();
+}
+
+void vanilla_retrieve_sps_pps_data(void *data, size_t *size)
+{
+    if (data != NULL) {
+        memcpy(data, sps_pps_params, MIN(*size, sizeof(sps_pps_params)));
+    }
+    *size = sizeof(sps_pps_params);
 }

lib/vanilla.h

@@ -128,6 +128,19 @@ void vanilla_log_no_newline_va(const char *format, va_list args);
  */
 void vanilla_install_logger(void (*logger)(const char *, va_list args));
 
+/**
+ * Request an IDR (instant decoder refresh) video frame from the console
+ */
+void vanilla_request_idr();
+
+/**
+ * Retrieve SPS/PPS data for H.264 encoding
+ * 
+ * If `data` is null, `*size` will be set to the number of bytes required.
+ * If `data` is not null, bytes will be copied up to `*size` or the total number of bytes.
+ */
+void vanilla_retrieve_sps_pps_data(void *data, size_t *size);
+
 #if defined(__cplusplus)
 }
 #endif

pipe/main.c

@@ -239,6 +239,11 @@ int main()
             }
             break;
         }
+        case VANILLA_PIPE_IN_REQ_IDR:
+        {
+            vanilla_request_idr();
+            break;
+        }
         case VANILLA_PIPE_IN_QUIT:
             m_quit = 1;
             break;

pipe/pipe.h

@@ -6,6 +6,7 @@
 #define VANILLA_PIPE_IN_CONNECT      0x02
 #define VANILLA_PIPE_IN_BUTTON       0x03
 #define VANILLA_PIPE_IN_TOUCH        0x04
+#define VANILLA_PIPE_IN_REQ_IDR      0x05
 #define VANILLA_PIPE_IN_INTERRUPT    0x1E
 #define VANILLA_PIPE_IN_QUIT         0x1F