pass-lang/src/io.c

#ifdef __unix__
#define _GNU_SOURCE
#endif

#include "io.h"

#include <errno.h>
#include <stdlib.h>
#include <string.h>

#ifdef __unix__
// This program can be trivially converted to work with only the C standard library
// at the cost of not being able to link the output file atomically.
#include <fcntl.h>
#include <libgen.h>
#include <sys/stat.h>
#include <unistd.h>
#endif

static const char* outfile_name;
FILE* infile;
FILE* outfile;

#ifdef __unix__
void open_files(const char* infile_name, const char* outfile_name_) {
    outfile_name = outfile_name_;
    // To avoid creating a corrupt or incomplete output file,
    // we operate on a temporary file and atomically link it only once compilation has succeeded.
    unlink(outfile_name);

    int infile_fd = open(infile_name, O_RDONLY);
    if (infile_fd == -1) {
        fprintf(stderr, "failed to open source file: %s\n", strerror(errno));
        exit(1);
    }
    off_t infile_len = lseek(infile_fd, 0, SEEK_END);
    if (infile_len == (off_t) -1) {
        fprintf(stderr, "failed to get length of source file: %s\n", strerror(errno));
        exit(1);
    }
    // There'll probably never be a source file large enough for this to make a difference,
    // and I *certainly* haven't profiled, but... I've always wanted to use these syscalls. :)
    posix_fadvise(infile_fd, 0, infile_len, POSIX_FADV_SEQUENTIAL);
    posix_fadvise(infile_fd, 0, infile_len, POSIX_FADV_NOREUSE);
    infile = fdopen(infile_fd, "rb");
    if (infile_fd == -1) {
        fprintf(stderr, "failed to open source file fd as file handle: %s\n", strerror(errno));
        exit(1);
    }

    int outfile_fd = open(dirname((char*) outfile_name), O_WRONLY | O_TMPFILE, S_IRWXU | S_IRWXG | S_IRWXO);
    if (outfile_fd == -1) {
        fprintf(stderr, "failed to create temporary output file: %s\n", strerror(errno));
        exit(1);
    }
    outfile = fdopen(outfile_fd, "wb");
    if (outfile == NULL) {
        fprintf(stderr, "failed to open output file fd as file handle: %s\n", strerror(errno));
        exit(1);
    }
}

void close_files(void) {
    if (fflush(outfile) != 0) {
        fprintf(stderr, "failed to flush output file: %s\n", strerror(errno));
        exit(1);
    }

    char outfile_tempname[20];
    snprintf(outfile_tempname, 20, "/proc/self/fd/%d", fileno(outfile));
    if (linkat(AT_FDCWD, outfile_tempname, AT_FDCWD, outfile_name, AT_SYMLINK_FOLLOW) == -1) {
        fprintf(stderr, "failed to link output file into file system: %s\n", strerror(errno));
        exit(1);
    }

    fclose(outfile);
    fclose(infile);
}
#else
void open_files(const char* infile_name, const char* outfile_name) {
    infile = fopen(infile_name, "rb");
    if (infile == NULL) {
        fprintf(stderr, "failed to open source file: %s\n", strerror(errno));
        exit(1);
    }

    // There is no way for us to mark the file as executable.
    // Then again, if it's not Unix, that probably doesn't matter.
    outfile = fopen(outfile_name, "wb");
    if (outfile == NULL) {
        fprintf(stderr, "failed to open output file: %s\n", strerror(errno));
        exit(1);
    }
}

void close_files(void) {
    if (fclose(outfile) != 0) {
        fprintf(stderr, "failed to close output file: %s\n", strerror(errno));
        // NOTE: ideally we'd do this on any dirty exit
        // TODO: use portable tempfiles and then just copy the entire file at the end?
        if (remove(outfile_name) != 0) {
            fprintf(stderr, "failed to remove output file, if it exists, it is corrupt: %s\n", strerror(errno));
        }
        exit(1);
    }
    fclose(infile);
}
#endif

void emit(const void* restrict ptr, size_t count) {
    fwrite(ptr, 1, count, outfile);
    if (ferror(outfile)) {
        fprintf(stderr, "failed to write to output file\n");
        exit(1);
    }
}

void emit_u8(uint8_t x) {
    emit(&x, sizeof(uint8_t));
}

void emit_u32(uint32_t x) {
    emit(&x, sizeof(uint32_t));
}

void emit_u64(uint64_t x) {
    emit(&x, sizeof(uint64_t));
}

void patch(size_t off, const void* ptr, size_t count) {
    fpos_t save;
    if (fgetpos(outfile, &save) != 0) {
        fprintf(stderr, "failed to save file position before patch: %s\n", strerror(errno));
        exit(1);
    }
    if (fseek(outfile, (long) off, SEEK_SET) != 0) {
        fprintf(stderr, "failed to set file position for patch: %s\n", strerror(errno));
        exit(1);
    }
    fwrite(ptr, 1, count, outfile);
    if (ferror(outfile) != 0) {
        fprintf(stderr, "failed to patch output file: %s\n", strerror(errno));
        exit(1);
    }
    if (fsetpos(outfile, &save) != 0) {
        fprintf(stderr, "failed to restore file position after patch: %s\n", strerror(errno));
        exit(1);
    }
}

void patch_u32(size_t off, uint32_t x) {
    patch(off, &x, sizeof(uint32_t));
}

void patch_i32(size_t off, int32_t x) {
    patch_u32(off, (uint32_t) x);
}