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