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