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 int *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 		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 
65 int 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 }
244