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 <ctype.h>
40 #include <errno.h>
41 #include <util.h>
42 #include <err.h>
43 
44 static int diskcount (void)
45 {
46     const int mib[2] = { CTL_HW, HW_DISKCOUNT };
47     int diskcount;
48     size_t len = sizeof diskcount;
49 
50     if (sysctl (mib, 2, &diskcount, &len, NULL, 0) == -1)
51         err (1, "sysctl(hw.diskcount)");
52 
53     return diskcount;
54 }
55 
56 static char *disknames (void)
57 {
58     const int num = diskcount ();
59     const int mib[2] = { CTL_HW, HW_DISKNAMES };
60     size_t len = 32 * num;
61     char *buffer = malloc (len);
62 
63     if (sysctl (mib, 2, buffer, &len, NULL, 0) == -1)
64         err (1, "sysctl(hw.disknames)");
65 
66     return buffer;
67 }
68 
69 static char *stripdisk (char *n)
70 {
71     const char sufx[] = " disk";
72     const size_t ln = strnlen (n, 16);
73     const size_t ls = sizeof sufx - 1;
74 
75     if (!memcmp (n + ln - ls, sufx, ls)) {
76         n[ln - ls] = '\0';
77     }
78 
79     return n;
80 }
81 
82 static void print_size (uint64_t sz)
83 {
84     struct unit {
85         char sym;
86         uint64_t factor;
87     };
88 
89     const struct unit units[] = {
90         { 'P', 1ull << 50 },
91         { 'T', 1ull << 40 },
92         { 'G', 1ull << 30 },
93         { 'M', 1ull << 20 },
94         { 'K', 1ull << 10 },
95         { '0', 0 },
96     };
97 
98     char sym = 'B';
99     uint64_t factor = 1;
100 
101     for (const struct unit *u = &units[0]; u->factor; ++u) {
102         if (sz >= (u->factor * 9 / 10)) {
103             sym = u->sym;
104             factor = u->factor;
105             break;
106         }
107     }
108 
109     const unsigned scaled10 = sz * 10 / factor;
110     const unsigned scaled = sz / factor;
111     if (scaled10 >= 1000) {
112         printf ("%u", scaled);
113     } else if (scaled10 >= 100) {
114         printf (" %u", scaled);
115     } else {
116         printf ("%u.%u", scaled, scaled10 % 10);
117     }
118 
119     putchar (sym);
120     putchar (' ');
121 }
122 
123 enum {
124     FIELD_NAME      = 0x01,
125     FIELD_DUID      = 0x02,
126     FIELD_SIZE      = 0x04,
127     FIELD_USED      = 0x08,
128     FIELD_FREE      = 0x10,
129     FIELD_TYPE      = 0x20,
130     FIELD_COMMENT   = 0x40,
131 
132     FIELD_DEFAULT   = FIELD_NAME | FIELD_SIZE | FIELD_TYPE | FIELD_COMMENT,
133 };
134 
135 enum {
136     OPT_NOHEADER    = 0x01,
137     OPT_NOUNICODE   = 0x02,
138     OPT_NOBIO       = 0x04,
139 };
140 
141 struct my_diskinfo;
142 
143 struct my_partinfo {
144     char letter;
145     uint64_t size;
146     uint64_t fssize;
147     uint64_t free;
148     const char *fstype;
149     const char *mount;
150     const struct my_diskinfo *sub; // If this is part of a RAID
151     const char *raidstatus; // Only available if (sub != NULL)
152 };
153 
154 struct my_diskinfo {
155     char type[16];
156     char label[16];
157     char name[8];
158     uint64_t size;
159     uint64_t used;
160     u_char duid[8];
161     uint8_t num_parts;
162     struct my_partinfo parts[MAXPARTITIONS];
163     const char *raidstatus; // If this is a RAID device
164 };
165 
166 struct padding {
167     int name;
168     /* duid = 8 */
169     /* size = 4 */
170     /* used = 4 */
171     /* free = 4 */
172     int type;
173     int comment;
174 };
175 
176 static void print_header (int fields, const struct padding *p)
177 {
178     if (fields & FIELD_NAME)
179         printf ("%-*s ", p->name, "NAME");
180 
181     if (fields & FIELD_DUID)
182         printf ("%-18s ", "DUID");
183 
184     if (fields & FIELD_SIZE)
185         printf ("%-4s ", "SIZE");
186 
187     if (fields & FIELD_USED)
188         printf ("%-4s ", "USED");
189 
190     if (fields & FIELD_FREE)
191         printf ("%-4s ", "FREE");
192 
193     if (fields & FIELD_TYPE)
194         printf ("%-*s ", p->type, "TYPE");
195 
196     if (fields & FIELD_COMMENT)
197         printf ("%-*s ", p->comment, "COMMENT");
198 
199 #if WSDEBUG
200     putchar ('X');
201 #endif
202 
203     putchar ('\n');
204 }
205 
206 static void print_duid (const u_char *duid)
207 {
208     for (size_t i = 0; i < 8; ++i) {
209         printf ("%02x", duid[i]);
210     }
211 }
212 
213 static void print_disk (const struct my_diskinfo *, int, int, const char *, const struct padding *);
214 static void print_part (
215     const struct my_diskinfo *disk,
216     const struct my_partinfo *part,
217     int fields,
218     int options,
219     bool last,
220     const struct padding *p
221 ) {
222     if (fields & FIELD_NAME) {
223         const char *prefix = (options & OPT_NOUNICODE) ? "  " : (last ? "└─" : "├─");
224         printf (
225             "%s%s%c%-*s ",
226             prefix,
227             disk->name,
228             part->letter,
229             p->name - 2 - (int)strlen (disk->name) - 1,
230             ""
231         );
232     }
233 
234     if (fields & FIELD_DUID) {
235         print_duid (disk->duid);
236         printf (".%c ", part->letter);
237     }
238 
239     if (fields & FIELD_SIZE)
240         print_size (part->size);
241 
242     if (fields & FIELD_USED) {
243         if (part->fssize) {
244             print_size (part->fssize - part->free);
245         } else {
246             printf (" N/A ");
247         }
248     }
249 
250     if (fields & FIELD_FREE) {
251         if (part->fssize) {
252             print_size (part->free);
253         } else {
254             printf (" N/A ");
255         }
256     }
257 
258     if (fields & FIELD_TYPE)
259         printf ("%-*s ", p->type, part->fstype);
260 
261     if (fields & FIELD_COMMENT)
262         printf ("%-*s ", p->comment, part->mount ? part->mount : "");
263 
264 #if WSDEBUG
265     putchar ('X');
266 #endif
267 
268     putchar ('\n');
269 
270     if (part->sub) {
271         print_disk (part->sub, fields, options, part->raidstatus, p);
272     }
273 }
274 
275 static void print_disk (
276     const struct my_diskinfo *disk,
277     int fields,
278     int options,
279     const char *raidstatus,
280     const struct padding *p
281 ) {
282     if (fields & FIELD_NAME) {
283         const char *prefix = raidstatus ? "│ └─" : "";
284 
285         printf (
286             "%s%-*s ",
287             prefix,
288             p->name - (int)strlen (prefix),
289             disk->name
290         );
291     }
292 
293     if (fields & FIELD_DUID) {
294         print_duid (disk->duid);
295         printf ("   ");
296     }
297 
298     if (fields & FIELD_SIZE)
299         print_size (disk->size);
300 
301     if (fields & FIELD_USED)
302         print_size (disk->used);
303 
304     if (fields & FIELD_FREE)
305         print_size (disk->size - disk->used);
306 
307     if (fields & FIELD_TYPE)
308         printf ("%-*.16s ", p->type, disk->type);
309 
310     if (fields & FIELD_COMMENT) {
311         if (raidstatus) {
312             printf ("%-*s ", p->comment, raidstatus);
313         } else if (disk->raidstatus) {
314             printf (
315                 "%.16s (%s)%*s ",
316                 disk->label,
317                 disk->raidstatus,
318                 p->comment - (int)strnlen (disk->label, sizeof disk->label) - (int)strlen (disk->raidstatus) - 3,
319                 ""
320             );
321         } else {
322             printf ("%-*.16s ", p->comment, disk->label);
323         }
324     }
325 
326 #if WSDEBUG
327     putchar ('X');
328 #endif
329 
330     putchar ('\n');
331 
332     if (!raidstatus) {
333         for (uint8_t i = 0; i < disk->num_parts; ++i)
334             print_part (disk, &disk->parts[i], fields, options, i == (disk->num_parts - 1), p);
335     }
336 }
337 
338 static const struct statfs *find_mount (const char *dev)
339 {
340     static struct statfs *mounts = NULL;
341     static int n_mounts;
342 
343     if (!mounts) {
344         n_mounts = getmntinfo (&mounts, MNT_NOWAIT);
345         if (n_mounts == 0)
346             err (1, "getmntinfo()");
347     }
348 
349     for (int i = 0; i < n_mounts; ++i) {
350         if (!strcmp (dev, mounts[i].f_mntfromname))
351             return &mounts[i];
352     }
353 
354     return NULL;
355 }
356 
357 static struct my_diskinfo read_disk (const char *name)
358 {
359     struct my_diskinfo disk;
360     struct disklabel label;
361     char *ppath, *letter;
362 
363     bzero (&disk, sizeof disk);
364 
365     {   // Read disklabel.
366         size_t len;
367         int fd;
368 
369         fd = opendev (name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, &ppath);
370         if (fd < 0)
371             err (1, "opendev(%s)", name);
372 
373         if (ioctl (fd, DIOCGDINFO, &label) < 0)
374             err (1, "ioctl(%s, DIOCGDINFO)", name);
375         close (fd);
376 
377         len = strlen (ppath);
378         letter = ppath + len - 1;
379     }
380 
381     strlcpy (disk.name, name, sizeof disk.name);
382     disk.size = DL_GETDSIZE (&label) * label.d_secsize;
383     memcpy (disk.type, label.d_typename, sizeof disk.type);
384     stripdisk (disk.type);
385     memcpy (disk.label, label.d_packname, sizeof disk.label);
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 || len_vendor > 8) {
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             char *last = disk + strlen (disk) - 1;
648             if (isalpha (*last)) {
649                 warnx ("%s: specifying a partition is not supported, using the drive.", disk);
650                 *last = '\0';
651             }
652             disks[num_disks++] = read_disk (disk);
653         }
654     }
655 
656     free (names);
657 
658     mergesort (disks, num_disks, sizeof *disks, compare_disk);
659 
660     if (!(options & OPT_NOBIO)) {
661         for (size_t i = 0; i < num_disks; ++i) {
662             read_raid (&disks[i], disks, num_disks);
663         }
664     }
665 
666     struct padding p = {
667         .name = strlen ("NAME"),
668         .type = strlen ("TYPE"),
669         .comment = strlen ("COMMENT"),
670     };
671 
672     for (size_t i = 0; i < num_disks; ++i)
673         pad_disk (&p, &disks[i]);
674 
675     if (!(options & OPT_NOHEADER))
676         print_header (fields, &p);
677 
678     for (size_t i = 0; i < num_disks; ++i) {
679         print_disk (&disks[i], fields, options, NULL, &p);
680     }
681 
682     return 0;
683 }