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