Aren
/
rulimine
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

limine-deploy: Add undeploy data save and undeploy mode. Addresses #197

This commit is contained in:
mintsuki 2022-07-14 14:27:19 +02:00
parent c50e02ec2e
commit dc0d9f8be8
1 changed files with 206 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

209
host/limine-deploy.c
View File

@ -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)