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 #include <stddef.h>
20 #include <dev/biovar.h>
21 #include <sys/cdefs.h>
22 #include <sys/types.h>
23 #include <sys/disklabel.h>
24 #include <sys/sysctl.h>
25 #include <sys/ioctl.h>
26 #include <sys/mount.h>
27 #include <sys/dkio.h>
28 #include <inttypes.h>
29 #include <stdbool.h>
30 #include <string.h>
31 #include <stdarg.h>
32 #include <stdint.h>
33 #include <stdlib.h>
34 #include <unistd.h>
35 #include <libgen.h>
36 #include <stdio.h>
37 #include <fcntl.h>
38 #include <errno.h>
39 #include <util.h>
40 #include <err.h>
41 
42 static int diskcount (void)
43 {
44     const int mib[2] = { CTL_HW, HW_DISKCOUNT };
45     int diskcount;
46     size_t len = sizeof diskcount;
47 
48     if (sysctl (mib, 2, &diskcount, &len, NULL, 0) == -1)
49         err (1, "sysctl(hw.diskcount)");
50 
51     return diskcount;
52 }
53 
54 static char *disknames (void)
55 {
56     const int num = diskcount ();
57     const int mib[2] = { CTL_HW, HW_DISKNAMES };
58     size_t len = 32 * num;
59     char *buffer = malloc (len);
60 
61     if (sysctl (mib, 2, buffer, &len, NULL, 0) == -1)
62         err (1, "sysctl(hw.disknames)");
63 
64     return buffer;
65 }
66 
67 static char *stripdisk (char *n)
68 {
69     const char sufx[] = " disk";
70     const size_t ln = strnlen (n, 16);
71     const size_t ls = sizeof sufx - 1;
72 
73     if (!memcmp (n + ln - ls, sufx, ls)) {
74         n[ln - ls] = '\0';
75     }
76 
77     return n;
78 }
79 
80 static void print_size (uint64_t sz)
81 {
82     struct unit {
83         char sym;
84         uint64_t factor;
85     };
86 
87     const struct unit 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[4];
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 static void print_header (int fields)
165 {
166     if (fields & FIELD_NAME)
167         printf ("%-7s ", "NAME");
168 
169     if (fields & FIELD_DUID)
170         printf ("%-18s ", "DUID");
171 
172     if (fields & FIELD_SIZE)
173         printf ("%-4s ", "SIZE");
174 
175     if (fields & FIELD_USED)
176         printf ("%-4s ", "USED");
177 
178     if (fields & FIELD_FREE)
179         printf ("%-4s ", "FREE");
180 
181     if (fields & FIELD_TYPE)
182         printf ("%-8s ", "TYPE");
183 
184     if (fields & FIELD_COMMENT)
185         printf ("COMMENT ");
186 
187     putchar ('\n');
188 }
189 
190 static void print_duid (const u_char *duid)
191 {
192     for (size_t i = 0; i < 8; ++i) {
193         printf ("%02x", duid[i]);
194     }
195 }
196 
197 static void print_disk (const struct my_diskinfo *disk, int fields, int options, const char *raidstatus);
198 static void print_part (
199     const struct my_diskinfo *disk,
200     const struct my_partinfo *part,
201     int fields,
202     int options,
203     bool last
204 ) {
205     if (fields & FIELD_NAME) {
206         const char *prefix = (options & OPT_NOUNICODE) ? "  " : (last ? "└─" : "├─");
207         printf ("%s%s%c  ", prefix, disk->name, part->letter);
208     }
209 
210     if (fields & FIELD_DUID) {
211         print_duid (disk->duid);
212         printf (".%c ", part->letter);
213     }
214 
215     if (fields & FIELD_SIZE)
216         print_size (part->size);
217 
218     if (fields & FIELD_USED) {
219         if (part->fssize) {
220             print_size (part->fssize - part->free);
221         } else {
222             printf (" N/A ");
223         }
224     }
225 
226     if (fields & FIELD_FREE) {
227         if (part->fssize) {
228             print_size (part->free);
229         } else {
230             printf (" N/A ");
231         }
232     }
233 
234     if (fields & FIELD_TYPE)
235         printf ("%-8.16s ", part->fstype);
236 
237     if (fields & FIELD_COMMENT && part->mount)
238         printf ("%s ", part->mount);
239 
240     putchar ('\n');
241 
242     if (part->sub) {
243         printf ("│ └─");
244         print_disk (part->sub, fields, options, part->raidstatus);
245     }
246 }
247 
248 static void print_disk (const struct my_diskinfo *disk, int fields, int options, const char *raidstatus)
249 {
250     if (fields & FIELD_NAME) {
251         printf ("%s ", disk->name);
252         if (!raidstatus)
253             printf ("    ");
254     }
255 
256     if (fields & FIELD_DUID) {
257         print_duid (disk->duid);
258         printf ("   ");
259     }
260 
261     if (fields & FIELD_SIZE)
262         print_size (disk->size);
263 
264     if (fields & FIELD_USED)
265         print_size (disk->used);
266 
267     if (fields & FIELD_FREE)
268         print_size (disk->size - disk->used);
269 
270     // Pad only upto 8 characters because most disk types are <=8 bytes long.
271     if (fields & FIELD_TYPE)
272         printf ("%-8.16s ", disk->type);
273 
274     if (fields & FIELD_COMMENT) {
275         if (raidstatus) {
276             printf ("%s ", raidstatus);
277         } else if (disk->raidstatus) {
278             printf ("%.16s (%s) ", disk->label, disk->raidstatus);
279         } else {
280             printf ("%.16s ", disk->label);
281         }
282     }
283 
284     putchar ('\n');
285 
286     if (!raidstatus) {
287         for (uint8_t i = 0; i < disk->num_parts; ++i)
288             print_part (disk, &disk->parts[i], fields, options, i == (disk->num_parts - 1));
289     }
290 }
291 
292 static const struct statfs *find_mount (const char *dev)
293 {
294     static struct statfs *mounts = NULL;
295     static int n_mounts;
296 
297     if (!mounts) {
298         n_mounts = getmntinfo (&mounts, MNT_NOWAIT);
299         if (n_mounts == 0)
300             err (1, "getmntinfo()");
301     }
302 
303     for (int i = 0; i < n_mounts; ++i) {
304         if (!strcmp (dev, mounts[i].f_mntfromname))
305             return &mounts[i];
306     }
307 
308     return NULL;
309 }
310 
311 static struct my_diskinfo read_disk (const char *name)
312 {
313     struct my_diskinfo disk;
314     struct disklabel label;
315     char *ppath, *letter;
316 
317     bzero (&disk, sizeof disk);
318 
319     {   // Read disklabel.
320         size_t len;
321         int fd;
322 
323         fd = opendev (name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, &ppath);
324         if (fd < 0)
325             err (1, "opendev(%s)", name);
326 
327         if (ioctl (fd, DIOCGDINFO, &label) < 0)
328             err (1, "ioctl(%s, DIOCGDINFO)", name);
329         close (fd);
330 
331         len = strlen (ppath);
332         letter = ppath + len - 1;
333     }
334 
335     memcpy (disk.name, name, 3);
336     disk.name[3] = '\0';
337     disk.size = DL_GETDSIZE (&label) * label.d_secsize;
338     memcpy (disk.type, label.d_typename, 16);
339     stripdisk (disk.type);
340     memcpy (disk.label, label.d_packname, 16);
341     memcpy (disk.duid, label.d_uid, sizeof disk.duid);
342     disk.num_parts = 0;
343 
344     for (uint16_t i = 0; i < label.d_npartitions; ++i) {
345         const struct partition *p = &label.d_partitions[i];
346         const struct statfs *mnt;
347         struct my_partinfo part;
348 
349         bzero (&part, sizeof part);
350 
351         part.size = DL_GETPSIZE (p) * label.d_secsize;
352 
353         if (!part.size)
354             continue;
355 
356         part.letter = 'a' + i;
357         *letter = part.letter;
358         mnt = find_mount (ppath);
359 
360         if (mnt) {
361             const uint64_t bs = mnt->f_bsize;
362             part.mount = mnt->f_mntonname;
363             part.fssize = mnt->f_blocks * bs;
364             part.free = mnt->f_bfree * bs;
365         }
366 
367         part.fstype = fstypenames[p->p_fstype];
368         if (i != 2)
369             disk.used += part.size;
370         disk.parts[disk.num_parts++] = part;
371     }
372 
373     return disk;
374 }
375 
376 static const char *bd_statusstr (int status) {
377     switch (status) {
378     case BIOC_SDONLINE:     return BIOC_SDONLINE_S;
379     case BIOC_SDOFFLINE:    return BIOC_SDOFFLINE_S;
380     case BIOC_SDFAILED:     return BIOC_SDFAILED_S;
381     case BIOC_SDREBUILD:    return BIOC_SDREBUILD_S;
382     case BIOC_SDHOTSPARE:   return BIOC_SDHOTSPARE_S;
383     case BIOC_SDUNUSED:     return BIOC_SDUNUSED_S;
384     case BIOC_SDSCRUB:      return BIOC_SDSCRUB_S;
385     case BIOC_SDINVALID:    return BIOC_SDINVALID_S;
386     default:                return "Unknown";
387     }
388 }
389 
390 static const char *bv_statusstr (int status) {
391     switch (status) {
392     case BIOC_SVONLINE:     return BIOC_SVONLINE_S;
393     case BIOC_SVOFFLINE:    return BIOC_SVOFFLINE_S;
394     case BIOC_SVDEGRADED:   return BIOC_SVDEGRADED_S;
395     case BIOC_SVBUILDING:   return BIOC_SVBUILDING_S;
396     case BIOC_SVSCRUB:      return BIOC_SVSCRUB_S;
397     case BIOC_SVREBUILD:    return BIOC_SVREBUILD_S;
398     case BIOC_SVINVALID:    return BIOC_SVINVALID_S;
399     default:                return "Unknown";
400     }
401 }
402 
403 static void read_raid (
404     struct my_diskinfo *disk,
405     struct my_diskinfo *disks,
406     size_t num_disks
407 ) {
408     struct bioc_inq bi;
409     int fd;
410 
411     fd = opendev (disk->name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, NULL);
412     if (fd < 0) {
413         warn ("read_raid(): opendev(%s)", disk->name);
414         return;
415     }
416 
417     bzero (&bi, sizeof bi);
418 
419     if (ioctl (fd, BIOCINQ, &bi) == -1)
420         goto ret;
421 
422     for (int i = 0; i < bi.bi_novol; ++i) {
423         struct bioc_vol bv;
424 
425         bzero (&bv, sizeof bv);
426         memcpy (&bv.bv_bio, &bi.bi_bio, sizeof bi.bi_bio);
427         bv.bv_volid = i;
428 
429         if (ioctl (fd, BIOCVOL, &bv) == -1) {
430             warn ("read_raid(%s): BIOCVOL(%d)", disk->name, i);
431             continue;
432         }
433 
434         if (strcmp (disk->name, bv.bv_dev) != 0)
435             continue;
436 
437         disk->raidstatus = bv_statusstr (bv.bv_status);
438 
439         for (int j = 0; j < bv.bv_nodisk; ++j) {
440             struct bioc_disk bd;
441             size_t len_vendor;
442             char letter;
443 
444             bzero (&bd, sizeof bd);
445             memcpy (&bd.bd_bio, &bi.bi_bio, sizeof bi.bi_bio);
446             bd.bd_volid = i;
447             bd.bd_diskid = j;
448 
449             if (ioctl (fd, BIOCDISK, &bd) == -1) {
450                 warn ("read_raid(%s): BIOCDISK(%d, %d)", disk->name, i, j);
451                 continue;
452             }
453 
454             len_vendor = strlen (bd.bd_vendor);
455             if (len_vendor != 4) {
456                 warnx ("read_raid(%s): unexpected vendor string: %.32s", disk->name, bd.bd_vendor);
457                 continue;
458             }
459             letter = bd.bd_vendor[len_vendor - 1];
460 
461             for (size_t k = 0; k < num_disks; ++k) {
462                 if (!memcmp (bd.bd_vendor, disks[k].name, 3)) {
463                     for (size_t l = 0; l < disks[k].num_parts; ++l) {
464                         if (letter == disks[k].parts[l].letter) {
465                             disks[k].parts[l].sub = disk;
466                             disks[k].parts[l].raidstatus = bd_statusstr(bd.bd_status);
467                             goto found;
468                         }
469                     }
470                 }
471             }
472         found:;
473         }
474     }
475 ret:
476     close (fd);
477 }
478 
479 static int usage (void)
480 {
481     fputs ("Usage: lsblk [-abinUuV] [disk...]\n", stderr);
482     return 1;
483 }
484 
485 int main (int argc, char *argv[])
486 {
487     int option;
488     int fields = FIELD_DEFAULT;
489     int options = 0;
490 
491     if (unveil ("/dev", "r") == -1)
492         err (1, "unveil(/dev)");
493 
494     if (unveil (NULL, NULL) == -1)
495         err (1, "unveil()");
496 
497     while ((option = getopt (argc, argv, ":abinUuV")) != -1) {
498         switch (option) {
499         case 'a':
500             fields = -1;
501             break;
502         case 'b':
503             options |= OPT_NOBIO;
504             break;
505         case 'i':
506             options |= OPT_NOUNICODE;
507             break;
508         case 'n':
509             options |= OPT_NOHEADER;
510             break;
511         case 'U':
512             fields |= FIELD_DUID;
513             break;
514         case 'u':
515             fields |= FIELD_USED | FIELD_FREE;
516             break;
517         case 'V':
518             puts ("lsblk-" VERSION);
519             return 0;
520         default:
521             return usage ();
522         }
523     }
524 
525     argv += optind;
526     argc -= optind;
527 
528     char *names = argc == 0 ? disknames () : NULL;
529 
530     if (pledge ("stdio rpath disklabel", NULL) == -1)
531         err (1, "pledge()");
532 
533     size_t cap_disks;
534     if (argc == 0) {
535         const char *s = names;
536         cap_disks = 0;
537         while ((s = strchr (s, ',')) != NULL)
538             ++cap_disks, ++s;
539     } else {
540         cap_disks = argc;
541     }
542 
543     struct my_diskinfo *disks = calloc (cap_disks, sizeof (struct my_diskinfo));
544     size_t num_disks = 0;
545 
546     if (argc == 0) {
547         for (char *disk; (disk = strsep (&names, ",")) != NULL; ) {
548             char *colon = strchr (disk, ':');
549             if (colon)
550                 *colon = '\0';
551             disks[num_disks++] = read_disk (disk);
552             if (colon)
553                 *colon = ':';
554         }
555     } else {
556         for (int i = 0; i < argc; ++i) {
557             char *disk = basename (argv[i]);
558             disks[num_disks++] = read_disk (disk);
559         }
560     }
561 
562     free (names);
563 
564     if (!(options & OPT_NOBIO)) {
565         for (size_t i = 0; i < num_disks; ++i) {
566             read_raid (&disks[i], disks, num_disks);
567         }
568     }
569 
570     if (!(options & OPT_NOHEADER))
571         print_header (fields);
572 
573     for (size_t i = 0; i < num_disks; ++i) {
574         print_disk (&disks[i], fields, options, false);
575     }
576 
577     return 0;
578 }