Module Name:    src
Committed By:   christos
Date:           Wed Dec  2 22:32:04 UTC 2015

Modified Files:
        src/sbin/gpt: restore.c

Log Message:
- factor out into smaller separate functions
- fix signed/unsigned confusion
- do proper write checks
- fix some memory leaks


To generate a diff of this commit:
cvs rdiff -u -r1.13 -r1.14 src/sbin/gpt/restore.c

Please note that diffs are not public domain; they are subject to the
copyright notices on the relevant files.

Modified files:

Index: src/sbin/gpt/restore.c
diff -u src/sbin/gpt/restore.c:1.13 src/sbin/gpt/restore.c:1.14
--- src/sbin/gpt/restore.c:1.13	Wed Dec  2 07:36:53 2015
+++ src/sbin/gpt/restore.c	Wed Dec  2 17:32:04 2015
@@ -33,7 +33,7 @@
 __FBSDID("$FreeBSD: src/sbin/gpt/create.c,v 1.11 2005/08/31 01:47:19 marcel Exp $");
 #endif
 #ifdef __RCSID
-__RCSID("$NetBSD: restore.c,v 1.13 2015/12/02 12:36:53 christos Exp $");
+__RCSID("$NetBSD: restore.c,v 1.14 2015/12/02 22:32:04 christos Exp $");
 #endif
 
 #include <sys/types.h>
@@ -76,6 +76,124 @@ struct gpt_cmd c_restore = {
 	return -1;				\
 }
 
