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