sbin/vquota/vquota.c

   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                 }
 529                 else if (prop_dictionary_get_uint32(item, "gid", &id)) {
 530                         print_group(id);
 531                 }
 532                 else
 533                         printf("total:");
 534                 if (flag_humanize) {
 535                         humanize_number(hbuf, sizeof(hbuf), space, "", HN_AUTOSCALE, HN_NOSPACE);
 536                         printf(" %s\n", hbuf);
 537                 } else {
 538                         printf(" %" PRIu64 "\n", space);
 539                 }
 540         }
 541         prop_object_iterator_release(iter);
 542
 543 end:
 544         prop_object_release(args);
 545         prop_object_release(res);
 546         return (rv == true);
 547 }
 548
 549 /* sync the in-kernel counters to the actual file system usage */
 550 static int cmd_sync(char *dirname)
 551 {
 552         prop_dictionary_t res, item;
 553         prop_array_t args;
 554         struct ac_unode *unp;
 555         struct ac_gnode *gnp;
 556         int rv = 0, i;
 557
 558         args = prop_array_create();
 559         if (args == NULL)
 560                 printf("cmd_sync(): couldn't create args dictionary\n");
 561         res  = prop_dictionary_create();
 562         if (res == NULL)
 563                 printf("cmd_sync(): couldn't create res dictionary\n");
 564
 565         rv = get_dirsize(dirname);
 566
 567         item = prop_dictionary_create();
 568         if (item == NULL)
 569                 printf("cmd_sync(): couldn't create item dictionary\n");
 570         (void) prop_dictionary_set_uint64(item, "space used", global_size);
 571         prop_array_add_and_rel(args, item);
 572
 573         RB_FOREACH(unp, ac_utree, &ac_uroot) {
 574             for (i=0; i<ACCT_CHUNK_NIDS; i++) {
 575                 if (unp->uid_chunk[i] != 0) {
 576                     item = prop_dictionary_create();
 577                     (void) prop_dictionary_set_uint32(item, "uid",
 578                                 (unp->left_bits << ACCT_CHUNK_BITS) + i);
 579                     (void) prop_dictionary_set_uint64(item, "space used",
 580                                 unp->uid_chunk[i]);
 581                     prop_array_add_and_rel(args, item);
 582                 }
 583             }
 584         }
 585         RB_FOREACH(gnp, ac_gtree, &ac_groot) {
 586             for (i=0; i<ACCT_CHUNK_NIDS; i++) {
 587                 if (gnp->gid_chunk[i] != 0) {
 588                     item = prop_dictionary_create();
 589                     (void) prop_dictionary_set_uint32(item, "gid",
 590                                 (gnp->left_bits << ACCT_CHUNK_BITS) + i);
 591                     (void) prop_dictionary_set_uint64(item, "space used",
 592                                 gnp->gid_chunk[i]);
 593                     prop_array_add_and_rel(args, item);
 594                 }
 595             }
 596         }
 597
 598         if (send_command(dirname, "set usage all", args, &res) == false) {
 599                 printf("Failed to send message to kernel\n");
 600                 rv = 1;
 601         }
 602
 603         prop_object_release(args);
 604         prop_object_release(res);
 605
 606         return rv;
 607 }
 608
 609 int
 610 main(int argc, char **argv)
 611 {
 612         int ch;
 613
 614         while ((ch = getopt(argc, argv, "Dhn")) != -1) {
 615                 switch(ch) {
 616                 case 'D':
 617                         flag_debug = 1;
 618                         break;
 619                 case 'h':
 620                         flag_humanize = 1;
 621                         break;
 622                 case 'n':
 623                         flag_resolve_ids = 0;
 624                         break;
 625                 }
 626         }
 627         argc -= optind;
 628         argv += optind;
 629         if (argc < 1)
 630                 usage(1);
 631
 632         if (strcmp(argv[0], "check") == 0) {
 633                 if (argc != 2)
 634                         usage(1);
 635                 return cmd_check(argv[1]);
 636         }
 637         if (strcmp(argv[0], "lsfs") == 0) {
 638                 return get_fslist();
 639         }
 640         if (strcmp(argv[0], "show") == 0) {
 641                 if (argc != 2)
 642                         usage(1);
 643                 return show_mp(argv[1]);
 644         }
 645         if (strcmp(argv[0], "sync") == 0) {
 646                 if (argc != 2)
 647                         usage(1);
 648                 return cmd_sync(argv[1]);
 649         }
 650
 651         usage(0);
 652 }