samples/nitro_enclaves/ne_ioctl_sample.c

   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  */
 148 struct 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  */
 164 static 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  */
 196 void *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  */
 255 static 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  */
 285 static 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  */
 408 static 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  */
 497 static 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  */
 517 static 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  */
 580 static 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  */
 650 static 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
 760 out:
 761         close(server_vsock_fd);
 762
 763         return rc;
 764 }
 765
 766 int 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
 877 release_enclave_fd:
 878         close(enclave_fd);
 879         ne_free_mem_regions(ne_user_mem_regions);
 880
 881         exit(EXIT_FAILURE);
 882 }