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. */
21 enum method {
22 	HUGETLBFS,
23 	MMAP_MAP_HUGETLB,
24 	SHM,
25 	MAX_METHOD
26 };
27 
28 
29 /* Global variables. */
30 static const char *self;
31 static char *shmaddr;
32 static int shmid;
33 
34 /*
35  * Show usage and exit.
36  */
37 static 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 
46 void 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 
63 int 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