1 /*
2  * Copyright (c) 2023 Benjamin Stürz <benni@stuerz.xyz>
3  *
4  * Permission to use, copy, modify, and distribute this software for any
5  * purpose with or without fee is hereby granted, provided that the above
6  * copyright notice and this permission notice appear in all copies.
7  *
8  * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9  * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10  * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11  * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12  * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13  * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14  * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
15  */
16 #define _XOPEN_SOURCE 700
17 #define _BSD_SOURCE 1
18 #define DKTYPENAMES
19 #define WSDEBUG 0
20 #include <stddef.h>
21 #include <dev/biovar.h>
22 #include <sys/cdefs.h>
23 #include <sys/types.h>
24 #include <sys/disklabel.h>
25 #include <sys/sysctl.h>
26 #include <sys/ioctl.h>
27 #include <sys/mount.h>
28 #include <sys/dkio.h>
29 #include <inttypes.h>
30 #include <stdbool.h>
31 #include <string.h>
32 #include <stdarg.h>
33 #include <stdint.h>
34 #include <stdlib.h>
35 #include <unistd.h>
36 #include <libgen.h>
37 #include <stdio.h>
38 #include <fcntl.h>
39 #include <errno.h>
40 #include <util.h>
41 #include <err.h>
42 
43 static int diskcount (void)
44 {
45     const int mib[2] = { CTL_HW, HW_DISKCOUNT };
46     int diskcount;
47     size_t len = sizeof diskcount;
48 
49     if (sysctl (mib, 2, &diskcount, &len, NULL, 0) == -1)
50         err (1, "sysctl(hw.diskcount)");
51 
52     return diskcount;
53 }
54 
55 static char *disknames (void)
56 {
57     const int num = diskcount ();
58     const int mib[2] = { CTL_HW, HW_DISKNAMES };
59     size_t len = 32 * num;
60     char *buffer = malloc (len);
61 
62     if (sysctl (mib, 2, buffer, &len, NULL, 0) == -1)
63         err (1, "sysctl(hw.disknames)");
64 
65     return buffer;
66 }
67 
68 static char *stripdisk (char *n)
69 {
70     const char sufx[] = " disk";
71     const size_t ln = strnlen (n, 16);
72     const size_t ls = sizeof sufx - 1;
73 
74     if (!memcmp (n + ln - ls, sufx, ls)) {
75         n[ln - ls] = '\0';
76     }
77 
78     return n;
79 }
80 
81 static void print_size (uint64_t sz)
82 {
83     struct unit {
84         char sym;
85         uint64_t factor;
86     };
87 
88     const struct unit units[] = {
89         { 'P', 1ull << 50 },
90         { 'T', 1ull << 40 },
91         { 'G', 1ull << 30 },
92         { 'M', 1ull << 20 },
93         { 'K', 1ull << 10 },
94         { '0', 0 },
95     };
96 
97     char sym = 'B';
98     uint64_t factor = 1;
99 
100     for (const struct unit *u = &units[0]; u->factor; ++u) {
101         if (sz >= (u->factor * 9 / 10)) {
102             sym = u->sym;
103             factor = u->factor;
104             break;
105         }
106     }
107 
108     const unsigned scaled10 = sz * 10 / factor;
109     const unsigned scaled = sz / factor;
110     if (scaled10 >= 1000) {
111         printf ("%u", scaled);
112     } else if (scaled10 >= 100) {
113         printf (" %u", scaled);
114     } else {
115         printf ("%u.%u", scaled, scaled10 % 10);
116     }
117 
118     putchar (sym);
119     putchar (' ');
120 }
121 
122 enum {
123     FIELD_NAME      = 0x01,
124     FIELD_DUID      = 0x02,
125     FIELD_SIZE      = 0x04,
126     FIELD_USED      = 0x08,
127     FIELD_FREE      = 0x10,
128     FIELD_TYPE      = 0x20,
129     FIELD_COMMENT   = 0x40,
130 
131     FIELD_DEFAULT   = FIELD_NAME | FIELD_SIZE | FIELD_TYPE | FIELD_COMMENT,
132 };
133 
134 enum {
135     OPT_NOHEADER    = 0x01,
136     OPT_NOUNICODE   = 0x02,
137     OPT_NOBIO       = 0x04,
138 };
139 
140 struct my_diskinfo;
141 
142 struct my_partinfo {
143     char letter;
144     uint64_t size;
145     uint64_t fssize;
146     uint64_t free;
147     const char *fstype;
148     const char *mount;
149     const struct my_diskinfo *sub; // If this is part of a RAID
150     const char *raidstatus; // Only available if (sub != NULL)
151 };
152 
153 struct my_diskinfo {
154     char type[16];
155     char label[16];
156     char name[4];
157     uint64_t size;
158     uint64_t used;
159     u_char duid[8];
160     uint8_t num_parts;
161     struct my_partinfo parts[MAXPARTITIONS];
162     const char *raidstatus; // If this is a RAID device
163 };
164 
165 struct padding {
166     int name;
167     /* duid = 8 */
168     /* size = 4 */
169     /* used = 4 */
170     /* free = 4 */
171     int type;
172     int comment;
173 };
174 
175 static void print_header (int fields, const struct padding *p)
176 {
177     if (fields & FIELD_NAME)
178         printf ("%-*s ", p->name, "NAME");
179 
180     if (fields & FIELD_DUID)
181         printf ("%-18s ", "DUID");
182 
183     if (fields & FIELD_SIZE)
184         printf ("%-4s ", "SIZE");
185 
186     if (fields & FIELD_USED)
187         printf ("%-4s ", "USED");
188 
189     if (fields & FIELD_FREE)
190         printf ("%-4s ", "FREE");
191 
192     if (fields & FIELD_TYPE)
193         printf ("%-*s ", p->type, "TYPE");
194 
195     if (fields & FIELD_COMMENT)
196         printf ("%-*s ", p->comment, "COMMENT");
197 
198 #if WSDEBUG
199     putchar ('X');
200 #endif
201 
202     putchar ('\n');
203 }
204 
205 static void print_duid (const u_char *duid)
206 {
207     for (size_t i = 0; i < 8; ++i) {
208         printf ("%02x", duid[i]);
209     }
210 }
211 
212 static void print_disk (const struct my_diskinfo *, int, int, const char *, const struct padding *);
213 static void print_part (
214     const struct my_diskinfo *disk,
215     const struct my_partinfo *part,
216     int fields,
217     int options,
218     bool last,
219     const struct padding *p
220 ) {
221     if (fields & FIELD_NAME) {
222         const char *prefix = (options & OPT_NOUNICODE) ? "  " : (last ? "└─" : "├─");
223         printf (
224             "%s%s%c%-*s ",
225             prefix,
226             disk->name,
227             part->letter,
228             p->name - 2 - (int)strlen (disk->name) - 1,
229             ""
230         );
231     }
232 
233     if (fields & FIELD_DUID) {
234         print_duid (disk->duid);
235         printf (".%c ", part->letter);
236     }
237 
238     if (fields & FIELD_SIZE)
239         print_size (part->size);
240 
241     if (fields & FIELD_USED) {
242         if (part->fssize) {
243             print_size (part->fssize - part->free);
244         } else {
245             printf (" N/A ");
246         }
247     }
248 
249     if (fields & FIELD_FREE) {
250         if (part->fssize) {
251             print_size (part->free);
252         } else {
253             printf (" N/A ");
254         }
255     }
256 
257     if (fields & FIELD_TYPE)
258         printf ("%-*s ", p->type, part->fstype);
259 
260     if (fields & FIELD_COMMENT)
261         printf ("%-*s ", p->comment, part->mount ? part->mount : "");
262 
263 #if WSDEBUG
264     putchar ('X');
265 #endif
266 
267     putchar ('\n');
268 
269     if (part->sub) {
270         print_disk (part->sub, fields, options, part->raidstatus, p);
271     }
272 }
273 
274 static void print_disk (
275     const struct my_diskinfo *disk,
276     int fields,
277     int options,
278     const char *raidstatus,
279     const struct padding *p
280 ) {
281     if (fields & FIELD_NAME) {
282         const char *prefix = raidstatus ? "│ └─" : "";
283 
284         printf (
285             "%s%-*s ",
286             prefix,
287             p->name - (int)strlen (prefix),
288             disk->name
289         );
290     }
291 
292     if (fields & FIELD_DUID) {
293         print_duid (disk->duid);
294         printf ("   ");
295     }
296 
297     if (fields & FIELD_SIZE)
298         print_size (disk->size);
299 
300     if (fields & FIELD_USED)
301         print_size (disk->used);
302 
303     if (fields & FIELD_FREE)
304         print_size (disk->size - disk->used);
305 
306     if (fields & FIELD_TYPE)
307         printf ("%-*.16s ", p->type, disk->type);
308 
309     if (fields & FIELD_COMMENT) {
310         if (raidstatus) {
311             printf ("%-*s ", p->comment, raidstatus);
312         } else if (disk->raidstatus) {
313             printf (
314                 "%.16s (%s)%*s ",
315                 disk->label,
316                 disk->raidstatus,
317                 p->comment - (int)strnlen (disk->label, sizeof disk->label) - (int)strlen (disk->raidstatus) - 3,
318                 ""
319             );
320         } else {
321             printf ("%-*.16s ", p->comment, disk->label);
322         }
323     }
324 
325 #if WSDEBUG
326     putchar ('X');
327 #endif
328 
329     putchar ('\n');
330 
331     if (!raidstatus) {
332         for (uint8_t i = 0; i < disk->num_parts; ++i)
333             print_part (disk, &disk->parts[i], fields, options, i == (disk->num_parts - 1), p);
334     }
335 }
336 
337 static const struct statfs *find_mount (const char *dev)
338 {
339     static struct statfs *mounts = NULL;
340     static int n_mounts;
341 
342     if (!mounts) {
343         n_mounts = getmntinfo (&mounts, MNT_NOWAIT);
344         if (n_mounts == 0)
345             err (1, "getmntinfo()");
346     }
347 
348     for (int i = 0; i < n_mounts; ++i) {
349         if (!strcmp (dev, mounts[i].f_mntfromname))
350             return &mounts[i];
351     }
352 
353     return NULL;
354 }
355 
356 static struct my_diskinfo read_disk (const char *name)
357 {
358     struct my_diskinfo disk;
359     struct disklabel label;
360     char *ppath, *letter;
361 
362     bzero (&disk, sizeof disk);
363 
364     {   // Read disklabel.
365         size_t len;
366         int fd;
367 
368         fd = opendev (name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, &ppath);
369         if (fd < 0)
370             err (1, "opendev(%s)", name);
371 
372         if (ioctl (fd, DIOCGDINFO, &label) < 0)
373             err (1, "ioctl(%s, DIOCGDINFO)", name);
374         close (fd);
375 
376         len = strlen (ppath);
377         letter = ppath + len - 1;
378     }
379 
380     memcpy (disk.name, name, 3);
381     disk.name[3] = '\0';
382     disk.size = DL_GETDSIZE (&label) * label.d_secsize;
383     memcpy (disk.type, label.d_typename, 16);
384     stripdisk (disk.type);
385     memcpy (disk.label, label.d_packname, 16);
386     memcpy (disk.duid, label.d_uid, sizeof disk.duid);
387     disk.num_parts = 0;
388 
389     for (uint16_t i = 0; i < label.d_npartitions; ++i) {
390         const struct partition *p = &label.d_partitions[i];
391         const struct statfs *mnt;
392         struct my_partinfo part;
393 
394         bzero (&part, sizeof part);
395 
396         part.size = DL_GETPSIZE (p) * label.d_secsize;
397 
398         if (!part.size)
399             continue;
400 
401         part.letter = 'a' + i;
402         *letter = part.letter;
403         mnt = find_mount (ppath);
404 
405         if (mnt) {
406             const uint64_t bs = mnt->f_bsize;
407             part.mount = mnt->f_mntonname;
408             part.fssize = mnt->f_blocks * bs;
409             part.free = mnt->f_bfree * bs;
410         }
411 
412         part.fstype = fstypenames[p->p_fstype];
413         if (i != 2)
414             disk.used += part.size;
415         disk.parts[disk.num_parts++] = part;
416     }
417 
418     return disk;
419 }
420 
421 static const char *bd_statusstr (int status) {
422     switch (status) {
423     case BIOC_SDONLINE:     return BIOC_SDONLINE_S;
424     case BIOC_SDOFFLINE:    return BIOC_SDOFFLINE_S;
425     case BIOC_SDFAILED:     return BIOC_SDFAILED_S;
426     case BIOC_SDREBUILD:    return BIOC_SDREBUILD_S;
427     case BIOC_SDHOTSPARE:   return BIOC_SDHOTSPARE_S;
428     case BIOC_SDUNUSED:     return BIOC_SDUNUSED_S;
429     case BIOC_SDSCRUB:      return BIOC_SDSCRUB_S;
430     case BIOC_SDINVALID:    return BIOC_SDINVALID_S;
431     default:                return "Unknown";
432     }
433 }
434 
435 static const char *bv_statusstr (int status) {
436     switch (status) {
437     case BIOC_SVONLINE:     return BIOC_SVONLINE_S;
438     case BIOC_SVOFFLINE:    return BIOC_SVOFFLINE_S;
439     case BIOC_SVDEGRADED:   return BIOC_SVDEGRADED_S;
440     case BIOC_SVBUILDING:   return BIOC_SVBUILDING_S;
441     case BIOC_SVSCRUB:      return BIOC_SVSCRUB_S;
442     case BIOC_SVREBUILD:    return BIOC_SVREBUILD_S;
443     case BIOC_SVINVALID:    return BIOC_SVINVALID_S;
444     default:                return "Unknown";
445     }
446 }
447 
448 static void read_raid (
449     struct my_diskinfo *disk,
450     struct my_diskinfo *disks,
451     size_t num_disks
452 ) {
453     struct bioc_inq bi;
454     int fd;
455 
456     fd = opendev (disk->name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, NULL);
457     if (fd < 0) {
458         warn ("read_raid(): opendev(%s)", disk->name);
459         return;
460     }
461 
462     bzero (&bi, sizeof bi);
463 
464     if (ioctl (fd, BIOCINQ, &bi) == -1)
465         goto ret;
466 
467     for (int i = 0; i < bi.bi_novol; ++i) {
468         struct bioc_vol bv;
469 
470         bzero (&bv, sizeof bv);
471         memcpy (&bv.bv_bio, &bi.bi_bio, sizeof bi.bi_bio);
472         bv.bv_volid = i;
473 
474         if (ioctl (fd, BIOCVOL, &bv) == -1) {
475             warn ("read_raid(%s): BIOCVOL(%d)", disk->name, i);
476             continue;
477         }
478 
479         if (strcmp (disk->name, bv.bv_dev) != 0)
480             continue;
481 
482         disk->raidstatus = bv_statusstr (bv.bv_status);
483 
484         for (int j = 0; j < bv.bv_nodisk; ++j) {
485             struct bioc_disk bd;
486             size_t len_vendor;
487             char letter;
488 
489             bzero (&bd, sizeof bd);
490             memcpy (&bd.bd_bio, &bi.bi_bio, sizeof bi.bi_bio);
491             bd.bd_volid = i;
492             bd.bd_diskid = j;
493 
494             if (ioctl (fd, BIOCDISK, &bd) == -1) {
495                 warn ("read_raid(%s): BIOCDISK(%d, %d)", disk->name, i, j);
496                 continue;
497             }
498 
499             len_vendor = strlen (bd.bd_vendor);
500             if (len_vendor != 4) {
501                 warnx ("read_raid(%s): unexpected vendor string: %.32s", disk->name, bd.bd_vendor);
502                 continue;
503             }
504             letter = bd.bd_vendor[len_vendor - 1];
505 
506             for (size_t k = 0; k < num_disks; ++k) {
507                 if (!memcmp (bd.bd_vendor, disks[k].name, 3)) {
508                     for (size_t l = 0; l < disks[k].num_parts; ++l) {
509                         if (letter == disks[k].parts[l].letter) {
510                             disks[k].parts[l].sub = disk;
511                             disks[k].parts[l].raidstatus = bd_statusstr(bd.bd_status);
512                             goto found;
513                         }
514                     }
515                 }
516             }
517         found:;
518         }
519     }
520 ret:
521     close (fd);
522 }
523 
524 static int usage (void)
525 {
526     fputs ("Usage: lsblk [-abinUuV] [disk...]\n", stderr);
527     return 1;
528 }
529 
530 static void pad_update (int *pad, int newval)
531 {
532     if (newval > *pad)
533         *pad = newval;
534 }
535 
536 static void pad_disk (struct padding *p, const struct my_diskinfo *disk)
537 {
538     size_t len_disk;
539     int comment;
540 
541     len_disk = strnlen (disk->name, sizeof disk->name);
542     comment = strnlen (disk->label, sizeof disk->label);
543     if (disk->raidstatus)
544         comment += strlen (disk->raidstatus) + 3;
545 
546     pad_update (&p->name, len_disk);
547     pad_update (&p->type, strnlen (disk->type, sizeof disk->type));
548     pad_update (&p->comment, comment);
549 
550     for (int i = 0; i < disk->num_parts; ++i) {
551         const struct my_partinfo *part = &disk->parts[i];
552 
553         pad_update (&p->name, len_disk + 3);
554         pad_update (&p->type, strlen (part->fstype));
555 
556         if (part->sub) {
557             pad_update (&p->name, strnlen (part->sub->name, sizeof part->sub->name) + 4);
558             pad_update (&p->comment, strlen (part->raidstatus));
559         } else if (part->mount) {
560             pad_update (&p->comment, strlen (part->mount));
561         }
562     }
563 }
564 
565 
566 static int compare_disk (const void *p1, const void *p2)
567 {
568     const struct my_diskinfo *d1 = p1;
569     const struct my_diskinfo *d2 = p2;
570 
571     return strcmp (d1->name, d2->name);
572 }
573 
574 int main (int argc, char *argv[])
575 {
576     int option;
577     int fields = FIELD_DEFAULT;
578     int options = 0;
579 
580     if (unveil ("/dev", "r") == -1)
581         err (1, "unveil(/dev)");
582 
583     if (unveil (NULL, NULL) == -1)
584         err (1, "unveil()");
585 
586     while ((option = getopt (argc, argv, ":abinUuV")) != -1) {
587         switch (option) {
588         case 'a':
589             fields = -1;
590             break;
591         case 'b':
592             options |= OPT_NOBIO;
593             break;
594         case 'i':
595             options |= OPT_NOUNICODE;
596             break;
597         case 'n':
598             options |= OPT_NOHEADER;
599             break;
600         case 'U':
601             fields |= FIELD_DUID;
602             break;
603         case 'u':
604             fields |= FIELD_USED | FIELD_FREE;
605             break;
606         case 'V':
607             puts ("lsblk-" VERSION);
608             return 0;
609         default:
610             return usage ();
611         }
612     }
613 
614     argv += optind;
615     argc -= optind;
616 
617     char *names = argc == 0 ? disknames () : NULL;
618 
619     if (pledge ("stdio rpath disklabel", NULL) == -1)
620         err (1, "pledge()");
621 
622     size_t cap_disks;
623     if (argc == 0) {
624         const char *s = names;
625         cap_disks = 0;
626         while ((s = strchr (s, ',')) != NULL)
627             ++cap_disks, ++s;
628     } else {
629         cap_disks = argc;
630     }
631 
632     struct my_diskinfo *disks = calloc (cap_disks, sizeof (struct my_diskinfo));
633     size_t num_disks = 0;
634 
635     if (argc == 0) {
636         for (char *disk; (disk = strsep (&names, ",")) != NULL; ) {
637             char *colon = strchr (disk, ':');
638             if (colon)
639                 *colon = '\0';
640             disks[num_disks++] = read_disk (disk);
641             if (colon)
642                 *colon = ':';
643         }
644     } else {
645         for (int i = 0; i < argc; ++i) {
646             char *disk = basename (argv[i]);
647             disks[num_disks++] = read_disk (disk);
648         }
649     }
650 
651     free (names);
652 
653     mergesort (disks, num_disks, sizeof *disks, compare_disk);
654 
655     if (!(options & OPT_NOBIO)) {
656         for (size_t i = 0; i < num_disks; ++i) {
657             read_raid (&disks[i], disks, num_disks);
658         }
659     }
660 
661     struct padding p = {
662         .name = strlen ("NAME"),
663         .type = strlen ("TYPE"),
664         .comment = strlen ("COMMENT"),
665     };
666 
667     for (size_t i = 0; i < num_disks; ++i)
668         pad_disk (&p, &disks[i]);
669 
670     if (!(options & OPT_NOHEADER))
671         print_header (fields, &p);
672 
673     for (size_t i = 0; i < num_disks; ++i) {
674         print_disk (&disks[i], fields, options, NULL, &p);
675     }
676 
677     return 0;
678 }