1// SPDX-License-Identifier: GPL-2.0
  2/*
  3 * This program reserves and uses hugetlb memory, supporting a bunch of
  4 * scenarios needed by the charged_reserved_hugetlb.sh test.
  5 */
  6
  7#include <err.h>
  8#include <errno.h>
  9#include <signal.h>
 10#include <stdio.h>
 11#include <stdlib.h>
 12#include <string.h>
 13#include <unistd.h>
 14#include <fcntl.h>
 15#include <sys/types.h>
 16#include <sys/shm.h>
 17#include <sys/stat.h>
 18#include <sys/mman.h>
 19
 20/* Global definitions. */
 21enum method {
 22	HUGETLBFS,
 23	MMAP_MAP_HUGETLB,
 24	SHM,
 25	MAX_METHOD
 26};
 27
 28
 29/* Global variables. */
 30static const char *self;
 31static int *shmaddr;
 32static int shmid;
 33
 34/*
 35 * Show usage and exit.
 36 */
 37static void exit_usage(void)
 38{
 39	printf("Usage: %s -p <path to hugetlbfs file> -s <size to map> "
 40	       "[-m <0=hugetlbfs | 1=mmap(MAP_HUGETLB)>] [-l] [-r] "
 41	       "[-o] [-w] [-n]\n",
 42	       self);
 43	exit(EXIT_FAILURE);
 44}
 45
 46void sig_handler(int signo)
 47{
 48	printf("Received %d.\n", signo);
 49	if (signo == SIGINT) {
 50		if (shmaddr) {
 51			printf("Deleting the memory\n");
 52			if (shmdt((const void *)shmaddr) != 0) {
 53				perror("Detach failure");
 54				shmctl(shmid, IPC_RMID, NULL);
 55				exit(4);
 56			}
 57
 58			shmctl(shmid, IPC_RMID, NULL);
 59			printf("Done deleting the memory\n");
 60		}
 61	}
 62	exit(2);
 63}
 64
 65int main(int argc, char **argv)
 66{
 67	int fd = 0;
 68	int key = 0;
 69	int *ptr = NULL;
 70	int c = 0;
 71	int size = 0;
 72	char path[256] = "";
 73	enum method method = MAX_METHOD;
 74	int want_sleep = 0, private = 0;
 75	int populate = 0;
 76	int write = 0;
 77	int reserve = 1;
 78
 79	if (signal(SIGINT, sig_handler) == SIG_ERR)
 80		err(1, "\ncan't catch SIGINT\n");
 81
 82	/* Parse command-line arguments. */
 83	setvbuf(stdout, NULL, _IONBF, 0);
 84	self = argv[0];
 85
 86	while ((c = getopt(argc, argv, "s:p:m:owlrn")) != -1) {
 87		switch (c) {
 88		case 's':
 89			size = atoi(optarg);
 90			break;
 91		case 'p':
 92			strncpy(path, optarg, sizeof(path));
 93			break;
 94		case 'm':
 95			if (atoi(optarg) >= MAX_METHOD) {
 96				errno = EINVAL;
 97				perror("Invalid -m.");
 98				exit_usage();
 99			}
100			method = atoi(optarg);
101			break;
102		case 'o':
103			populate = 1;
104			break;
105		case 'w':
106			write = 1;
107			break;
108		case 'l':
109			want_sleep = 1;
110			break;
111		case 'r':
112		    private
113			= 1;
114			break;
115		case 'n':
116			reserve = 0;
117			break;
118		default:
119			errno = EINVAL;
120			perror("Invalid arg");
121			exit_usage();
122		}
123	}
124
125	if (strncmp(path, "", sizeof(path)) != 0) {
126		printf("Writing to this path: %s\n", path);
127	} else {
128		errno = EINVAL;
129		perror("path not found");
130		exit_usage();
131	}
132
133	if (size != 0) {
134		printf("Writing this size: %d\n", size);
135	} else {
136		errno = EINVAL;
137		perror("size not found");
138		exit_usage();
139	}
140
141	if (!populate)
142		printf("Not populating.\n");
143	else
144		printf("Populating.\n");
145
146	if (!write)
147		printf("Not writing to memory.\n");
148
149	if (method == MAX_METHOD) {
150		errno = EINVAL;
151		perror("-m Invalid");
152		exit_usage();
153	} else
154		printf("Using method=%d\n", method);
155
156	if (!private)
157		printf("Shared mapping.\n");
158	else
159		printf("Private mapping.\n");
160
161	if (!reserve)
162		printf("NO_RESERVE mapping.\n");
163	else
164		printf("RESERVE mapping.\n");
165
166	switch (method) {
167	case HUGETLBFS:
168		printf("Allocating using HUGETLBFS.\n");
169		fd = open(path, O_CREAT | O_RDWR, 0777);
170		if (fd == -1)
171			err(1, "Failed to open file.");
172
173		ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
174			   (private ? MAP_PRIVATE : MAP_SHARED) |
175				   (populate ? MAP_POPULATE : 0) |
176				   (reserve ? 0 : MAP_NORESERVE),
177			   fd, 0);
178
179		if (ptr == MAP_FAILED) {
180			close(fd);
181			err(1, "Error mapping the file");
182		}
183		break;
184	case MMAP_MAP_HUGETLB:
185		printf("Allocating using MAP_HUGETLB.\n");
186		ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
187			   (private ? (MAP_PRIVATE | MAP_ANONYMOUS) :
188				      MAP_SHARED) |
189				   MAP_HUGETLB | (populate ? MAP_POPULATE : 0) |
190				   (reserve ? 0 : MAP_NORESERVE),
191			   -1, 0);
192
193		if (ptr == MAP_FAILED)
194			err(1, "mmap");
195
196		printf("Returned address is %p\n", ptr);
197		break;
198	case SHM:
199		printf("Allocating using SHM.\n");
200		shmid = shmget(key, size,
201			       SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
202		if (shmid < 0) {
203			shmid = shmget(++key, size,
204				       SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
205			if (shmid < 0)
206				err(1, "shmget");
207		}
208		printf("shmid: 0x%x, shmget key:%d\n", shmid, key);
209
210		ptr = shmat(shmid, NULL, 0);
211		if (ptr == (int *)-1) {
212			perror("Shared memory attach failure");
213			shmctl(shmid, IPC_RMID, NULL);
214			exit(2);
215		}
216		shmaddr = ptr;
217		printf("shmaddr: %p\n", shmaddr);
218
219		break;
220	default:
221		errno = EINVAL;
222		err(1, "Invalid method.");
223	}
224
225	if (write) {
226		printf("Writing to memory.\n");
227		memset(ptr, 1, size);
228	}
229
230	if (want_sleep) {
231		/* Signal to caller that we're done. */
232		printf("DONE\n");
233
234		/* Hold memory until external kill signal is delivered. */
235		while (1)
236			sleep(100);
237	}
238
239	if (method == HUGETLBFS)
240		close(fd);
241
242	return 0;
243}