From fff17cae7738d6b3d65426a0ce7df76e8034ad1c Mon Sep 17 00:00:00 2001
From: Andri Yngvason <andri@yngvason.is>
Date: Wed, 24 Jul 2024 20:50:14 +0000
Subject: [PATCH] H264Libav: Use AVFrame for pixel format conversion buffer

This ensures that the buffer is allocated with the correct alignment and
padding for use with sws_scale.

This fixes out-of-bounds writes which would in some cases cause
segmentation faults and/or heap corruption.
---
 common/rfb/H264LibavDecoderContext.cxx | 27 ++++++++++++++++----------
 common/rfb/H264LibavDecoderContext.h   |  2 +-
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/common/rfb/H264LibavDecoderContext.cxx b/common/rfb/H264LibavDecoderContext.cxx
index 5bd1dbbf..2bd7d343 100644
--- a/common/rfb/H264LibavDecoderContext.cxx
+++ b/common/rfb/H264LibavDecoderContext.cxx
@@ -43,7 +43,6 @@ bool H264LibavDecoderContext::initCodec() {
   os::AutoMutex lock(&mutex);
 
   sws = nullptr;
-  swsBuffer = nullptr;
   h264WorkBuffer = nullptr;
   h264WorkBufferLength = 0;
 
@@ -87,9 +86,6 @@ bool H264LibavDecoderContext::initCodec() {
     return false;
   }
 
-  int numBytes = av_image_get_buffer_size(AV_PIX_FMT_RGB32, rect.width(), rect.height(), 1);
-  swsBuffer = new uint8_t[numBytes];
-
   initialized = true;
   return true;
 }
@@ -101,8 +97,8 @@ void H264LibavDecoderContext::freeCodec() {
     return;
   av_parser_close(parser);
   avcodec_free_context(&avctx);
+  av_frame_free(&rgbFrame);
   av_frame_free(&frame);
-  delete[] swsBuffer;
   free(h264WorkBuffer);
   initialized = false;
 }
@@ -220,11 +216,22 @@ void H264LibavDecoderContext::decode(const uint8_t* h264_in_buffer,
                              frame->width, frame->height, AV_PIX_FMT_RGB32,
                              0, nullptr, nullptr, nullptr);
 
-  int stride;
-  pb->getBuffer(rect, &stride);
-  int dst_linesize = stride * pb->getPF().bpp/8;  // stride is in pixels, linesize is in bytes (stride x4). We need bytes
+  if (rgbFrame && (rgbFrame->width != frame->width || rgbFrame->height != frame->height)) {
+    av_frame_free(&rgbFrame);
+
+  }
+
+  if (!rgbFrame) {
+    rgbFrame = av_frame_alloc();
+    // TODO: Can we really assume that the pixel format will always be RGB32?
+    rgbFrame->format = AV_PIX_FMT_RGB32;
+    rgbFrame->width = frame->width;
+    rgbFrame->height = frame->height;
+    av_frame_get_buffer(rgbFrame, 0);
+  }
 
-  sws_scale(sws, frame->data, frame->linesize, 0, frame->height, &swsBuffer, &dst_linesize);
+  sws_scale(sws, frame->data, frame->linesize, 0, frame->height, rgbFrame->data,
+            rgbFrame->linesize);
 
-  pb->imageRect(rect, swsBuffer, stride);
+  pb->imageRect(rect, rgbFrame->data[0], rgbFrame->linesize[0] / 4);
 }
diff --git a/common/rfb/H264LibavDecoderContext.h b/common/rfb/H264LibavDecoderContext.h
index f399b3cc..96558bee 100644
--- a/common/rfb/H264LibavDecoderContext.h
+++ b/common/rfb/H264LibavDecoderContext.h
@@ -47,8 +47,8 @@ namespace rfb {
       AVCodecContext *avctx;
       AVCodecParserContext *parser;
       AVFrame* frame;
+      AVFrame* rgbFrame = nullptr;
       SwsContext* sws;
-      uint8_t* swsBuffer;
       uint8_t* h264WorkBuffer;
       uint32_t h264WorkBufferLength;
   };
-- 
2.39.5