X-Git-Url: http://git.ozlabs.org/?p=petitboot;a=blobdiff_plain;f=discover%2Fgrub2%2Fenv.c;h=7eda0953fd4f5ce7d1cd6b6f20964a1f62f16bed;hp=90e9c69667abe47db82ee324c3f52052bba99648;hb=8f09986340e602a69398b2866a265ea3f16609c6;hpb=ccacb83a030da564c59c8e0f081fe967fba5a3eb

diff --git a/discover/grub2/env.c b/discover/grub2/env.c
index 90e9c69..7eda095 100644
--- a/discover/grub2/env.c
+++ b/discover/grub2/env.c
@@ -3,12 +3,12 @@
 #include <string.h>
 
 #include <log/log.h>
+#include <file/file.h>
 #include <types/types.h>
 #include <talloc/talloc.h>
-#include <array-size/array-size.h>
+#include <util/util.h>
 
 #include <discover/parser.h>
-#include <discover/file.h>
 
 #include "grub2.h"
 
@@ -17,7 +17,7 @@ static const char *signature = "# GRUB Environment Block\n";
 
 static int parse_buf_to_env(struct grub2_script *script, void *buf, int len)
 {
-	char *tmp, *line, *sep;
+	char *tmp = NULL, *line, *sep;
 	int siglen;
 
 	siglen = strlen(signature);
@@ -84,10 +84,10 @@ int builtin_load_env(struct grub2_script *script,
 
 	rc = parser_request_file(script->ctx, dev, envpath, &buf, &len);
 
-	if (!rc)
+	if (!rc) {
 		rc = parse_buf_to_env(script, buf, len);
-
-	talloc_free(buf);
+		talloc_free(buf);
+	}
 
 	return 0;
 }
@@ -214,6 +214,7 @@ int builtin_save_env(struct grub2_script *script,
 	int i, rc, len, siglen;
 	char *buf, *envpath;
 	const char *envfile;
+	bool using_dash_f = false;
 
 	/* we only support local filesystems */
 	if (!dev->mounted) {
@@ -222,9 +223,15 @@ int builtin_save_env(struct grub2_script *script,
 		return -1;
 	}
 
-	if (argc == 3 && !strcmp(argv[1], "-f"))
+	if (argc >= 2 && !strcmp(argv[1], "-f")) {
+		if (argc < 3) {
+			pb_log("save_env: for -f, need argument\n");
+			return -1;
+		}
+
 		envfile = argv[2];
-	else
+		using_dash_f = true;
+	} else
 		envfile = default_envfile;
 
 	envpath = talloc_asprintf(script, "%s/%s",
@@ -241,11 +248,19 @@ int builtin_save_env(struct grub2_script *script,
 	if (rc || len < siglen || memcmp(buf, signature, siglen))
 		goto err;
 
-	for (i = 1; i < argc; i++) {
+	/* For "-f", skip the "-f <file>" arguments in picking the
+	 * variables to save. */
+	i = (using_dash_f ? 3 : 1);
+
+	for (; i < argc; i++) {
 		const char *name, *value;
 
 		name = argv[i];
 		value = script_env_get(script, name);
+		if (!value) {
+			pb_log("Saved unset environment variable %s!\n", name);
+			value = "";
+		}
 
 		update_env(buf + siglen, len - siglen, name, value);
 	}