+#define prop_uint(a) prop_number_unsigned_integer_value(a)
+
+static int
+restore_mbr(gpt_t gpt, struct mbr *mbr, prop_dictionary_t mbr_dict, off_t last)
+{
+	unsigned int i;
+	prop_number_t propnum;
+	struct mbr_part *part;
+
+	propnum = prop_dictionary_get(mbr_dict, "index");
+	PROP_ERR(propnum);
+
+	i = prop_number_integer_value(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "flag");
+	PROP_ERR(propnum);
+	part = &mbr->mbr_part[i];
+
+	part->part_flag = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "start_head");
+	PROP_ERR(propnum);
+	part->part_shd = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "start_sector");
+	PROP_ERR(propnum);
+	part->part_ssect = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "start_cylinder");
+	PROP_ERR(propnum);
+	part->part_scyl = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "type");
+	PROP_ERR(propnum);
+	part->part_typ = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "end_head");
+	PROP_ERR(propnum);
+	part->part_ehd = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "end_sector");
+	PROP_ERR(propnum);
+	part->part_esect = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "end_cylinder");
+	PROP_ERR(propnum);
+	part->part_ecyl = prop_uint(propnum);
+	propnum = prop_dictionary_get(mbr_dict, "lba_start_low");
+	PROP_ERR(propnum);
+	part->part_start_lo = htole16(prop_uint(propnum));
+	propnum = prop_dictionary_get(mbr_dict, "lba_start_high");
+	PROP_ERR(propnum);
+	part->part_start_hi = htole16(prop_uint(propnum));
+	/* adjust PMBR size to size of device */
+	if (part->part_typ == MBR_PTYPE_PMBR) {
+		if (last > 0xffffffff) {
+			mbr->mbr_part[0].part_size_lo = htole16(0xffff);
+			mbr->mbr_part[0].part_size_hi = htole16(0xffff);
+		} else {
+			mbr->mbr_part[0].part_size_lo = htole16(last);
+			mbr->mbr_part[0].part_size_hi = htole16(last >> 16);
+		}
+	} else {
+		propnum = prop_dictionary_get(mbr_dict, "lba_size_low");
+		PROP_ERR(propnum);
+		part->part_size_lo = htole16(prop_uint(propnum));
+		propnum = prop_dictionary_get(mbr_dict, "lba_size_high");
+		PROP_ERR(propnum);
+		part->part_size_hi = htole16(prop_uint(propnum));
+	}
+	return 0;
+}
+
+static int
+restore_ent(gpt_t gpt, prop_dictionary_t gpt_dict, void *secbuf, u_int gpt_size,
+    u_int entries)
+{
+	unsigned int i;
+	struct gpt_ent ent;
+	const char *s;
+	prop_string_t propstr;
+	prop_number_t propnum;
+
+	memset(&ent, 0, sizeof(ent));
+	propstr = prop_dictionary_get(gpt_dict, "type");
+	PROP_ERR(propstr);
+	s = prop_string_cstring_nocopy(propstr);
+	if (gpt_uuid_parse(s, ent.ent_type) != 0) {
+		gpt_warnx(gpt, "%s: not able to convert to an UUID", s);
+		return -1;
+	}
+	propstr = prop_dictionary_get(gpt_dict, "guid");
+	PROP_ERR(propstr);
+	s = prop_string_cstring_nocopy(propstr);
+	if (gpt_uuid_parse(s, ent.ent_guid) != 0) {
+		gpt_warnx(gpt, "%s: not able to convert to an UUID", s);
+		return -1;
+	}
+	propnum = prop_dictionary_get(gpt_dict, "start");
+	PROP_ERR(propnum);
+	ent.ent_lba_start = htole64(prop_uint(propnum));
+	propnum = prop_dictionary_get(gpt_dict, "end");
+	PROP_ERR(propnum);
+	ent.ent_lba_end = htole64(prop_uint(propnum));
+	propnum = prop_dictionary_get(gpt_dict, "attributes");
+	PROP_ERR(propnum);
+	ent.ent_attr = htole64(prop_uint(propnum));
+	propstr = prop_dictionary_get(gpt_dict, "name");
+	if (propstr != NULL) {
+		s = prop_string_cstring_nocopy(propstr);
+		utf8_to_utf16((const uint8_t *)s, ent.ent_name,
+			__arraycount(ent.ent_name));
+	}
+	propnum = prop_dictionary_get(gpt_dict, "index");
+	PROP_ERR(propnum);
+	i = prop_number_integer_value(propnum);
+	if (i > entries) {
+		gpt_warnx(gpt, "Entity index out of bounds %u > %u\n",
+		    i, entries);
+		return -1;
+	}
+	memcpy((char *)secbuf + gpt->secsz + ((i - 1) * sizeof(ent)),
+	    &ent, sizeof(ent));
+	return 0;
+}
+
 static int
 restore(gpt_t gpt)
 {
@@ -84,17 +202,17 @@ restore(gpt_t gpt)
 	map_t map;
 	struct mbr *mbr;
 	struct gpt_hdr *hdr;
-	struct gpt_ent ent;
-	unsigned int i;
+	unsigned int i, gpt_size;
 	prop_dictionary_t props, gpt_dict, mbr_dict, type_dict;
 	prop_object_iterator_t propiter;
 	prop_data_t propdata;
 	prop_array_t mbr_array, gpt_array;
 	prop_number_t propnum;
 	prop_string_t propstr;
-	int entries, gpt_size;
+	unsigned int entries;
 	const char *s;
-	void *secbuf;
+	void *secbuf = NULL;;
+	int rv = -1;
 
 	last = gpt->mediasz / gpt->secsz - 1LL;
 
@@ -143,7 +261,7 @@ restore(gpt_t gpt)
 
 	propnum = prop_dictionary_get(gpt_dict, "entries");
 	PROP_ERR(propnum);
-	entries = prop_number_integer_value(propnum);
+	entries = prop_uint(propnum);
 	gpt_size = entries * sizeof(struct gpt_ent) / gpt->secsz;
 	if (gpt_size * sizeof(struct gpt_ent) % gpt->secsz)
 		gpt_size++;
@@ -153,8 +271,7 @@ restore(gpt_t gpt)
 	s = prop_string_cstring_nocopy(propstr);
 	if (gpt_uuid_parse(s, gpt_guid) != 0) {
 		gpt_warnx(gpt, "%s: not able to convert to an UUID", s);
-		// XXX: leak
-		return -1;
+		goto out;
 	}
 	firstdata = gpt_size + 2;		/* PMBR and GPT header */
 	lastdata = last - gpt_size - 1;		/* alt. GPT table and header */
@@ -171,51 +288,51 @@ restore(gpt_t gpt)
 		s = prop_string_cstring_nocopy(propstr);
 		if (gpt_uuid_parse(s, uuid) != 0) {
 			gpt_warnx(gpt, "%s: not able to convert to an UUID", s);
-			return -1;
+			goto out;
 		}
 		if (gpt_uuid_is_nil(uuid))
 			continue;
 		propnum = prop_dictionary_get(gpt_dict, "start");
 		PROP_ERR(propnum);
-		gpe_start = prop_number_unsigned_integer_value(propnum);
+		gpe_start = prop_uint(propnum);
 		propnum = prop_dictionary_get(gpt_dict, "end");
 		PROP_ERR(propnum);
-		gpe_end = prop_number_unsigned_integer_value(propnum);
+		gpe_end = prop_uint(propnum);
 		if (gpe_start < firstdata || gpe_end > lastdata) {
 			gpt_warnx(gpt, "Backup GPT doesn't fit");
-			return -1;
+			goto out;
 		}
 	}
 	prop_object_iterator_release(propiter);
 
