1 /* 2 * Copyright (c) 2011 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 50 static bool flag_debug = 0; 51 static bool flag_humanize = 0; 52 53 static void usage(int); 54 static int get_dirsize(char *); 55 static int get_fslist(void); 56 57 static void 58 usage(int retcode) 59 { 60 fprintf(stderr, "usage: vquota [-Dh] check directory\n"); 61 fprintf(stderr, " vquota [-Dh] lsfs\n"); 62 fprintf(stderr, " vquota [-Dh] show mount_point\n"); 63 fprintf(stderr, " vquota [-Dh] sync mount_point\n"); 64 exit(retcode); 65 } 66 67 /* 68 * Inode numbers with more than one hard link often come in groups; 69 * use linear arrays of 1024 ones as the basic unit of allocation. 70 * We only need to check if the inodes have been previously processed, 71 * bit arrays are perfect for that purpose. 72 */ 73 #define HL_CHUNK_BITS 10 74 #define HL_CHUNK_ENTRIES (1<<HL_CHUNK_BITS) 75 #define HL_CHUNK_MASK (HL_CHUNK_ENTRIES - 1) 76 #define BA_UINT64_BITS 6 77 #define BA_UINT64_ENTRIES (1<<BA_UINT64_BITS) 78 #define BA_UINT64_MASK (BA_UINT64_ENTRIES - 1) 79 80 struct hl_node { 81 RB_ENTRY(hl_node) rb_entry; 82 ino_t ino_left_bits; 83 uint64_t hl_chunk[HL_CHUNK_ENTRIES/64]; 84 }; 85 86 RB_HEAD(hl_tree,hl_node) hl_root; 87 88 RB_PROTOTYPE(hl_tree, hl_node, rb_entry, rb_hl_node_cmp); 89 90 static int 91 rb_hl_node_cmp(struct hl_node *a, struct hl_node *b); 92 93 RB_GENERATE(hl_tree, hl_node, rb_entry, rb_hl_node_cmp); 94 95 struct hl_node* hl_node_insert(ino_t); 96 97 98 static int 99 rb_hl_node_cmp(struct hl_node *a, struct hl_node *b) 100 { 101 if (a->ino_left_bits < b->ino_left_bits) 102 return(-1); 103 else if (a->ino_left_bits > b->ino_left_bits) 104 return(1); 105 return(0); 106 } 107 108 struct hl_node* hl_node_insert(ino_t inode) 109 { 110 struct hl_node *hlp, *res; 111 112 hlp = malloc(sizeof(struct hl_node)); 113 if (hlp == NULL) { 114 /* shouldn't happen */ 115 printf("hl_node_insert(): malloc failed\n"); 116 exit(ENOMEM); 117 } 118 bzero(hlp, sizeof(struct hl_node)); 119 120 hlp->ino_left_bits = (inode >> HL_CHUNK_BITS); 121 res = RB_INSERT(hl_tree, &hl_root, hlp); 122 123 if (res != NULL) /* shouldn't happen */ 124 printf("hl_node_insert(): RB_INSERT didn't return NULL\n"); 125 126 return hlp; 127 } 128 129 /* 130 * hl_register: register an inode number in a rb-tree of bit arrays 131 * returns: 132 * - true if the inode was already processed 133 * - false otherwise 134 */ 135 static bool 136 hl_register(ino_t inode) 137 { 138 struct hl_node hl_find, *hlp; 139 uint64_t ino_right_bits, ba_index, ba_offset; 140 uint64_t bitmask, bitval; 141 bool retval = false; 142 143 /* calculate the different addresses of the wanted bit */ 144 hl_find.ino_left_bits = (inode >> HL_CHUNK_BITS); 145 146 ino_right_bits = inode & HL_CHUNK_MASK; 147 ba_index = ino_right_bits >> BA_UINT64_BITS; 148 ba_offset = ino_right_bits & BA_UINT64_MASK; 149 150 /* no existing node? create and initialize it */ 151 if ((hlp = RB_FIND(hl_tree, &hl_root, &hl_find)) == NULL) { 152 hlp = hl_node_insert(inode); 153 } 154 155 /* node was found, check the bit value */ 156 bitmask = 1 << ba_offset; 157 bitval = hlp->hl_chunk[ba_index] & bitmask; 158 if (bitval != 0) { 159 retval = true; 160 } 161 162 /* set the bit */ 163 hlp->hl_chunk[ba_index] |= bitmask; 164 165 return retval; 166 } 167 168 /* global variable used by get_dir_size() */ 169 uint64_t global_size; 170 171 /* storage for collected id numbers */ 172 /* FIXME: same data structures used in kernel, should find a way to 173 * deduplicate this code */ 174 175 static int 176 rb_ac_unode_cmp(struct ac_unode*, struct ac_unode*); 177 static int 178 rb_ac_gnode_cmp(struct ac_gnode*, struct ac_gnode*); 179 180 RB_HEAD(ac_utree,ac_unode) ac_uroot; 181 RB_HEAD(ac_gtree,ac_gnode) ac_groot; 182 RB_PROTOTYPE(ac_utree, ac_unode, rb_entry, rb_ac_unode_cmp); 183 RB_PROTOTYPE(ac_gtree, ac_gnode, rb_entry, rb_ac_gnode_cmp); 184 RB_GENERATE(ac_utree, ac_unode, rb_entry, rb_ac_unode_cmp); 185 RB_GENERATE(ac_gtree, ac_gnode, rb_entry, rb_ac_gnode_cmp); 186 187 static int 188 rb_ac_unode_cmp(struct ac_unode *a, struct ac_unode *b) 189 { 190 if (a->left_bits < b->left_bits) 191 return(-1); 192 else if (a->left_bits > b->left_bits) 193 return(1); 194 return(0); 195 } 196 197 static int 198 rb_ac_gnode_cmp(struct ac_gnode *a, struct ac_gnode *b) 199 { 200 if (a->left_bits < b->left_bits) 201 return(-1); 202 else if (a->left_bits > b->left_bits) 203 return(1); 204 return(0); 205 } 206 207 static struct ac_unode* 208 unode_insert(uid_t uid) 209 { 210 struct ac_unode *unp, *res; 211 212 unp = malloc(sizeof(struct ac_unode)); 213 if (unp == NULL) { 214 printf("unode_insert(): malloc failed\n"); 215 exit(ENOMEM); 216 } 217 bzero(unp, sizeof(struct ac_unode)); 218 219 unp->left_bits = (uid >> ACCT_CHUNK_BITS); 220 res = RB_INSERT(ac_utree, &ac_uroot, unp); 221 222 if (res != NULL) /* shouldn't happen */ 223 printf("unode_insert(): RB_INSERT didn't return NULL\n"); 224 225 return unp; 226 } 227 228 static struct ac_gnode* 229 gnode_insert(gid_t gid) 230 { 231 struct ac_gnode *gnp, *res; 232 233 gnp = malloc(sizeof(struct ac_gnode)); 234 if (gnp == NULL) { 235 printf("gnode_insert(): malloc failed\n"); 236 exit(ENOMEM); 237 } 238 bzero(gnp, sizeof(struct ac_gnode)); 239 240 gnp->left_bits = (gid >> ACCT_CHUNK_BITS); 241 res = RB_INSERT(ac_gtree, &ac_groot, gnp); 242 243 if (res != NULL) /* shouldn't happen */ 244 printf("gnode_insert(): RB_INSERT didn't return NULL\n"); 245 246 return gnp; 247 } 248 249 /* 250 * get_dirsize(): walks a directory tree in the same filesystem 251 * output: 252 * - global rb-trees ac_uroot and ac_groot 253 * - global variable global_size 254 */ 255 static int 256 get_dirsize(char* dirname) 257 { 258 FTS *fts; 259 FTSENT *p; 260 char* fts_args[2]; 261 int retval = 0; 262 263 /* what we need */ 264 ino_t file_inode; 265 off_t file_size; 266 uid_t file_uid; 267 gid_t file_gid; 268 269 struct ac_unode *unp, ufind; 270 struct ac_gnode *gnp, gfind; 271 272 /* TODO: check directory name sanity */ 273 fts_args[0] = dirname; 274 fts_args[1] = NULL; 275 276 if ((fts = fts_open(fts_args, FTS_PHYSICAL|FTS_XDEV, NULL)) == NULL) 277 err(1, "fts_open() failed"); 278 279 while ((p = fts_read(fts)) != NULL) { 280 switch (p->fts_info) { 281 /* directories, ignore them */ 282 case FTS_D: 283 case FTS_DC: 284 case FTS_DP: 285 break; 286 /* read errors, warn, continue and flag */ 287 case FTS_DNR: 288 case FTS_ERR: 289 case FTS_NS: 290 warnx("%s: %s", p->fts_path, strerror(p->fts_errno)); 291 retval = 1; 292 break; 293 default: 294 file_inode = p->fts_statp->st_ino; 295 file_size = p->fts_statp->st_size; 296 file_uid = p->fts_statp->st_uid; 297 file_gid = p->fts_statp->st_gid; 298 299 /* files with more than one hard link: */ 300 /* process them only once */ 301 if (p->fts_statp->st_nlink > 1) 302 if (hl_register(file_inode) == false) 303 break; 304 305 global_size += file_size; 306 ufind.left_bits = (file_uid >> ACCT_CHUNK_BITS); 307 gfind.left_bits = (file_gid >> ACCT_CHUNK_BITS); 308 if ((unp = RB_FIND(ac_utree, &ac_uroot, &ufind)) == NULL) 309 unp = unode_insert(file_uid); 310 if ((gnp = RB_FIND(ac_gtree, &ac_groot, &gfind)) == NULL) 311 gnp = gnode_insert(file_gid); 312 unp->uid_chunk[(file_uid & ACCT_CHUNK_MASK)] += file_size; 313 gnp->gid_chunk[(file_gid & ACCT_CHUNK_MASK)] += file_size; 314 } 315 } 316 fts_close(fts); 317 318 return retval; 319 } 320 321 static int 322 cmd_check(char* dirname) 323 { 324 int32_t uid, gid; 325 char hbuf[5]; 326 struct ac_unode *unp; 327 struct ac_gnode *gnp; 328 int rv, i; 329 330 rv = get_dirsize(dirname); 331 332 if (flag_humanize) { 333 humanize_number(hbuf, sizeof(hbuf), global_size, "", 334 HN_AUTOSCALE, HN_NOSPACE); 335 printf("total: %s\n", hbuf); 336 } else { 337 printf("total: %"PRIu64"\n", global_size); 338 } 339 RB_FOREACH(unp, ac_utree, &ac_uroot) { 340 for (i=0; i<ACCT_CHUNK_NIDS; i++) { 341 if (unp->uid_chunk[i] != 0) { 342 uid = (unp->left_bits << ACCT_CHUNK_BITS) + i; 343 if (flag_humanize) { 344 humanize_number(hbuf, sizeof(hbuf), 345 unp->uid_chunk[i], "", HN_AUTOSCALE, HN_NOSPACE); 346 printf("uid %"PRIu32": %s\n", uid, hbuf); 347 } else { 348 printf("uid %"PRIu32": %"PRIu64"\n", uid, unp->uid_chunk[i]); 349 } 350 } 351 } 352 } 353 RB_FOREACH(gnp, ac_gtree, &ac_groot) { 354 for (i=0; i<ACCT_CHUNK_NIDS; i++) { 355 if (gnp->gid_chunk[i] != 0) { 356 gid = (gnp->left_bits << ACCT_CHUNK_BITS) + i; 357 if (flag_humanize) { 358 humanize_number(hbuf, sizeof(hbuf), 359 gnp->gid_chunk[i], "", HN_AUTOSCALE, HN_NOSPACE); 360 printf("gid %"PRIu32": %s\n", gid, hbuf); 361 } else { 362 printf("gid %"PRIu32": %"PRIu64"\n", gid, gnp->gid_chunk[i]); 363 } 364 } 365 } 366 } 367 368 return rv; 369 } 370 371 /* print a list of filesystems with accounting enabled */ 372 static int 373 get_fslist(void) 374 { 375 struct statfs *mntbufp; 376 int nloc, i; 377 378 /* read mount table from kernel */ 379 nloc = getmntinfo(&mntbufp, MNT_NOWAIT|MNT_LOCAL); 380 if (nloc <= 0) { 381 perror("getmntinfo"); 382 exit(1); 383 } 384 385 /* iterate mounted filesystems */ 386 for (i=0; i<nloc; i++) { 387 /* vfs accounting enabled on this one ? */ 388 if (mntbufp[i].f_flags & MNT_ACCOUNTING) 389 printf("%s on %s\n", mntbufp[i].f_mntfromname, 390 mntbufp[i].f_mntonname); 391 } 392 393 return 0; 394 } 395 396 static bool 397 send_command(const char *path, const char *cmd, 398 prop_object_t args, prop_dictionary_t *res) 399 { 400 prop_dictionary_t dict; 401 struct plistref pref; 402 403 bool rv; 404 int error; 405 406 dict = prop_dictionary_create(); 407 408 if (dict == NULL) { 409 printf("send_command(): couldn't create dictionary\n"); 410 return false; 411 } 412 413 rv = prop_dictionary_set_cstring(dict, "command", cmd); 414 if (rv== false) { 415 printf("send_command(): couldn't initialize dictionary\n"); 416 return false; 417 } 418 419 rv = prop_dictionary_set(dict, "arguments", args); 420 if (rv == false) { 421 printf("prop_dictionary_set() failed\n"); 422 return false; 423 } 424 425 error = prop_dictionary_send_syscall(dict, &pref); 426 if (error != 0) { 427 printf("prop_dictionary_send_syscall() failed\n"); 428 prop_object_release(dict); 429 return false; 430 } 431 432 if (flag_debug) 433 printf("Message to kernel:\n%s\n", 434 prop_dictionary_externalize(dict)); 435 436 error = vquotactl(path, &pref); 437 if (error != 0) { 438 printf("send_command: vquotactl = %d\n", error); 439 return false; 440 } 441 442 error = prop_dictionary_recv_syscall(&pref, res); 443 if (error != 0) { 444 printf("prop_dictionary_recv_syscall() failed\n"); 445 } 446 447 if (flag_debug) 448 printf("Message from kernel:\n%s\n", 449 prop_dictionary_externalize(*res)); 450 451 return true; 452 } 453 454 /* show collected statistics on mount point */ 455 static int 456 show_mp(char *path) 457 { 458 prop_dictionary_t args, res; 459 prop_array_t reslist; 460 bool rv; 461 prop_object_iterator_t iter; 462 prop_dictionary_t item; 463 uint32_t id; 464 uint64_t space; 465 char hbuf[5]; 466 467 args = prop_dictionary_create(); 468 res = prop_dictionary_create(); 469 if (args == NULL) 470 printf("show_mp(): couldn't create args dictionary\n"); 471 res = prop_dictionary_create(); 472 if (res == NULL) 473 printf("show_mp(): couldn't create res dictionary\n"); 474 475 rv = send_command(path, "get usage all", args, &res); 476 if (rv == false) { 477 printf("show-mp(): failed to send message to kernel\n"); 478 goto end; 479 } 480 481 reslist = prop_dictionary_get(res, "returned data"); 482 if (reslist == NULL) { 483 printf("show_mp(): failed to get array of results"); 484 rv = false; 485 goto end; 486 } 487 488 iter = prop_array_iterator(reslist); 489 if (iter == NULL) { 490 printf("show_mp(): failed to create iterator\n"); 491 rv = false; 492 goto end; 493 } 494 495 while ((item = prop_object_iterator_next(iter)) != NULL) { 496 rv = prop_dictionary_get_uint64(item, "space used", &space); 497 if (prop_dictionary_get_uint32(item, "uid", &id)) 498 printf("uid %u:", id); 499 else if (prop_dictionary_get_uint32(item, "gid", &id)) 500 printf("gid %u:", id); 501 else 502 printf("total:"); 503 if (flag_humanize) { 504 humanize_number(hbuf, sizeof(hbuf), space, "", HN_AUTOSCALE, HN_NOSPACE); 505 printf(" %s\n", hbuf); 506 } else { 507 printf(" %" PRIu64 "\n", space); 508 } 509 } 510 prop_object_iterator_release(iter); 511 512 end: 513 prop_object_release(args); 514 prop_object_release(res); 515 return (rv == true); 516 } 517 518 /* sync the in-kernel counters to the actual file system usage */ 519 static int cmd_sync(char *dirname) 520 { 521 prop_dictionary_t res, item; 522 prop_array_t args; 523 struct ac_unode *unp; 524 struct ac_gnode *gnp; 525 int rv = 0, i; 526 527 args = prop_array_create(); 528 if (args == NULL) 529 printf("cmd_sync(): couldn't create args dictionary\n"); 530 res = prop_dictionary_create(); 531 if (res == NULL) 532 printf("cmd_sync(): couldn't create res dictionary\n"); 533 534 rv = get_dirsize(dirname); 535 536 item = prop_dictionary_create(); 537 if (item == NULL) 538 printf("cmd_sync(): couldn't create item dictionary\n"); 539 (void) prop_dictionary_set_uint64(item, "space used", global_size); 540 prop_array_add_and_rel(args, item); 541 542 RB_FOREACH(unp, ac_utree, &ac_uroot) { 543 for (i=0; i<ACCT_CHUNK_NIDS; i++) { 544 if (unp->uid_chunk[i] != 0) { 545 item = prop_dictionary_create(); 546 (void) prop_dictionary_set_uint32(item, "uid", 547 (unp->left_bits << ACCT_CHUNK_BITS) + i); 548 (void) prop_dictionary_set_uint64(item, "space used", 549 unp->uid_chunk[i]); 550 prop_array_add_and_rel(args, item); 551 } 552 } 553 } 554 RB_FOREACH(gnp, ac_gtree, &ac_groot) { 555 for (i=0; i<ACCT_CHUNK_NIDS; i++) { 556 if (gnp->gid_chunk[i] != 0) { 557 item = prop_dictionary_create(); 558 (void) prop_dictionary_set_uint32(item, "gid", 559 (gnp->left_bits << ACCT_CHUNK_BITS) + i); 560 (void) prop_dictionary_set_uint64(item, "space used", 561 gnp->gid_chunk[i]); 562 prop_array_add_and_rel(args, item); 563 } 564 } 565 } 566 567 if (send_command(dirname, "set usage all", args, &res) == false) { 568 printf("Failed to send message to kernel\n"); 569 rv = 1; 570 } 571 572 prop_object_release(args); 573 prop_object_release(res); 574 575 return rv; 576 } 577 578 int 579 main(int argc, char **argv) 580 { 581 int ch; 582 583 while ((ch = getopt(argc, argv, "Dh")) != -1) { 584 switch(ch) { 585 case 'D': 586 flag_debug = 1; 587 break; 588 case 'h': 589 flag_humanize = 1; 590 break; 591 } 592 } 593 argc -= optind; 594 argv += optind; 595 if (argc < 1) 596 usage(1); 597 598 if (strcmp(argv[0], "check") == 0) { 599 if (argc != 2) 600 usage(1); 601 return cmd_check(argv[1]); 602 } 603 if (strcmp(argv[0], "lsfs") == 0) { 604 return get_fslist(); 605 } 606 if (strcmp(argv[0], "show") == 0) { 607 if (argc != 2) 608 usage(1); 609 return show_mp(argv[1]); 610 } 611 if (strcmp(argv[0], "sync") == 0) { 612 if (argc != 2) 613 usage(1); 614 return cmd_sync(argv[1]); 615 } 616 617 usage(0); 618 } 619