From 633c1abf44bd4d1ef37623571ac6b4597b8b9ef6 Mon Sep 17 00:00:00 2001
From: Ely Ronnen <elyronnen@gmail.com>
Date: Tue, 25 Jan 2022 17:35:10 +0200
Subject: [PATCH] handling small input buffer case

---
 libfreerdp/codec/h264.c            | 14 ++++----
 libfreerdp/codec/h264_mediacodec.c | 53 +++++++++++++++++-------------
 2 files changed, 37 insertions(+), 30 deletions(-)

diff --git a/libfreerdp/codec/h264.c b/libfreerdp/codec/h264.c
index 3ef91092c..a52cf0dc5 100644
--- a/libfreerdp/codec/h264.c
+++ b/libfreerdp/codec/h264.c
@@ -582,13 +582,6 @@ static BOOL CALLBACK h264_register_subsystems(PINIT_ONCE once, PVOID param, PVOI
 		i++;
 	}
 #endif
-#ifdef WITH_MEDIACODEC
-	{
-		extern H264_CONTEXT_SUBSYSTEM g_Subsystem_mediacodec;
-		subSystems[i] = &g_Subsystem_mediacodec;
-		i++;
-	}
-#endif
 #ifdef WITH_OPENH264
 	{
 		extern H264_CONTEXT_SUBSYSTEM g_Subsystem_OpenH264;
@@ -602,6 +595,13 @@ static BOOL CALLBACK h264_register_subsystems(PINIT_ONCE once, PVOID param, PVOI
 		subSystems[i] = &g_Subsystem_libavcodec;
 		i++;
 	}
+#endif
+#ifdef WITH_MEDIACODEC
+	{
+		extern H264_CONTEXT_SUBSYSTEM g_Subsystem_mediacodec;
+		subSystems[i] = &g_Subsystem_mediacodec;
+		i++;
+	}
 #endif
 	return i > 0;
 }
diff --git a/libfreerdp/codec/h264_mediacodec.c b/libfreerdp/codec/h264_mediacodec.c
index 3cecc2680..dd718768b 100644
--- a/libfreerdp/codec/h264_mediacodec.c
+++ b/libfreerdp/codec/h264_mediacodec.c
@@ -248,33 +248,40 @@ static int mediacodec_decompress(H264_CONTEXT* h264, const BYTE* pSrcData, UINT3
 
     while (true)
     {
-        inputBufferId = sys->fnAMediaCodec_dequeueInputBuffer(sys->decoder, -1);
-        if (inputBufferId < 0)
+        UINT32 inputBufferCurrnetOffset = 0;
+        while (inputBufferCurrnetOffset < SrcSize)
         {
-            WLog_Print(h264->log, WLOG_ERROR, "AMediaCodec_dequeueInputBuffer failed [%d]", inputBufferId);
-            // TODO: sleep?
-            continue;
-        }
+            UINT32 numberOfBytesToCopy = SrcSize - inputBufferCurrnetOffset;
+            inputBufferId = sys->fnAMediaCodec_dequeueInputBuffer(sys->decoder, -1);
+            if (inputBufferId < 0)
+            {
+                WLog_Print(h264->log, WLOG_ERROR, "AMediaCodec_dequeueInputBuffer failed [%d]", inputBufferId);
+                // TODO: sleep?
+                continue;
+            }
 
-        inputBuffer = sys->fnAMediaCodec_getInputBuffer(sys->decoder, inputBufferId, &inputBufferSize);
-        if (inputBuffer == NULL)
-        {
-            WLog_Print(h264->log, WLOG_ERROR, "AMediaCodec_getInputBuffer failed");
-            return -1;
-        }
+            inputBuffer = sys->fnAMediaCodec_getInputBuffer(sys->decoder, inputBufferId, &inputBufferSize);
+            if (inputBuffer == NULL)
+            {
+                WLog_Print(h264->log, WLOG_ERROR, "AMediaCodec_getInputBuffer failed");
+                return -1;
+            }
 
-        if (SrcSize > inputBufferSize)
-        {
-            WLog_Print(h264->log, WLOG_ERROR, "MediaCodec input buffer size is too small: got [%d] but wanted [%d]", inputBufferSize, SrcSize);
-            return -1;
-        }
+            if (numberOfBytesToCopy > inputBufferSize)
+            {
+                WLog_Print(h264->log, WLOG_WARN, "MediaCodec inputBufferSize: got [%d] but wanted [%d]", inputBufferSize, numberOfBytesToCopy);
+                numberOfBytesToCopy = inputBufferSize;
+            }
 
-        memcpy(inputBuffer, pSrcData, SrcSize);
-        status = sys->fnAMediaCodec_queueInputBuffer(sys->decoder, inputBufferId, 0, SrcSize, 0, 0);
-        if (status != AMEDIA_OK)
-        {
-            WLog_Print(h264->log, WLOG_ERROR, "Error AMediaCodec_queueInputBuffer %d", status);
-            return -1;
+            memcpy(inputBuffer, pSrcData + inputBufferCurrnetOffset, numberOfBytesToCopy);
+            inputBufferCurrnetOffset += numberOfBytesToCopy;
+
+            status = sys->fnAMediaCodec_queueInputBuffer(sys->decoder, inputBufferId, 0, numberOfBytesToCopy, 0, 0);
+            if (status != AMEDIA_OK)
+            {
+                WLog_Print(h264->log, WLOG_ERROR, "Error AMediaCodec_queueInputBuffer %d", status);
+                return -1;
+            }
         }
 
         while (true)