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 <sys/cdefs.h>
21 #include <sys/types.h>
22 #include <sys/disklabel.h>
23 #include <sys/sysctl.h>
24 #include <sys/ioctl.h>
25 #include <sys/mount.h>
26 #include <sys/dkio.h>
27 #include <inttypes.h>
28 #include <stdbool.h>
29 #include <string.h>
30 #include <stdarg.h>
31 #include <stdint.h>
32 #include <stdlib.h>
33 #include <unistd.h>
34 #include <libgen.h>
35 #include <stdio.h>
36 #include <fcntl.h>
37 #include <errno.h>
38 
39 static __dead void die (const char *fmt, ...)
40 {
41     va_list ap;
42 
43     va_start (ap, fmt);
44     fputs ("Error: ", stderr);
45     vfprintf (stderr, fmt, ap);
46 
47     if (errno != 0) {
48         fprintf (stderr, ": %s\n", strerror (errno));
49     } else {
50         fputc ('\n', stderr);
51     }
52 
53     exit (1);
54 }
55 
56 static int diskcount (void)
57 {
58     const int mib[2] = { CTL_HW, HW_DISKCOUNT };
59     int diskcount;
60     size_t len = sizeof diskcount;
61 
62     if (sysctl (mib, 2, &diskcount, &len, NULL, 0) == -1)
63         die ("sysctl(hw.diskcount)");
64 
65     return diskcount;
66 }
67 
68 static char *disknames (void)
69 {
70     const int num = diskcount ();
71     const int mib[2] = { CTL_HW, HW_DISKNAMES };
72     size_t len = 32 * num;
73     char *buffer = malloc (len);
74 
75     if (sysctl (mib, 2, buffer, &len, NULL, 0) == -1)
76         die ("sysctl(hw.disknames)");
77 
78     return buffer;
79 }
80 
81 static char *stripdisk (char *n)
82 {
83     const char sufx[] = " disk";
84     const size_t ln = strnlen (n, 16);
85     const size_t ls = sizeof sufx - 1;
86 
87     if (!memcmp (n + ln - ls, sufx, ls)) {
88         n[ln - ls] = '\0';
89     }
90 
91     return n;
92 }
93 
94 static void print_size (uint64_t sz)
95 {
96     struct unit {
97         char sym;
98         uint64_t factor;
99     };
100 
101     const struct unit units[] = {
102         { 'P', 1ull << 50 },
103         { 'T', 1ull << 40 },
104         { 'G', 1ull << 30 },
105         { 'M', 1ull << 20 },
106         { 'K', 1ull << 10 },
107         { '0', 0 },
108     };
109 
110     char sym = 'B';
111     uint64_t factor = 1;
112 
113     for (const struct unit *u = &units[0]; u->factor; ++u) {
114         if (sz >= (u->factor * 9 / 10)) {
115             sym = u->sym;
116             factor = u->factor;
117             break;
118         }
119     }
120 
121     const unsigned scaled10 = sz * 10 / factor;
122     const unsigned scaled = sz / factor;
123     if (scaled10 >= 1000) {
124         printf ("%u", scaled);
125     } else if (scaled10 >= 100) {
126         printf (" %u", scaled);
127     } else {
128         printf ("%u.%u", scaled, scaled10 % 10);
129     }
130 
131     putchar (sym);
132     putchar (' ');
133 }
134 
135 enum {
136     FIELD_NAME      = 0x01,
137     FIELD_SIZE      = 0x02,
138     FIELD_USED      = 0x04,
139     FIELD_FREE      = 0x08,
140     FIELD_TYPE      = 0x10,
141     FIELD_LMNT      = 0x20,
142 
143     FIELD_DEFAULT   = FIELD_NAME | FIELD_SIZE | FIELD_TYPE | FIELD_LMNT,
144 };
145 
146 enum {
147     OPT_NOHEADER    = 0x01,
148     OPT_NOUNICODE   = 0x02,
149 };
150 
151 struct my_partinfo {
152     char name[5];
153     uint64_t size;
154     uint64_t fssize;
155     uint64_t free;
156     const char *fstype;
157     const char *mount;
158 };
159 
160 struct my_diskinfo {
161     char type[16];
162     char label[16];
163     char name[4];
164     uint64_t size;
165     uint64_t used;
166     uint8_t num_parts;
167     struct my_partinfo parts[MAXPARTITIONS];
168 };
169 
170 static void print_header (int fields)
171 {
172     if (fields & FIELD_NAME)
173         printf ("%-6s ", "NAME");
174 
175     if (fields & FIELD_SIZE)
176         printf ("%-4s ", "SIZE");
177 
178     if (fields & FIELD_USED)
179         printf ("%-4s ", "USED");
180 
181     if (fields & FIELD_FREE)
182         printf ("%-4s ", "FREE");
183 
184     if (fields & FIELD_TYPE)
185         printf ("%-8s ", "TYPE");
186 
187     if (fields & FIELD_LMNT)
188         printf ("LABEL/MOUNT ");
189 
190     putchar ('\n');
191 }
192 
193 static void print_part (const struct my_partinfo *part, int fields, int options, bool last)
194 {
195     if (fields & FIELD_NAME) {
196         const char *prefix = (options & OPT_NOUNICODE) ? "  " : (last ? "└─" : "├─");
197         printf ("%s%s ", prefix, part->name);
198     }
199 
200     if (fields & FIELD_SIZE)
201         print_size (part->size);
202 
203     if (fields & FIELD_USED) {
204         if (part->fssize) {
205             print_size (part->fssize - part->free);
206         } else {
207             printf (" N/A ");
208         }
209     }
210 
211     if (fields & FIELD_FREE) {
212         if (part->fssize) {
213             print_size (part->free);
214         } else {
215             printf (" N/A ");
216         }
217     }
218 
219     if (fields & FIELD_TYPE)
220         printf ("%-8.16s ", part->fstype);
221 
222     if (fields & FIELD_LMNT && part->mount)
223         printf ("%s ", part->mount);
224 
225     putchar ('\n');
226 }
227 
228 static void print_disk (const struct my_diskinfo *disk, int fields, int options)
229 {
230     if (fields & FIELD_NAME)
231         printf ("%s    ", disk->name);
232 
233     if (fields & FIELD_SIZE)
234         print_size (disk->size);
235 
236     if (fields & FIELD_USED)
237         print_size (disk->used);
238 
239     if (fields & FIELD_FREE)
240         print_size (disk->size - disk->used);
241 
242     // Pad only upto 8 characters because most disk types are <=8 bytes long.
243     if (fields & FIELD_TYPE)
244         printf ("%-8.16s ", disk->type);
245 
246     if (fields & FIELD_LMNT)
247         printf ("%.16s ", disk->label);
248 
249     putchar ('\n');
250 
251     for (uint8_t i = 0; i < disk->num_parts; ++i)
252         print_part (&disk->parts[i], fields, options, i == (disk->num_parts - 1));
253 }
254 
255 static const struct statfs *find_mount (const char *dev)
256 {
257     static struct statfs *mounts = NULL;
258     static int n_mounts;
259 
260     if (!mounts) {
261         n_mounts = getmntinfo (&mounts, MNT_NOWAIT);
262         if (n_mounts == 0)
263             die ("getmntinfo()");
264     }
265 
266     for (int i = 0; i < n_mounts; ++i) {
267         if (!strcmp (dev, mounts[i].f_mntfromname))
268             return &mounts[i];
269     }
270 
271     return NULL;
272 }
273 
274 static struct my_diskinfo read_disk (const char *name)
275 {
276     struct my_diskinfo disk;
277     struct disklabel label;
278 
279     bzero (&disk, sizeof disk);
280 
281     {   // Read disklabel.
282         char *path = NULL;
283         int fd;
284 
285         asprintf (&path, "/dev/%sc", name);
286         fd = open (path, O_RDONLY);
287         if (fd < 0)
288             die ("open(%s)", path);
289         if (ioctl (fd, DIOCGDINFO, &label) < 0)
290             die ("ioctl(%s, DIOCGDINFO)", path);
291         close (fd);
292         free (path);
293     }
294 
295 
296     memcpy (disk.name, name, 3);
297     disk.name[3] = '\0';
298     disk.size = DL_GETDSIZE (&label) * label.d_secsize;
299     memcpy (disk.type, label.d_typename, 16);
300     stripdisk (disk.type);
301     memcpy (disk.label, label.d_packname, 16);
302     disk.num_parts = 0;
303 
304     for (uint16_t i = 0; i < label.d_npartitions; ++i) {
305         const struct partition *p = &label.d_partitions[i];
306         const struct statfs *mnt;
307         struct my_partinfo part;
308         char path[5 + 4 + 1];
309 
310         bzero (&part, sizeof part);
311 
312         part.size = DL_GETPSIZE (p) * label.d_secsize;
313 
314         if (!part.size)
315             continue;
316 
317         memcpy (part.name, disk.name, 3);
318         part.name[3] = 'a' + i;
319         part.name[4] = '\0';
320 
321         snprintf (path, sizeof path, "/dev/%s", part.name);
322         mnt = find_mount (path);
323 
324         if (mnt) {
325             const uint64_t bs = mnt->f_bsize;
326             part.mount = mnt->f_mntonname;
327             part.fssize = mnt->f_blocks * bs;
328             part.free = mnt->f_bfree * bs;
329         }
330 
331         part.fstype = fstypenames[p->p_fstype];
332         if (i != 2)
333             disk.used += part.size;
334         disk.parts[disk.num_parts++] = part;
335     }
336 
337     return disk;
338 }
339 
340 static int usage (void)
341 {
342     fputs ("Usage: lsblk [-Vainu] [disk...]\n", stderr);
343     return 1;
344 }
345 
346 int main (int argc, char *argv[])
347 {
348     int option;
349     int fields = FIELD_DEFAULT;
350     int options = 0;
351 
352     if (unveil ("/dev", "r") == -1)
353         die ("unveil(/dev)");
354 
355     if (unveil (NULL, NULL) == -1)
356         die ("unveil()");
357 
358     while ((option = getopt (argc, argv, ":Vainu")) != -1) {
359         switch (option) {
360         case 'V':
361             puts ("lsblk-" VERSION);
362             return 0;
363         case 'a':
364             fields = -1;
365             break;
366         case 'i':
367             options |= OPT_NOUNICODE;
368             break;
369         case 'n':
370             options |= OPT_NOHEADER;
371             break;
372         case 'u':
373             fields |= FIELD_USED | FIELD_FREE;
374             break;
375         default:
376             return usage ();
377         }
378     }
379 
380     argv += optind;
381     argc -= optind;
382 
383     char *names = argc == 0 ? disknames () : NULL;
384 
385     if (pledge ("stdio rpath disklabel", NULL) == -1)
386         die ("pledge()");
387 
388     size_t cap_disks;
389     if (argc == 0) {
390         const char *s = names;
391         cap_disks = 0;
392         while ((s = strchr (s, ',')) != NULL)
393             ++cap_disks, ++s;
394     } else {
395         cap_disks = argc;
396     }
397 
398     struct my_diskinfo *disks = calloc (cap_disks, sizeof (struct my_diskinfo));
399     size_t num_disks = 0;
400 
401     if (argc == 0) {
402         for (char *disk; (disk = strsep (&names, ",")) != NULL; ) {
403             char *colon = strchr (disk, ':');
404             if (colon)
405                 *colon = '\0';
406             disks[num_disks++] = read_disk (disk);
407             if (colon)
408                 *colon = ':';
409         }
410     } else {
411         for (int i = 0; i < argc; ++i) {
412             char *disk = basename (argv[i]);
413             disks[num_disks++] = read_disk (disk);
414         }
415     }
416 
417     free (names);
418 
419     if (!(options & OPT_NOHEADER))
420         print_header (fields);
421 
422     for (size_t i = 0; i < num_disks; ++i) {
423         print_disk (&disks[i], fields, options);
424     }
425 
426     return 0;
427 }