-	secbuf = calloc(gpt_size + 1, gpt->secsz);	/* GPT TABLE + GPT HEADER */
-	if (secbuf == NULL) {
-		gpt_warnx(gpt, "not enough memory to create a sector buffer");
-		return -1;
+	/* GPT TABLE + GPT HEADER */
+	if ((secbuf = calloc(gpt_size + 1, gpt->secsz)) == NULL) {
+		gpt_warn(gpt, "not enough memory to create a sector buffer");
+		goto out;
 	}
 
 	if (lseek(gpt->fd, 0LL, SEEK_SET) == -1) {
-		gpt_warnx(gpt, "Can't seek to beginning");
-		return -1;
+		gpt_warn(gpt, "Can't seek to beginning");
+		goto out;
 	}
 	for (i = 0; i < firstdata; i++) {
-		if (write(gpt->fd, secbuf, gpt->secsz) == -1) {
-			gpt_warnx(gpt, "Error writing");
-			return -1;
+		if (write(gpt->fd, secbuf, gpt->secsz) != (ssize_t)gpt->secsz) {
+			gpt_warn(gpt, "Error writing");
+			goto out;
 		}
 	}
 	if (lseek(gpt->fd, (lastdata + 1) * gpt->secsz, SEEK_SET) == -1) {
-		gpt_warnx(gpt, "Can't seek to end");
-		return -1;
+		gpt_warn(gpt, "Can't seek to end");
+		goto out;
 	}
 	for (i = lastdata + 1; i <= last; i++) {
-		if (write(gpt->fd, secbuf, gpt->secsz) == -1) {
-			gpt_warnx(gpt, "Error writing");
-			return -1;
+		if (write(gpt->fd, secbuf, gpt->secsz) != (ssize_t)gpt->secsz) {
+			gpt_warn(gpt, "Error writing");
+			goto out;
 		}
 	}
 
-	mbr = (struct mbr *)secbuf;
+	mbr = secbuf;
 	type_dict = prop_dictionary_get(props, "MBR");
 	PROP_ERR(type_dict);
 	propdata = prop_dictionary_get(type_dict, "code");
@@ -227,137 +344,42 @@ restore(gpt_t gpt)
 	propiter = prop_array_iterator(mbr_array);
 	PROP_ERR(propiter);
 	while ((mbr_dict = prop_object_iterator_next(propiter)) != NULL) {
-		propnum = prop_dictionary_get(mbr_dict, "index");
-		PROP_ERR(propnum);
-		i = prop_number_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "flag");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_flag =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "start_head");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_shd =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "start_sector");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_ssect =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "start_cylinder");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_scyl =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "type");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_typ =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "end_head");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_ehd =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "end_sector");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_esect =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "end_cylinder");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_ecyl =
-		    prop_number_unsigned_integer_value(propnum);
-		propnum = prop_dictionary_get(mbr_dict, "lba_start_low");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_start_lo =
-		    htole16(prop_number_unsigned_integer_value(propnum));
-		propnum = prop_dictionary_get(mbr_dict, "lba_start_high");
-		PROP_ERR(propnum);
-		mbr->mbr_part[i].part_start_hi =
-		    htole16(prop_number_unsigned_integer_value(propnum));
-		/* adjust PMBR size to size of device */
-                if (mbr->mbr_part[i].part_typ == MBR_PTYPE_PMBR) {
-			if (last > 0xffffffff) {
-				mbr->mbr_part[0].part_size_lo = htole16(0xffff);
-				mbr->mbr_part[0].part_size_hi = htole16(0xffff);
-			} else {
-				mbr->mbr_part[0].part_size_lo = htole16(last);
-				mbr->mbr_part[0].part_size_hi =
-				    htole16(last >> 16);
-			}
-		} else {
-			propnum = prop_dictionary_get(mbr_dict, "lba_size_low");
-			PROP_ERR(propnum);
-			mbr->mbr_part[i].part_size_lo =
-			    htole16(prop_number_unsigned_integer_value(propnum));
-			propnum =
-			    prop_dictionary_get(mbr_dict, "lba_size_high");
-			PROP_ERR(propnum);
-			mbr->mbr_part[i].part_size_hi =
-			    htole16(prop_number_unsigned_integer_value(propnum));
-		}
+		if (restore_mbr(gpt, mbr, mbr_dict, last) == -1)
+			goto out;
 	}
