xref: /dragonfly/sbin/vquota/vquota.c (revision 2020c8fe)
1 /*
2  * Copyright (c) 2011, 2012 François Tigeot <ftigeot@wolfpond.org>
3  * All rights reserved.
4  *
5  * Redistribution and use in source and binary forms, with or without
6  * modification, are permitted provided that the following conditions
7  * are met:
8  *
9  * 1. Redistributions of source code must retain the above copyright
10  *    notice, this list of conditions and the following disclaimer.
11  * 2. Redistributions in binary form must reproduce the above copyright
12  *    notice, this list of conditions and the following disclaimer in
13  *    the documentation and/or other materials provided with the
14  *    distribution.
15  * 3. Neither the name of The DragonFly Project nor the names of its
16  *    contributors may be used to endorse or promote products derived
17  *    from this software without specific, prior written permission.
18  *
19  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20  * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21  * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
22  * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
23  * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
24  * INCIDENTAL, SPECIAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES (INCLUDING,
25  * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
26  * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
27  * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28  * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
29  * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
30  * SUCH DAMAGE.
31  */
32 
33 #include <sys/stat.h>
34 #include <sys/mount.h>
35 #include <sys/vfs_quota.h>
36 
37 #include <stdio.h>
38 #include <stdlib.h>
39 #include <err.h>
40 #include <string.h>
41 #include <fts.h>
42 #include <libprop/proplib.h>
43 #include <unistd.h>
44 #include <sys/tree.h>
45 #include <errno.h>
46 #include <inttypes.h>
47 #include <sys/types.h>
48 #include <libutil.h>
49 #include <pwd.h>
50 #include <grp.h>
51 
52 static bool flag_debug = 0;
53 static bool flag_humanize = 0;
54 static bool flag_resolve_ids = 1;
55 
56 static void usage(int);
57 static int get_dirsize(char *);
58 static int get_fslist(void);
59 
60 static void
61 usage(int retcode)
62 {
63 	fprintf(stderr, "usage: vquota [-Dhn] check directory\n");
64 	fprintf(stderr, "       vquota [-Dhn] lsfs\n");
65 	fprintf(stderr, "       vquota [-Dhn] show mount_point\n");
66 	fprintf(stderr, "       vquota [-Dhn] sync mount_point\n");
67 	exit(retcode);
68 }
69 
70 /*
71  * Inode numbers with more than one hard link often come in groups;
72  * use linear arrays of 1024 ones as the basic unit of allocation.
73  * We only need to check if the inodes have been previously processed,
74  * bit arrays are perfect for that purpose.
75  */
76 #define HL_CHUNK_BITS		10
77 #define HL_CHUNK_ENTRIES	(1<<HL_CHUNK_BITS)
78 #define HL_CHUNK_MASK		(HL_CHUNK_ENTRIES - 1)
79 #define BA_UINT64_BITS		6
80 #define BA_UINT64_ENTRIES	(1<<BA_UINT64_BITS)
81 #define BA_UINT64_MASK		(BA_UINT64_ENTRIES - 1)
82 
83 struct hl_node {
84 	RB_ENTRY(hl_node)	rb_entry;
85 	ino_t			ino_left_bits;
86 	uint64_t		hl_chunk[HL_CHUNK_ENTRIES/64];
87 };
88 
89 RB_HEAD(hl_tree,hl_node)	hl_root;
90 
91 RB_PROTOTYPE(hl_tree, hl_node, rb_entry, rb_hl_node_cmp);
92 
93 static int
94 rb_hl_node_cmp(struct hl_node *a, struct hl_node *b);
95 
96 RB_GENERATE(hl_tree, hl_node, rb_entry, rb_hl_node_cmp);
97 
98 struct hl_node* hl_node_insert(ino_t);
99 
100 
101 static int
102 rb_hl_node_cmp(struct hl_node *a, struct hl_node *b)
103 {
104 	if (a->ino_left_bits < b->ino_left_bits)
105 		return(-1);
106 	else if (a->ino_left_bits > b->ino_left_bits)
107 		return(1);
108 	return(0);
109 }
110 
111 struct hl_node* hl_node_insert(ino_t inode)
112 {
113 	struct hl_node *hlp, *res;
114 
115 	hlp = malloc(sizeof(struct hl_node));
116 	if (hlp == NULL) {
117 		/* shouldn't happen */
118 		printf("hl_node_insert(): malloc failed\n");
119 		exit(ENOMEM);
120 	}
121 	bzero(hlp, sizeof(struct hl_node));
122 
123 	hlp->ino_left_bits = (inode >> HL_CHUNK_BITS);
124 	res = RB_INSERT(hl_tree, &hl_root, hlp);
125 
126 	if (res != NULL)	/* shouldn't happen */
127 		printf("hl_node_insert(): RB_INSERT didn't return NULL\n");
128 
129 	return hlp;
130 }
131 
132 /*
133  * hl_register: register an inode number in a rb-tree of bit arrays
134  * returns:
135  * - true if the inode was already processed
136  * - false otherwise
137  */
138 static bool
139 hl_register(ino_t inode)
140 {
141 	struct hl_node hl_find, *hlp;
142 	uint64_t ino_right_bits, ba_index, ba_offset;
143 	uint64_t bitmask, bitval;
144 	bool retval = false;
145 
146 	/* calculate the different addresses of the wanted bit */
147 	hl_find.ino_left_bits = (inode >> HL_CHUNK_BITS);
148 
149 	ino_right_bits = inode & HL_CHUNK_MASK;
150 	ba_index  = ino_right_bits >> BA_UINT64_BITS;
151 	ba_offset = ino_right_bits & BA_UINT64_MASK;
152 
153 	/* no existing node? create and initialize it */
154 	if ((hlp = RB_FIND(hl_tree, &hl_root, &hl_find)) == NULL) {
155 		hlp = hl_node_insert(inode);
156 	}
157 
158 	/* node was found, check the bit value */
159 	bitmask = 1 << ba_offset;
160 	bitval = hlp->hl_chunk[ba_index] & bitmask;
161 	if (bitval != 0) {
162 		retval = true;
163 	}
164 
165 	/* set the bit */
166 	hlp->hl_chunk[ba_index] |= bitmask;
167 
168 	return retval;
169 }
170 
171 /* global variable used by get_dir_size() */
172 uint64_t global_size;
173 
174 /* storage for collected id numbers */
175 /* FIXME: same data structures used in kernel, should find a way to
176  * deduplicate this code */
177 
178 static int
179 rb_ac_unode_cmp(struct ac_unode*, struct ac_unode*);
180 static int
181 rb_ac_gnode_cmp(struct ac_gnode*, struct ac_gnode*);
182 
183 RB_HEAD(ac_utree,ac_unode) ac_uroot;
184 RB_HEAD(ac_gtree,ac_gnode) ac_groot;
185 RB_PROTOTYPE(ac_utree, ac_unode, rb_entry, rb_ac_unode_cmp);
186 RB_PROTOTYPE(ac_gtree, ac_gnode, rb_entry, rb_ac_gnode_cmp);
187 RB_GENERATE(ac_utree, ac_unode, rb_entry, rb_ac_unode_cmp);
188 RB_GENERATE(ac_gtree, ac_gnode, rb_entry, rb_ac_gnode_cmp);
189 
190 static int
191 rb_ac_unode_cmp(struct ac_unode *a, struct ac_unode *b)
192 {
193 	if (a->left_bits < b->left_bits)
194 		return(-1);
195 	else if (a->left_bits > b->left_bits)
196 		return(1);
197 	return(0);
198 }
199 
200 static int
201 rb_ac_gnode_cmp(struct ac_gnode *a, struct ac_gnode *b)
202 {
203 	if (a->left_bits < b->left_bits)
204 		return(-1);
205 	else if (a->left_bits > b->left_bits)
206 		return(1);
207 	return(0);
208 }
209 
210 static struct ac_unode*
211 unode_insert(uid_t uid)
212 {
213 	struct ac_unode *unp, *res;
214 
215 	unp = malloc(sizeof(struct ac_unode));
216 	if (unp == NULL) {
217 		printf("unode_insert(): malloc failed\n");
218 		exit(ENOMEM);
219 	}
220 	bzero(unp, sizeof(struct ac_unode));
221 
222 	unp->left_bits = (uid >> ACCT_CHUNK_BITS);
223 	res = RB_INSERT(ac_utree, &ac_uroot, unp);
224 
225 	if (res != NULL)	/* shouldn't happen */
226 		printf("unode_insert(): RB_INSERT didn't return NULL\n");
227 
228 	return unp;
229 }
230 
231 static struct ac_gnode*
232 gnode_insert(gid_t gid)
233 {
234 	struct ac_gnode *gnp, *res;
235 
236 	gnp = malloc(sizeof(struct ac_gnode));
237 	if (gnp == NULL) {
238 		printf("gnode_insert(): malloc failed\n");
239 		exit(ENOMEM);
240 	}
241 	bzero(gnp, sizeof(struct ac_gnode));
242 
243 	gnp->left_bits = (gid >> ACCT_CHUNK_BITS);
244 	res = RB_INSERT(ac_gtree, &ac_groot, gnp);
245 
246 	if (res != NULL)	/* shouldn't happen */
247 		printf("gnode_insert(): RB_INSERT didn't return NULL\n");
248 
249 	return gnp;
250 }
251 
252 /*
253  * get_dirsize(): walks a directory tree in the same filesystem
254  * output:
255  * - global rb-trees ac_uroot and ac_groot
256  * - global variable global_size
257  */
258 static int
259 get_dirsize(char* dirname)
260 {
261 	FTS		*fts;
262 	FTSENT		*p;
263 	char*		fts_args[2];
264 	int		retval = 0;
265 
266 	/* what we need */
267 	ino_t		file_inode;
268 	off_t		file_size;
269 	uid_t		file_uid;
270 	gid_t		file_gid;
271 
272 	struct ac_unode *unp, ufind;
273 	struct ac_gnode *gnp, gfind;
274 
275 	/* TODO: check directory name sanity */
276 	fts_args[0] = dirname;
277 	fts_args[1] = NULL;
278 
279 	if ((fts = fts_open(fts_args, FTS_PHYSICAL|FTS_XDEV, NULL)) == NULL)
280 		err(1, "fts_open() failed");
281 
282 	while ((p = fts_read(fts)) != NULL) {
283 		switch (p->fts_info) {
284 		/* directories, ignore them */
285 		case FTS_D:
286 		case FTS_DC:
287 		case FTS_DP:
288 			break;
289 		/* read errors, warn, continue and flag */
290 		case FTS_DNR:
291 		case FTS_ERR:
292 		case FTS_NS:
293 			warnx("%s: %s", p->fts_path, strerror(p->fts_errno));
294 			retval = 1;
295 			break;
296 		default:
297 			file_inode = p->fts_statp->st_ino;
298 			file_size = p->fts_statp->st_size;
299 			file_uid = p->fts_statp->st_uid;
300 			file_gid = p->fts_statp->st_gid;
301 
302 			/* files with more than one hard link: */
303 			/* process them only once */
304 			if (p->fts_statp->st_nlink > 1)
305 				if (hl_register(file_inode) == false)
306 					break;
307 
308 			global_size += file_size;
309 			ufind.left_bits = (file_uid >> ACCT_CHUNK_BITS);
310 			gfind.left_bits = (file_gid >> ACCT_CHUNK_BITS);
311 			if ((unp = RB_FIND(ac_utree, &ac_uroot, &ufind)) == NULL)
312 				unp = unode_insert(file_uid);
313 			if ((gnp = RB_FIND(ac_gtree, &ac_groot, &gfind)) == NULL)
314 				gnp = gnode_insert(file_gid);
315 			unp->uid_chunk[(file_uid & ACCT_CHUNK_MASK)] += file_size;
316 			gnp->gid_chunk[(file_gid & ACCT_CHUNK_MASK)] += file_size;
317 		}
318 	}
319 	fts_close(fts);
320 
321 	return retval;
322 }
323 
324 static void
325 print_user(uid_t uid)
326 {
327 	struct passwd *pw;
328 
329 	if (flag_resolve_ids && ((pw = getpwuid(uid)) != NULL)) {
330 		printf("user %s:", pw->pw_name);
331 	} else {
332 		printf("uid %u:", uid);
333 	}
334 }
335 
336 static void
337 print_group(gid_t gid)
338 {
339 	struct group *gr;
340 
341 	if (flag_resolve_ids && ((gr = getgrgid(gid)) != NULL)) {
342 		printf("group %s:", gr->gr_name);
343 	} else {
344 		printf("gid %u:", gid);
345 	}
346 }
347 
348 static int
349 cmd_check(char* dirname)
350 {
351 	int32_t uid, gid;
352 	char	hbuf[5];
353 	struct  ac_unode *unp;
354 	struct  ac_gnode *gnp;
355 	int	rv, i;
356 
357 	rv = get_dirsize(dirname);
358 
359 	if (flag_humanize) {
360 		humanize_number(hbuf, sizeof(hbuf), global_size, "",
361 		    HN_AUTOSCALE, HN_NOSPACE);
362 		printf("total: %s\n", hbuf);
363 	} else {
364 		printf("total: %"PRIu64"\n", global_size);
365 	}
366 	RB_FOREACH(unp, ac_utree, &ac_uroot) {
367 	    for (i=0; i<ACCT_CHUNK_NIDS; i++) {
368 		if (unp->uid_chunk[i] != 0) {
369 		    uid = (unp->left_bits << ACCT_CHUNK_BITS) + i;
370 		    print_user(uid);
371 		    if (flag_humanize) {
372 			humanize_number(hbuf, sizeof(hbuf),
373 			unp->uid_chunk[i], "", HN_AUTOSCALE, HN_NOSPACE);
374 			printf(" %s\n", hbuf);
375 		    } else {
376 			printf(" %" PRIu64 "\n", unp->uid_chunk[i]);
377 		    }
378 		}
379 	    }
380 	}
381 	RB_FOREACH(gnp, ac_gtree, &ac_groot) {
382 	    for (i=0; i<ACCT_CHUNK_NIDS; i++) {
383 		if (gnp->gid_chunk[i] != 0) {
384 		    gid = (gnp->left_bits << ACCT_CHUNK_BITS) + i;
385 		    print_group(gid);
386 		    if (flag_humanize) {
387 			humanize_number(hbuf, sizeof(hbuf),
388 			gnp->gid_chunk[i], "", HN_AUTOSCALE, HN_NOSPACE);
389 			printf(" %s\n", hbuf);
390 		    } else {
391 			printf(" %" PRIu64 "\n", gnp->gid_chunk[i]);
392 		    }
393 		}
394 	    }
395 	}
396 
397 	return rv;
398 }
399 
400 /* print a list of filesystems with accounting enabled */
401 static int
402 get_fslist(void)
403 {
404 	struct statfs *mntbufp;
405 	int nloc, i;
406 
407 	/* read mount table from kernel */
408 	nloc = getmntinfo(&mntbufp, MNT_NOWAIT|MNT_LOCAL);
409 	if (nloc <= 0) {
410 		perror("getmntinfo");
411 		exit(1);
412 	}
413 
414 	/* iterate mounted filesystems */
415 	for (i=0; i<nloc; i++) {
416 	    /* vfs accounting enabled on this one ? */
417 	    if (mntbufp[i].f_flags & MNT_ACCOUNTING)
418 		printf("%s on %s\n", mntbufp[i].f_mntfromname,
419 		    mntbufp[i].f_mntonname);
420 	}
421 
422 	return 0;
423 }
424 
425 static bool
426 send_command(const char *path, const char *cmd,
427 		prop_object_t args, prop_dictionary_t *res)
428 {
429 	prop_dictionary_t dict;
430 	struct plistref pref;
431 
432 	bool rv;
433 	int error;
434 
435 	dict = prop_dictionary_create();
436 
437 	if (dict == NULL) {
438 		printf("send_command(): couldn't create dictionary\n");
439 		return false;
440 	}
441 
442 	rv = prop_dictionary_set_cstring(dict, "command", cmd);
443 	if (rv== false) {
444 		printf("send_command(): couldn't initialize dictionary\n");
445 		return false;
446 	}
447 
448 	rv = prop_dictionary_set(dict, "arguments", args);
449 	if (rv == false) {
450 		printf("prop_dictionary_set() failed\n");
451 		return false;
452 	}
453 
454 	error = prop_dictionary_send_syscall(dict, &pref);
455 	if (error != 0) {
456 		printf("prop_dictionary_send_syscall() failed\n");
457 		prop_object_release(dict);
458 		return false;
459 	}
460 
461 	if (flag_debug)
462 		printf("Message to kernel:\n%s\n",
463 		    prop_dictionary_externalize(dict));
464 
465 	error = vquotactl(path, &pref);
466 	if (error != 0) {
467 		printf("send_command: vquotactl = %d\n", error);
468 		return false;
469 	}
470 
471 	error = prop_dictionary_recv_syscall(&pref, res);
472 	if (error != 0) {
473 		printf("prop_dictionary_recv_syscall() failed\n");
474 	}
475 
476 	if (flag_debug)
477 		printf("Message from kernel:\n%s\n",
478 		    prop_dictionary_externalize(*res));
479 
480 	return true;
481 }
482 
483 /* show collected statistics on mount point */
484 static int
485 show_mp(char *path)
486 {
487 	prop_dictionary_t args, res;
488 	prop_array_t reslist;
489 	bool rv;
490 	prop_object_iterator_t	iter;
491 	prop_dictionary_t item;
492 	uint32_t id;
493 	uint64_t space;
494 	char hbuf[5];
495 
496 	args = prop_dictionary_create();
497 	res  = prop_dictionary_create();
498 	if (args == NULL)
499 		printf("show_mp(): couldn't create args dictionary\n");
500 	res  = prop_dictionary_create();
501 	if (res == NULL)
502 		printf("show_mp(): couldn't create res dictionary\n");
503 
504 	rv = send_command(path, "get usage all", args, &res);
505 	if (rv == false) {
506 		printf("show-mp(): failed to send message to kernel\n");
507 		goto end;
508 	}
509 
510 	reslist = prop_dictionary_get(res, "returned data");
511 	if (reslist == NULL) {
512 		printf("show_mp(): failed to get array of results");
513 		rv = false;
514 		goto end;
515 	}
516 
517 	iter = prop_array_iterator(reslist);
518 	if (iter == NULL) {
519 		printf("show_mp(): failed to create iterator\n");
520 		rv = false;
521 		goto end;
522 	}
523 
524 	while ((item = prop_object_iterator_next(iter)) != NULL) {
525 		rv = prop_dictionary_get_uint64(item, "space used", &space);
526 		if (prop_dictionary_get_uint32(item, "uid", &id))
527 			print_user(id);
528 		else if (prop_dictionary_get_uint32(item, "gid", &id))
529 			print_group(id);
530 		else
531 			printf("total:");
532 		if (flag_humanize) {
533 			humanize_number(hbuf, sizeof(hbuf), space, "", HN_AUTOSCALE, HN_NOSPACE);
534 			printf(" %s\n", hbuf);
535 		} else {
536 			printf(" %" PRIu64 "\n", space);
537 		}
538 	}
539 	prop_object_iterator_release(iter);
540 
541 end:
542 	prop_object_release(args);
543 	prop_object_release(res);
544 	return (rv == true);
545 }
546 
547 /* sync the in-kernel counters to the actual file system usage */
548 static int cmd_sync(char *dirname)
549 {
550 	prop_dictionary_t res, item;
551 	prop_array_t args;
552 	struct ac_unode *unp;
553 	struct ac_gnode *gnp;
554 	int rv = 0, i;
555 
556 	args = prop_array_create();
557 	if (args == NULL)
558 		printf("cmd_sync(): couldn't create args dictionary\n");
559 	res  = prop_dictionary_create();
560 	if (res == NULL)
561 		printf("cmd_sync(): couldn't create res dictionary\n");
562 
563 	rv = get_dirsize(dirname);
564 
565 	item = prop_dictionary_create();
566 	if (item == NULL)
567 		printf("cmd_sync(): couldn't create item dictionary\n");
568 	(void) prop_dictionary_set_uint64(item, "space used", global_size);
569 	prop_array_add_and_rel(args, item);
570 
571 	RB_FOREACH(unp, ac_utree, &ac_uroot) {
572 	    for (i=0; i<ACCT_CHUNK_NIDS; i++) {
573 		if (unp->uid_chunk[i] != 0) {
574 		    item = prop_dictionary_create();
575 		    (void) prop_dictionary_set_uint32(item, "uid",
576 				(unp->left_bits << ACCT_CHUNK_BITS) + i);
577 		    (void) prop_dictionary_set_uint64(item, "space used",
578 				unp->uid_chunk[i]);
579 		    prop_array_add_and_rel(args, item);
580 		}
581 	    }
582 	}
583 	RB_FOREACH(gnp, ac_gtree, &ac_groot) {
584 	    for (i=0; i<ACCT_CHUNK_NIDS; i++) {
585 		if (gnp->gid_chunk[i] != 0) {
586 		    item = prop_dictionary_create();
587 		    (void) prop_dictionary_set_uint32(item, "gid",
588 				(gnp->left_bits << ACCT_CHUNK_BITS) + i);
589 		    (void) prop_dictionary_set_uint64(item, "space used",
590 				gnp->gid_chunk[i]);
591 		    prop_array_add_and_rel(args, item);
592 		}
593 	    }
594 	}
595 
596 	if (send_command(dirname, "set usage all", args, &res) == false) {
597 		printf("Failed to send message to kernel\n");
598 		rv = 1;
599 	}
600 
601 	prop_object_release(args);
602 	prop_object_release(res);
603 
604 	return rv;
605 }
606 
607 int
608 main(int argc, char **argv)
609 {
610 	int ch;
611 
612 	while ((ch = getopt(argc, argv, "Dhn")) != -1) {
613 		switch(ch) {
614 		case 'D':
615 			flag_debug = 1;
616 			break;
617 		case 'h':
618 			flag_humanize = 1;
619 			break;
620 		case 'n':
621 			flag_resolve_ids = 0;
622 			break;
623 		}
624 	}
625 	argc -= optind;
626 	argv += optind;
627 	if (argc < 1)
628 		usage(1);
629 
630 	if (strcmp(argv[0], "check") == 0) {
631 		if (argc != 2)
632 			usage(1);
633 		return cmd_check(argv[1]);
634 	}
635 	if (strcmp(argv[0], "lsfs") == 0) {
636 		return get_fslist();
637 	}
638 	if (strcmp(argv[0], "show") == 0) {
639 		if (argc != 2)
640 			usage(1);
641 		return show_mp(argv[1]);
642 	}
643 	if (strcmp(argv[0], "sync") == 0) {
644 		if (argc != 2)
645 			usage(1);
646 		return cmd_sync(argv[1]);
647 	}
648 
649 	usage(0);
650 }
651