cachepc-linux

Fork of AMDESE/linux with modifications for CachePC side-channel attack
git clone https://git.sinitax.com/sinitax/cachepc-linux
Log | Files | Refs | README | LICENSE | sfeed.txt

ne_ioctl_sample.c (20900B)


      1// SPDX-License-Identifier: GPL-2.0
      2/*
      3 * Copyright 2020-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
      4 */
      5
      6/**
      7 * DOC: Sample flow of using the ioctl interface provided by the Nitro Enclaves (NE)
      8 * kernel driver.
      9 *
     10 * Usage
     11 * -----
     12 *
     13 * Load the nitro_enclaves module, setting also the enclave CPU pool. The
     14 * enclave CPUs need to be full cores from the same NUMA node. CPU 0 and its
     15 * siblings have to remain available for the primary / parent VM, so they
     16 * cannot be included in the enclave CPU pool.
     17 *
     18 * See the cpu list section from the kernel documentation.
     19 * https://www.kernel.org/doc/html/latest/admin-guide/kernel-parameters.html#cpu-lists
     20 *
     21 *	insmod drivers/virt/nitro_enclaves/nitro_enclaves.ko
     22 *	lsmod
     23 *
     24 *	The CPU pool can be set at runtime, after the kernel module is loaded.
     25 *
     26 *	echo <cpu-list> > /sys/module/nitro_enclaves/parameters/ne_cpus
     27 *
     28 *	NUMA and CPU siblings information can be found using:
     29 *
     30 *	lscpu
     31 *	/proc/cpuinfo
     32 *
     33 * Check the online / offline CPU list. The CPUs from the pool should be
     34 * offlined.
     35 *
     36 *	lscpu
     37 *
     38 * Check dmesg for any warnings / errors through the NE driver lifetime / usage.
     39 * The NE logs contain the "nitro_enclaves" or "pci 0000:00:02.0" pattern.
     40 *
     41 *	dmesg
     42 *
     43 * Setup hugetlbfs huge pages. The memory needs to be from the same NUMA node as
     44 * the enclave CPUs.
     45 *
     46 * https://www.kernel.org/doc/html/latest/admin-guide/mm/hugetlbpage.html
     47 *
     48 * By default, the allocation of hugetlb pages are distributed on all possible
     49 * NUMA nodes. Use the following configuration files to set the number of huge
     50 * pages from a NUMA node:
     51 *
     52 *	/sys/devices/system/node/node<X>/hugepages/hugepages-2048kB/nr_hugepages
     53 *	/sys/devices/system/node/node<X>/hugepages/hugepages-1048576kB/nr_hugepages
     54 *
     55 *	or, if not on a system with multiple NUMA nodes, can also set the number
     56 *	of 2 MiB / 1 GiB huge pages using
     57 *
     58 *	/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages
     59 *	/sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages
     60 *
     61 *	In this example 256 hugepages of 2 MiB are used.
     62 *
     63 * Build and run the NE sample.
     64 *
     65 *	make -C samples/nitro_enclaves clean
     66 *	make -C samples/nitro_enclaves
     67 *	./samples/nitro_enclaves/ne_ioctl_sample <path_to_enclave_image>
     68 *
     69 * Unload the nitro_enclaves module.
     70 *
     71 *	rmmod nitro_enclaves
     72 *	lsmod
     73 */
     74
     75#include <stdio.h>
     76#include <stdlib.h>
     77#include <errno.h>
     78#include <fcntl.h>
     79#include <limits.h>
     80#include <poll.h>
     81#include <pthread.h>
     82#include <string.h>
     83#include <sys/eventfd.h>
     84#include <sys/ioctl.h>
     85#include <sys/mman.h>
     86#include <sys/socket.h>
     87#include <sys/stat.h>
     88#include <sys/types.h>
     89#include <unistd.h>
     90
     91#include <linux/mman.h>
     92#include <linux/nitro_enclaves.h>
     93#include <linux/vm_sockets.h>
     94
     95/**
     96 * NE_DEV_NAME - Nitro Enclaves (NE) misc device that provides the ioctl interface.
     97 */
     98#define NE_DEV_NAME			"/dev/nitro_enclaves"
     99
    100/**
    101 * NE_POLL_WAIT_TIME - Timeout in seconds for each poll event.
    102 */
    103#define NE_POLL_WAIT_TIME		(60)
    104/**
    105 * NE_POLL_WAIT_TIME_MS - Timeout in milliseconds for each poll event.
    106 */
    107#define NE_POLL_WAIT_TIME_MS		(NE_POLL_WAIT_TIME * 1000)
    108
    109/**
    110 * NE_SLEEP_TIME - Amount of time in seconds for the process to keep the enclave alive.
    111 */
    112#define NE_SLEEP_TIME			(300)
    113
    114/**
    115 * NE_DEFAULT_NR_VCPUS - Default number of vCPUs set for an enclave.
    116 */
    117#define NE_DEFAULT_NR_VCPUS		(2)
    118
    119/**
    120 * NE_MIN_MEM_REGION_SIZE - Minimum size of a memory region - 2 MiB.
    121 */
    122#define NE_MIN_MEM_REGION_SIZE		(2 * 1024 * 1024)
    123
    124/**
    125 * NE_DEFAULT_NR_MEM_REGIONS - Default number of memory regions of 2 MiB set for
    126 *			       an enclave.
    127 */
    128#define NE_DEFAULT_NR_MEM_REGIONS	(256)
    129
    130/**
    131 * NE_IMAGE_LOAD_HEARTBEAT_CID - Vsock CID for enclave image loading heartbeat logic.
    132 */
    133#define NE_IMAGE_LOAD_HEARTBEAT_CID	(3)
    134/**
    135 * NE_IMAGE_LOAD_HEARTBEAT_PORT - Vsock port for enclave image loading heartbeat logic.
    136 */
    137#define NE_IMAGE_LOAD_HEARTBEAT_PORT	(9000)
    138/**
    139 * NE_IMAGE_LOAD_HEARTBEAT_VALUE - Heartbeat value for enclave image loading.
    140 */
    141#define NE_IMAGE_LOAD_HEARTBEAT_VALUE	(0xb7)
    142
    143/**
    144 * struct ne_user_mem_region - User space memory region set for an enclave.
    145 * @userspace_addr:	Address of the user space memory region.
    146 * @memory_size:	Size of the user space memory region.
    147 */
    148struct ne_user_mem_region {
    149	void	*userspace_addr;
    150	size_t	memory_size;
    151};
    152
    153/**
    154 * ne_create_vm() - Create a slot for the enclave VM.
    155 * @ne_dev_fd:		The file descriptor of the NE misc device.
    156 * @slot_uid:		The generated slot uid for the enclave.
    157 * @enclave_fd :	The generated file descriptor for the enclave.
    158 *
    159 * Context: Process context.
    160 * Return:
    161 * * 0 on success.
    162 * * Negative return value on failure.
    163 */
    164static int ne_create_vm(int ne_dev_fd, unsigned long *slot_uid, int *enclave_fd)
    165{
    166	int rc = -EINVAL;
    167	*enclave_fd = ioctl(ne_dev_fd, NE_CREATE_VM, slot_uid);
    168
    169	if (*enclave_fd < 0) {
    170		rc = *enclave_fd;
    171		switch (errno) {
    172		case NE_ERR_NO_CPUS_AVAIL_IN_POOL: {
    173			printf("Error in create VM, no CPUs available in the NE CPU pool\n");
    174
    175			break;
    176		}
    177
    178		default:
    179			printf("Error in create VM [%m]\n");
    180		}
    181
    182		return rc;
    183	}
    184
    185	return 0;
    186}
    187
    188/**
    189 * ne_poll_enclave_fd() - Thread function for polling the enclave fd.
    190 * @data:	Argument provided for the polling function.
    191 *
    192 * Context: Process context.
    193 * Return:
    194 * * NULL on success / failure.
    195 */
    196void *ne_poll_enclave_fd(void *data)
    197{
    198	int enclave_fd = *(int *)data;
    199	struct pollfd fds[1] = {};
    200	int i = 0;
    201	int rc = -EINVAL;
    202
    203	printf("Running from poll thread, enclave fd %d\n", enclave_fd);
    204
    205	fds[0].fd = enclave_fd;
    206	fds[0].events = POLLIN | POLLERR | POLLHUP;
    207
    208	/* Keep on polling until the current process is terminated. */
    209	while (1) {
    210		printf("[iter %d] Polling ...\n", i);
    211
    212		rc = poll(fds, 1, NE_POLL_WAIT_TIME_MS);
    213		if (rc < 0) {
    214			printf("Error in poll [%m]\n");
    215
    216			return NULL;
    217		}
    218
    219		i++;
    220
    221		if (!rc) {
    222			printf("Poll: %d seconds elapsed\n",
    223			       i * NE_POLL_WAIT_TIME);
    224
    225			continue;
    226		}
    227
    228		printf("Poll received value 0x%x\n", fds[0].revents);
    229
    230		if (fds[0].revents & POLLHUP) {
    231			printf("Received POLLHUP\n");
    232
    233			return NULL;
    234		}
    235
    236		if (fds[0].revents & POLLNVAL) {
    237			printf("Received POLLNVAL\n");
    238
    239			return NULL;
    240		}
    241	}
    242
    243	return NULL;
    244}
    245
    246/**
    247 * ne_alloc_user_mem_region() - Allocate a user space memory region for an enclave.
    248 * @ne_user_mem_region:	User space memory region allocated using hugetlbfs.
    249 *
    250 * Context: Process context.
    251 * Return:
    252 * * 0 on success.
    253 * * Negative return value on failure.
    254 */
    255static int ne_alloc_user_mem_region(struct ne_user_mem_region *ne_user_mem_region)
    256{
    257	/**
    258	 * Check available hugetlb encodings for different huge page sizes in
    259	 * include/uapi/linux/mman.h.
    260	 */
    261	ne_user_mem_region->userspace_addr = mmap(NULL, ne_user_mem_region->memory_size,
    262						  PROT_READ | PROT_WRITE,
    263						  MAP_PRIVATE | MAP_ANONYMOUS |
    264						  MAP_HUGETLB | MAP_HUGE_2MB, -1, 0);
    265	if (ne_user_mem_region->userspace_addr == MAP_FAILED) {
    266		printf("Error in mmap memory [%m]\n");
    267
    268		return -1;
    269	}
    270
    271	return 0;
    272}
    273
    274/**
    275 * ne_load_enclave_image() - Place the enclave image in the enclave memory.
    276 * @enclave_fd :		The file descriptor associated with the enclave.
    277 * @ne_user_mem_regions:	User space memory regions allocated for the enclave.
    278 * @enclave_image_path :	The file path of the enclave image.
    279 *
    280 * Context: Process context.
    281 * Return:
    282 * * 0 on success.
    283 * * Negative return value on failure.
    284 */
    285static int ne_load_enclave_image(int enclave_fd, struct ne_user_mem_region ne_user_mem_regions[],
    286				 char *enclave_image_path)
    287{
    288	unsigned char *enclave_image = NULL;
    289	int enclave_image_fd = -1;
    290	size_t enclave_image_size = 0;
    291	size_t enclave_memory_size = 0;
    292	unsigned long i = 0;
    293	size_t image_written_bytes = 0;
    294	struct ne_image_load_info image_load_info = {
    295		.flags = NE_EIF_IMAGE,
    296	};
    297	struct stat image_stat_buf = {};
    298	int rc = -EINVAL;
    299	size_t temp_image_offset = 0;
    300
    301	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++)
    302		enclave_memory_size += ne_user_mem_regions[i].memory_size;
    303
    304	rc = stat(enclave_image_path, &image_stat_buf);
    305	if (rc < 0) {
    306		printf("Error in get image stat info [%m]\n");
    307
    308		return rc;
    309	}
    310
    311	enclave_image_size = image_stat_buf.st_size;
    312
    313	if (enclave_memory_size < enclave_image_size) {
    314		printf("The enclave memory is smaller than the enclave image size\n");
    315
    316		return -ENOMEM;
    317	}
    318
    319	rc = ioctl(enclave_fd, NE_GET_IMAGE_LOAD_INFO, &image_load_info);
    320	if (rc < 0) {
    321		switch (errno) {
    322		case NE_ERR_NOT_IN_INIT_STATE: {
    323			printf("Error in get image load info, enclave not in init state\n");
    324
    325			break;
    326		}
    327
    328		case NE_ERR_INVALID_FLAG_VALUE: {
    329			printf("Error in get image load info, provided invalid flag\n");
    330
    331			break;
    332		}
    333
    334		default:
    335			printf("Error in get image load info [%m]\n");
    336		}
    337
    338		return rc;
    339	}
    340
    341	printf("Enclave image offset in enclave memory is %lld\n",
    342	       image_load_info.memory_offset);
    343
    344	enclave_image_fd = open(enclave_image_path, O_RDONLY);
    345	if (enclave_image_fd < 0) {
    346		printf("Error in open enclave image file [%m]\n");
    347
    348		return enclave_image_fd;
    349	}
    350
    351	enclave_image = mmap(NULL, enclave_image_size, PROT_READ,
    352			     MAP_PRIVATE, enclave_image_fd, 0);
    353	if (enclave_image == MAP_FAILED) {
    354		printf("Error in mmap enclave image [%m]\n");
    355
    356		return -1;
    357	}
    358
    359	temp_image_offset = image_load_info.memory_offset;
    360
    361	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++) {
    362		size_t bytes_to_write = 0;
    363		size_t memory_offset = 0;
    364		size_t memory_size = ne_user_mem_regions[i].memory_size;
    365		size_t remaining_bytes = 0;
    366		void *userspace_addr = ne_user_mem_regions[i].userspace_addr;
    367
    368		if (temp_image_offset >= memory_size) {
    369			temp_image_offset -= memory_size;
    370
    371			continue;
    372		} else if (temp_image_offset != 0) {
    373			memory_offset = temp_image_offset;
    374			memory_size -= temp_image_offset;
    375			temp_image_offset = 0;
    376		}
    377
    378		remaining_bytes = enclave_image_size - image_written_bytes;
    379		bytes_to_write = memory_size < remaining_bytes ?
    380				 memory_size : remaining_bytes;
    381
    382		memcpy(userspace_addr + memory_offset,
    383		       enclave_image + image_written_bytes, bytes_to_write);
    384
    385		image_written_bytes += bytes_to_write;
    386
    387		if (image_written_bytes == enclave_image_size)
    388			break;
    389	}
    390
    391	munmap(enclave_image, enclave_image_size);
    392
    393	close(enclave_image_fd);
    394
    395	return 0;
    396}
    397
    398/**
    399 * ne_set_user_mem_region() - Set a user space memory region for the given enclave.
    400 * @enclave_fd :		The file descriptor associated with the enclave.
    401 * @ne_user_mem_region :	User space memory region to be set for the enclave.
    402 *
    403 * Context: Process context.
    404 * Return:
    405 * * 0 on success.
    406 * * Negative return value on failure.
    407 */
    408static int ne_set_user_mem_region(int enclave_fd, struct ne_user_mem_region ne_user_mem_region)
    409{
    410	struct ne_user_memory_region mem_region = {
    411		.flags = NE_DEFAULT_MEMORY_REGION,
    412		.memory_size = ne_user_mem_region.memory_size,
    413		.userspace_addr = (__u64)ne_user_mem_region.userspace_addr,
    414	};
    415	int rc = -EINVAL;
    416
    417	rc = ioctl(enclave_fd, NE_SET_USER_MEMORY_REGION, &mem_region);
    418	if (rc < 0) {
    419		switch (errno) {
    420		case NE_ERR_NOT_IN_INIT_STATE: {
    421			printf("Error in set user memory region, enclave not in init state\n");
    422
    423			break;
    424		}
    425
    426		case NE_ERR_INVALID_MEM_REGION_SIZE: {
    427			printf("Error in set user memory region, mem size not multiple of 2 MiB\n");
    428
    429			break;
    430		}
    431
    432		case NE_ERR_INVALID_MEM_REGION_ADDR: {
    433			printf("Error in set user memory region, invalid user space address\n");
    434
    435			break;
    436		}
    437
    438		case NE_ERR_UNALIGNED_MEM_REGION_ADDR: {
    439			printf("Error in set user memory region, unaligned user space address\n");
    440
    441			break;
    442		}
    443
    444		case NE_ERR_MEM_REGION_ALREADY_USED: {
    445			printf("Error in set user memory region, memory region already used\n");
    446
    447			break;
    448		}
    449
    450		case NE_ERR_MEM_NOT_HUGE_PAGE: {
    451			printf("Error in set user memory region, not backed by huge pages\n");
    452
    453			break;
    454		}
    455
    456		case NE_ERR_MEM_DIFFERENT_NUMA_NODE: {
    457			printf("Error in set user memory region, different NUMA node than CPUs\n");
    458
    459			break;
    460		}
    461
    462		case NE_ERR_MEM_MAX_REGIONS: {
    463			printf("Error in set user memory region, max memory regions reached\n");
    464
    465			break;
    466		}
    467
    468		case NE_ERR_INVALID_PAGE_SIZE: {
    469			printf("Error in set user memory region, has page not multiple of 2 MiB\n");
    470
    471			break;
    472		}
    473
    474		case NE_ERR_INVALID_FLAG_VALUE: {
    475			printf("Error in set user memory region, provided invalid flag\n");
    476
    477			break;
    478		}
    479
    480		default:
    481			printf("Error in set user memory region [%m]\n");
    482		}
    483
    484		return rc;
    485	}
    486
    487	return 0;
    488}
    489
    490/**
    491 * ne_free_mem_regions() - Unmap all the user space memory regions that were set
    492 *			   aside for the enclave.
    493 * @ne_user_mem_regions:	The user space memory regions associated with an enclave.
    494 *
    495 * Context: Process context.
    496 */
    497static void ne_free_mem_regions(struct ne_user_mem_region ne_user_mem_regions[])
    498{
    499	unsigned int i = 0;
    500
    501	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++)
    502		munmap(ne_user_mem_regions[i].userspace_addr,
    503		       ne_user_mem_regions[i].memory_size);
    504}
    505
    506/**
    507 * ne_add_vcpu() - Add a vCPU to the given enclave.
    508 * @enclave_fd :	The file descriptor associated with the enclave.
    509 * @vcpu_id:		vCPU id to be set for the enclave, either provided or
    510 *			auto-generated (if provided vCPU id is 0).
    511 *
    512 * Context: Process context.
    513 * Return:
    514 * * 0 on success.
    515 * * Negative return value on failure.
    516 */
    517static int ne_add_vcpu(int enclave_fd, unsigned int *vcpu_id)
    518{
    519	int rc = -EINVAL;
    520
    521	rc = ioctl(enclave_fd, NE_ADD_VCPU, vcpu_id);
    522	if (rc < 0) {
    523		switch (errno) {
    524		case NE_ERR_NO_CPUS_AVAIL_IN_POOL: {
    525			printf("Error in add vcpu, no CPUs available in the NE CPU pool\n");
    526
    527			break;
    528		}
    529
    530		case NE_ERR_VCPU_ALREADY_USED: {
    531			printf("Error in add vcpu, the provided vCPU is already used\n");
    532
    533			break;
    534		}
    535
    536		case NE_ERR_VCPU_NOT_IN_CPU_POOL: {
    537			printf("Error in add vcpu, the provided vCPU is not in the NE CPU pool\n");
    538
    539			break;
    540		}
    541
    542		case NE_ERR_VCPU_INVALID_CPU_CORE: {
    543			printf("Error in add vcpu, the core id of the provided vCPU is invalid\n");
    544
    545			break;
    546		}
    547
    548		case NE_ERR_NOT_IN_INIT_STATE: {
    549			printf("Error in add vcpu, enclave not in init state\n");
    550
    551			break;
    552		}
    553
    554		case NE_ERR_INVALID_VCPU: {
    555			printf("Error in add vcpu, the provided vCPU is out of avail CPUs range\n");
    556
    557			break;
    558		}
    559
    560		default:
    561			printf("Error in add vcpu [%m]\n");
    562		}
    563
    564		return rc;
    565	}
    566
    567	return 0;
    568}
    569
    570/**
    571 * ne_start_enclave() - Start the given enclave.
    572 * @enclave_fd :		The file descriptor associated with the enclave.
    573 * @enclave_start_info :	Enclave metadata used for starting e.g. vsock CID.
    574 *
    575 * Context: Process context.
    576 * Return:
    577 * * 0 on success.
    578 * * Negative return value on failure.
    579 */
    580static int ne_start_enclave(int enclave_fd,  struct ne_enclave_start_info *enclave_start_info)
    581{
    582	int rc = -EINVAL;
    583
    584	rc = ioctl(enclave_fd, NE_START_ENCLAVE, enclave_start_info);
    585	if (rc < 0) {
    586		switch (errno) {
    587		case NE_ERR_NOT_IN_INIT_STATE: {
    588			printf("Error in start enclave, enclave not in init state\n");
    589
    590			break;
    591		}
    592
    593		case NE_ERR_NO_MEM_REGIONS_ADDED: {
    594			printf("Error in start enclave, no memory regions have been added\n");
    595
    596			break;
    597		}
    598
    599		case NE_ERR_NO_VCPUS_ADDED: {
    600			printf("Error in start enclave, no vCPUs have been added\n");
    601
    602			break;
    603		}
    604
    605		case NE_ERR_FULL_CORES_NOT_USED: {
    606			printf("Error in start enclave, enclave has no full cores set\n");
    607
    608			break;
    609		}
    610
    611		case NE_ERR_ENCLAVE_MEM_MIN_SIZE: {
    612			printf("Error in start enclave, enclave memory is less than min size\n");
    613
    614			break;
    615		}
    616
    617		case NE_ERR_INVALID_FLAG_VALUE: {
    618			printf("Error in start enclave, provided invalid flag\n");
    619
    620			break;
    621		}
    622
    623		case NE_ERR_INVALID_ENCLAVE_CID: {
    624			printf("Error in start enclave, provided invalid enclave CID\n");
    625
    626			break;
    627		}
    628
    629		default:
    630			printf("Error in start enclave [%m]\n");
    631		}
    632
    633		return rc;
    634	}
    635
    636	return 0;
    637}
    638
    639/**
    640 * ne_start_enclave_check_booted() - Start the enclave and wait for a heartbeat
    641 *				     from it, on a newly created vsock channel,
    642 *				     to check it has booted.
    643 * @enclave_fd :	The file descriptor associated with the enclave.
    644 *
    645 * Context: Process context.
    646 * Return:
    647 * * 0 on success.
    648 * * Negative return value on failure.
    649 */
    650static int ne_start_enclave_check_booted(int enclave_fd)
    651{
    652	struct sockaddr_vm client_vsock_addr = {};
    653	int client_vsock_fd = -1;
    654	socklen_t client_vsock_len = sizeof(client_vsock_addr);
    655	struct ne_enclave_start_info enclave_start_info = {};
    656	struct pollfd fds[1] = {};
    657	int rc = -EINVAL;
    658	unsigned char recv_buf = 0;
    659	struct sockaddr_vm server_vsock_addr = {
    660		.svm_family = AF_VSOCK,
    661		.svm_cid = NE_IMAGE_LOAD_HEARTBEAT_CID,
    662		.svm_port = NE_IMAGE_LOAD_HEARTBEAT_PORT,
    663	};
    664	int server_vsock_fd = -1;
    665
    666	server_vsock_fd = socket(AF_VSOCK, SOCK_STREAM, 0);
    667	if (server_vsock_fd < 0) {
    668		rc = server_vsock_fd;
    669
    670		printf("Error in socket [%m]\n");
    671
    672		return rc;
    673	}
    674
    675	rc = bind(server_vsock_fd, (struct sockaddr *)&server_vsock_addr,
    676		  sizeof(server_vsock_addr));
    677	if (rc < 0) {
    678		printf("Error in bind [%m]\n");
    679
    680		goto out;
    681	}
    682
    683	rc = listen(server_vsock_fd, 1);
    684	if (rc < 0) {
    685		printf("Error in listen [%m]\n");
    686
    687		goto out;
    688	}
    689
    690	rc = ne_start_enclave(enclave_fd, &enclave_start_info);
    691	if (rc < 0)
    692		goto out;
    693
    694	printf("Enclave started, CID %llu\n", enclave_start_info.enclave_cid);
    695
    696	fds[0].fd = server_vsock_fd;
    697	fds[0].events = POLLIN;
    698
    699	rc = poll(fds, 1, NE_POLL_WAIT_TIME_MS);
    700	if (rc < 0) {
    701		printf("Error in poll [%m]\n");
    702
    703		goto out;
    704	}
    705
    706	if (!rc) {
    707		printf("Poll timeout, %d seconds elapsed\n", NE_POLL_WAIT_TIME);
    708
    709		rc = -ETIMEDOUT;
    710
    711		goto out;
    712	}
    713
    714	if ((fds[0].revents & POLLIN) == 0) {
    715		printf("Poll received value %d\n", fds[0].revents);
    716
    717		rc = -EINVAL;
    718
    719		goto out;
    720	}
    721
    722	rc = accept(server_vsock_fd, (struct sockaddr *)&client_vsock_addr,
    723		    &client_vsock_len);
    724	if (rc < 0) {
    725		printf("Error in accept [%m]\n");
    726
    727		goto out;
    728	}
    729
    730	client_vsock_fd = rc;
    731
    732	/*
    733	 * Read the heartbeat value that the init process in the enclave sends
    734	 * after vsock connect.
    735	 */
    736	rc = read(client_vsock_fd, &recv_buf, sizeof(recv_buf));
    737	if (rc < 0) {
    738		printf("Error in read [%m]\n");
    739
    740		goto out;
    741	}
    742
    743	if (rc != sizeof(recv_buf) || recv_buf != NE_IMAGE_LOAD_HEARTBEAT_VALUE) {
    744		printf("Read %d instead of %d\n", recv_buf,
    745		       NE_IMAGE_LOAD_HEARTBEAT_VALUE);
    746
    747		goto out;
    748	}
    749
    750	/* Write the heartbeat value back. */
    751	rc = write(client_vsock_fd, &recv_buf, sizeof(recv_buf));
    752	if (rc < 0) {
    753		printf("Error in write [%m]\n");
    754
    755		goto out;
    756	}
    757
    758	rc = 0;
    759
    760out:
    761	close(server_vsock_fd);
    762
    763	return rc;
    764}
    765
    766int main(int argc, char *argv[])
    767{
    768	int enclave_fd = -1;
    769	unsigned int i = 0;
    770	int ne_dev_fd = -1;
    771	struct ne_user_mem_region ne_user_mem_regions[NE_DEFAULT_NR_MEM_REGIONS] = {};
    772	unsigned int ne_vcpus[NE_DEFAULT_NR_VCPUS] = {};
    773	int rc = -EINVAL;
    774	pthread_t thread_id = 0;
    775	unsigned long slot_uid = 0;
    776
    777	if (argc != 2) {
    778		printf("Usage: %s <path_to_enclave_image>\n", argv[0]);
    779
    780		exit(EXIT_FAILURE);
    781	}
    782
    783	if (strlen(argv[1]) >= PATH_MAX) {
    784		printf("The size of the path to enclave image is higher than max path\n");
    785
    786		exit(EXIT_FAILURE);
    787	}
    788
    789	ne_dev_fd = open(NE_DEV_NAME, O_RDWR | O_CLOEXEC);
    790	if (ne_dev_fd < 0) {
    791		printf("Error in open NE device [%m]\n");
    792
    793		exit(EXIT_FAILURE);
    794	}
    795
    796	printf("Creating enclave slot ...\n");
    797
    798	rc = ne_create_vm(ne_dev_fd, &slot_uid, &enclave_fd);
    799
    800	close(ne_dev_fd);
    801
    802	if (rc < 0)
    803		exit(EXIT_FAILURE);
    804
    805	printf("Enclave fd %d\n", enclave_fd);
    806
    807	rc = pthread_create(&thread_id, NULL, ne_poll_enclave_fd, (void *)&enclave_fd);
    808	if (rc < 0) {
    809		printf("Error in thread create [%m]\n");
    810
    811		close(enclave_fd);
    812
    813		exit(EXIT_FAILURE);
    814	}
    815
    816	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++) {
    817		ne_user_mem_regions[i].memory_size = NE_MIN_MEM_REGION_SIZE;
    818
    819		rc = ne_alloc_user_mem_region(&ne_user_mem_regions[i]);
    820		if (rc < 0) {
    821			printf("Error in alloc userspace memory region, iter %d\n", i);
    822
    823			goto release_enclave_fd;
    824		}
    825	}
    826
    827	rc = ne_load_enclave_image(enclave_fd, ne_user_mem_regions, argv[1]);
    828	if (rc < 0)
    829		goto release_enclave_fd;
    830
    831	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++) {
    832		rc = ne_set_user_mem_region(enclave_fd, ne_user_mem_regions[i]);
    833		if (rc < 0) {
    834			printf("Error in set memory region, iter %d\n", i);
    835
    836			goto release_enclave_fd;
    837		}
    838	}
    839
    840	printf("Enclave memory regions were added\n");
    841
    842	for (i = 0; i < NE_DEFAULT_NR_VCPUS; i++) {
    843		/*
    844		 * The vCPU is chosen from the enclave vCPU pool, if the value
    845		 * of the vcpu_id is 0.
    846		 */
    847		ne_vcpus[i] = 0;
    848		rc = ne_add_vcpu(enclave_fd, &ne_vcpus[i]);
    849		if (rc < 0) {
    850			printf("Error in add vcpu, iter %d\n", i);
    851
    852			goto release_enclave_fd;
    853		}
    854
    855		printf("Added vCPU %d to the enclave\n", ne_vcpus[i]);
    856	}
    857
    858	printf("Enclave vCPUs were added\n");
    859
    860	rc = ne_start_enclave_check_booted(enclave_fd);
    861	if (rc < 0) {
    862		printf("Error in the enclave start / image loading heartbeat logic [rc=%d]\n", rc);
    863
    864		goto release_enclave_fd;
    865	}
    866
    867	printf("Entering sleep for %d seconds ...\n", NE_SLEEP_TIME);
    868
    869	sleep(NE_SLEEP_TIME);
    870
    871	close(enclave_fd);
    872
    873	ne_free_mem_regions(ne_user_mem_regions);
    874
    875	exit(EXIT_SUCCESS);
    876
    877release_enclave_fd:
    878	close(enclave_fd);
    879	ne_free_mem_regions(ne_user_mem_regions);
    880
    881	exit(EXIT_FAILURE);
    882}