vquota(8): Remove unneeded braces
[dragonfly.git] / 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                 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 }