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 ? (options & OPT_NOUNICODE ? "    " : "│ └─") : "";
284 
285         printf (
286             "%s%-*s ",
287             prefix,
288             p->name - (raidstatus ? 4 : 0),
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 int read_disk (const char *name, struct my_diskinfo *disk)
358 {
359     struct disklabel label;
360     char *ppath, *letter;
361 
362 	memset (disk, 0, 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             warn ("read_disk(): opendev(%s)", name);
371 			return -1;
372 		}
373 
374         if (ioctl (fd, DIOCGDINFO, &label) < 0) {
375             warn ("read_disk(): ioctl(%s, DIOCGDINFO)", name);
376 			close (fd);
377 			return -1;
378 		}
379         close (fd);
380 
381         len = strlen (ppath);
382         letter = ppath + len - 1;
383     }
384 
385     strlcpy (disk->name, name, sizeof disk->name);
386     disk->size = DL_GETDSIZE (&label) * label.d_secsize;
387     memcpy (disk->type, label.d_typename, sizeof disk->type);
388     stripdisk (disk->type);
389     memcpy (disk->label, label.d_packname, sizeof disk->label);
390     memcpy (disk->duid, label.d_uid, sizeof disk->duid);
391     disk->num_parts = 0;
392 
393     for (uint16_t i = 0; i < label.d_npartitions; ++i) {
394         const struct partition *p = &label.d_partitions[i];
395         const struct statfs *mnt;
396         struct my_partinfo part;
397 
398 		memset (&part, 0, sizeof part);
399 
400         part.size = DL_GETPSIZE (p) * label.d_secsize;
401 
402         if (!part.size)
403             continue;
404 
405         part.letter = 'a' + i;
406         *letter = part.letter;
407         mnt = find_mount (ppath);
408 
409         if (mnt) {
410             const uint64_t bs = mnt->f_bsize;
411             part.mount = mnt->f_mntonname;
412             part.fssize = mnt->f_blocks * bs;
413             part.free = mnt->f_bfree * bs;
414         }
415 
416         part.fstype = fstypenames[p->p_fstype];
417         if (i != 2)
418             disk->used += part.size;
419         disk->parts[disk->num_parts++] = part;
420     }
421 
422     return 0;
423 }
424 
425 static const char *bd_statusstr (int status) {
426     switch (status) {
427     case BIOC_SDONLINE:     return BIOC_SDONLINE_S;
428     case BIOC_SDOFFLINE:    return BIOC_SDOFFLINE_S;
429     case BIOC_SDFAILED:     return BIOC_SDFAILED_S;
430     case BIOC_SDREBUILD:    return BIOC_SDREBUILD_S;
431     case BIOC_SDHOTSPARE:   return BIOC_SDHOTSPARE_S;
432     case BIOC_SDUNUSED:     return BIOC_SDUNUSED_S;
433     case BIOC_SDSCRUB:      return BIOC_SDSCRUB_S;
434     case BIOC_SDINVALID:    return BIOC_SDINVALID_S;
435     default:                return "Unknown";
436     }
437 }
438 
439 static const char *bv_statusstr (int status) {
440     switch (status) {
441     case BIOC_SVONLINE:     return BIOC_SVONLINE_S;
442     case BIOC_SVOFFLINE:    return BIOC_SVOFFLINE_S;
443     case BIOC_SVDEGRADED:   return BIOC_SVDEGRADED_S;
444     case BIOC_SVBUILDING:   return BIOC_SVBUILDING_S;
445     case BIOC_SVSCRUB:      return BIOC_SVSCRUB_S;
446     case BIOC_SVREBUILD:    return BIOC_SVREBUILD_S;
447     case BIOC_SVINVALID:    return BIOC_SVINVALID_S;
448     default:                return "Unknown";
449     }
450 }
451 
452 static void read_raid (
453     struct my_diskinfo *disk,
454     struct my_diskinfo *disks,
455     size_t num_disks
456 ) {
457     struct bioc_inq bi;
458     int fd;
459 
460     fd = opendev (disk->name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, NULL);
461     if (fd < 0) {
462         warn ("read_raid(): opendev(%s)", disk->name);
463         return;
464     }
465 
466 	memset (&bi, 0, sizeof bi);
467 
468     if (ioctl (fd, BIOCINQ, &bi) == -1)
469         goto ret;
470 
471     for (int i = 0; i < bi.bi_novol; ++i) {
472         struct bioc_vol bv;
473 
474 		memset (&bv, 0, sizeof bv);
475         memcpy (&bv.bv_bio, &bi.bi_bio, sizeof bi.bi_bio);
476         bv.bv_volid = i;
477 
478         if (ioctl (fd, BIOCVOL, &bv) == -1) {
479             warn ("read_raid(%s): BIOCVOL(%d)", disk->name, i);
480             continue;
481         }
482 
483         if (strcmp (disk->name, bv.bv_dev) != 0)
484             continue;
485 
486         disk->raidstatus = bv_statusstr (bv.bv_status);
487 
488         for (int j = 0; j < bv.bv_nodisk; ++j) {
489             struct bioc_disk bd;
490             size_t len_vendor;
491             char letter;
492 
493 			memset (&bd, 0, sizeof bd);
494             memcpy (&bd.bd_bio, &bi.bi_bio, sizeof bi.bi_bio);
495             bd.bd_volid = i;
496             bd.bd_diskid = j;
497 
498             if (ioctl (fd, BIOCDISK, &bd) == -1) {
499                 warn ("read_raid(%s): BIOCDISK(%d, %d)", disk->name, i, j);
500                 continue;
501             }
502 
503             len_vendor = strlen (bd.bd_vendor);
504             if (len_vendor < 4 || len_vendor > 8) {
505                 warnx ("read_raid(%s): unexpected vendor string: %.32s", disk->name, bd.bd_vendor);
506                 continue;
507             }
508             letter = bd.bd_vendor[len_vendor - 1];
509 
510             for (size_t k = 0; k < num_disks; ++k) {
511                 if (!memcmp (bd.bd_vendor, disks[k].name, 3)) {
512                     for (size_t l = 0; l < disks[k].num_parts; ++l) {
513                         if (letter == disks[k].parts[l].letter) {
514                             disks[k].parts[l].sub = disk;
515                             disks[k].parts[l].raidstatus = bd_statusstr(bd.bd_status);
516                             goto found;
517                         }
518                     }
519                 }
520             }
521         found:;
522         }
523     }
524 ret:
525     close (fd);
526 }
527 
528 static int usage (void)
529 {
530     fputs ("Usage: lsblk [-abinUuV] [disk...]\n", stderr);
531     return 1;
532 }
533 
534 static void pad_update (int *pad, int newval)
535 {
536     if (newval > *pad)
537         *pad = newval;
538 }
539 
540 static void pad_disk (struct padding *p, const struct my_diskinfo *disk)
541 {
542     size_t len_disk;
543     int comment;
544 
545     len_disk = strnlen (disk->name, sizeof disk->name);
546     comment = strnlen (disk->label, sizeof disk->label);
547     if (disk->raidstatus)
548         comment += strlen (disk->raidstatus) + 3;
549 
550     pad_update (&p->name, len_disk);
551     pad_update (&p->type, strnlen (disk->type, sizeof disk->type));
552     pad_update (&p->comment, comment);
553 
554     for (int i = 0; i < disk->num_parts; ++i) {
555         const struct my_partinfo *part = &disk->parts[i];
556 
557         pad_update (&p->name, len_disk + 3);
558         pad_update (&p->type, strlen (part->fstype));
559 
560         if (part->sub) {
561             pad_update (&p->name, strnlen (part->sub->name, sizeof part->sub->name) + 4);
562             pad_update (&p->comment, strlen (part->raidstatus));
563         } else if (part->mount) {
564             pad_update (&p->comment, strlen (part->mount));
565         }
566     }
567 }
568 
569 
570 static int compare_disk (const void *p1, const void *p2)
571 {
572     const struct my_diskinfo *d1 = p1;
573     const struct my_diskinfo *d2 = p2;
574 
575     return strcmp (d1->name, d2->name);
576 }
577 
578 int main (int argc, char *argv[])
579 {
580     int option;
581     int fields = FIELD_DEFAULT;
582     int options = 0;
583 
584     if (unveil ("/dev", "r") == -1)
585         err (1, "unveil(/dev)");
586 
587     if (unveil (NULL, NULL) == -1)
588         err (1, "unveil()");
589 
590     while ((option = getopt (argc, argv, ":abinUuV")) != -1) {
591         switch (option) {
592         case 'a':
593             fields = -1;
594             break;
595         case 'b':
596             options |= OPT_NOBIO;
597             break;
598         case 'i':
599             options |= OPT_NOUNICODE;
600             break;
601         case 'n':
602             options |= OPT_NOHEADER;
603             break;
604         case 'U':
605             fields |= FIELD_DUID;
606             break;
607         case 'u':
608             fields |= FIELD_USED | FIELD_FREE;
609             break;
610         case 'V':
611             puts ("lsblk-" VERSION);
612             return 0;
613         default:
614             return usage ();
615         }
616     }
617 
618     argv += optind;
619     argc -= optind;
620 
621     char *names = argc == 0 ? disknames () : NULL;
622 
623     if (pledge ("stdio rpath disklabel", NULL) == -1)
624         err (1, "pledge()");
625 
626     size_t cap_disks;
627     if (argc == 0) {
628         const char *s = names;
629         cap_disks = 0;
630         while ((s = strchr (s, ',')) != NULL)
631             ++cap_disks, ++s;
632     } else {
633         cap_disks = argc;
634     }
635 
636     struct my_diskinfo *disks = calloc (cap_disks, sizeof (struct my_diskinfo));
637     size_t num_disks = 0;
638 
639     if (argc == 0) {
640         for (char *name; (name = strsep (&names, ",")) != NULL; ) {
641             char *colon = strchr (name, ':');
642 			struct my_diskinfo disk;
643 
644             if (colon)
645                 *colon = '\0';
646 			if (read_disk (name, &disk) == 0) {
647 				disks[num_disks++] = disk;
648 			}
649 
650             if (colon)
651                 *colon = ':';
652         }
653     } else {
654         for (int i = 0; i < argc; ++i) {
655             char *name = basename (argv[i]);
656             char *last = name + strlen (name) - 1;
657 			struct my_diskinfo disk;
658 
659             if (isalpha (*last)) {
660                 warnx ("%s: specifying a partition is not supported, using the drive.", name);
661                 *last = '\0';
662             }
663 			if (read_disk (name, &disk) == 0) {
664 				disks[num_disks++] = disk;
665 			}
666         }
667     }
668 
669     free (names);
670 
671     mergesort (disks, num_disks, sizeof *disks, compare_disk);
672 
673     if (!(options & OPT_NOBIO)) {
674         for (size_t i = 0; i < num_disks; ++i) {
675             read_raid (&disks[i], disks, num_disks);
676         }
677     }
678 
679     struct padding p = {
680         .name = strlen ("NAME"),
681         .type = strlen ("TYPE"),
682         .comment = strlen ("COMMENT"),
683     };
684 
685     for (size_t i = 0; i < num_disks; ++i)
686         pad_disk (&p, &disks[i]);
687 
688     if (!(options & OPT_NOHEADER))
689         print_header (fields, &p);
690 
691     for (size_t i = 0; i < num_disks; ++i) {
692         print_disk (&disks[i], fields, options, NULL, &p);
693     }
694 
695     return 0;
696 }