sbin/vquota/vquota.c

   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 }