writeimg/writeimg.c

#include "crc32.h"
#include <assert.h>
#include <fcntl.h>
#include <getopt.h>
#include <libgen.h>
#include <limits.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#if __linux__
#include <linux/fs.h>  /* IOTCL flush number */
#include <sys/ioctl.h> /* IOCTL */
#endif

#ifndef GITREV
#define GITREV "unknown"
#endif

#ifndef VERSION
#define VERSION "unknown"
#endif

#ifndef CR_YEAR
#define CR_YEAR "2026"
#endif

#ifndef BLOCKSIZE
#define BLOCKSIZE (1024 * 1024)
#endif

#define WI_VERIFY (1 << 0)
#define WI_WRITE  (1 << 1)
#define WI_ASK    (1 << 2)

#define BYTES_TO_MIB(bts) ((double)bts / (1024 * 1024))

#define BAR_WIDTH 50

void print_progress(int current, int total) {
    static int last = INT_MAX;
    float      fraction = (float)current / total;
    int        filled = (int)(fraction * BAR_WIDTH);

    if (filled == last)
        return; /* Avoid unnecessary io/flushes */

    last = filled;

    printf("\r[");
    for (int i = 0; i < BAR_WIDTH; i++) {
        if (i < filled)
            putchar('#');
        else
            putchar(' ');
    }
    printf("] %3d%%", (int)(fraction * 100));
    fflush(stdout);
}

// clang-format off
const char help[] =
    "Usage:\n"
    "   writeimg [-v] -d <device> <file.img>\n"
    "\n"
    "Args:\n"
    "   <file.img>          Binary image file\n"
    "   -v                  Verify only\n"
    "   -d device           Target block device\n"
    "   -h, --help          Print this help message\n"
    "   -n, --noconfirm     Do not ask for permission\n"
    "   -V, --version       Print version\n"
    "\0";
// clang-format on

const char copyright[] = "Copyright (C) %s Imbus, BSD-2-Clause\n";

struct write_job {
    char  *iname;
    char  *oname;
    char  *buffer;
    char  *buffer2; /* For memcmp integrity checks */
    size_t bufsize;
    size_t block_size;
    size_t total_bytes;
    char   flags;
} wjob = {0};

typedef struct write_job write_job_t;

void int_handler(int signum) {
    fprintf(stderr, "\nInterrupted: %s\n", strsignal(signum));

    if (wjob.buffer)
        free(wjob.buffer);

    if (wjob.buffer2)
        free(wjob.buffer2);

    exit(EXIT_FAILURE);
}

int perform_write(write_job_t *job) {
    int block_fd = open(job->oname, O_RDWR);
    int file_fd = open(job->iname, O_RDONLY);

    assert(block_fd >= 0);
    assert(file_fd >= 0);

    /* TODO: Checks */
    assert(job->bufsize >= job->block_size);

    int    crc = crc32_init();
    size_t b_written = 0;

    while (42) {
        fsync(block_fd);
        ssize_t read_bytes = read(file_fd, job->buffer, job->block_size);
        assert(read_bytes >= 0);

        if (read_bytes == 0) {
            crc = crc32_finalize(crc);

            if (job->flags & WI_WRITE) {
                printf("\nWriting done...\n");
                assert(job->total_bytes == b_written);
            } else
                assert(0 == b_written);

            break; /* Finished */
        }

        if (read_bytes < 0) {
            fprintf(stderr, "%s: Read error\n", job->iname);
            perror("Read");
            exit(EXIT_FAILURE);
        }

        crc = crc32_update(crc, job->buffer, read_bytes);

        if (job->flags & WI_WRITE) {
            ssize_t written_bytes = write(block_fd, job->buffer, read_bytes);
            if (written_bytes < 0) {
                fprintf(stderr, "%s: Write error\n", job->oname);
                perror("Write");
                exit(EXIT_FAILURE);
            }
            print_progress(b_written += written_bytes, job->total_bytes);
        } /* Else maybe give some helpful insights? */
    }

    lseek(block_fd, 0, SEEK_SET);
    lseek(file_fd, 0, SEEK_SET);

    memset(job->buffer, 0, BLOCKSIZE);
    memset(job->buffer2, 0, BLOCKSIZE);

    /* This is essentially $ blockdev --flushbufs */
    ioctl(block_fd, BLKFLSBUF);

    int crc_back = crc32_init();
    b_written = 0;

    while (1) {
        ssize_t read_file = read(file_fd, job->buffer2, job->block_size);
        ssize_t read_block = read(block_fd, job->buffer, read_file);

        assert(read_block == read_file);
        assert(read_block >= 0);
        assert(read_file >= 0);

        if (read_file == 0) {
            crc_back = crc32_finalize(crc_back);
            assert(crc_back == crc);
            assert(job->total_bytes == b_written);
            printf("\nVerification done!\n");
            break;
        }

        crc_back = crc32_update(crc_back, job->buffer, read_file);
        print_progress(b_written += read_block, job->total_bytes);

        int cmp = memcmp((const void *)job->buffer2, (const void *)job->buffer, (size_t)read_file);
        if (0 != cmp) {
            fflush(stdout);
            fprintf(stderr, "\nWARNING: File did not verify correctly!\n");
            exit(EXIT_FAILURE);
        }
    }

    close(block_fd);
    close(file_fd);
    return 1;
}