+
 	prop_object_iterator_release(propiter);
 	mbr->mbr_sig = htole16(MBR_SIG);
 	if (lseek(gpt->fd, 0LL, SEEK_SET) == -1 ||
-	    write(gpt->fd, mbr, gpt->secsz) == -1) {
-		gpt_warnx(gpt, "Unable to write MBR");
+	    write(gpt->fd, mbr, gpt->secsz) != (ssize_t)gpt->secsz) {
+		gpt_warn(gpt, "Unable to seek/write MBR");
 		return -1;
 	}
 	
 	propiter = prop_array_iterator(gpt_array);
 	PROP_ERR(propiter);
+
 	while ((gpt_dict = prop_object_iterator_next(propiter)) != NULL) {
-		memset(&ent, 0, sizeof(ent));
-		propstr = prop_dictionary_get(gpt_dict, "type");
-		PROP_ERR(propstr);
-		s = prop_string_cstring_nocopy(propstr);
-		if (gpt_uuid_parse(s, ent.ent_type) != 0) {
-			gpt_warnx(gpt, "%s: not able to convert to an UUID", s);
-			return -1;
-		}
-		propstr = prop_dictionary_get(gpt_dict, "guid");
-		PROP_ERR(propstr);
-		s = prop_string_cstring_nocopy(propstr);
-		if (gpt_uuid_parse(s, ent.ent_guid) != 0) {
-			gpt_warnx(gpt, "%s: not able to convert to an UUID", s);
-			return -1;
-		}
-		propnum = prop_dictionary_get(gpt_dict, "start");
-		PROP_ERR(propnum);
-		ent.ent_lba_start =
-		    htole64(prop_number_unsigned_integer_value(propnum));
-		propnum = prop_dictionary_get(gpt_dict, "end");
-		PROP_ERR(propnum);
-		ent.ent_lba_end =
-		    htole64(prop_number_unsigned_integer_value(propnum));
-		propnum = prop_dictionary_get(gpt_dict, "attributes");
-		PROP_ERR(propnum);
-		ent.ent_attr =
-		    htole64(prop_number_unsigned_integer_value(propnum));
-		propstr = prop_dictionary_get(gpt_dict, "name");
-		if (propstr != NULL) {
-			s = prop_string_cstring_nocopy(propstr);
-			utf8_to_utf16((const uint8_t *)s, ent.ent_name,
-				__arraycount(ent.ent_name));
-		}
-		propnum = prop_dictionary_get(gpt_dict, "index");
-		PROP_ERR(propnum);
-		i = prop_number_integer_value(propnum);
-		memcpy((char *)secbuf + gpt->secsz + ((i - 1) * sizeof(ent)),
-		    &ent, sizeof(ent));
+		if (restore_ent(gpt, gpt_dict, secbuf, gpt_size, entries) == -1)
+			goto out;
 	}
 	prop_object_iterator_release(propiter);
