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/types.h>
21 #include <sys/disklabel.h>
22 #include <sys/sysctl.h>
23 #include <sys/ioctl.h>
24 #include <sys/mount.h>
25 #include <sys/dkio.h>
26 #include <inttypes.h>
27 #include <string.h>
28 #include <stdarg.h>
29 #include <stdint.h>
30 #include <stdlib.h>
31 #include <unistd.h>
32 #include <stdio.h>
33 #include <fcntl.h>
34 #include <errno.h>
35 
36 static void die (const char *fmt, ...)
37 {
38     va_list ap;
39 
40     va_start (ap, fmt);
41     fputs ("Error: ", stderr);
42     vfprintf (stderr, fmt, ap);
43 
44     if (errno != 0) {
45         fprintf (stderr, ": %s\n", strerror (errno));
46     } else {
47         fputc ('\n', stderr);
48     }
49 
50     exit (1);
51 }
52 
53 static int diskcount (void)
54 {
55     const int mib[2] = { CTL_HW, HW_DISKCOUNT };
56     int diskcount;
57     size_t len = sizeof diskcount;
58 
59     if (sysctl (mib, 2, &diskcount, &len, NULL, 0) == -1)
60         die ("sysctl(hw.diskcount)");
61 
62     return diskcount;
63 }
64 
65 static char *disknames (void)
66 {
67     const int num = diskcount ();
68     const int mib[2] = { CTL_HW, HW_DISKNAMES };
69     size_t len = 32 * num;
70     char *buffer = malloc (len);
71 
72     if (sysctl (mib, 2, buffer, &len, NULL, 0) == -1)
73         die ("sysctl(hw.disknames)");
74 
75     return buffer;
76 }
77 
78 static void printsize (uint64_t sz)
79 {
80     struct unit {
81         char sym;
82         uint64_t factor;
83     };
84 
85     const struct unit units[] = {
86         { 'P', 1ull << 50 },
87         { 'T', 1ull << 40 },
88         { 'G', 1ull << 30 },
89         { 'M', 1ull << 20 },
90         { 'K', 1ull << 10 },
91         { '0', 0 },
92     };
93 
94     char sym = 'B';
95     uint64_t factor = 1;
96 
97     for (const struct unit *u = &units[0]; u->factor; ++u) {
98         if (sz >= (u->factor * 9 / 10)) {
99             sym = u->sym;
100             factor = u->factor;
101             break;
102         }
103     }
104 
105     const unsigned scaled10 = sz * 10 / factor;
106     const unsigned scaled = sz / factor;
107     if (scaled10 >= 1000) {
108         printf (" %u", scaled);
109     } else if (scaled10 >= 100) {
110         printf ("%u.%u", scaled, scaled10 % 10);
111     } else {
112         printf (" %u.%u", scaled, scaled10 % 10);
113     }
114 
115     putchar (sym);
116 }
117 
118 static int usage (void)
119 {
120     fputs ("Usage: lsblk [-V]\n", stderr);
121     return 1;
122 }
123 
124 int main (int argc, char *argv[])
125 {
126     int option;
127     while ((option = getopt (argc, argv, ":V")) != -1) {
128         switch (option) {
129         case 'V':
130             puts ("lsblk-" VERSION);
131             return 0;
132         default:
133             return usage ();
134         }
135     }
136 
137     if (argc != optind)
138         return usage ();
139 
140     struct statfs *mounts;
141     char *names = disknames ();
142 
143     if (pledge ("stdio rpath disklabel", NULL) == -1)
144         die ("pledge()");
145 
146     const int n_mounts = getmntinfo (&mounts, MNT_NOWAIT);
147     if (n_mounts == 0)
148         die ("getmntinfo()");
149 
150     for (char *disk; (disk = strsep (&names, ",")) != NULL;) {
151         struct disklabel label;
152         char *colon, *name = NULL;
153         int fd;
154 
155         colon = strchr (disk, ':');
156         if (colon)
157             *colon = '\0';
158 
159         asprintf (&name, "/dev/%sc", disk);
160 
161         fd = open (name, O_RDONLY);
162         if (fd < 0)
163             die ("open(%s)", name);
164 
165         if (ioctl (fd, DIOCGDINFO, &label) < 0)
166             die ("ioctl(%s, IOCGDINFO)", name);
167 
168         close (fd);
169 
170         const uint32_t ss = label.d_secsize;
171 
172         char *letter = name + strlen (name) - 1;
173 
174         printf ("%s    ", disk);
175         printsize ((((uint64_t)label.d_secperunith << 32) | label.d_secperunit) * ss);
176         printf (" %s\n", label.d_typename);
177         for (uint16_t i = 0; i < label.d_npartitions; ++i) {
178             const struct partition *p = &label.d_partitions[i];
179             const uint64_t size = (((uint64_t)p->p_sizeh << 32) | p->p_size) * ss;
180             if (size == 0)
181                 continue;
182 
183             *letter = i + 'a';
184             const char *mp = "";
185             for (int i = 0; i < n_mounts; ++i) {
186                 const struct statfs *m = &mounts[i];
187                 if (!strcmp (name, m->f_mntfromname)) {
188                     mp = m->f_mntonname;
189                     break;
190                 }
191             }
192 
193             printf ("  %s%c ", disk, *letter);
194             printsize (size);
195             printf (" %s %s\n", fstypenames[p->p_fstype], mp);
196         }
197     }
198 
199     free (names);
200     return 0;
201 }