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