toolbox: enable mkswap to work on block devices

Change-Id: I9b71ca3fa09ffbe12dcc4d04229edab80851db4d
Signed-off-by: Andrew Boie <andrew.p.boie@intel.com>

1 files changed, 39 insertions, 41 deletions

diff --git a/toolbox/mkswap.c b/toolbox/mkswap.c
index 0904152..ad66353 100644
--- a/toolbox/mkswap.c
+++ b/toolbox/mkswap.c
@@ -1,11 +1,13 @@
+#include <fcntl.h>
+#include <linux/fs.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <unistd.h>
+#include <sys/stat.h>
 #include <sys/swap.h>
 #include <sys/types.h>
-#include <sys/stat.h>
-#include <fcntl.h>
 
-/* XXX This needs to be obtained from kernel headers. See b/9336527 */
+/* This is not in a uapi header. */
 struct linux_swap_header {
     char            bootbits[1024]; /* Space for disklabel etc. */
     uint32_t        version;
@@ -23,71 +25,67 @@ struct linux_swap_header {
 
 int mkswap_main(int argc, char **argv)
 {
-    int err = 0;
-    int fd;
-    ssize_t len;
-    off_t swap_size;
-    int pagesize;
-    struct linux_swap_header sw_hdr;
-
     if (argc != 2) {
         fprintf(stderr, "Usage: %s <filename>\n", argv[0]);
-        return -EINVAL;
+        return EXIT_FAILURE;
     }
 
-    fd = open(argv[1], O_WRONLY);
+    int fd = open(argv[1], O_RDWR);
     if (fd < 0) {
-        err = errno;
-        fprintf(stderr, "Cannot open %s\n", argv[1]);
-        return err;
+        fprintf(stderr, "Cannot open %s: %s\n", argv[1], strerror(errno));
+        return EXIT_FAILURE;
     }
 
-    pagesize = getpagesize();
     /* Determine the length of the swap file */
-    swap_size = lseek(fd, 0, SEEK_END);
-    if (swap_size < MIN_PAGES * pagesize) {
-        fprintf(stderr, "Swap file needs to be at least %dkB\n",
-            (MIN_PAGES * pagesize) >> 10);
-        err = -ENOSPC;
-        goto err;
+    off64_t swap_size;
+    struct stat sb;
+    if (fstat(fd, &sb)) {
+        fprintf(stderr, "Couldn't fstat file: %s\n", strerror(errno));
+        return EXIT_FAILURE;
+    }
+    if (S_ISBLK(sb.st_mode)) {
+        if (ioctl(fd, BLKGETSIZE64, &swap_size) < 0) {
+            fprintf(stderr, "Couldn't determine block device size: %s\n", strerror(errno));
+            return EXIT_FAILURE;
+        }
+    } else {
+        swap_size = sb.st_size;
     }
-    if (lseek(fd, 0, SEEK_SET)) {
-        err = errno;
-        fprintf(stderr, "Can't seek to the beginning of the file\n");
-        goto err;
+
+    int pagesize = getpagesize();
+    if (swap_size < MIN_PAGES * pagesize) {
+        fprintf(stderr, "Swap file needs to be at least %d KiB\n", (MIN_PAGES * pagesize) >> 10);
+        return EXIT_FAILURE;
     }
 
+    struct linux_swap_header sw_hdr;
     memset(&sw_hdr, 0, sizeof(sw_hdr));
     sw_hdr.version = 1;
     sw_hdr.last_page = (swap_size / pagesize) - 1;
 
-    len = write(fd, &sw_hdr, sizeof(sw_hdr));
+    ssize_t len = write(fd, &sw_hdr, sizeof(sw_hdr));
     if (len != sizeof(sw_hdr)) {
-        err = errno;
-        fprintf(stderr, "Failed to write swap header into %s\n", argv[1]);
-        goto err;
+        fprintf(stderr, "Failed to write swap header into %s: %s\n", argv[1], strerror(errno));
+        return EXIT_FAILURE;
     }
 
     /* Write the magic header */
     if (lseek(fd, pagesize - MAGIC_SWAP_HEADER_LEN, SEEK_SET) < 0) {
-        err = errno;
-        fprintf(stderr, "Failed to seek into %s\n", argv[1]);
-        goto err;
+        fprintf(stderr, "Failed to seek into %s: %s\n", argv[1], strerror(errno));
+        return EXIT_FAILURE;
     }
 
     len = write(fd, MAGIC_SWAP_HEADER, MAGIC_SWAP_HEADER_LEN);
     if (len != MAGIC_SWAP_HEADER_LEN) {
-        err = errno;
-        fprintf(stderr, "Failed to write magic swap header into %s\n", argv[1]);
-        goto err;
+        fprintf(stderr, "Failed to write magic swap header into %s: %s\n", argv[1], strerror(errno));
+        return EXIT_FAILURE;
     }
 
     if (fsync(fd) < 0) {
-        err = errno;
-        fprintf(stderr, "Failed to sync %s\n", argv[1]);
-        goto err;
+        fprintf(stderr, "Failed to sync %s: %s\n", argv[1], strerror(errno));
+        return EXIT_FAILURE;
     }
-err:
+
     close(fd);
-    return err;
+    return EXIT_SUCCESS;
 }