+
+	size_t len = gpt_size * gpt->secsz;
 	if (lseek(gpt->fd, 2 * gpt->secsz, SEEK_SET) == -1 ||
-	    write(gpt->fd, (char *)secbuf + 1 * gpt->secsz,
-	    gpt_size * gpt->secsz) == -1) {
-		gpt_warnx(gpt, "Unable to write primary GPT");
-		return -1;
+	    write(gpt->fd, (char *)secbuf + gpt->secsz, len) != (ssize_t) len) {
+		gpt_warn(gpt, "Unable to write primary GPT");
+		goto out;
 	}
+
 	if (lseek(gpt->fd, (lastdata + 1) * gpt->secsz, SEEK_SET) == -1 ||
-	    write(gpt->fd, (char *)secbuf + 1 * gpt->secsz,
-	    gpt_size * gpt->secsz) == -1) {
-		gpt_warnx(gpt, "Unable to write secondary GPT");
-		return -1;
+	    write(gpt->fd, (char *)secbuf + gpt->secsz, len) != (ssize_t) len) {
+		gpt_warn(gpt, "Unable to write secondary GPT");
+		goto out;
 	}
 
 	memset(secbuf, 0, gpt->secsz);
-	hdr = (struct gpt_hdr *)secbuf;
+	hdr = secbuf;
 	memcpy(hdr->hdr_sig, GPT_HDR_SIG, sizeof(hdr->hdr_sig));
 	hdr->hdr_revision = htole32(GPT_HDR_REVISION);
 	hdr->hdr_size = htole32(GPT_HDR_SIZE);
@@ -369,13 +391,12 @@ restore(gpt_t gpt)
 	hdr->hdr_lba_table = htole64(2);
 	hdr->hdr_entries = htole32(entries);
 	hdr->hdr_entsz = htole32(sizeof(struct gpt_ent));
-	hdr->hdr_crc_table = htole32(crc32((char *)secbuf + 1 * gpt->secsz,
-	    gpt_size * gpt->secsz));
+	hdr->hdr_crc_table = htole32(crc32((char *)secbuf + gpt->secsz, len));
 	hdr->hdr_crc_self = htole32(crc32(hdr, GPT_HDR_SIZE));
-	if (lseek(gpt->fd, 1 * gpt->secsz, SEEK_SET) == -1 ||
-	    write(gpt->fd, hdr, gpt->secsz) == -1) {
-		gpt_warnx(gpt, "Unable to write primary header");
-		return -1;
+	if (lseek(gpt->fd, gpt->secsz, SEEK_SET) == -1 ||
+	    write(gpt->fd, hdr, gpt->secsz) != (ssize_t)gpt->secsz) {
+		gpt_warn(gpt, "Unable to write primary header");
+		goto out;
 	}
 
 	hdr->hdr_lba_self = htole64(last);
@@ -384,13 +405,16 @@ restore(gpt_t gpt)
 	hdr->hdr_crc_self = 0;
 	hdr->hdr_crc_self = htole32(crc32(hdr, GPT_HDR_SIZE));
 	if (lseek(gpt->fd, last * gpt->secsz, SEEK_SET) == -1 ||
-	    write(gpt->fd, hdr, gpt->secsz) == -1) {
-		gpt_warnx(gpt, "Unable to write secondary header");
-		return -1;
+	    write(gpt->fd, hdr, gpt->secsz) != (ssize_t)gpt->secsz) {
+		gpt_warn(gpt, "Unable to write secondary header");
+		goto out;
 	}
+	rv = 0;
 
+out:
+	free(secbuf);
 	prop_object_release(props);
-	return 0;
+	return rv;
 }
 
 static int

Reply via email to