static const struct option longopts[] = {
    {"help", 0, 0, 'h'},
    {"version", 0, 0, 'V'},
    {"device", 1, 0, 'd'},
    {"noconfirm", 0, 0, 'n'},
    {NULL, 0, 0, 0},
};

int main(int argc, char *argv[]) {
    printf("%s %s, Rev. %s\n", "WriteIMG", VERSION, GITREV);

    /* Line buffering, system allocated */
    setvbuf(stdout, NULL, _IOLBF, 0);
    setvbuf(stderr, NULL, _IOLBF, 0);

    /* Hang these up to alert the user that we were interrupted */
    signal(SIGINT, int_handler);
    signal(SIGHUP, int_handler);
    signal(SIGTERM, int_handler);

    wjob.flags = WI_VERIFY | WI_WRITE | WI_ASK;

    int c = {0};
    while ((c = getopt_long(argc, argv, "vd:hnV", longopts, 0)) != -1) {
        switch (c) {
        case 'v':
            wjob.flags |= WI_VERIFY;
            wjob.flags &= ~WI_WRITE;
            continue;
        case 'd': wjob.oname = optarg; continue;
        case 'h': break;
        case 'n': wjob.flags &= ~WI_ASK; continue;
        case 'V': exit(EXIT_SUCCESS);
        }
        printf("In honor of SwePwnage - the OG disk destroyer\n");
        printf(copyright, CR_YEAR);
#ifdef BLDDATE
        printf("Build date: %s\n", BLDDATE);
#endif
        printf("\n");
        printf("%s\n", help);
        exit(EXIT_SUCCESS);
    }

    argc -= optind;
    argv += optind;

    /* argv[0] should now be the image file*/

    if (argc != 1) {
        printf("You need to specify an image file.\n");
        exit(EXIT_FAILURE);
    }

    wjob.iname = argv[0];
    struct stat file_stat = {0};
    if (0 != stat(wjob.iname, &file_stat)) {
        printf("File does not exist...\n");
        exit(EXIT_FAILURE);
    }

    if (NULL == wjob.oname) {
        printf("You need to specify a device.\n");
        exit(EXIT_FAILURE);
    }

    if (0 != strncmp(wjob.oname, "/dev/", 4)) {
        printf("\"%s\" does not appear to be a block device...\n", wjob.oname);
        exit(EXIT_FAILURE);
    }

    uint64_t device_size;
    {
        /* Seems to be the cleanest way to check for write perm on a blockdev */
        int fd = open(wjob.oname, O_WRONLY);
        if (fd < 0) {
            printf("Cannot write to \"%s\", do you have write permissions?\n", wjob.oname);
            close(fd);
            exit(EXIT_FAILURE);
        }

#ifdef __GLIBC__
        if (ioctl(fd, BLKGETSIZE64, &device_size) < 0) {
#else /* With musl, ioctl's take ints, kheaders provide unsigned longs. Passes tests. */
        if (ioctl(fd, (int)BLKGETSIZE64, &device_size) < 0) {
#endif
            perror("ioctl");
            close(fd);
            exit(EXIT_FAILURE);
        }

        close(fd);
    }

    wjob.total_bytes = file_stat.st_size;
    assert(file_stat.st_size >= 0);

    if (wjob.flags & WI_WRITE)
        printf("Writing \"%s\" (%.1f MiB) to \"%s\" (%.1f MiB)\n",
               basename(wjob.iname),
               BYTES_TO_MIB(wjob.total_bytes),
               wjob.oname,
               BYTES_TO_MIB(device_size));

    if ((wjob.flags & WI_ASK) && (wjob.flags & WI_WRITE)) {
        printf("Is this okay? (y/N): ");
        fflush(stdout);
        if ('y' != getchar()) {
            printf("Aborting...\n");
            exit(EXIT_SUCCESS);
        }
    }

    wjob.buffer = malloc(BLOCKSIZE);
    wjob.buffer2 = malloc(BLOCKSIZE);
    assert(wjob.buffer);
    assert(wjob.buffer2);

    wjob.bufsize = BLOCKSIZE;
    wjob.block_size = BLOCKSIZE;

    perform_write(&wjob);

    if (wjob.buffer)
        free(wjob.buffer);
    if (wjob.buffer2)
        free(wjob.buffer2);

    printf("\n%.1f MiB's verified.\nAll good!\n", BYTES_TO_MIB(wjob.total_bytes));

    exit(EXIT_SUCCESS);
}