1// SPDX-License-Identifier: GPL-2.0-or-later
2/*
3 * Copyright (c) 2009 FUJITSU LIMITED
4 * Author: Li Zefan <lizf@cn.fujitsu.com>
5 */
6
7#include <sys/types.h>
8#include <sys/mman.h>
9#include <sys/shm.h>
10#include <sys/stat.h>
11#include <err.h>
12#include <errno.h>
13#include <fcntl.h>
14#include <getopt.h>
15#include <limits.h>
16#include <signal.h>
17#include <stdio.h>
18#include <stdlib.h>
19#include <string.h>
20#include <unistd.h>
21#define TST_NO_DEFAULT_MAIN
22#include "tst_test.h"
23
24static int fd;
25
26static volatile int flag_exit;
27static volatile int flag_allocated;
28
29static int opt_mmap_anon;
30static int opt_mmap_file;
31static int opt_mmap_lock1;
32static int opt_mmap_lock2;
33static int opt_shm;
34static int opt_hugepage;
35
36static int key_id;			/* used with opt_shm */
37static unsigned long memsize;
38
39#define FILE_HUGEPAGE	"/hugetlb/hugepagefile"
40
41#define MMAP_ANON	(SCHAR_MAX + 1)
42#define MMAP_FILE	(SCHAR_MAX + 2)
43#define MMAP_LOCK1	(SCHAR_MAX + 3)
44#define MMAP_LOCK2	(SCHAR_MAX + 4)
45#define SHM		(SCHAR_MAX + 5)
46#define HUGEPAGE	(SCHAR_MAX + 6)
47
48static const struct option long_opts[] = {
49	{"mmap-anon", 0, NULL, MMAP_ANON},
50	{"mmap-file", 0, NULL, MMAP_FILE},
51	{"mmap-lock1", 0, NULL, MMAP_LOCK1},
52	{"mmap-lock2", 0, NULL, MMAP_LOCK2},
53	{"shm", 0, NULL, SHM},
54	{"hugepage", 0, NULL, HUGEPAGE},
55	{"size", 1, NULL, 's'},
56	{"key", 1, NULL, 'k'},
57	{NULL, 0, NULL, 0},
58};
59
60/*
61 * process_options: read options from user input
62 */
63static void process_options(int argc, char *argv[])
64{
65	int c;
66	char *end;
67
68	while ((c = getopt_long(argc, argv, "k:s:", long_opts, NULL)) != -1) {
69		switch (c) {
70		case 'k':
71			key_id = atoi(optarg);
72			break;
73		case 's':
74			memsize = strtoul(optarg, &end, 10);
75			if (*end != '\0')
76				errx(1, "wrong -s argument!");
77			break;
78		case MMAP_ANON:
79			opt_mmap_anon = 1;
80			break;
81		case MMAP_FILE:
82			opt_mmap_file = 1;
83			break;
84		case MMAP_LOCK1:
85			opt_mmap_lock1 = 1;
86			break;
87		case MMAP_LOCK2:
88			opt_mmap_lock2 = 1;
89			break;
90		case SHM:
91			opt_shm = 1;
92			break;
93		case HUGEPAGE:
94			opt_hugepage = 1;
95			break;
96		default:
97			errx(1, "unknown option: %c", c);
98			break;
99		}
100	}
101}
102
103/*
104 * touch_memory: force allocating phy memory
105 */
106static void touch_memory(char *p, int size)
107{
108	int i;
109	int pagesize = getpagesize();
110
111	for (i = 0; i < size; i += pagesize)
112		p[i] = 0xef;
113}
114
115static void mmap_anon(void)
116{
117	static char *p;
118
119	if (!flag_allocated) {
120		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
121			 MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
122		if (p == MAP_FAILED)
123			err(1, "mmap(anonymous) failed");
124		touch_memory(p, memsize);
125	} else {
126		if (munmap(p, memsize) == -1)
127			err(1, "munmap(anonymous) failed");
128	}
129}
130
131static void mmap_file(void)
132{
133	static char *p;
134	static int fd_hugepage;
135	int fd_tmp;
136
137	if (!flag_allocated) {
138		if (opt_hugepage) {
139			fd_hugepage = open(FILE_HUGEPAGE,
140					   O_CREAT | O_RDWR, 0755);
141			if (fd_hugepage < 0)
142				err(1, "open hugepage file failed");
143			fd_tmp = fd_hugepage;
144		} else
145			fd_tmp = fd;
146
147		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
148			 MAP_SHARED, fd_tmp, 0);
149		if (p == MAP_FAILED) {
150			if (opt_hugepage)
151				unlink(FILE_HUGEPAGE);
152			err(1, "mmap(file) failed");
153		}
154		touch_memory(p, memsize);
155	} else {
156		if (munmap(p, memsize) == -1)
157			err(1, "munmap(file) failed");
158
159		if (opt_hugepage) {
160			close(fd_hugepage);
161			unlink(FILE_HUGEPAGE);
162		}
163	}
164}
165
166static void mmap_lock1(void)
167{
168	static char *p;
169
170	if (!flag_allocated) {
171		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
172			 MAP_PRIVATE | MAP_ANONYMOUS | MAP_LOCKED, 0, 0);
173		if (p == MAP_FAILED)
174			err(1, "mmap(lock) failed");
175		touch_memory(p, memsize);
176	} else {
177		if (munmap(p, memsize) == -1)
178			err(1, "munmap(lock) failed");
179	}
180}
181
182static void mmap_lock2(void)
183{
184	static char *p;
185
186	if (!flag_allocated) {
187		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
188			 MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
189		if (p == MAP_FAILED)
190			err(1, "mmap failed");
191
192		if (mlock(p, memsize) == -1) {
193			if (errno == EAGAIN)
194				exit(2);
195			else
196				err(1, "mlock failed");
197		}
198	} else {
199		if (munmap(p, memsize) == -1)
200			err(1, "munmap failed");
201	}
202}
203
204static void shm(void)
205{
206	static char *p;
207	static int shmid;
208	unsigned long flag;
209
210	key_t key;
211
212	if (!flag_allocated) {
213		flag = IPC_CREAT | SHM_R | SHM_W;
214		if (opt_hugepage)
215			flag |= SHM_HUGETLB;
216
217		key = ftok("/dev/null", key_id);
218		if (key == -1)
219			err(1, "ftok() failed");
220
221		shmid = shmget(key, memsize, flag);
222		if (shmid == -1)
223			err(1, "shmget() failed");
224
225		p = shmat(shmid, NULL, 0);
226		if (p == (void *)-1) {
227			shmctl(shmid, IPC_RMID, NULL);
228			err(1, "shmat() failed");
229		}
230
231		if (shmctl(shmid, IPC_RMID, NULL) == -1)
232			err(1, "shmctl() failed");
233
234		touch_memory(p, memsize);
235	} else {
236		if (shmdt(p) == -1)
237			err(1, "shmdt() failed");
238	}
239}
240
241/*
242 * sigint_handler: handle SIGINT by set the exit flag.
243 */
244static void sigint_handler(int __attribute__ ((unused)) signo)
245{
246	flag_exit = 1;
247}
248
249/*
250 * sigusr_handler: handler SIGUSR
251 *
252 * When we receive SIGUSR, we allocate some memory according
253 * to the user input when the process started.
254 *
255 * When we receive SIGUSR again, we will free all the allocated
256 * memory.
257 */
258static void sigusr_handler(int __attribute__ ((unused)) signo)
259{
260	if (opt_mmap_anon)
261		mmap_anon();
262
263	if (opt_mmap_file)
264		mmap_file();
265
266	if (opt_mmap_lock1)
267		mmap_lock1();
268
269	if (opt_mmap_lock2)
270		mmap_lock2();
271
272	if (opt_shm)
273		shm();
274
275	flag_allocated = !flag_allocated;
276}
277
278int main(int argc, char *argv[])
279{
280	struct sigaction sigint_action;
281	struct sigaction sigusr_action;
282
283	if ((fd = open("/dev/zero", O_RDWR)) == -1)
284		err(1, "open /dev/zero failed");
285
286	memset(&sigint_action, 0, sizeof(sigint_action));
287	memset(&sigusr_action, 0, sizeof(sigusr_action));
288
289	sigemptyset(&sigint_action.sa_mask);
290	sigint_action.sa_handler = &sigint_handler;
291	if (sigaction(SIGINT, &sigint_action, NULL))
292		err(1, "sigaction(SIGINT)");
293
294	sigemptyset(&sigusr_action.sa_mask);
295	sigusr_action.sa_handler = &sigusr_handler;
296	if (sigaction(SIGUSR1, &sigusr_action, NULL))
297		err(1, "sigaction(SIGUSR1)");
298
299	process_options(argc, argv);
300
301	tst_reinit();
302
303	TST_CHECKPOINT_WAKE(0);
304
305	while (!flag_exit)
306		sleep(1);
307
308	close(fd);
309
310	return 0;
311}
312