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
25int 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
46int 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("cann't 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 */
85int 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 */
138int 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 */
215int 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
262int 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