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 char *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		printf("Deleting the memory\n");
51		if (shmdt((const void *)shmaddr) != 0) {
52			perror("Detach failure");
53			shmctl(shmid, IPC_RMID, NULL);
54			exit(4);
55		}
56
57		shmctl(shmid, IPC_RMID, NULL);
58		printf("Done deleting the memory\n");
59	}
60	exit(2);
61}
62
63int main(int argc, char **argv)
64{
65	int fd = 0;
66	int key = 0;
67	int *ptr = NULL;
68	int c = 0;
69	int size = 0;
70	char path[256] = "";
71	enum method method = MAX_METHOD;
72	int want_sleep = 0, private = 0;
73	int populate = 0;
74	int write = 0;
75	int reserve = 1;
76
77	if (signal(SIGINT, sig_handler) == SIG_ERR)
78		err(1, "\ncan't catch SIGINT\n");
79
80	/* Parse command-line arguments. */
81	setvbuf(stdout, NULL, _IONBF, 0);
82	self = argv[0];
83
84	while ((c = getopt(argc, argv, "s:p:m:owlrn")) != -1) {
85		switch (c) {
86		case 's':
87			size = atoi(optarg);
88			break;
89		case 'p':
90			strncpy(path, optarg, sizeof(path));
91			break;
92		case 'm':
93			if (atoi(optarg) >= MAX_METHOD) {
94				errno = EINVAL;
95				perror("Invalid -m.");
96				exit_usage();
97			}
98			method = atoi(optarg);
99			break;
100		case 'o':
101			populate = 1;
102			break;
103		case 'w':
104			write = 1;
105			break;
106		case 'l':
107			want_sleep = 1;
108			break;
109		case 'r':
110		    private
111			= 1;
112			break;
113		case 'n':
114			reserve = 0;
115			break;
116		default:
117			errno = EINVAL;
118			perror("Invalid arg");
119			exit_usage();
120		}
121	}
122
123	if (strncmp(path, "", sizeof(path)) != 0) {
124		printf("Writing to this path: %s\n", path);
125	} else {
126		errno = EINVAL;
127		perror("path not found");
128		exit_usage();
129	}
130
131	if (size != 0) {
132		printf("Writing this size: %d\n", size);
133	} else {
134		errno = EINVAL;
135		perror("size not found");
136		exit_usage();
137	}
138
139	if (!populate)
140		printf("Not populating.\n");
141	else
142		printf("Populating.\n");
143
144	if (!write)
145		printf("Not writing to memory.\n");
146
147	if (method == MAX_METHOD) {
148		errno = EINVAL;
149		perror("-m Invalid");
150		exit_usage();
151	} else
152		printf("Using method=%d\n", method);
153
154	if (!private)
155		printf("Shared mapping.\n");
156	else
157		printf("Private mapping.\n");
158
159	if (!reserve)
160		printf("NO_RESERVE mapping.\n");
161	else
162		printf("RESERVE mapping.\n");
163
164	switch (method) {
165	case HUGETLBFS:
166		printf("Allocating using HUGETLBFS.\n");
167		fd = open(path, O_CREAT | O_RDWR, 0777);
168		if (fd == -1)
169			err(1, "Failed to open file.");
170
171		ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
172			   (private ? MAP_PRIVATE : MAP_SHARED) |
173				   (populate ? MAP_POPULATE : 0) |
174				   (reserve ? 0 : MAP_NORESERVE),
175			   fd, 0);
176
177		if (ptr == MAP_FAILED) {
178			close(fd);
179			err(1, "Error mapping the file");
180		}
181		break;
182	case MMAP_MAP_HUGETLB:
183		printf("Allocating using MAP_HUGETLB.\n");
184		ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
185			   (private ? (MAP_PRIVATE | MAP_ANONYMOUS) :
186				      MAP_SHARED) |
187				   MAP_HUGETLB | (populate ? MAP_POPULATE : 0) |
188				   (reserve ? 0 : MAP_NORESERVE),
189			   -1, 0);
190
191		if (ptr == MAP_FAILED)
192			err(1, "mmap");
193
194		printf("Returned address is %p\n", ptr);
195		break;
196	case SHM:
197		printf("Allocating using SHM.\n");
198		shmid = shmget(key, size,
199			       SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
200		if (shmid < 0) {
201			shmid = shmget(++key, size,
202				       SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
203			if (shmid < 0)
204				err(1, "shmget");
205		}
206		printf("shmid: 0x%x, shmget key:%d\n", shmid, key);
207
208		ptr = shmat(shmid, NULL, 0);
209		if (ptr == (int *)-1) {
210			perror("Shared memory attach failure");
211			shmctl(shmid, IPC_RMID, NULL);
212			exit(2);
213		}
214		printf("shmaddr: %p\n", ptr);
215
216		break;
217	default:
218		errno = EINVAL;
219		err(1, "Invalid method.");
220	}
221
222	if (write) {
223		printf("Writing to memory.\n");
224		memset(ptr, 1, size);
225	}
226
227	if (want_sleep) {
228		/* Signal to caller that we're done. */
229		printf("DONE\n");
230
231		/* Hold memory until external kill signal is delivered. */
232		while (1)
233			sleep(100);
234	}
235
236	if (method == HUGETLBFS)
237		close(fd);
238
239	return 0;
240}
241