1 // SPDX-License-Identifier: GPL-2.0
2 /*
3  * It tests the mlock/mlock2() when they are invoked
4  * on randomly memory region.
5  */
6 #include <unistd.h>
7 #include <sys/resource.h>
8 #include <sys/capability.h>
9 #include <sys/mman.h>
10 #include <linux/mman.h>
11 #include <fcntl.h>
12 #include <string.h>
13 #include <sys/ipc.h>
14 #include <sys/shm.h>
15 #include <time.h>
16 #include "mlock2.h"
17 
18 #define CHUNK_UNIT (128 * 1024)
19 #define MLOCK_RLIMIT_SIZE (CHUNK_UNIT * 2)
20 #define MLOCK_WITHIN_LIMIT_SIZE CHUNK_UNIT
21 #define MLOCK_OUTOF_LIMIT_SIZE (CHUNK_UNIT * 3)
22 
23 #define TEST_LOOP 100
24 #define PAGE_ALIGN(size, ps) (((size) + ((ps) - 1)) & ~((ps) - 1))
25 
26 int set_cap_limits(rlim_t max)
27 {
28 	struct rlimit new;
29 	cap_t cap = cap_init();
30 
31 	new.rlim_cur = max;
32 	new.rlim_max = max;
33 	if (setrlimit(RLIMIT_MEMLOCK, &new)) {
34 		perror("setrlimit() returns error\n");
35 		return -1;
36 	}
37 
38 	/* drop capabilities including CAP_IPC_LOCK */
39 	if (cap_set_proc(cap)) {
40 		perror("cap_set_proc() returns error\n");
41 		return -2;
42 	}
43 
44 	return 0;
45 }
46 
47 int get_proc_locked_vm_size(void)
48 {
49 	FILE *f;
50 	int ret = -1;
51 	char line[1024] = {0};
52 	unsigned long lock_size = 0;
53 
54 	f = fopen("/proc/self/status", "r");
55 	if (!f) {
56 		perror("fopen");
57 		return -1;
58 	}
59 
60 	while (fgets(line, 1024, f)) {
61 		if (strstr(line, "VmLck")) {
62 			ret = sscanf(line, "VmLck:\t%8lu kB", &lock_size);
63 			if (ret <= 0) {
64 				printf("sscanf() on VmLck error: %s: %d\n",
65 						line, ret);
66 				fclose(f);
67 				return -1;
68 			}
69 			fclose(f);
70 			return (int)(lock_size << 10);
71 		}
72 	}
73 
74 	perror("cannot parse VmLck in /proc/self/status\n");
75 	fclose(f);
76 	return -1;
77 }
78 
79 /*
80  * Get the MMUPageSize of the memory region including input
81  * address from proc file.
82  *
83  * return value: on error case, 0 will be returned.
84  * Otherwise the page size(in bytes) is returned.
85  */
86 int get_proc_page_size(unsigned long addr)
87 {
88 	FILE *smaps;
89 	char *line;
90 	unsigned long mmupage_size = 0;
91 	size_t size;
92 
93 	smaps = seek_to_smaps_entry(addr);
94 	if (!smaps) {
95 		printf("Unable to parse /proc/self/smaps\n");
96 		return 0;
97 	}
98 
99 	while (getline(&line, &size, smaps) > 0) {
100 		if (!strstr(line, "MMUPageSize")) {
101 			free(line);
102 			line = NULL;
103 			size = 0;
104 			continue;
105 		}
106 
107 		/* found the MMUPageSize of this section */
108 		if (sscanf(line, "MMUPageSize:    %8lu kB",
109 					&mmupage_size) < 1) {
110 			printf("Unable to parse smaps entry for Size:%s\n",
111 					line);
112 			break;
113 		}
114 
115 	}
116 	free(line);
117 	if (smaps)
118 		fclose(smaps);
119 	return mmupage_size << 10;
120 }
121 
122 /*
123  * Test mlock/mlock2() on provided memory chunk.
124  * It expects the mlock/mlock2() to be successful (within rlimit)
125  *
126  * With allocated memory chunk [p, p + alloc_size), this
127  * test will choose start/len randomly to perform mlock/mlock2
128  * [start, start +  len] memory range. The range is within range
129  * of the allocated chunk.
130  *
131  * The memory region size alloc_size is within the rlimit.
132  * So we always expect a success of mlock/mlock2.
133  *
134  * VmLck is assumed to be 0 before this test.
135  *
136  *    return value: 0 - success
137  *    else: failure
138  */
139 int test_mlock_within_limit(char *p, int alloc_size)
140 {
141 	int i;
142 	int ret = 0;
143 	int locked_vm_size = 0;
144 	struct rlimit cur;
145 	int page_size = 0;
146 
147 	getrlimit(RLIMIT_MEMLOCK, &cur);
148 	if (cur.rlim_cur < alloc_size) {
149 		printf("alloc_size[%d] < %u rlimit,lead to mlock failure\n",
150 				alloc_size, (unsigned int)cur.rlim_cur);
151 		return -1;
152 	}
153 
154 	srand(time(NULL));
155 	for (i = 0; i < TEST_LOOP; i++) {
156 		/*
157 		 * - choose mlock/mlock2 randomly
158 		 * - choose lock_size randomly but lock_size < alloc_size
159 		 * - choose start_offset randomly but p+start_offset+lock_size
160 		 *   < p+alloc_size
161 		 */
162 		int is_mlock = !!(rand() % 2);
163 		int lock_size = rand() % alloc_size;
164 		int start_offset = rand() % (alloc_size - lock_size);
165 
166 		if (is_mlock)
167 			ret = mlock(p + start_offset, lock_size);
168 		else
169 			ret = mlock2_(p + start_offset, lock_size,
170 				       MLOCK_ONFAULT);
171 
172 		if (ret) {
173 			printf("%s() failure at |%p(%d)| mlock:|%p(%d)|\n",
174 					is_mlock ? "mlock" : "mlock2",
175 					p, alloc_size,
176 					p + start_offset, lock_size);
177 			return ret;
178 		}
179 	}
180 
181 	/*
182 	 * Check VmLck left by the tests.
183 	 */
184 	locked_vm_size = get_proc_locked_vm_size();
185 	page_size = get_proc_page_size((unsigned long)p);
186 	if (page_size == 0) {
187 		printf("cannot get proc MMUPageSize\n");
188 		return -1;
189 	}
190 
191 	if (locked_vm_size > PAGE_ALIGN(alloc_size, page_size) + page_size) {
192 		printf("test_mlock_within_limit() left VmLck:%d on %d chunk\n",
193 				locked_vm_size, alloc_size);
194 		return -1;
195 	}
196 
197 	return 0;
198 }
199 
200 
201 /*
202  * We expect the mlock/mlock2() to be fail (outof limitation)
203  *
204  * With allocated memory chunk [p, p + alloc_size), this
205  * test will randomly choose start/len and perform mlock/mlock2
206  * on [start, start+len] range.
207  *
208  * The memory region size alloc_size is above the rlimit.
209  * And the len to be locked is higher than rlimit.
210  * So we always expect a failure of mlock/mlock2.
211  * No locked page number should be increased as a side effect.
212  *
213  *    return value: 0 - success
214  *    else: failure
215  */
216 int test_mlock_outof_limit(char *p, int alloc_size)
217 {
218 	int i;
219 	int ret = 0;
220 	int locked_vm_size = 0, old_locked_vm_size = 0;
221 	struct rlimit cur;
222 
223 	getrlimit(RLIMIT_MEMLOCK, &cur);
224 	if (cur.rlim_cur >= alloc_size) {
225 		printf("alloc_size[%d] >%u rlimit, violates test condition\n",
226 				alloc_size, (unsigned int)cur.rlim_cur);
227 		return -1;
228 	}
229 
230 	old_locked_vm_size = get_proc_locked_vm_size();
231 	srand(time(NULL));
232 	for (i = 0; i < TEST_LOOP; i++) {
233 		int is_mlock = !!(rand() % 2);
234 		int lock_size = (rand() % (alloc_size - cur.rlim_cur))
235 			+ cur.rlim_cur;
236 		int start_offset = rand() % (alloc_size - lock_size);
237 
238 		if (is_mlock)
239 			ret = mlock(p + start_offset, lock_size);
240 		else
241 			ret = mlock2_(p + start_offset, lock_size,
242 					MLOCK_ONFAULT);
243 		if (ret == 0) {
244 			printf("%s() succeeds? on %p(%d) mlock%p(%d)\n",
245 					is_mlock ? "mlock" : "mlock2",
246 					p, alloc_size,
247 					p + start_offset, lock_size);
248 			return -1;
249 		}
250 	}
251 
252 	locked_vm_size = get_proc_locked_vm_size();
253 	if (locked_vm_size != old_locked_vm_size) {
254 		printf("tests leads to new mlocked page: old[%d], new[%d]\n",
255 				old_locked_vm_size,
256 				locked_vm_size);
257 		return -1;
258 	}
259 
260 	return 0;
261 }
262 
263 int main(int argc, char **argv)
264 {
265 	char *p = NULL;
266 	int ret = 0;
267 
268 	if (set_cap_limits(MLOCK_RLIMIT_SIZE))
269 		return -1;
270 
271 	p = malloc(MLOCK_WITHIN_LIMIT_SIZE);
272 	if (p == NULL) {
273 		perror("malloc() failure\n");
274 		return -1;
275 	}
276 	ret = test_mlock_within_limit(p, MLOCK_WITHIN_LIMIT_SIZE);
277 	if (ret)
278 		return ret;
279 	munlock(p, MLOCK_WITHIN_LIMIT_SIZE);
280 	free(p);
281 
282 
283 	p = malloc(MLOCK_OUTOF_LIMIT_SIZE);
284 	if (p == NULL) {
285 		perror("malloc() failure\n");
286 		return -1;
287 	}
288 	ret = test_mlock_outof_limit(p, MLOCK_OUTOF_LIMIT_SIZE);
289 	if (ret)
290 		return ret;
291 	munlock(p, MLOCK_OUTOF_LIMIT_SIZE);
292 	free(p);
293 
294 	return 0;
295 }
296