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
26int 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
47int 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 */
86int 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 */
139int 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 */
216int 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
263int 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