CCodec: handle input image for input buffers as well

Bug: 126743446
Test: atest CtsMediaTestCases:VideoEncoderTest#testOtherH264Flex480p
Change-Id: Idb33f27e43bd579ea4b0e7877581db79ec3c3223
diff --git a/media/codec2/sfplugin/CCodecBufferChannel.cpp b/media/codec2/sfplugin/CCodecBufferChannel.cpp
index 7a444a3..fb6af93 100644
--- a/media/codec2/sfplugin/CCodecBufferChannel.cpp
+++ b/media/codec2/sfplugin/CCodecBufferChannel.cpp
@@ -99,6 +99,34 @@
      */
     virtual size_t numClientBuffers() const = 0;
 
+    void handleImageData(const sp<Codec2Buffer> &buffer) {
+        sp<ABuffer> imageDataCandidate = buffer->getImageData();
+        if (imageDataCandidate == nullptr) {
+            return;
+        }
+        sp<ABuffer> imageData;
+        if (!mFormat->findBuffer("image-data", &imageData)
+                || imageDataCandidate->size() != imageData->size()
+                || memcmp(imageDataCandidate->data(), imageData->data(), imageData->size()) != 0) {
+            ALOGD("[%s] updating image-data", mName);
+            sp<AMessage> newFormat = dupFormat();
+            newFormat->setBuffer("image-data", imageDataCandidate);
+            MediaImage2 *img = (MediaImage2*)imageDataCandidate->data();
+            if (img->mNumPlanes > 0 && img->mType != img->MEDIA_IMAGE_TYPE_UNKNOWN) {
+                int32_t stride = img->mPlane[0].mRowInc;
+                newFormat->setInt32(KEY_STRIDE, stride);
+                ALOGD("[%s] updating stride = %d", mName, stride);
+                if (img->mNumPlanes > 1 && stride > 0) {
+                    int32_t vstride = (img->mPlane[1].mOffset - img->mPlane[0].mOffset) / stride;
+                    newFormat->setInt32(KEY_SLICE_HEIGHT, vstride);
+                    ALOGD("[%s] updating vstride = %d", mName, vstride);
+                }
+            }
+            setFormat(newFormat);
+            buffer->setFormat(newFormat);
+        }
+    }
+
 protected:
     std::string mComponentName; ///< name of component for debugging
     std::string mChannelName; ///< name of channel for debugging
@@ -255,34 +283,6 @@
         mSkipCutBuffer = scb;
     }
 
-    void handleImageData(const sp<Codec2Buffer> &buffer) {
-        sp<ABuffer> imageDataCandidate = buffer->getImageData();
-        if (imageDataCandidate == nullptr) {
-            return;
-        }
-        sp<ABuffer> imageData;
-        if (!mFormat->findBuffer("image-data", &imageData)
-                || imageDataCandidate->size() != imageData->size()
-                || memcmp(imageDataCandidate->data(), imageData->data(), imageData->size()) != 0) {
-            ALOGD("[%s] updating image-data", mName);
-            sp<AMessage> newFormat = dupFormat();
-            newFormat->setBuffer("image-data", imageDataCandidate);
-            MediaImage2 *img = (MediaImage2*)imageDataCandidate->data();
-            if (img->mNumPlanes > 0 && img->mType != img->MEDIA_IMAGE_TYPE_UNKNOWN) {
-                int32_t stride = img->mPlane[0].mRowInc;
-                newFormat->setInt32(KEY_STRIDE, stride);
-                ALOGD("[%s] updating stride = %d", mName, stride);
-                if (img->mNumPlanes > 1 && stride > 0) {
-                    int32_t vstride = (img->mPlane[1].mOffset - img->mPlane[0].mOffset) / stride;
-                    newFormat->setInt32(KEY_SLICE_HEIGHT, vstride);
-                    ALOGD("[%s] updating vstride = %d", mName, vstride);
-                }
-            }
-            setFormat(newFormat);
-            buffer->setFormat(newFormat);
-        }
-    }
-
 protected:
     sp<SkipCutBuffer> mSkipCutBuffer;
 
@@ -783,6 +783,7 @@
         status_t err = mImpl.grabBuffer(index, &c2Buffer);
         if (err == OK) {
             c2Buffer->setFormat(mFormat);
+            handleImageData(c2Buffer);
             *buffer = c2Buffer;
             return true;
         }
@@ -1053,6 +1054,7 @@
             return false;
         }
         *index = mImpl.assignSlot(newBuffer);
+        handleImageData(newBuffer);
         *buffer = newBuffer;
         return true;
     }