diff --git a/host/limine-deploy.c b/host/limine-deploy.c index a62ece59..f171dca9 100644 --- a/host/limine-deploy.c +++ b/host/limine-deploy.c @@ -257,6 +257,108 @@ static bool device_cache_block(uint64_t block) { return true; } +struct undeploy_data { + void *data; + uint64_t loc; + uint64_t count; +}; + +#define UNDEPLOY_DATA_MAX 256 + +static bool undeploying = false; +static struct undeploy_data undeploy_data[UNDEPLOY_DATA_MAX]; +static struct undeploy_data undeploy_data_rev[UNDEPLOY_DATA_MAX]; +static uint64_t undeploy_data_i = 0; +static const char *undeploy_file = NULL; + +static void reverse_undeploy_data(void) { + for (size_t i = 0, j = undeploy_data_i - 1; i < undeploy_data_i; i++, j--) { + undeploy_data_rev[j] = undeploy_data[i]; + } + + memcpy(undeploy_data, undeploy_data_rev, undeploy_data_i * sizeof(struct undeploy_data)); +} + +static void free_undeploy_data(void) { + for (size_t i = 0; i < undeploy_data_i; i++) { + free(undeploy_data[i].data); + } +} + +static bool store_undeploy_data(const char *filename) { + fprintf(stderr, "Storing undeploy data to file: `%s`...\n", filename); + + FILE *udfile = fopen(filename, "wb"); + if (udfile == NULL) { + goto error; + } + + if (fwrite(&undeploy_data_i, sizeof(uint64_t), 1, udfile) != 1) { + goto error; + } + + for (size_t i = 0; i < undeploy_data_i; i++) { + if (fwrite(&undeploy_data[i].loc, sizeof(uint64_t), 1, udfile) != 1) { + goto error; + } + if (fwrite(&undeploy_data[i].count, sizeof(uint64_t), 1, udfile) != 1) { + goto error; + } + if (fwrite(undeploy_data[i].data, undeploy_data[i].count, 1, udfile) != 1) { + goto error; + } + } + + fclose(udfile); + return true; + +error: + perror("ERROR"); + if (udfile != NULL) { + fclose(udfile); + } + return false; +} + +static bool load_undeploy_data(const char *filename) { + fprintf(stderr, "Loading undeploy data from file: `%s`...\n", filename); + + FILE *udfile = fopen(filename, "rb"); + if (udfile == NULL) { + goto error; + } + + if (fread(&undeploy_data_i, sizeof(uint64_t), 1, udfile) != 1) { + goto error; + } + + for (size_t i = 0; i < undeploy_data_i; i++) { + if (fread(&undeploy_data[i].loc, sizeof(uint64_t), 1, udfile) != 1) { + goto error; + } + if (fread(&undeploy_data[i].count, sizeof(uint64_t), 1, udfile) != 1) { + goto error; + } + undeploy_data[i].data = malloc(undeploy_data[i].count); + if (undeploy_data[i].data == NULL) { + goto error; + } + if (fread(undeploy_data[i].data, undeploy_data[i].count, 1, udfile) != 1) { + goto error; + } + } + + fclose(udfile); + return true; + +error: + perror("ERROR"); + if (udfile != NULL) { + fclose(udfile); + } + return false; +} + static bool _device_read(void *_buffer, uint64_t loc, size_t count) { uint8_t *buffer = _buffer; uint64_t progress = 0; @@ -281,6 +383,32 @@ static bool _device_read(void *_buffer, uint64_t loc, size_t count) { } static bool _device_write(const void *_buffer, uint64_t loc, size_t count) { + if (undeploying) { + goto skip_save; + } + + if (undeploy_data_i >= UNDEPLOY_DATA_MAX) { + fprintf(stderr, "Internal error: Too many undeploy data entries!\n"); + return false; + } + + struct undeploy_data *ud = &undeploy_data[undeploy_data_i]; + + ud->data = malloc(count); + if (ud->data == NULL) { + fprintf(stderr, "ERROR: Memory allocation failure.\n"); + return false; + } + + if (!_device_read(ud->data, loc, count)) { + fprintf(stderr, "ERROR: Device read failure.\n"); + return false; + } + + ud->loc = loc; + ud->count = count; + +skip_save:; const uint8_t *buffer = _buffer; uint64_t progress = 0; while (progress < count) { @@ -301,9 +429,43 @@ static bool _device_write(const void *_buffer, uint64_t loc, size_t count) { progress += chunk; } + if (!undeploying) { + undeploy_data_i++; + } return true; } +static void undeploy(void) { + undeploying = true; + + cache_state = CACHE_CLEAN; + cached_block = (uint64_t)-1; + + for (size_t i = 0; i < undeploy_data_i; i++) { + struct undeploy_data *ud = &undeploy_data[i]; + bool retry = false; + while (!_device_write(ud->data, ud->loc, ud->count)) { + if (retry) { + fprintf(stderr, "ERROR: Undeploy data index %zu failed to write. Undeploy may be incomplete!\n", i); + break; + } + fprintf(stderr, "Warning: Undeploy data index %zu failed to write, retrying...\n", i); + if (!device_flush_cache()) { + fprintf(stderr, "ERROR: Device cache flush failure. Undeploy may be incomplete!\n"); + } + cache_state = CACHE_CLEAN; + cached_block = (uint64_t)-1; + retry = true; + } + } + + if (!device_flush_cache()) { + fprintf(stderr, "ERROR: Device cache flush failure. Undeploy may be incomplete!\n"); + } + + fprintf(stderr, "Undeploy data restored successfully. Limine undeployed!\n"); +} + #define device_read(BUFFER, LOC, COUNT) \ do { \ if (!_device_read(BUFFER, LOC, COUNT)) \ @@ -326,6 +488,7 @@ static void usage(const char *name) { int main(int argc, char *argv[]) { int ok = EXIT_FAILURE; int force_mbr = 0; + bool undeploy_mode = false; const uint8_t *bootloader_img = binary_limine_hdd_bin_data; size_t bootloader_file_size = sizeof(binary_limine_hdd_bin_data); uint8_t orig_mbr[70], timestamp[6]; @@ -343,9 +506,24 @@ int main(int argc, char *argv[]) { for (int i = 1; i < argc; i++) { if (strcmp(argv[i], "--force-mbr") == 0) { // TODO: add to usage if (force_mbr) { - puts("Warning: --force-mbr already set"); + fprintf(stderr, "Warning: --force-mbr already set.\n"); } force_mbr = 1; + } else if (strcmp(argv[i], "--undeploy") == 0) { + if (undeploy_mode) { + fprintf(stderr, "Warning: --undeploy already set.\n"); + } + undeploy_mode = true; + } else if (memcmp(argv[i], "--undeploy-data-file=", 21) == 0) { + if (undeploy_file != NULL) { + fprintf(stderr, "Warning: --undeploy-data-file already set. Overriding...\n"); + } + undeploy_file = argv[i] + 21; + if (strlen(undeploy_file) == 0) { + fprintf(stderr, "ERROR: Undeploy data file has a zero-length name!\n"); + undeploy_file = NULL; + goto cleanup; + } } else { if (device != NULL) { // [GPT partition index] part_ndx = argv[i]; // TODO: Make this non-positional? @@ -365,6 +543,22 @@ int main(int argc, char *argv[]) { if (!device_init()) goto cleanup; + if (undeploy_mode) { + if (undeploy_file == NULL) { + fprintf(stderr, "ERROR: Undeploy mode set but no --undeploy-data-file=... passed.\n"); + goto undeploy_mode_cleanup; + } + + if (!load_undeploy_data(undeploy_file)) { + goto undeploy_mode_cleanup; + } + + undeploy(); + + ok = EXIT_SUCCESS; + goto undeploy_mode_cleanup; + } + // Probe for GPT and logical block size int gpt = 0; struct gpt_table_header gpt_header; @@ -379,8 +573,8 @@ int main(int argc, char *argv[]) { fprintf(stderr, "Deploying to GPT. Logical block size of %" PRIu64 " bytes.\n", lb_guesses[i]); } else { - memset(&gpt_header, 0, sizeof(struct gpt_table_header)); - device_write(&gpt_header, lb_guesses[i], sizeof(struct gpt_table_header)); + fprintf(stderr, "Device has a valid GPT, refusing to force MBR.\n"); + goto cleanup; } break; } @@ -737,6 +931,15 @@ int main(int argc, char *argv[]) { ok = EXIT_SUCCESS; cleanup: + reverse_undeploy_data(); + if (ok != EXIT_SUCCESS) { + // If we failed, attempt to reverse deploy process + undeploy(); + } else if (undeploy_file != NULL) { + store_undeploy_data(undeploy_file); + } +undeploy_mode_cleanup: + free_undeploy_data(); if (cache) free(cache); if (device != NULL)