writeimg/writeimg.c

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

#ifdef GITREV
#define VERSION GITREV
#else
#define VERSION "1.0.0"
#endif

// 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"
    "   -V, --version   Print version\n"
    "\0";
// clang-format on

const char copyright[] = "Copyright (C) %s Imbus\n";

#define BUFSIZE (4 * 1024 * 1024)

typedef struct write_job {
    char  *filename;
    char  *dev_name;
    char  *buffer;
    char  *buffer2; /* For memcmp integrity checks */
    size_t bufsize;
    size_t block_size;
    char   verify_only;
} write_job_t;

write_job_t wjob = {0};

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->dev_name, O_RDWR);
    int file_fd = open(job->filename, O_RDONLY);

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

    while (!job->verify_only) {
        ssize_t read_bytes = read(file_fd, job->buffer, job->block_size);
        if (read_bytes == 0) {
            fsync(block_fd);
            break; /* Finished */
        }

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

        ssize_t written_bytes = write(block_fd, job->buffer, read_bytes);
        if (written_bytes < 0) {
            fprintf(stderr, "%s: Write error\n", job->dev_name);
            perror("Write");
            exit(EXIT_FAILURE);
        }
    }

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

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

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

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

        if (read_file == 0)
            break;

        int cmp = memcmp((const void *)job->buffer2, (const void *)job->buffer, (size_t)read_file);
        if (0 != cmp) {
            fprintf(stderr, "WARNING: 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'},
    {NULL, 0, 0, 0},
};

int main(int argc, char *argv[]) {
    printf("%s version %s\n", basename(argv[0]), VERSION);
    printf("Build date: %s\n", BLDDATE);
    fprintf(stdout, copyright, CR_YEAR);
    printf("\n");

    /* 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);

    int c = {0};
    while ((c = getopt_long(argc, argv, "vd:hV", longopts, 0)) != -1) {
        switch (c) {
        case 'v': ++wjob.verify_only; continue;
        case 'd': wjob.dev_name = optarg; continue;
        case 'h': break;
        case 'V': exit(EXIT_SUCCESS);
        }
        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");
        // printf("%s\n", help);
        exit(EXIT_FAILURE);
    }

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

    if (NULL == wjob.dev_name) {
        printf("Device required...\n");
        exit(EXIT_FAILURE);
    }

    printf("Writing %s to %s\n", wjob.filename, wjob.dev_name);

    wjob.buffer = malloc(BUFSIZE);
    wjob.buffer2 = malloc(BUFSIZE);

    wjob.bufsize = BUFSIZE;
    wjob.block_size = BUFSIZE;

    perform_write(&wjob);

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

    printf("\nOkay!\n");
    exit(0);
    return 0;
}