- Add GPT header/table recovery command

	- Minor related cleanup in add command

Approved by:	marcel

2 files changed, 92 insertions, 4 deletions

diff --git a/sbin/gpt/add.c b/sbin/gpt/add.c
index ee6454d..f0ce9a8 100644
--- a/sbin/gpt/add.c
+++ b/sbin/gpt/add.c
@@ -65,13 +65,14 @@ add(int fd)
 
 	gpt = map_find(MAP_TYPE_PRI_GPT_HDR);
 	if (gpt == NULL) {
-		warnx("%s: error: device does not contain a GPT", device_name);
+		warnx("%s: error: no primary GPT header; run create or recover",
+		    device_name);
 		return;
 	}
 
 	tpg = map_find(MAP_TYPE_SEC_GPT_HDR);
 	if (tpg == NULL) {
-		warnx("%s: error: no secundary table; run recover",
+		warnx("%s: error: no secondary GPT header; run recover",
 		    device_name);
 		return;
 	}
diff --git a/sbin/gpt/recover.c b/sbin/gpt/recover.c
index 040e0c6..d4b8418 100644
--- a/sbin/gpt/recover.c
+++ b/sbin/gpt/recover.c
@@ -51,10 +51,97 @@ usage_recover(void)
 }
 
 static void
-recover(int fd __unused)
+recover(int fd)
 {
+	off_t last;
+	map_t *gpt, *tpg;
+	map_t *tbl, *lbt;
+	struct gpt_hdr *hdr;
 
-	/* TODO */
+	if (map_find(MAP_TYPE_MBR) != NULL) {
+		warnx("%s: error: device contains a MBR", device_name);
+		return;
+	}
+
+	gpt = map_find(MAP_TYPE_PRI_GPT_HDR);
+	tpg = map_find(MAP_TYPE_SEC_GPT_HDR);
+	tbl = map_find(MAP_TYPE_PRI_GPT_TBL);
+	lbt = map_find(MAP_TYPE_SEC_GPT_TBL);
+
+	if (gpt == NULL && tpg == NULL) {
+		warnx("%s: no primary or secondary GPT headers, can't recover",
+		    device_name);
+		return;
+	}
+	if (tbl == NULL && lbt == NULL) {
+		warnx("%s: no primary or secondary GPT tables, can't recover",
+		    device_name);
+		return;
+	}
+
+	last = mediasz / secsz - 1LL;
+
+	if (tbl != NULL && lbt == NULL) {
+		lbt = map_add(last - tbl->map_size, tbl->map_size,
+		    MAP_TYPE_SEC_GPT_TBL, tbl->map_data);
+		if (lbt == NULL) {
+			warnx("%s: adding secondary GPT table failed",
+			    device_name);
+			return;
+		}
+		gpt_write(fd, lbt);
+		warnx("%s: recovered secondary GPT table from primary",
+		    device_name);
+	} else if (tbl == NULL && lbt != NULL) {
+		tbl = map_add(2LL, lbt->map_size, MAP_TYPE_PRI_GPT_TBL,
+		    lbt->map_data);
+		if (tbl == NULL) {
+			warnx("%s: adding primary GPT table failed",
+			    device_name);
+			return;
+		}
+		gpt_write(fd, tbl);
+		warnx("%s: recovered primary GPT table from secondary",
+		    device_name);
+	}
+
+	if (gpt != NULL && tpg == NULL) {
+		tpg = map_add(last, 1LL, MAP_TYPE_SEC_GPT_HDR,
+		    calloc(1, secsz));
+		if (tpg == NULL) {
+			warnx("%s: adding secondary GPT header failed",
+			    device_name);
+			return;
+		}
+		memcpy(tpg->map_data, gpt->map_data, secsz);
+		hdr = tpg->map_data;
+		hdr->hdr_lba_self = tpg->map_start;
+		hdr->hdr_lba_alt = gpt->map_start;
+		hdr->hdr_lba_table = lbt->map_start;
+		hdr->hdr_crc_self = 0;
+		hdr->hdr_crc_self = crc32(hdr, hdr->hdr_size);
+		gpt_write(fd, tpg);
+		warnx("%s: recovered secondary GPT header from primary",
+		    device_name);
+	} else if (gpt == NULL && tpg != NULL) {
+		gpt = map_add(1LL, 1LL, MAP_TYPE_PRI_GPT_HDR,
+		    calloc(1, secsz));
+		if (gpt == NULL) {
+			warnx("%s: adding primary GPT header failed",
+			    device_name);
+			return;
+		}
+		memcpy(gpt->map_data, tpg->map_data, secsz);
+		hdr = gpt->map_data;
+		hdr->hdr_lba_self = gpt->map_start;
+		hdr->hdr_lba_alt = tpg->map_start;
+		hdr->hdr_lba_table = tbl->map_start;
+		hdr->hdr_crc_self = 0;
+		hdr->hdr_crc_self = crc32(hdr, hdr->hdr_size);
+		gpt_write(fd, gpt);
+		warnx("%s: recovered primary GPT header from secondary",
+		    device_name);
+	}
 }
 
 int