4 files changed, 115 insertions, 19 deletions

diff --git a/usr.sbin/bsdinstall/partedit/gpart_ops.c b/usr.sbin/bsdinstall/partedit/gpart_ops.c
index 37b172e..48bbb68 100644
--- a/usr.sbin/bsdinstall/partedit/gpart_ops.c
+++ b/usr.sbin/bsdinstall/partedit/gpart_ops.c
@@ -206,12 +206,11 @@ newfs_command(const char *fstype, char *command, int use_default)
 	}
 }
 
-int
-gpart_partition(const char *lg_name, const char *scheme)
+const char *
+choose_part_type(const char *def_scheme)
 {
 	int cancel, choice;
-	struct gctl_req *r;
-	const char *errstr;
+	const char *scheme = NULL;
 
 	DIALOG_LISTITEM items[] = {
 		{"APM", "Apple Partition Map",
@@ -228,30 +227,61 @@ gpart_partition(const char *lg_name, const char *scheme)
 		    "Bootable on Sun SPARC systems", 0 },
 	};
 
+parttypemenu:
+	dialog_vars.default_item = __DECONST(char *, def_scheme);
+	cancel = dlg_menu("Partition Scheme",
+	    "Select a partition scheme for this volume:", 0, 0, 0,
+	    sizeof(items) / sizeof(items[0]), items, &choice, NULL);
+	dialog_vars.default_item = NULL;
+
+	if (cancel)
+		return NULL;
+
+	if (!is_scheme_bootable(items[choice].name)) {
+		char message[512];
+		sprintf(message, "This partition scheme (%s) is not "
+		    "bootable on this platform. Are you sure you want "
+		    "to proceed?", items[choice].name);
+		dialog_vars.defaultno = TRUE;
+		cancel = dialog_yesno("Warning", message, 0, 0);
+		dialog_vars.defaultno = FALSE;
+		if (cancel) /* cancel */
+			goto parttypemenu;
+	}
+
+	scheme = items[choice].name;
+
+	return scheme;
+}
+
+int
+gpart_partition(const char *lg_name, const char *scheme)
+{
+	int cancel;
+	struct gctl_req *r;
+	const char *errstr;
+
 schememenu:
 	if (scheme == NULL) {
-		dialog_vars.default_item = __DECONST(char *, default_scheme());
-		cancel = dlg_menu("Partition Scheme",
-		    "Select a partition scheme for this volume:", 0, 0, 0,
-		    sizeof(items) / sizeof(items[0]), items, &choice, NULL);
-		dialog_vars.default_item = NULL;
+		scheme = choose_part_type(default_scheme());
 
-		if (cancel)
+		if (scheme == NULL)
 			return (-1);
 
-		if (!is_scheme_bootable(items[choice].name)) {
+		if (!is_scheme_bootable(scheme)) {
 			char message[512];
 			sprintf(message, "This partition scheme (%s) is not "
 			    "bootable on this platform. Are you sure you want "
-			    "to proceed?", items[choice].name);
+			    "to proceed?", scheme);
 			dialog_vars.defaultno = TRUE;
 			cancel = dialog_yesno("Warning", message, 0, 0);
 			dialog_vars.defaultno = FALSE;
-			if (cancel) /* cancel */
+			if (cancel) { /* cancel */
+				/* Reset scheme so user can choose another */
+				scheme = NULL;
 				goto schememenu;
+			}
 		}
-
-		scheme = items[choice].name;
 	}
 
 	r = gctl_get_handle();
@@ -322,6 +352,26 @@ gpart_activate(struct gprovider *pp)
 	gctl_free(r);
 }
 
+void
+gpart_set_root(const char *lg_name, const char *attribute)
+{
+	struct gctl_req *r;
+	const char *errstr;
+
+	r = gctl_get_handle();
+	gctl_ro_param(r, "class", -1, "PART");
+	gctl_ro_param(r, "arg0", -1, lg_name);
+	gctl_ro_param(r, "flags", -1, "C");
+	gctl_ro_param(r, "verb", -1, "set");
+	gctl_ro_param(r, "attrib", -1, attribute);
+
+	errstr = gctl_issue(r);
+	if (errstr != NULL && errstr[0] != '\0') 
+		gpart_show_error("Error", "Error setting parameter on disk:",
+		    errstr);
+	gctl_free(r);
+}
+
 static void
 gpart_bootcode(struct ggeom *gp)
 {
diff --git a/usr.sbin/bsdinstall/partedit/part_wizard.c b/usr.sbin/bsdinstall/partedit/part_wizard.c
index 3f7ccd5..1e9c899 100644
--- a/usr.sbin/bsdinstall/partedit/part_wizard.c
+++ b/usr.sbin/bsdinstall/partedit/part_wizard.c
@@ -257,8 +257,10 @@ query:
 			goto query;
 
 		gpart_destroy(gpart);
-		gpart_partition(disk, default_scheme());
-		scheme = default_scheme();
+		scheme = choose_part_type(default_scheme());
+		if (scheme == NULL)
+			return NULL;
+		gpart_partition(disk, scheme);
 	}
 
 	if (scheme == NULL || choice == 0) {
@@ -272,8 +274,10 @@ query:
 			gpart_destroy(gpart);
 		}
 
