|
| 1 | +/* |
| 2 | + * Copyright (C) 2024 Intel Corporation |
| 3 | + * |
| 4 | + * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT. |
| 5 | + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 6 | + */ |
| 7 | + |
| 8 | +#include <dlfcn.h> |
| 9 | +#include <fcntl.h> |
| 10 | +#include <stdio.h> |
| 11 | +#include <stdlib.h> |
| 12 | +#include <string.h> |
| 13 | +#include <sys/socket.h> |
| 14 | +#include <sys/stat.h> |
| 15 | +#include <sys/types.h> |
| 16 | +#include <unistd.h> |
| 17 | + |
| 18 | +#include <umf/ipc.h> |
| 19 | + |
| 20 | +#include "ipc_common.h" |
| 21 | +#include "utils_load_library.h" |
| 22 | + |
| 23 | +umf_result_t (*pfnGetIPCHandle)(const void *ptr, umf_ipc_handle_t *umfIPCHandle, |
| 24 | + size_t *size); |
| 25 | +umf_result_t (*pfnPutIPCHandle)(umf_ipc_handle_t umfIPCHandle); |
| 26 | + |
| 27 | +// This is a test for a scenario where a user process is started using the |
| 28 | +// LD_PRELOAD with the UMF Proxy Lib and this process uses UMF by loading |
| 29 | +// libumf.so at runtime. |
| 30 | +// In this test, we expect that all allocations made by the process will be |
| 31 | +// handled by UMF in the Proxy Lib and added to the UMF tracker so that they |
| 32 | +// can be used later in the UMF IPC API. |
| 33 | +int main(int argc, char *argv[]) { |
| 34 | + int ret = 0; |
| 35 | + umf_result_t umf_result = UMF_RESULT_ERROR_UNKNOWN; |
| 36 | + int producer_socket = -1; |
| 37 | + const size_t MSG_SIZE = 2048; |
| 38 | + char consumer_message[MSG_SIZE]; |
| 39 | + |
| 40 | + if (argc < 2) { |
| 41 | + fprintf(stderr, "usage: %s <port> [shm_name]\n", argv[0]); |
| 42 | + return -1; |
| 43 | + } |
| 44 | + |
| 45 | + int port = atoi(argv[1]); |
| 46 | + |
| 47 | + int fd = open("/proc/self/maps", O_RDONLY); |
| 48 | + if (fd == -1) { |
| 49 | + return -1; |
| 50 | + } |
| 51 | + |
| 52 | + // read the "/proc/self/maps" file until the "libumf_proxy.so" of the maps |
| 53 | + // is found or EOF is reached. |
| 54 | + const size_t size_buf = 8192; |
| 55 | + char buf[size_buf]; |
| 56 | + size_t nbytes = 1; |
| 57 | + char *found = NULL; |
| 58 | + while (nbytes > 0 && found == NULL) { |
| 59 | + memset(buf, 0, nbytes); // erase previous data |
| 60 | + nbytes = read(fd, buf, size_buf); |
| 61 | + found = strstr(buf, "libumf_proxy.so"); |
| 62 | + } |
| 63 | + (void)close(fd); |
| 64 | + |
| 65 | + if (found == NULL) { |
| 66 | + fprintf( |
| 67 | + stderr, |
| 68 | + "test binary not run under LD_PRELOAD with \"libumf_proxy.so\"\n"); |
| 69 | + return -1; |
| 70 | + } |
| 71 | + |
| 72 | + // open the UMF library and get umfGetIPCHandle() function |
| 73 | + const char *umf_lib_name = "libumf.so"; |
| 74 | + void *umf_lib_handle = utils_open_library(umf_lib_name, 0); |
| 75 | + if (umf_lib_handle == NULL) { |
| 76 | + fprintf(stderr, "utils_open_library: UMF library not found (%s)\n", |
| 77 | + umf_lib_name); |
| 78 | + return -1; |
| 79 | + } |
| 80 | + |
| 81 | + *(void **)&pfnGetIPCHandle = |
| 82 | + utils_get_symbol_addr(umf_lib_handle, "umfGetIPCHandle", umf_lib_name); |
| 83 | + if (pfnGetIPCHandle == NULL) { |
| 84 | + ret = -1; |
| 85 | + goto err_close_lib; |
| 86 | + } |
| 87 | + |
| 88 | + *(void **)&pfnPutIPCHandle = |
| 89 | + utils_get_symbol_addr(umf_lib_handle, "umfPutIPCHandle", umf_lib_name); |
| 90 | + if (pfnPutIPCHandle == NULL) { |
| 91 | + ret = -1; |
| 92 | + goto err_close_lib; |
| 93 | + } |
| 94 | + |
| 95 | + // create simple allocation - it should be added to the UMF tracker if the |
| 96 | + // process was launched under UMF Proxy Lib |
| 97 | + size_t size = 2137; |
| 98 | + void *ptr = malloc(size); |
| 99 | + if (ptr == NULL) { |
| 100 | + ret = -1; |
| 101 | + goto err_close_lib; |
| 102 | + } |
| 103 | + |
| 104 | + fprintf(stderr, "Allocated memory - %zu\n", size); |
| 105 | + unsigned long long val = 144; |
| 106 | + unsigned long long expected_val = val / 2; |
| 107 | + *(unsigned long long *)ptr = val; |
| 108 | + |
| 109 | + // get IPC handle of the allocation |
| 110 | + umf_ipc_handle_t ipc_handle = NULL; |
| 111 | + size_t ipc_handle_size = 0; |
| 112 | + umf_result_t res = pfnGetIPCHandle(ptr, &ipc_handle, &ipc_handle_size); |
| 113 | + if (res != UMF_RESULT_SUCCESS) { |
| 114 | + fprintf(stderr, "pfnGetIPCHandle() failed!\n"); |
| 115 | + ret = -1; |
| 116 | + goto err_free_mem; |
| 117 | + } |
| 118 | + |
| 119 | + // check if we got valid data |
| 120 | + if (ipc_handle == NULL || ipc_handle_size == 0) { |
| 121 | + fprintf(stderr, "pfnGetIPCHandle() couldn't find the handle data!\n"); |
| 122 | + ret = -1; |
| 123 | + goto err_free_mem; |
| 124 | + } |
| 125 | + |
| 126 | + fprintf(stderr, "Got IPCHandle for memory - %p | size - %zu\n", |
| 127 | + (void *)ipc_handle, ipc_handle_size); |
| 128 | + |
| 129 | + producer_socket = producer_connect(port); |
| 130 | + if (producer_socket < 0) { |
| 131 | + goto err_PutIPCHandle; |
| 132 | + } |
| 133 | + |
| 134 | + // send the ipc_handle_size to the consumer |
| 135 | + ssize_t len = |
| 136 | + send(producer_socket, &ipc_handle_size, sizeof(ipc_handle_size), 0); |
| 137 | + if (len < 0) { |
| 138 | + fprintf(stderr, "[producer] ERROR: unable to send the message\n"); |
| 139 | + goto err_close_producer_socket; |
| 140 | + } |
| 141 | + |
| 142 | + fprintf(stderr, |
| 143 | + "[producer] Sent the size of the IPC handle (%zu) to the consumer " |
| 144 | + "(sent %zu bytes)\n", |
| 145 | + ipc_handle_size, len); |
| 146 | + |
| 147 | + // zero the consumer_message buffer |
| 148 | + memset(consumer_message, 0, sizeof(consumer_message)); |
| 149 | + |
| 150 | + // receive the consumer's confirmation - IPC handle size |
| 151 | + len = recv(producer_socket, consumer_message, sizeof(consumer_message), 0); |
| 152 | + if (len < 0) { |
| 153 | + fprintf(stderr, "[producer] ERROR: error while receiving the " |
| 154 | + "confirmation from the consumer\n"); |
| 155 | + goto err_close_producer_socket; |
| 156 | + } |
| 157 | + |
| 158 | + size_t conf_IPC_handle_size = *(size_t *)consumer_message; |
| 159 | + if (conf_IPC_handle_size == ipc_handle_size) { |
| 160 | + fprintf(stderr, |
| 161 | + "[producer] Received the correct confirmation (%zu) from the " |
| 162 | + "consumer (%zu bytes)\n", |
| 163 | + conf_IPC_handle_size, len); |
| 164 | + } else { |
| 165 | + fprintf(stderr, |
| 166 | + "[producer] Received an INCORRECT confirmation (%zu) from the " |
| 167 | + "consumer (%zu bytes)\n", |
| 168 | + conf_IPC_handle_size, len); |
| 169 | + goto err_close_producer_socket; |
| 170 | + } |
| 171 | + |
| 172 | + // send the ipc_handle of ipc_handle_size to the consumer |
| 173 | + if (send(producer_socket, ipc_handle, ipc_handle_size, 0) < 0) { |
| 174 | + fprintf(stderr, "[producer] ERROR: unable to send the message\n"); |
| 175 | + goto err_close_producer_socket; |
| 176 | + } |
| 177 | + |
| 178 | + fprintf(stderr, |
| 179 | + "[producer] Sent the IPC handle to the consumer (%zu bytes)\n", |
| 180 | + ipc_handle_size); |
| 181 | + |
| 182 | + // zero the consumer_message buffer |
| 183 | + memset(consumer_message, 0, sizeof(consumer_message)); |
| 184 | + |
| 185 | + // receive the consumer's response |
| 186 | + if (recv(producer_socket, consumer_message, sizeof(consumer_message) - 1, |
| 187 | + 0) < 0) { |
| 188 | + fprintf( |
| 189 | + stderr, |
| 190 | + "[producer] ERROR: error while receiving the consumer's message\n"); |
| 191 | + goto err_close_producer_socket; |
| 192 | + } |
| 193 | + |
| 194 | + fprintf(stderr, "[producer] Received the consumer's response: \"%s\"\n", |
| 195 | + consumer_message); |
| 196 | + |
| 197 | + if (strncmp(consumer_message, "SKIP", 5 /* length of "SKIP" + 1 */) == 0) { |
| 198 | + fprintf(stderr, "[producer] SKIP: received the 'SKIP' response from " |
| 199 | + "consumer, skipping ...\n"); |
| 200 | + ret = 1; |
| 201 | + goto err_close_producer_socket; |
| 202 | + } |
| 203 | + |
| 204 | + // read a new value - the expected correct value val / 2 |
| 205 | + volatile unsigned long long new_val = *(unsigned long long *)ptr; |
| 206 | + if (new_val == expected_val) { |
| 207 | + ret = 0; // got the correct value - success! |
| 208 | + fprintf( |
| 209 | + stderr, |
| 210 | + "[producer] The consumer wrote the correct value (the old one / 2) " |
| 211 | + "to my shared memory: %llu\n", |
| 212 | + new_val); |
| 213 | + } else { |
| 214 | + fprintf( |
| 215 | + stderr, |
| 216 | + "[producer] ERROR: The consumer did NOT write the correct value " |
| 217 | + "(the old one / 2 = %llu) to my shared memory: %llu\n", |
| 218 | + expected_val, new_val); |
| 219 | + } |
| 220 | + |
| 221 | +err_close_producer_socket: |
| 222 | + close(producer_socket); |
| 223 | + |
| 224 | +err_PutIPCHandle: |
| 225 | + umf_result = pfnPutIPCHandle(ipc_handle); |
| 226 | + if (umf_result != UMF_RESULT_SUCCESS) { |
| 227 | + fprintf(stderr, "[producer] ERROR: putting the IPC handle failed\n"); |
| 228 | + } |
| 229 | + |
| 230 | + fprintf(stderr, "[producer] Put the IPC handle\n"); |
| 231 | + |
| 232 | + if (ret == 0) { |
| 233 | + fprintf(stderr, "[producer] Shutting down (status OK) ...\n"); |
| 234 | + } else if (ret == 1) { |
| 235 | + fprintf(stderr, "[producer] Shutting down (status SKIP) ...\n"); |
| 236 | + ret = 0; |
| 237 | + } else { |
| 238 | + fprintf(stderr, "[producer] Shutting down (status ERROR) ...\n"); |
| 239 | + } |
| 240 | + |
| 241 | + return ret; |
| 242 | + |
| 243 | +err_free_mem: |
| 244 | + free(ptr); |
| 245 | + |
| 246 | +err_close_lib: |
| 247 | + utils_close_library(umf_lib_handle); |
| 248 | + |
| 249 | + return ret; |
| 250 | +} |
0 commit comments