/*
 * Copyright (C) 2025 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.
 */

#define private public

#include <binder/Parcel.h>
#include <binder/RpcSession.h>

#include <vector>

#include "../includes/common.h"
#include "binary_loader.h"

using namespace android;

typedef status_t (*truncate_rpc_objects_func)(void*, size_t);

int main(int /* argc */, char* argv[]) {
    // Get the path to the shared library.
    BinaryLoader binaryLoader(argv[1]);

    // Get the function address of 'Parcel::truncateRpcObjects()' from loaded library.
    const uintptr_t functionOffset = strtoul(argv[2], nullptr, 0);
    const uintptr_t truncate_rpc_objects_addr = binaryLoader.getFunctionAddress(functionOffset);
    FAIL_CHECK(truncate_rpc_objects_addr);

    // Create the function pointer to 'Parcel::truncateRpcObjects()'.
    truncate_rpc_objects_func truncate_rpc_objects_ptr =
            (truncate_rpc_objects_func)truncate_rpc_objects_addr;

    // Create an object for 'Parcel' and assign 'session' of type 'RpcSession' to
    // 'Parcel::mVariantFields' by invoking 'Parcel::markForRpc()'.
    Parcel parcel;
    const sp<RpcSession> session = sp<RpcSession>(RpcSession::make());
    parcel.markForRpc(session);

    // Create buffer of size 1 to avoid segmentation fault and assign it to 'Parcel::mData'.
    std::vector<uint8_t> buffer(1);
    parcel.mData = buffer.data();
    parcel.mDataSize = 0;

    // Configure the argument newObjectSize passed to 'Parcel::truncateRpcObjects()' to a value
    // greater than '0' to reach the vulnerable condition. Resize the vector
    // 'Parcel::RpcFields::mObjectPositions' greater than newObjectSize to check if
    // 'Parcel::RpcFields::mObjectPositions' is modified.
    int newObjectSize = 1;
    int initialObjectPositionsSize = newObjectSize + 1;
    auto* rpcFields = parcel.maybeRpcFields();
    rpcFields->mObjectPositions.resize(initialObjectPositionsSize);
    rpcFields->mObjectPositions.insert(rpcFields->mObjectPositions.end(), {0, 0});

    // Adding a check to ensure the fix condition added in truncateRpcObjects() is satisfied
    uint32_t pos = rpcFields->mObjectPositions.back(), minObjectEnd;
    if (!(__builtin_add_overflow(pos, sizeof(Parcel::RpcFields::ObjectType), &minObjectEnd) ||
          minObjectEnd > parcel.mDataSize)) {
        return EXIT_FAILURE;
    }

    // Without fix, 'Parcel::truncateRpcObjects()' tries to delete 'Parcel::RpcFields' even if
    // 'Parcel::mDataSize' indicates no data. With fix, the checks prevent this behavior.
    status_t status = truncate_rpc_objects_ptr(&parcel, newObjectSize);
    parcel.mData = nullptr;

    // Fail the test when vulnerable function returns 'OK' and 'Parcel::RpcFields::mObjectPositions'
    // is modified.
    if (status == OK && rpcFields->mObjectPositions.size() != initialObjectPositionsSize) {
        return EXIT_VULNERABLE;
    }
    return EXIT_SUCCESS;
}