-		gpart_partition(disk, default_scheme());
-		scheme = default_scheme();
+		scheme = choose_part_type(default_scheme());
+		if (scheme == NULL)
+			return NULL;
+		gpart_partition(disk, scheme);
 	}
 
 	if (strcmp(scheme, "PC98") == 0 || strcmp(scheme, "MBR") == 0) {
diff --git a/usr.sbin/bsdinstall/partedit/partedit.c b/usr.sbin/bsdinstall/partedit/partedit.c
index ac3cd8a..f022044 100644
--- a/usr.sbin/bsdinstall/partedit/partedit.c
+++ b/usr.sbin/bsdinstall/partedit/partedit.c
@@ -44,6 +44,7 @@ struct pmetadata_head part_metadata;
 static int sade_mode = 0;
 
 static int apply_changes(struct gmesh *mesh);
+static void apply_workaround(struct gmesh *mesh);
 static struct partedit_item *read_geom_mesh(struct gmesh *mesh, int *nitems);
 static void add_geom_children(struct ggeom *gp, int recurse,
     struct partedit_item **items, int *nitems);
@@ -189,6 +190,8 @@ main(int argc, const char **argv)
 
 			if (op == 0 && validate_setup()) { /* Save */
 				error = apply_changes(&mesh);
+				if (!error)
+					apply_workaround(&mesh);
 				break;
 			} else if (op == 3) { /* Quit */
 				gpart_revert_all(&mesh);
@@ -390,6 +393,43 @@ apply_changes(struct gmesh *mesh)
 	return (0);
 }
 
+static void
+apply_workaround(struct gmesh *mesh)
+{
+	struct gclass *classp;
+	struct ggeom *gp;
+	struct gconfig *gc;
+	const char *scheme = NULL, *modified = NULL;
+
+	LIST_FOREACH(classp, &mesh->lg_class, lg_class) {
+		if (strcmp(classp->lg_name, "PART") == 0)
+			break;
+	}
+
+	if (strcmp(classp->lg_name, "PART") != 0) {
+		dialog_msgbox("Error", "gpart not found!", 0, 0, TRUE);
+		return;
+	}
+
+	LIST_FOREACH(gp, &classp->lg_geom, lg_geom) {
+		LIST_FOREACH(gc, &gp->lg_config, lg_config) {
+			if (strcmp(gc->lg_name, "scheme") == 0) {
+				scheme = gc->lg_val;
+			} else if (strcmp(gc->lg_name, "modified") == 0) {
+				modified = gc->lg_val;
+			}
+		}
+
+		if (scheme && strcmp(scheme, "GPT") == 0 &&
+		    modified && strcmp(modified, "true") == 0) {
+			if (getenv("WORKAROUND_LENOVO"))
+				gpart_set_root(gp->lg_name, "lenovofix");
+			if (getenv("WORKAROUND_GPTACTIVE"))
+				gpart_set_root(gp->lg_name, "active");
+		}
+	}
+}
+
 static struct partedit_item *
 read_geom_mesh(struct gmesh *mesh, int *nitems)
 {
diff --git a/usr.sbin/bsdinstall/partedit/partedit.h b/usr.sbin/bsdinstall/partedit/partedit.h
index b6f7258..3eb9173 100644
--- a/usr.sbin/bsdinstall/partedit/partedit.h
+++ b/usr.sbin/bsdinstall/partedit/partedit.h
@@ -72,6 +72,8 @@ void gpart_commit(struct gmesh *mesh);
 int gpart_partition(const char *lg_name, const char *scheme);
 void set_default_part_metadata(const char *name, const char *scheme,
     const char *type, const char *mountpoint, const char *newfs);
+void gpart_set_root(const char *lg_name, const char *attribute);
+const char *choose_part_type(const char *def_scheme);
 
 /* machine-dependent bootability checks */
 const char *default_scheme(void);