162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0
262306a36Sopenharmony_ci/*
362306a36Sopenharmony_ci * Copyright 2020-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
462306a36Sopenharmony_ci */
562306a36Sopenharmony_ci
662306a36Sopenharmony_ci/**
762306a36Sopenharmony_ci * DOC: Sample flow of using the ioctl interface provided by the Nitro Enclaves (NE)
862306a36Sopenharmony_ci * kernel driver.
962306a36Sopenharmony_ci *
1062306a36Sopenharmony_ci * Usage
1162306a36Sopenharmony_ci * -----
1262306a36Sopenharmony_ci *
1362306a36Sopenharmony_ci * Load the nitro_enclaves module, setting also the enclave CPU pool. The
1462306a36Sopenharmony_ci * enclave CPUs need to be full cores from the same NUMA node. CPU 0 and its
1562306a36Sopenharmony_ci * siblings have to remain available for the primary / parent VM, so they
1662306a36Sopenharmony_ci * cannot be included in the enclave CPU pool.
1762306a36Sopenharmony_ci *
1862306a36Sopenharmony_ci * See the cpu list section from the kernel documentation.
1962306a36Sopenharmony_ci * https://www.kernel.org/doc/html/latest/admin-guide/kernel-parameters.html#cpu-lists
2062306a36Sopenharmony_ci *
2162306a36Sopenharmony_ci *	insmod drivers/virt/nitro_enclaves/nitro_enclaves.ko
2262306a36Sopenharmony_ci *	lsmod
2362306a36Sopenharmony_ci *
2462306a36Sopenharmony_ci *	The CPU pool can be set at runtime, after the kernel module is loaded.
2562306a36Sopenharmony_ci *
2662306a36Sopenharmony_ci *	echo <cpu-list> > /sys/module/nitro_enclaves/parameters/ne_cpus
2762306a36Sopenharmony_ci *
2862306a36Sopenharmony_ci *	NUMA and CPU siblings information can be found using:
2962306a36Sopenharmony_ci *
3062306a36Sopenharmony_ci *	lscpu
3162306a36Sopenharmony_ci *	/proc/cpuinfo
3262306a36Sopenharmony_ci *
3362306a36Sopenharmony_ci * Check the online / offline CPU list. The CPUs from the pool should be
3462306a36Sopenharmony_ci * offlined.
3562306a36Sopenharmony_ci *
3662306a36Sopenharmony_ci *	lscpu
3762306a36Sopenharmony_ci *
3862306a36Sopenharmony_ci * Check dmesg for any warnings / errors through the NE driver lifetime / usage.
3962306a36Sopenharmony_ci * The NE logs contain the "nitro_enclaves" or "pci 0000:00:02.0" pattern.
4062306a36Sopenharmony_ci *
4162306a36Sopenharmony_ci *	dmesg
4262306a36Sopenharmony_ci *
4362306a36Sopenharmony_ci * Setup hugetlbfs huge pages. The memory needs to be from the same NUMA node as
4462306a36Sopenharmony_ci * the enclave CPUs.
4562306a36Sopenharmony_ci *
4662306a36Sopenharmony_ci * https://www.kernel.org/doc/html/latest/admin-guide/mm/hugetlbpage.html
4762306a36Sopenharmony_ci *
4862306a36Sopenharmony_ci * By default, the allocation of hugetlb pages are distributed on all possible
4962306a36Sopenharmony_ci * NUMA nodes. Use the following configuration files to set the number of huge
5062306a36Sopenharmony_ci * pages from a NUMA node:
5162306a36Sopenharmony_ci *
5262306a36Sopenharmony_ci *	/sys/devices/system/node/node<X>/hugepages/hugepages-2048kB/nr_hugepages
5362306a36Sopenharmony_ci *	/sys/devices/system/node/node<X>/hugepages/hugepages-1048576kB/nr_hugepages
5462306a36Sopenharmony_ci *
5562306a36Sopenharmony_ci *	or, if not on a system with multiple NUMA nodes, can also set the number
5662306a36Sopenharmony_ci *	of 2 MiB / 1 GiB huge pages using
5762306a36Sopenharmony_ci *
5862306a36Sopenharmony_ci *	/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages
5962306a36Sopenharmony_ci *	/sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages
6062306a36Sopenharmony_ci *
6162306a36Sopenharmony_ci *	In this example 256 hugepages of 2 MiB are used.
6262306a36Sopenharmony_ci *
6362306a36Sopenharmony_ci * Build and run the NE sample.
6462306a36Sopenharmony_ci *
6562306a36Sopenharmony_ci *	make -C samples/nitro_enclaves clean
6662306a36Sopenharmony_ci *	make -C samples/nitro_enclaves
6762306a36Sopenharmony_ci *	./samples/nitro_enclaves/ne_ioctl_sample <path_to_enclave_image>
6862306a36Sopenharmony_ci *
6962306a36Sopenharmony_ci * Unload the nitro_enclaves module.
7062306a36Sopenharmony_ci *
7162306a36Sopenharmony_ci *	rmmod nitro_enclaves
7262306a36Sopenharmony_ci *	lsmod
7362306a36Sopenharmony_ci */
7462306a36Sopenharmony_ci
7562306a36Sopenharmony_ci#include <stdio.h>
7662306a36Sopenharmony_ci#include <stdlib.h>
7762306a36Sopenharmony_ci#include <errno.h>
7862306a36Sopenharmony_ci#include <fcntl.h>
7962306a36Sopenharmony_ci#include <limits.h>
8062306a36Sopenharmony_ci#include <poll.h>
8162306a36Sopenharmony_ci#include <pthread.h>
8262306a36Sopenharmony_ci#include <string.h>
8362306a36Sopenharmony_ci#include <sys/eventfd.h>
8462306a36Sopenharmony_ci#include <sys/ioctl.h>
8562306a36Sopenharmony_ci#include <sys/mman.h>
8662306a36Sopenharmony_ci#include <sys/socket.h>
8762306a36Sopenharmony_ci#include <sys/stat.h>
8862306a36Sopenharmony_ci#include <sys/types.h>
8962306a36Sopenharmony_ci#include <unistd.h>
9062306a36Sopenharmony_ci
9162306a36Sopenharmony_ci#include <linux/mman.h>
9262306a36Sopenharmony_ci#include <linux/nitro_enclaves.h>
9362306a36Sopenharmony_ci#include <linux/vm_sockets.h>
9462306a36Sopenharmony_ci
9562306a36Sopenharmony_ci/**
9662306a36Sopenharmony_ci * NE_DEV_NAME - Nitro Enclaves (NE) misc device that provides the ioctl interface.
9762306a36Sopenharmony_ci */
9862306a36Sopenharmony_ci#define NE_DEV_NAME			"/dev/nitro_enclaves"
9962306a36Sopenharmony_ci
10062306a36Sopenharmony_ci/**
10162306a36Sopenharmony_ci * NE_POLL_WAIT_TIME - Timeout in seconds for each poll event.
10262306a36Sopenharmony_ci */
10362306a36Sopenharmony_ci#define NE_POLL_WAIT_TIME		(60)
10462306a36Sopenharmony_ci/**
10562306a36Sopenharmony_ci * NE_POLL_WAIT_TIME_MS - Timeout in milliseconds for each poll event.
10662306a36Sopenharmony_ci */
10762306a36Sopenharmony_ci#define NE_POLL_WAIT_TIME_MS		(NE_POLL_WAIT_TIME * 1000)
10862306a36Sopenharmony_ci
10962306a36Sopenharmony_ci/**
11062306a36Sopenharmony_ci * NE_SLEEP_TIME - Amount of time in seconds for the process to keep the enclave alive.
11162306a36Sopenharmony_ci */
11262306a36Sopenharmony_ci#define NE_SLEEP_TIME			(300)
11362306a36Sopenharmony_ci
11462306a36Sopenharmony_ci/**
11562306a36Sopenharmony_ci * NE_DEFAULT_NR_VCPUS - Default number of vCPUs set for an enclave.
11662306a36Sopenharmony_ci */
11762306a36Sopenharmony_ci#define NE_DEFAULT_NR_VCPUS		(2)
11862306a36Sopenharmony_ci
11962306a36Sopenharmony_ci/**
12062306a36Sopenharmony_ci * NE_MIN_MEM_REGION_SIZE - Minimum size of a memory region - 2 MiB.
12162306a36Sopenharmony_ci */
12262306a36Sopenharmony_ci#define NE_MIN_MEM_REGION_SIZE		(2 * 1024 * 1024)
12362306a36Sopenharmony_ci
12462306a36Sopenharmony_ci/**
12562306a36Sopenharmony_ci * NE_DEFAULT_NR_MEM_REGIONS - Default number of memory regions of 2 MiB set for
12662306a36Sopenharmony_ci *			       an enclave.
12762306a36Sopenharmony_ci */
12862306a36Sopenharmony_ci#define NE_DEFAULT_NR_MEM_REGIONS	(256)
12962306a36Sopenharmony_ci
13062306a36Sopenharmony_ci/**
13162306a36Sopenharmony_ci * NE_IMAGE_LOAD_HEARTBEAT_CID - Vsock CID for enclave image loading heartbeat logic.
13262306a36Sopenharmony_ci */
13362306a36Sopenharmony_ci#define NE_IMAGE_LOAD_HEARTBEAT_CID	(3)
13462306a36Sopenharmony_ci/**
13562306a36Sopenharmony_ci * NE_IMAGE_LOAD_HEARTBEAT_PORT - Vsock port for enclave image loading heartbeat logic.
13662306a36Sopenharmony_ci */
13762306a36Sopenharmony_ci#define NE_IMAGE_LOAD_HEARTBEAT_PORT	(9000)
13862306a36Sopenharmony_ci/**
13962306a36Sopenharmony_ci * NE_IMAGE_LOAD_HEARTBEAT_VALUE - Heartbeat value for enclave image loading.
14062306a36Sopenharmony_ci */
14162306a36Sopenharmony_ci#define NE_IMAGE_LOAD_HEARTBEAT_VALUE	(0xb7)
14262306a36Sopenharmony_ci
14362306a36Sopenharmony_ci/**
14462306a36Sopenharmony_ci * struct ne_user_mem_region - User space memory region set for an enclave.
14562306a36Sopenharmony_ci * @userspace_addr:	Address of the user space memory region.
14662306a36Sopenharmony_ci * @memory_size:	Size of the user space memory region.
14762306a36Sopenharmony_ci */
14862306a36Sopenharmony_cistruct ne_user_mem_region {
14962306a36Sopenharmony_ci	void	*userspace_addr;
15062306a36Sopenharmony_ci	size_t	memory_size;
15162306a36Sopenharmony_ci};
15262306a36Sopenharmony_ci
15362306a36Sopenharmony_ci/**
15462306a36Sopenharmony_ci * ne_create_vm() - Create a slot for the enclave VM.
15562306a36Sopenharmony_ci * @ne_dev_fd:		The file descriptor of the NE misc device.
15662306a36Sopenharmony_ci * @slot_uid:		The generated slot uid for the enclave.
15762306a36Sopenharmony_ci * @enclave_fd :	The generated file descriptor for the enclave.
15862306a36Sopenharmony_ci *
15962306a36Sopenharmony_ci * Context: Process context.
16062306a36Sopenharmony_ci * Return:
16162306a36Sopenharmony_ci * * 0 on success.
16262306a36Sopenharmony_ci * * Negative return value on failure.
16362306a36Sopenharmony_ci */
16462306a36Sopenharmony_cistatic int ne_create_vm(int ne_dev_fd, unsigned long *slot_uid, int *enclave_fd)
16562306a36Sopenharmony_ci{
16662306a36Sopenharmony_ci	int rc = -EINVAL;
16762306a36Sopenharmony_ci	*enclave_fd = ioctl(ne_dev_fd, NE_CREATE_VM, slot_uid);
16862306a36Sopenharmony_ci
16962306a36Sopenharmony_ci	if (*enclave_fd < 0) {
17062306a36Sopenharmony_ci		rc = *enclave_fd;
17162306a36Sopenharmony_ci		switch (errno) {
17262306a36Sopenharmony_ci		case NE_ERR_NO_CPUS_AVAIL_IN_POOL: {
17362306a36Sopenharmony_ci			printf("Error in create VM, no CPUs available in the NE CPU pool\n");
17462306a36Sopenharmony_ci
17562306a36Sopenharmony_ci			break;
17662306a36Sopenharmony_ci		}
17762306a36Sopenharmony_ci
17862306a36Sopenharmony_ci		default:
17962306a36Sopenharmony_ci			printf("Error in create VM [%m]\n");
18062306a36Sopenharmony_ci		}
18162306a36Sopenharmony_ci
18262306a36Sopenharmony_ci		return rc;
18362306a36Sopenharmony_ci	}
18462306a36Sopenharmony_ci
18562306a36Sopenharmony_ci	return 0;
18662306a36Sopenharmony_ci}
18762306a36Sopenharmony_ci
18862306a36Sopenharmony_ci/**
18962306a36Sopenharmony_ci * ne_poll_enclave_fd() - Thread function for polling the enclave fd.
19062306a36Sopenharmony_ci * @data:	Argument provided for the polling function.
19162306a36Sopenharmony_ci *
19262306a36Sopenharmony_ci * Context: Process context.
19362306a36Sopenharmony_ci * Return:
19462306a36Sopenharmony_ci * * NULL on success / failure.
19562306a36Sopenharmony_ci */
19662306a36Sopenharmony_civoid *ne_poll_enclave_fd(void *data)
19762306a36Sopenharmony_ci{
19862306a36Sopenharmony_ci	int enclave_fd = *(int *)data;
19962306a36Sopenharmony_ci	struct pollfd fds[1] = {};
20062306a36Sopenharmony_ci	int i = 0;
20162306a36Sopenharmony_ci	int rc = -EINVAL;
20262306a36Sopenharmony_ci
20362306a36Sopenharmony_ci	printf("Running from poll thread, enclave fd %d\n", enclave_fd);
20462306a36Sopenharmony_ci
20562306a36Sopenharmony_ci	fds[0].fd = enclave_fd;
20662306a36Sopenharmony_ci	fds[0].events = POLLIN | POLLERR | POLLHUP;
20762306a36Sopenharmony_ci
20862306a36Sopenharmony_ci	/* Keep on polling until the current process is terminated. */
20962306a36Sopenharmony_ci	while (1) {
21062306a36Sopenharmony_ci		printf("[iter %d] Polling ...\n", i);
21162306a36Sopenharmony_ci
21262306a36Sopenharmony_ci		rc = poll(fds, 1, NE_POLL_WAIT_TIME_MS);
21362306a36Sopenharmony_ci		if (rc < 0) {
21462306a36Sopenharmony_ci			printf("Error in poll [%m]\n");
21562306a36Sopenharmony_ci
21662306a36Sopenharmony_ci			return NULL;
21762306a36Sopenharmony_ci		}
21862306a36Sopenharmony_ci
21962306a36Sopenharmony_ci		i++;
22062306a36Sopenharmony_ci
22162306a36Sopenharmony_ci		if (!rc) {
22262306a36Sopenharmony_ci			printf("Poll: %d seconds elapsed\n",
22362306a36Sopenharmony_ci			       i * NE_POLL_WAIT_TIME);
22462306a36Sopenharmony_ci
22562306a36Sopenharmony_ci			continue;
22662306a36Sopenharmony_ci		}
22762306a36Sopenharmony_ci
22862306a36Sopenharmony_ci		printf("Poll received value 0x%x\n", fds[0].revents);
22962306a36Sopenharmony_ci
23062306a36Sopenharmony_ci		if (fds[0].revents & POLLHUP) {
23162306a36Sopenharmony_ci			printf("Received POLLHUP\n");
23262306a36Sopenharmony_ci
23362306a36Sopenharmony_ci			return NULL;
23462306a36Sopenharmony_ci		}
23562306a36Sopenharmony_ci
23662306a36Sopenharmony_ci		if (fds[0].revents & POLLNVAL) {
23762306a36Sopenharmony_ci			printf("Received POLLNVAL\n");
23862306a36Sopenharmony_ci
23962306a36Sopenharmony_ci			return NULL;
24062306a36Sopenharmony_ci		}
24162306a36Sopenharmony_ci	}
24262306a36Sopenharmony_ci
24362306a36Sopenharmony_ci	return NULL;
24462306a36Sopenharmony_ci}
24562306a36Sopenharmony_ci
24662306a36Sopenharmony_ci/**
24762306a36Sopenharmony_ci * ne_alloc_user_mem_region() - Allocate a user space memory region for an enclave.
24862306a36Sopenharmony_ci * @ne_user_mem_region:	User space memory region allocated using hugetlbfs.
24962306a36Sopenharmony_ci *
25062306a36Sopenharmony_ci * Context: Process context.
25162306a36Sopenharmony_ci * Return:
25262306a36Sopenharmony_ci * * 0 on success.
25362306a36Sopenharmony_ci * * Negative return value on failure.
25462306a36Sopenharmony_ci */
25562306a36Sopenharmony_cistatic int ne_alloc_user_mem_region(struct ne_user_mem_region *ne_user_mem_region)
25662306a36Sopenharmony_ci{
25762306a36Sopenharmony_ci	/**
25862306a36Sopenharmony_ci	 * Check available hugetlb encodings for different huge page sizes in
25962306a36Sopenharmony_ci	 * include/uapi/linux/mman.h.
26062306a36Sopenharmony_ci	 */
26162306a36Sopenharmony_ci	ne_user_mem_region->userspace_addr = mmap(NULL, ne_user_mem_region->memory_size,
26262306a36Sopenharmony_ci						  PROT_READ | PROT_WRITE,
26362306a36Sopenharmony_ci						  MAP_PRIVATE | MAP_ANONYMOUS |
26462306a36Sopenharmony_ci						  MAP_HUGETLB | MAP_HUGE_2MB, -1, 0);
26562306a36Sopenharmony_ci	if (ne_user_mem_region->userspace_addr == MAP_FAILED) {
26662306a36Sopenharmony_ci		printf("Error in mmap memory [%m]\n");
26762306a36Sopenharmony_ci
26862306a36Sopenharmony_ci		return -1;
26962306a36Sopenharmony_ci	}
27062306a36Sopenharmony_ci
27162306a36Sopenharmony_ci	return 0;
27262306a36Sopenharmony_ci}
27362306a36Sopenharmony_ci
27462306a36Sopenharmony_ci/**
27562306a36Sopenharmony_ci * ne_load_enclave_image() - Place the enclave image in the enclave memory.
27662306a36Sopenharmony_ci * @enclave_fd :		The file descriptor associated with the enclave.
27762306a36Sopenharmony_ci * @ne_user_mem_regions:	User space memory regions allocated for the enclave.
27862306a36Sopenharmony_ci * @enclave_image_path :	The file path of the enclave image.
27962306a36Sopenharmony_ci *
28062306a36Sopenharmony_ci * Context: Process context.
28162306a36Sopenharmony_ci * Return:
28262306a36Sopenharmony_ci * * 0 on success.
28362306a36Sopenharmony_ci * * Negative return value on failure.
28462306a36Sopenharmony_ci */
28562306a36Sopenharmony_cistatic int ne_load_enclave_image(int enclave_fd, struct ne_user_mem_region ne_user_mem_regions[],
28662306a36Sopenharmony_ci				 char *enclave_image_path)
28762306a36Sopenharmony_ci{
28862306a36Sopenharmony_ci	unsigned char *enclave_image = NULL;
28962306a36Sopenharmony_ci	int enclave_image_fd = -1;
29062306a36Sopenharmony_ci	size_t enclave_image_size = 0;
29162306a36Sopenharmony_ci	size_t enclave_memory_size = 0;
29262306a36Sopenharmony_ci	unsigned long i = 0;
29362306a36Sopenharmony_ci	size_t image_written_bytes = 0;
29462306a36Sopenharmony_ci	struct ne_image_load_info image_load_info = {
29562306a36Sopenharmony_ci		.flags = NE_EIF_IMAGE,
29662306a36Sopenharmony_ci	};
29762306a36Sopenharmony_ci	struct stat image_stat_buf = {};
29862306a36Sopenharmony_ci	int rc = -EINVAL;
29962306a36Sopenharmony_ci	size_t temp_image_offset = 0;
30062306a36Sopenharmony_ci
30162306a36Sopenharmony_ci	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++)
30262306a36Sopenharmony_ci		enclave_memory_size += ne_user_mem_regions[i].memory_size;
30362306a36Sopenharmony_ci
30462306a36Sopenharmony_ci	rc = stat(enclave_image_path, &image_stat_buf);
30562306a36Sopenharmony_ci	if (rc < 0) {
30662306a36Sopenharmony_ci		printf("Error in get image stat info [%m]\n");
30762306a36Sopenharmony_ci
30862306a36Sopenharmony_ci		return rc;
30962306a36Sopenharmony_ci	}
31062306a36Sopenharmony_ci
31162306a36Sopenharmony_ci	enclave_image_size = image_stat_buf.st_size;
31262306a36Sopenharmony_ci
31362306a36Sopenharmony_ci	if (enclave_memory_size < enclave_image_size) {
31462306a36Sopenharmony_ci		printf("The enclave memory is smaller than the enclave image size\n");
31562306a36Sopenharmony_ci
31662306a36Sopenharmony_ci		return -ENOMEM;
31762306a36Sopenharmony_ci	}
31862306a36Sopenharmony_ci
31962306a36Sopenharmony_ci	rc = ioctl(enclave_fd, NE_GET_IMAGE_LOAD_INFO, &image_load_info);
32062306a36Sopenharmony_ci	if (rc < 0) {
32162306a36Sopenharmony_ci		switch (errno) {
32262306a36Sopenharmony_ci		case NE_ERR_NOT_IN_INIT_STATE: {
32362306a36Sopenharmony_ci			printf("Error in get image load info, enclave not in init state\n");
32462306a36Sopenharmony_ci
32562306a36Sopenharmony_ci			break;
32662306a36Sopenharmony_ci		}
32762306a36Sopenharmony_ci
32862306a36Sopenharmony_ci		case NE_ERR_INVALID_FLAG_VALUE: {
32962306a36Sopenharmony_ci			printf("Error in get image load info, provided invalid flag\n");
33062306a36Sopenharmony_ci
33162306a36Sopenharmony_ci			break;
33262306a36Sopenharmony_ci		}
33362306a36Sopenharmony_ci
33462306a36Sopenharmony_ci		default:
33562306a36Sopenharmony_ci			printf("Error in get image load info [%m]\n");
33662306a36Sopenharmony_ci		}
33762306a36Sopenharmony_ci
33862306a36Sopenharmony_ci		return rc;
33962306a36Sopenharmony_ci	}
34062306a36Sopenharmony_ci
34162306a36Sopenharmony_ci	printf("Enclave image offset in enclave memory is %lld\n",
34262306a36Sopenharmony_ci	       image_load_info.memory_offset);
34362306a36Sopenharmony_ci
34462306a36Sopenharmony_ci	enclave_image_fd = open(enclave_image_path, O_RDONLY);
34562306a36Sopenharmony_ci	if (enclave_image_fd < 0) {
34662306a36Sopenharmony_ci		printf("Error in open enclave image file [%m]\n");
34762306a36Sopenharmony_ci
34862306a36Sopenharmony_ci		return enclave_image_fd;
34962306a36Sopenharmony_ci	}
35062306a36Sopenharmony_ci
35162306a36Sopenharmony_ci	enclave_image = mmap(NULL, enclave_image_size, PROT_READ,
35262306a36Sopenharmony_ci			     MAP_PRIVATE, enclave_image_fd, 0);
35362306a36Sopenharmony_ci	if (enclave_image == MAP_FAILED) {
35462306a36Sopenharmony_ci		printf("Error in mmap enclave image [%m]\n");
35562306a36Sopenharmony_ci
35662306a36Sopenharmony_ci		return -1;
35762306a36Sopenharmony_ci	}
35862306a36Sopenharmony_ci
35962306a36Sopenharmony_ci	temp_image_offset = image_load_info.memory_offset;
36062306a36Sopenharmony_ci
36162306a36Sopenharmony_ci	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++) {
36262306a36Sopenharmony_ci		size_t bytes_to_write = 0;
36362306a36Sopenharmony_ci		size_t memory_offset = 0;
36462306a36Sopenharmony_ci		size_t memory_size = ne_user_mem_regions[i].memory_size;
36562306a36Sopenharmony_ci		size_t remaining_bytes = 0;
36662306a36Sopenharmony_ci		void *userspace_addr = ne_user_mem_regions[i].userspace_addr;
36762306a36Sopenharmony_ci
36862306a36Sopenharmony_ci		if (temp_image_offset >= memory_size) {
36962306a36Sopenharmony_ci			temp_image_offset -= memory_size;
37062306a36Sopenharmony_ci
37162306a36Sopenharmony_ci			continue;
37262306a36Sopenharmony_ci		} else if (temp_image_offset != 0) {
37362306a36Sopenharmony_ci			memory_offset = temp_image_offset;
37462306a36Sopenharmony_ci			memory_size -= temp_image_offset;
37562306a36Sopenharmony_ci			temp_image_offset = 0;
37662306a36Sopenharmony_ci		}
37762306a36Sopenharmony_ci
37862306a36Sopenharmony_ci		remaining_bytes = enclave_image_size - image_written_bytes;
37962306a36Sopenharmony_ci		bytes_to_write = memory_size < remaining_bytes ?
38062306a36Sopenharmony_ci				 memory_size : remaining_bytes;
38162306a36Sopenharmony_ci
38262306a36Sopenharmony_ci		memcpy(userspace_addr + memory_offset,
38362306a36Sopenharmony_ci		       enclave_image + image_written_bytes, bytes_to_write);
38462306a36Sopenharmony_ci
38562306a36Sopenharmony_ci		image_written_bytes += bytes_to_write;
38662306a36Sopenharmony_ci
38762306a36Sopenharmony_ci		if (image_written_bytes == enclave_image_size)
38862306a36Sopenharmony_ci			break;
38962306a36Sopenharmony_ci	}
39062306a36Sopenharmony_ci
39162306a36Sopenharmony_ci	munmap(enclave_image, enclave_image_size);
39262306a36Sopenharmony_ci
39362306a36Sopenharmony_ci	close(enclave_image_fd);
39462306a36Sopenharmony_ci
39562306a36Sopenharmony_ci	return 0;
39662306a36Sopenharmony_ci}
39762306a36Sopenharmony_ci
39862306a36Sopenharmony_ci/**
39962306a36Sopenharmony_ci * ne_set_user_mem_region() - Set a user space memory region for the given enclave.
40062306a36Sopenharmony_ci * @enclave_fd :		The file descriptor associated with the enclave.
40162306a36Sopenharmony_ci * @ne_user_mem_region :	User space memory region to be set for the enclave.
40262306a36Sopenharmony_ci *
40362306a36Sopenharmony_ci * Context: Process context.
40462306a36Sopenharmony_ci * Return:
40562306a36Sopenharmony_ci * * 0 on success.
40662306a36Sopenharmony_ci * * Negative return value on failure.
40762306a36Sopenharmony_ci */
40862306a36Sopenharmony_cistatic int ne_set_user_mem_region(int enclave_fd, struct ne_user_mem_region ne_user_mem_region)
40962306a36Sopenharmony_ci{
41062306a36Sopenharmony_ci	struct ne_user_memory_region mem_region = {
41162306a36Sopenharmony_ci		.flags = NE_DEFAULT_MEMORY_REGION,
41262306a36Sopenharmony_ci		.memory_size = ne_user_mem_region.memory_size,
41362306a36Sopenharmony_ci		.userspace_addr = (__u64)ne_user_mem_region.userspace_addr,
41462306a36Sopenharmony_ci	};
41562306a36Sopenharmony_ci	int rc = -EINVAL;
41662306a36Sopenharmony_ci
41762306a36Sopenharmony_ci	rc = ioctl(enclave_fd, NE_SET_USER_MEMORY_REGION, &mem_region);
41862306a36Sopenharmony_ci	if (rc < 0) {
41962306a36Sopenharmony_ci		switch (errno) {
42062306a36Sopenharmony_ci		case NE_ERR_NOT_IN_INIT_STATE: {
42162306a36Sopenharmony_ci			printf("Error in set user memory region, enclave not in init state\n");
42262306a36Sopenharmony_ci
42362306a36Sopenharmony_ci			break;
42462306a36Sopenharmony_ci		}
42562306a36Sopenharmony_ci
42662306a36Sopenharmony_ci		case NE_ERR_INVALID_MEM_REGION_SIZE: {
42762306a36Sopenharmony_ci			printf("Error in set user memory region, mem size not multiple of 2 MiB\n");
42862306a36Sopenharmony_ci
42962306a36Sopenharmony_ci			break;
43062306a36Sopenharmony_ci		}
43162306a36Sopenharmony_ci
43262306a36Sopenharmony_ci		case NE_ERR_INVALID_MEM_REGION_ADDR: {
43362306a36Sopenharmony_ci			printf("Error in set user memory region, invalid user space address\n");
43462306a36Sopenharmony_ci
43562306a36Sopenharmony_ci			break;
43662306a36Sopenharmony_ci		}
43762306a36Sopenharmony_ci
43862306a36Sopenharmony_ci		case NE_ERR_UNALIGNED_MEM_REGION_ADDR: {
43962306a36Sopenharmony_ci			printf("Error in set user memory region, unaligned user space address\n");
44062306a36Sopenharmony_ci
44162306a36Sopenharmony_ci			break;
44262306a36Sopenharmony_ci		}
44362306a36Sopenharmony_ci
44462306a36Sopenharmony_ci		case NE_ERR_MEM_REGION_ALREADY_USED: {
44562306a36Sopenharmony_ci			printf("Error in set user memory region, memory region already used\n");
44662306a36Sopenharmony_ci
44762306a36Sopenharmony_ci			break;
44862306a36Sopenharmony_ci		}
44962306a36Sopenharmony_ci
45062306a36Sopenharmony_ci		case NE_ERR_MEM_NOT_HUGE_PAGE: {
45162306a36Sopenharmony_ci			printf("Error in set user memory region, not backed by huge pages\n");
45262306a36Sopenharmony_ci
45362306a36Sopenharmony_ci			break;
45462306a36Sopenharmony_ci		}
45562306a36Sopenharmony_ci
45662306a36Sopenharmony_ci		case NE_ERR_MEM_DIFFERENT_NUMA_NODE: {
45762306a36Sopenharmony_ci			printf("Error in set user memory region, different NUMA node than CPUs\n");
45862306a36Sopenharmony_ci
45962306a36Sopenharmony_ci			break;
46062306a36Sopenharmony_ci		}
46162306a36Sopenharmony_ci
46262306a36Sopenharmony_ci		case NE_ERR_MEM_MAX_REGIONS: {
46362306a36Sopenharmony_ci			printf("Error in set user memory region, max memory regions reached\n");
46462306a36Sopenharmony_ci
46562306a36Sopenharmony_ci			break;
46662306a36Sopenharmony_ci		}
46762306a36Sopenharmony_ci
46862306a36Sopenharmony_ci		case NE_ERR_INVALID_PAGE_SIZE: {
46962306a36Sopenharmony_ci			printf("Error in set user memory region, has page not multiple of 2 MiB\n");
47062306a36Sopenharmony_ci
47162306a36Sopenharmony_ci			break;
47262306a36Sopenharmony_ci		}
47362306a36Sopenharmony_ci
47462306a36Sopenharmony_ci		case NE_ERR_INVALID_FLAG_VALUE: {
47562306a36Sopenharmony_ci			printf("Error in set user memory region, provided invalid flag\n");
47662306a36Sopenharmony_ci
47762306a36Sopenharmony_ci			break;
47862306a36Sopenharmony_ci		}
47962306a36Sopenharmony_ci
48062306a36Sopenharmony_ci		default:
48162306a36Sopenharmony_ci			printf("Error in set user memory region [%m]\n");
48262306a36Sopenharmony_ci		}
48362306a36Sopenharmony_ci
48462306a36Sopenharmony_ci		return rc;
48562306a36Sopenharmony_ci	}
48662306a36Sopenharmony_ci
48762306a36Sopenharmony_ci	return 0;
48862306a36Sopenharmony_ci}
48962306a36Sopenharmony_ci
49062306a36Sopenharmony_ci/**
49162306a36Sopenharmony_ci * ne_free_mem_regions() - Unmap all the user space memory regions that were set
49262306a36Sopenharmony_ci *			   aside for the enclave.
49362306a36Sopenharmony_ci * @ne_user_mem_regions:	The user space memory regions associated with an enclave.
49462306a36Sopenharmony_ci *
49562306a36Sopenharmony_ci * Context: Process context.
49662306a36Sopenharmony_ci */
49762306a36Sopenharmony_cistatic void ne_free_mem_regions(struct ne_user_mem_region ne_user_mem_regions[])
49862306a36Sopenharmony_ci{
49962306a36Sopenharmony_ci	unsigned int i = 0;
50062306a36Sopenharmony_ci
50162306a36Sopenharmony_ci	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++)
50262306a36Sopenharmony_ci		munmap(ne_user_mem_regions[i].userspace_addr,
50362306a36Sopenharmony_ci		       ne_user_mem_regions[i].memory_size);
50462306a36Sopenharmony_ci}
50562306a36Sopenharmony_ci
50662306a36Sopenharmony_ci/**
50762306a36Sopenharmony_ci * ne_add_vcpu() - Add a vCPU to the given enclave.
50862306a36Sopenharmony_ci * @enclave_fd :	The file descriptor associated with the enclave.
50962306a36Sopenharmony_ci * @vcpu_id:		vCPU id to be set for the enclave, either provided or
51062306a36Sopenharmony_ci *			auto-generated (if provided vCPU id is 0).
51162306a36Sopenharmony_ci *
51262306a36Sopenharmony_ci * Context: Process context.
51362306a36Sopenharmony_ci * Return:
51462306a36Sopenharmony_ci * * 0 on success.
51562306a36Sopenharmony_ci * * Negative return value on failure.
51662306a36Sopenharmony_ci */
51762306a36Sopenharmony_cistatic int ne_add_vcpu(int enclave_fd, unsigned int *vcpu_id)
51862306a36Sopenharmony_ci{
51962306a36Sopenharmony_ci	int rc = -EINVAL;
52062306a36Sopenharmony_ci
52162306a36Sopenharmony_ci	rc = ioctl(enclave_fd, NE_ADD_VCPU, vcpu_id);
52262306a36Sopenharmony_ci	if (rc < 0) {
52362306a36Sopenharmony_ci		switch (errno) {
52462306a36Sopenharmony_ci		case NE_ERR_NO_CPUS_AVAIL_IN_POOL: {
52562306a36Sopenharmony_ci			printf("Error in add vcpu, no CPUs available in the NE CPU pool\n");
52662306a36Sopenharmony_ci
52762306a36Sopenharmony_ci			break;
52862306a36Sopenharmony_ci		}
52962306a36Sopenharmony_ci
53062306a36Sopenharmony_ci		case NE_ERR_VCPU_ALREADY_USED: {
53162306a36Sopenharmony_ci			printf("Error in add vcpu, the provided vCPU is already used\n");
53262306a36Sopenharmony_ci
53362306a36Sopenharmony_ci			break;
53462306a36Sopenharmony_ci		}
53562306a36Sopenharmony_ci
53662306a36Sopenharmony_ci		case NE_ERR_VCPU_NOT_IN_CPU_POOL: {
53762306a36Sopenharmony_ci			printf("Error in add vcpu, the provided vCPU is not in the NE CPU pool\n");
53862306a36Sopenharmony_ci
53962306a36Sopenharmony_ci			break;
54062306a36Sopenharmony_ci		}
54162306a36Sopenharmony_ci
54262306a36Sopenharmony_ci		case NE_ERR_VCPU_INVALID_CPU_CORE: {
54362306a36Sopenharmony_ci			printf("Error in add vcpu, the core id of the provided vCPU is invalid\n");
54462306a36Sopenharmony_ci
54562306a36Sopenharmony_ci			break;
54662306a36Sopenharmony_ci		}
54762306a36Sopenharmony_ci
54862306a36Sopenharmony_ci		case NE_ERR_NOT_IN_INIT_STATE: {
54962306a36Sopenharmony_ci			printf("Error in add vcpu, enclave not in init state\n");
55062306a36Sopenharmony_ci
55162306a36Sopenharmony_ci			break;
55262306a36Sopenharmony_ci		}
55362306a36Sopenharmony_ci
55462306a36Sopenharmony_ci		case NE_ERR_INVALID_VCPU: {
55562306a36Sopenharmony_ci			printf("Error in add vcpu, the provided vCPU is out of avail CPUs range\n");
55662306a36Sopenharmony_ci
55762306a36Sopenharmony_ci			break;
55862306a36Sopenharmony_ci		}
55962306a36Sopenharmony_ci
56062306a36Sopenharmony_ci		default:
56162306a36Sopenharmony_ci			printf("Error in add vcpu [%m]\n");
56262306a36Sopenharmony_ci		}
56362306a36Sopenharmony_ci
56462306a36Sopenharmony_ci		return rc;
56562306a36Sopenharmony_ci	}
56662306a36Sopenharmony_ci
56762306a36Sopenharmony_ci	return 0;
56862306a36Sopenharmony_ci}
56962306a36Sopenharmony_ci
57062306a36Sopenharmony_ci/**
57162306a36Sopenharmony_ci * ne_start_enclave() - Start the given enclave.
57262306a36Sopenharmony_ci * @enclave_fd :		The file descriptor associated with the enclave.
57362306a36Sopenharmony_ci * @enclave_start_info :	Enclave metadata used for starting e.g. vsock CID.
57462306a36Sopenharmony_ci *
57562306a36Sopenharmony_ci * Context: Process context.
57662306a36Sopenharmony_ci * Return:
57762306a36Sopenharmony_ci * * 0 on success.
57862306a36Sopenharmony_ci * * Negative return value on failure.
57962306a36Sopenharmony_ci */
58062306a36Sopenharmony_cistatic int ne_start_enclave(int enclave_fd,  struct ne_enclave_start_info *enclave_start_info)
58162306a36Sopenharmony_ci{
58262306a36Sopenharmony_ci	int rc = -EINVAL;
58362306a36Sopenharmony_ci
58462306a36Sopenharmony_ci	rc = ioctl(enclave_fd, NE_START_ENCLAVE, enclave_start_info);
58562306a36Sopenharmony_ci	if (rc < 0) {
58662306a36Sopenharmony_ci		switch (errno) {
58762306a36Sopenharmony_ci		case NE_ERR_NOT_IN_INIT_STATE: {
58862306a36Sopenharmony_ci			printf("Error in start enclave, enclave not in init state\n");
58962306a36Sopenharmony_ci
59062306a36Sopenharmony_ci			break;
59162306a36Sopenharmony_ci		}
59262306a36Sopenharmony_ci
59362306a36Sopenharmony_ci		case NE_ERR_NO_MEM_REGIONS_ADDED: {
59462306a36Sopenharmony_ci			printf("Error in start enclave, no memory regions have been added\n");
59562306a36Sopenharmony_ci
59662306a36Sopenharmony_ci			break;
59762306a36Sopenharmony_ci		}
59862306a36Sopenharmony_ci
59962306a36Sopenharmony_ci		case NE_ERR_NO_VCPUS_ADDED: {
60062306a36Sopenharmony_ci			printf("Error in start enclave, no vCPUs have been added\n");
60162306a36Sopenharmony_ci
60262306a36Sopenharmony_ci			break;
60362306a36Sopenharmony_ci		}
60462306a36Sopenharmony_ci
60562306a36Sopenharmony_ci		case NE_ERR_FULL_CORES_NOT_USED: {
60662306a36Sopenharmony_ci			printf("Error in start enclave, enclave has no full cores set\n");
60762306a36Sopenharmony_ci
60862306a36Sopenharmony_ci			break;
60962306a36Sopenharmony_ci		}
61062306a36Sopenharmony_ci
61162306a36Sopenharmony_ci		case NE_ERR_ENCLAVE_MEM_MIN_SIZE: {
61262306a36Sopenharmony_ci			printf("Error in start enclave, enclave memory is less than min size\n");
61362306a36Sopenharmony_ci
61462306a36Sopenharmony_ci			break;
61562306a36Sopenharmony_ci		}
61662306a36Sopenharmony_ci
61762306a36Sopenharmony_ci		case NE_ERR_INVALID_FLAG_VALUE: {
61862306a36Sopenharmony_ci			printf("Error in start enclave, provided invalid flag\n");
61962306a36Sopenharmony_ci
62062306a36Sopenharmony_ci			break;
62162306a36Sopenharmony_ci		}
62262306a36Sopenharmony_ci
62362306a36Sopenharmony_ci		case NE_ERR_INVALID_ENCLAVE_CID: {
62462306a36Sopenharmony_ci			printf("Error in start enclave, provided invalid enclave CID\n");
62562306a36Sopenharmony_ci
62662306a36Sopenharmony_ci			break;
62762306a36Sopenharmony_ci		}
62862306a36Sopenharmony_ci
62962306a36Sopenharmony_ci		default:
63062306a36Sopenharmony_ci			printf("Error in start enclave [%m]\n");
63162306a36Sopenharmony_ci		}
63262306a36Sopenharmony_ci
63362306a36Sopenharmony_ci		return rc;
63462306a36Sopenharmony_ci	}
63562306a36Sopenharmony_ci
63662306a36Sopenharmony_ci	return 0;
63762306a36Sopenharmony_ci}
63862306a36Sopenharmony_ci
63962306a36Sopenharmony_ci/**
64062306a36Sopenharmony_ci * ne_start_enclave_check_booted() - Start the enclave and wait for a heartbeat
64162306a36Sopenharmony_ci *				     from it, on a newly created vsock channel,
64262306a36Sopenharmony_ci *				     to check it has booted.
64362306a36Sopenharmony_ci * @enclave_fd :	The file descriptor associated with the enclave.
64462306a36Sopenharmony_ci *
64562306a36Sopenharmony_ci * Context: Process context.
64662306a36Sopenharmony_ci * Return:
64762306a36Sopenharmony_ci * * 0 on success.
64862306a36Sopenharmony_ci * * Negative return value on failure.
64962306a36Sopenharmony_ci */
65062306a36Sopenharmony_cistatic int ne_start_enclave_check_booted(int enclave_fd)
65162306a36Sopenharmony_ci{
65262306a36Sopenharmony_ci	struct sockaddr_vm client_vsock_addr = {};
65362306a36Sopenharmony_ci	int client_vsock_fd = -1;
65462306a36Sopenharmony_ci	socklen_t client_vsock_len = sizeof(client_vsock_addr);
65562306a36Sopenharmony_ci	struct ne_enclave_start_info enclave_start_info = {};
65662306a36Sopenharmony_ci	struct pollfd fds[1] = {};
65762306a36Sopenharmony_ci	int rc = -EINVAL;
65862306a36Sopenharmony_ci	unsigned char recv_buf = 0;
65962306a36Sopenharmony_ci	struct sockaddr_vm server_vsock_addr = {
66062306a36Sopenharmony_ci		.svm_family = AF_VSOCK,
66162306a36Sopenharmony_ci		.svm_cid = NE_IMAGE_LOAD_HEARTBEAT_CID,
66262306a36Sopenharmony_ci		.svm_port = NE_IMAGE_LOAD_HEARTBEAT_PORT,
66362306a36Sopenharmony_ci	};
66462306a36Sopenharmony_ci	int server_vsock_fd = -1;
66562306a36Sopenharmony_ci
66662306a36Sopenharmony_ci	server_vsock_fd = socket(AF_VSOCK, SOCK_STREAM, 0);
66762306a36Sopenharmony_ci	if (server_vsock_fd < 0) {
66862306a36Sopenharmony_ci		rc = server_vsock_fd;
66962306a36Sopenharmony_ci
67062306a36Sopenharmony_ci		printf("Error in socket [%m]\n");
67162306a36Sopenharmony_ci
67262306a36Sopenharmony_ci		return rc;
67362306a36Sopenharmony_ci	}
67462306a36Sopenharmony_ci
67562306a36Sopenharmony_ci	rc = bind(server_vsock_fd, (struct sockaddr *)&server_vsock_addr,
67662306a36Sopenharmony_ci		  sizeof(server_vsock_addr));
67762306a36Sopenharmony_ci	if (rc < 0) {
67862306a36Sopenharmony_ci		printf("Error in bind [%m]\n");
67962306a36Sopenharmony_ci
68062306a36Sopenharmony_ci		goto out;
68162306a36Sopenharmony_ci	}
68262306a36Sopenharmony_ci
68362306a36Sopenharmony_ci	rc = listen(server_vsock_fd, 1);
68462306a36Sopenharmony_ci	if (rc < 0) {
68562306a36Sopenharmony_ci		printf("Error in listen [%m]\n");
68662306a36Sopenharmony_ci
68762306a36Sopenharmony_ci		goto out;
68862306a36Sopenharmony_ci	}
68962306a36Sopenharmony_ci
69062306a36Sopenharmony_ci	rc = ne_start_enclave(enclave_fd, &enclave_start_info);
69162306a36Sopenharmony_ci	if (rc < 0)
69262306a36Sopenharmony_ci		goto out;
69362306a36Sopenharmony_ci
69462306a36Sopenharmony_ci	printf("Enclave started, CID %llu\n", enclave_start_info.enclave_cid);
69562306a36Sopenharmony_ci
69662306a36Sopenharmony_ci	fds[0].fd = server_vsock_fd;
69762306a36Sopenharmony_ci	fds[0].events = POLLIN;
69862306a36Sopenharmony_ci
69962306a36Sopenharmony_ci	rc = poll(fds, 1, NE_POLL_WAIT_TIME_MS);
70062306a36Sopenharmony_ci	if (rc < 0) {
70162306a36Sopenharmony_ci		printf("Error in poll [%m]\n");
70262306a36Sopenharmony_ci
70362306a36Sopenharmony_ci		goto out;
70462306a36Sopenharmony_ci	}
70562306a36Sopenharmony_ci
70662306a36Sopenharmony_ci	if (!rc) {
70762306a36Sopenharmony_ci		printf("Poll timeout, %d seconds elapsed\n", NE_POLL_WAIT_TIME);
70862306a36Sopenharmony_ci
70962306a36Sopenharmony_ci		rc = -ETIMEDOUT;
71062306a36Sopenharmony_ci
71162306a36Sopenharmony_ci		goto out;
71262306a36Sopenharmony_ci	}
71362306a36Sopenharmony_ci
71462306a36Sopenharmony_ci	if ((fds[0].revents & POLLIN) == 0) {
71562306a36Sopenharmony_ci		printf("Poll received value %d\n", fds[0].revents);
71662306a36Sopenharmony_ci
71762306a36Sopenharmony_ci		rc = -EINVAL;
71862306a36Sopenharmony_ci
71962306a36Sopenharmony_ci		goto out;
72062306a36Sopenharmony_ci	}
72162306a36Sopenharmony_ci
72262306a36Sopenharmony_ci	rc = accept(server_vsock_fd, (struct sockaddr *)&client_vsock_addr,
72362306a36Sopenharmony_ci		    &client_vsock_len);
72462306a36Sopenharmony_ci	if (rc < 0) {
72562306a36Sopenharmony_ci		printf("Error in accept [%m]\n");
72662306a36Sopenharmony_ci
72762306a36Sopenharmony_ci		goto out;
72862306a36Sopenharmony_ci	}
72962306a36Sopenharmony_ci
73062306a36Sopenharmony_ci	client_vsock_fd = rc;
73162306a36Sopenharmony_ci
73262306a36Sopenharmony_ci	/*
73362306a36Sopenharmony_ci	 * Read the heartbeat value that the init process in the enclave sends
73462306a36Sopenharmony_ci	 * after vsock connect.
73562306a36Sopenharmony_ci	 */
73662306a36Sopenharmony_ci	rc = read(client_vsock_fd, &recv_buf, sizeof(recv_buf));
73762306a36Sopenharmony_ci	if (rc < 0) {
73862306a36Sopenharmony_ci		printf("Error in read [%m]\n");
73962306a36Sopenharmony_ci
74062306a36Sopenharmony_ci		goto out;
74162306a36Sopenharmony_ci	}
74262306a36Sopenharmony_ci
74362306a36Sopenharmony_ci	if (rc != sizeof(recv_buf) || recv_buf != NE_IMAGE_LOAD_HEARTBEAT_VALUE) {
74462306a36Sopenharmony_ci		printf("Read %d instead of %d\n", recv_buf,
74562306a36Sopenharmony_ci		       NE_IMAGE_LOAD_HEARTBEAT_VALUE);
74662306a36Sopenharmony_ci
74762306a36Sopenharmony_ci		goto out;
74862306a36Sopenharmony_ci	}
74962306a36Sopenharmony_ci
75062306a36Sopenharmony_ci	/* Write the heartbeat value back. */
75162306a36Sopenharmony_ci	rc = write(client_vsock_fd, &recv_buf, sizeof(recv_buf));
75262306a36Sopenharmony_ci	if (rc < 0) {
75362306a36Sopenharmony_ci		printf("Error in write [%m]\n");
75462306a36Sopenharmony_ci
75562306a36Sopenharmony_ci		goto out;
75662306a36Sopenharmony_ci	}
75762306a36Sopenharmony_ci
75862306a36Sopenharmony_ci	rc = 0;
75962306a36Sopenharmony_ci
76062306a36Sopenharmony_ciout:
76162306a36Sopenharmony_ci	close(server_vsock_fd);
76262306a36Sopenharmony_ci
76362306a36Sopenharmony_ci	return rc;
76462306a36Sopenharmony_ci}
76562306a36Sopenharmony_ci
76662306a36Sopenharmony_ciint main(int argc, char *argv[])
76762306a36Sopenharmony_ci{
76862306a36Sopenharmony_ci	int enclave_fd = -1;
76962306a36Sopenharmony_ci	unsigned int i = 0;
77062306a36Sopenharmony_ci	int ne_dev_fd = -1;
77162306a36Sopenharmony_ci	struct ne_user_mem_region ne_user_mem_regions[NE_DEFAULT_NR_MEM_REGIONS] = {};
77262306a36Sopenharmony_ci	unsigned int ne_vcpus[NE_DEFAULT_NR_VCPUS] = {};
77362306a36Sopenharmony_ci	int rc = -EINVAL;
77462306a36Sopenharmony_ci	pthread_t thread_id = 0;
77562306a36Sopenharmony_ci	unsigned long slot_uid = 0;
77662306a36Sopenharmony_ci
77762306a36Sopenharmony_ci	if (argc != 2) {
77862306a36Sopenharmony_ci		printf("Usage: %s <path_to_enclave_image>\n", argv[0]);
77962306a36Sopenharmony_ci
78062306a36Sopenharmony_ci		exit(EXIT_FAILURE);
78162306a36Sopenharmony_ci	}
78262306a36Sopenharmony_ci
78362306a36Sopenharmony_ci	if (strlen(argv[1]) >= PATH_MAX) {
78462306a36Sopenharmony_ci		printf("The size of the path to enclave image is higher than max path\n");
78562306a36Sopenharmony_ci
78662306a36Sopenharmony_ci		exit(EXIT_FAILURE);
78762306a36Sopenharmony_ci	}
78862306a36Sopenharmony_ci
78962306a36Sopenharmony_ci	ne_dev_fd = open(NE_DEV_NAME, O_RDWR | O_CLOEXEC);
79062306a36Sopenharmony_ci	if (ne_dev_fd < 0) {
79162306a36Sopenharmony_ci		printf("Error in open NE device [%m]\n");
79262306a36Sopenharmony_ci
79362306a36Sopenharmony_ci		exit(EXIT_FAILURE);
79462306a36Sopenharmony_ci	}
79562306a36Sopenharmony_ci
79662306a36Sopenharmony_ci	printf("Creating enclave slot ...\n");
79762306a36Sopenharmony_ci
79862306a36Sopenharmony_ci	rc = ne_create_vm(ne_dev_fd, &slot_uid, &enclave_fd);
79962306a36Sopenharmony_ci
80062306a36Sopenharmony_ci	close(ne_dev_fd);
80162306a36Sopenharmony_ci
80262306a36Sopenharmony_ci	if (rc < 0)
80362306a36Sopenharmony_ci		exit(EXIT_FAILURE);
80462306a36Sopenharmony_ci
80562306a36Sopenharmony_ci	printf("Enclave fd %d\n", enclave_fd);
80662306a36Sopenharmony_ci
80762306a36Sopenharmony_ci	rc = pthread_create(&thread_id, NULL, ne_poll_enclave_fd, (void *)&enclave_fd);
80862306a36Sopenharmony_ci	if (rc < 0) {
80962306a36Sopenharmony_ci		printf("Error in thread create [%m]\n");
81062306a36Sopenharmony_ci
81162306a36Sopenharmony_ci		close(enclave_fd);
81262306a36Sopenharmony_ci
81362306a36Sopenharmony_ci		exit(EXIT_FAILURE);
81462306a36Sopenharmony_ci	}
81562306a36Sopenharmony_ci
81662306a36Sopenharmony_ci	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++) {
81762306a36Sopenharmony_ci		ne_user_mem_regions[i].memory_size = NE_MIN_MEM_REGION_SIZE;
81862306a36Sopenharmony_ci
81962306a36Sopenharmony_ci		rc = ne_alloc_user_mem_region(&ne_user_mem_regions[i]);
82062306a36Sopenharmony_ci		if (rc < 0) {
82162306a36Sopenharmony_ci			printf("Error in alloc userspace memory region, iter %d\n", i);
82262306a36Sopenharmony_ci
82362306a36Sopenharmony_ci			goto release_enclave_fd;
82462306a36Sopenharmony_ci		}
82562306a36Sopenharmony_ci	}
82662306a36Sopenharmony_ci
82762306a36Sopenharmony_ci	rc = ne_load_enclave_image(enclave_fd, ne_user_mem_regions, argv[1]);
82862306a36Sopenharmony_ci	if (rc < 0)
82962306a36Sopenharmony_ci		goto release_enclave_fd;
83062306a36Sopenharmony_ci
83162306a36Sopenharmony_ci	for (i = 0; i < NE_DEFAULT_NR_MEM_REGIONS; i++) {
83262306a36Sopenharmony_ci		rc = ne_set_user_mem_region(enclave_fd, ne_user_mem_regions[i]);
83362306a36Sopenharmony_ci		if (rc < 0) {
83462306a36Sopenharmony_ci			printf("Error in set memory region, iter %d\n", i);
83562306a36Sopenharmony_ci
83662306a36Sopenharmony_ci			goto release_enclave_fd;
83762306a36Sopenharmony_ci		}
83862306a36Sopenharmony_ci	}
83962306a36Sopenharmony_ci
84062306a36Sopenharmony_ci	printf("Enclave memory regions were added\n");
84162306a36Sopenharmony_ci
84262306a36Sopenharmony_ci	for (i = 0; i < NE_DEFAULT_NR_VCPUS; i++) {
84362306a36Sopenharmony_ci		/*
84462306a36Sopenharmony_ci		 * The vCPU is chosen from the enclave vCPU pool, if the value
84562306a36Sopenharmony_ci		 * of the vcpu_id is 0.
84662306a36Sopenharmony_ci		 */
84762306a36Sopenharmony_ci		ne_vcpus[i] = 0;
84862306a36Sopenharmony_ci		rc = ne_add_vcpu(enclave_fd, &ne_vcpus[i]);
84962306a36Sopenharmony_ci		if (rc < 0) {
85062306a36Sopenharmony_ci			printf("Error in add vcpu, iter %d\n", i);
85162306a36Sopenharmony_ci
85262306a36Sopenharmony_ci			goto release_enclave_fd;
85362306a36Sopenharmony_ci		}
85462306a36Sopenharmony_ci
85562306a36Sopenharmony_ci		printf("Added vCPU %d to the enclave\n", ne_vcpus[i]);
85662306a36Sopenharmony_ci	}
85762306a36Sopenharmony_ci
85862306a36Sopenharmony_ci	printf("Enclave vCPUs were added\n");
85962306a36Sopenharmony_ci
86062306a36Sopenharmony_ci	rc = ne_start_enclave_check_booted(enclave_fd);
86162306a36Sopenharmony_ci	if (rc < 0) {
86262306a36Sopenharmony_ci		printf("Error in the enclave start / image loading heartbeat logic [rc=%d]\n", rc);
86362306a36Sopenharmony_ci
86462306a36Sopenharmony_ci		goto release_enclave_fd;
86562306a36Sopenharmony_ci	}
86662306a36Sopenharmony_ci
86762306a36Sopenharmony_ci	printf("Entering sleep for %d seconds ...\n", NE_SLEEP_TIME);
86862306a36Sopenharmony_ci
86962306a36Sopenharmony_ci	sleep(NE_SLEEP_TIME);
87062306a36Sopenharmony_ci
87162306a36Sopenharmony_ci	close(enclave_fd);
87262306a36Sopenharmony_ci
87362306a36Sopenharmony_ci	ne_free_mem_regions(ne_user_mem_regions);
87462306a36Sopenharmony_ci
87562306a36Sopenharmony_ci	exit(EXIT_SUCCESS);
87662306a36Sopenharmony_ci
87762306a36Sopenharmony_cirelease_enclave_fd:
87862306a36Sopenharmony_ci	close(enclave_fd);
87962306a36Sopenharmony_ci	ne_free_mem_regions(ne_user_mem_regions);
88062306a36Sopenharmony_ci
88162306a36Sopenharmony_ci	exit(EXIT_FAILURE);
88262306a36Sopenharmony_ci}
883