:: commit 71025a497ad9705c13dc17830afdcddc39909f9a
limine-deploy: Add undeploy data save and undeploy mode. Addresses #197
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)
