futility: add --type arg to show and sign commands

This allows the user to manually specify the type of an input
file, since not all file types can be reliably identified.

This also adds a test to ensure that futility doesn't coredump if
you give it the wrong type (although I'm sure it's not exhaustive).

BUG=chromium:231574
BRANCH=none
TEST=make runtests

Signed-off-by: Bill Richardson <wfrichar@chromium.org>
Change-Id: I9d909305d9989fe7299e744c585de380109cf8cd
Reviewed-on: https://chromium-review.googlesource.com/262895
Reviewed-by: Randall Spangler <rspangler@chromium.org>

7 files changed, 129 insertions, 17 deletions

diff --git a/futility/cmd_show.c b/futility/cmd_show.c
index 661aef2b..f00f2224 100644
--- a/futility/cmd_show.c
+++ b/futility/cmd_show.c
@@ -39,8 +39,10 @@ static struct local_data_s {
 	uint32_t padding;
 	int strict;
 	int t_flag;
+	enum futil_file_type type;
 } option = {
 	.padding = 65536,
+	.type = FILE_TYPE_UNKNOWN,
 };
 
 /* Stuff for BIOS images. */
@@ -615,6 +617,7 @@ int ft_show_kernel_preamble(const char *name, uint8_t *buf, uint32_t len,
 
 enum no_short_opts {
 	OPT_PADDING = 1000,
+	OPT_TYPE,
 	OPT_HELP,
 };
 
@@ -635,6 +638,8 @@ static const char usage[] = "\n"
 	"            Use this public key for validation\n"
 	"  -f|--fv          FILE            Verify this payload (FW_MAIN_A/B)\n"
 	"  --pad            NUM             Kernel vblock padding size\n"
+	"  --type           TYPE            Override the detected file type\n"
+	"                                     Use \"--type help\" for a list\n"
 	"%s"
 	"\n";
 
@@ -655,6 +660,7 @@ static const struct option long_opts[] = {
 	{"publickey",   1, 0, 'k'},
 	{"fv",          1, 0, 'f'},
 	{"pad",         1, NULL, OPT_PADDING},
+	{"type",        1, NULL, OPT_TYPE},
 	{"strict",      0, &option.strict, 1},
 	{"help",        0, NULL, OPT_HELP},
 	{NULL, 0, NULL, 0},
@@ -662,7 +668,7 @@ static const struct option long_opts[] = {
 static char *short_opts = ":f:k:t";
 
 
-static void show_type(char *filename)
+static int show_type(char *filename)
 {
 	enum futil_file_err err;
 	enum futil_file_type type;
@@ -670,7 +676,8 @@ static void show_type(char *filename)
 	switch (err) {
 	case FILE_ERR_NONE:
 		printf("%s:\t%s\n", filename, futil_file_type_name(type));
-		break;
+		/* Only our recognized types return success */
+		return 0;
 	case FILE_ERR_DIR:
 		printf("%s:\t%s\n", filename, "directory");
 		break;
@@ -686,6 +693,8 @@ static void show_type(char *filename)
 	default:
 		break;
 	}
+	/* Everything else is an error */
+	return 1;
 }
 
 static int do_show(int argc, char *argv[])
@@ -693,10 +702,11 @@ static int do_show(int argc, char *argv[])
 	char *infile = 0;
 	int ifd, i;
 	int errorcnt = 0;
-	enum futil_file_type type;
 	uint8_t *buf;
 	uint32_t len;
 	char *e = 0;
+	int type_override = 0;
+	enum futil_file_type type;
 
 	opterr = 0;		/* quiet, you */
 	while ((i = getopt_long(argc, argv, short_opts, long_opts, 0)) != -1) {
@@ -727,6 +737,16 @@ static int do_show(int argc, char *argv[])
 				errorcnt++;
 			}
 			break;
+		case OPT_TYPE:
+			if (!futil_str_to_file_type(optarg, &option.type)) {
+				if (!strcasecmp("help", optarg))
+					print_file_types_and_exit(errorcnt);
+				fprintf(stderr,
+					"Invalid --type \"%s\"\n", optarg);
+				errorcnt++;
+			}
+			type_override = 1;
+			break;
 		case OPT_HELP:
 			print_help(argc, argv);
 			return !!errorcnt;
@@ -763,7 +783,7 @@ static int do_show(int argc, char *argv[])
 
 	if (option.t_flag) {
 		for (i = optind; i < argc; i++)
-			show_type(argv[i]);
+			errorcnt += show_type(argv[i]);
 		goto done;
 	}
 
@@ -782,7 +802,11 @@ static int do_show(int argc, char *argv[])
 			goto boo;
 		}
 
-		type = futil_file_type_buf(buf, len);
+		/* Allow the user to override the type */
+		if (type_override)
+			type = option.type;
+		else
+			type = futil_file_type_buf(buf, len);
 
 		errorcnt += futil_file_type_show(type, infile, buf, len);
 
diff --git a/futility/cmd_sign.c b/futility/cmd_sign.c
index cae08b2e..98a197e0 100644
--- a/futility/cmd_sign.c
+++ b/futility/cmd_sign.c
@@ -57,11 +57,13 @@ static struct local_data_s {
 	int pem_algo_specified;
 	uint32_t pem_algo;
 	char *pem_external;
+	enum futil_file_type type;
 } option = {
 	.version = 1,
 	.arch = ARCH_UNSPECIFIED,
 	.kloadaddr = CROS_32BIT_ENTRY_ADDR,
 	.padding = 65536,
+	.type = FILE_TYPE_UNKNOWN,
 };
 
 
@@ -748,6 +750,7 @@ enum no_short_opts {
 	OPT_PEM_SIGNPRIV,
 	OPT_PEM_ALGO,
 	OPT_PEM_EXTERNAL,
+	OPT_TYPE,
 	OPT_HELP,
 };
 
@@ -775,6 +778,7 @@ static const struct option long_opts[] = {
 	{"pem_signpriv", 1, NULL, OPT_PEM_SIGNPRIV},
 	{"pem_algo",     1, NULL, OPT_PEM_ALGO},
 	{"pem_external", 1, NULL, OPT_PEM_EXTERNAL},
+	{"type",         1, NULL, OPT_TYPE},
 	{"vblockonly",   0, &option.vblockonly, 1},
 	{"help",         0, NULL, OPT_HELP},
 	{NULL,           0, NULL, 0},
@@ -790,7 +794,6 @@ static int do_sign(int argc, char *argv[])
 	uint8_t *buf;
 	uint32_t buf_len;
 	char *e = 0;
-	enum futil_file_type type;
 	int inout_file_count = 0;
 	int mapping;
 	int helpind = 0;
@@ -940,6 +943,15 @@ static int do_sign(int argc, char *argv[])
 		case OPT_PEM_EXTERNAL:
 			option.pem_external = optarg;
 			break;
+		case OPT_TYPE:
+			if (!futil_str_to_file_type(optarg, &option.type)) {
+				if (!strcasecmp("help", optarg))
+				    print_file_types_and_exit(errorcnt);
+				fprintf(stderr,
+					"Invalid --type \"%s\"\n", optarg);
+				errorcnt++;
+			}
+			break;
 		case OPT_HELP:
 			helpind = optind - 1;
 			break;
@@ -994,24 +1006,25 @@ static int do_sign(int argc, char *argv[])
 	}
 
 	/* What are we looking at? */
-	if (futil_file_type(infile, &type)) {
+	if (option.type == FILE_TYPE_UNKNOWN &&
+	    futil_file_type(infile, &option.type)) {
 		errorcnt++;
 		goto done;
 	}
 
 	/* We may be able to infer the type based on the other args */
-	if (type == FILE_TYPE_UNKNOWN) {
+	if (option.type == FILE_TYPE_UNKNOWN) {
 		if (option.bootloader_data || option.config_data
 		    || option.arch != ARCH_UNSPECIFIED)
-			type = FILE_TYPE_RAW_KERNEL;
+			option.type = FILE_TYPE_RAW_KERNEL;
 		else if (option.kernel_subkey || option.fv_specified)
-			type = FILE_TYPE_RAW_FIRMWARE;
+			option.type = FILE_TYPE_RAW_FIRMWARE;
 	}
 
-	Debug("type=%s\n", futil_file_type_name(type));
+	Debug("type=%s\n", futil_file_type_name(option.type));
 
 	/* Check the arguments for the type of thing we want to sign */
-	switch (type) {
+	switch (option.type) {
 	case FILE_TYPE_PUBKEY:
 		option.create_new_outfile = 1;
 		if (option.signprivate && option.pem_signpriv) {
@@ -1063,7 +1076,7 @@ static int do_sign(int argc, char *argv[])
 		break;
 	default:
 		fprintf(stderr, "Unable to sign type %s\n",
-			futil_file_type_name(type));
+			futil_file_type_name(option.type));
 		errorcnt++;
 	}
 
@@ -1124,7 +1137,7 @@ static int do_sign(int argc, char *argv[])
 		goto done;
 	}
 
-	errorcnt += futil_file_type_sign(type, infile, buf, buf_len);
+	errorcnt += futil_file_type_sign(option.type, infile, buf, buf_len);
 
 	errorcnt += futil_unmap_file(ifd, MAP_RW, buf, buf_len);
 
diff --git a/futility/cmd_vbutil_kernel.c b/futility/cmd_vbutil_kernel.c
index 8ddc5ef2..d13749bf 100644
--- a/futility/cmd_vbutil_kernel.c
+++ b/futility/cmd_vbutil_kernel.c
@@ -24,7 +24,6 @@
 #include "futility.h"
 #include "host_common.h"
 #include "kernel_blob.h"
-#include "traversal.h"
 #include "vb1_helper.h"
 
 static void Fatal(const char *format, ...)
diff --git a/futility/futility.c b/futility/futility.c
index 91b7cdc6..60cc9c56 100644
--- a/futility/futility.c
+++ b/futility/futility.c
@@ -319,7 +319,7 @@ int main(int argc, char *argv[], char *envp[])
 	cmd = find_command(progname);
 	if (cmd) {
 		/* Yep, just do that */
-		return run_command(cmd, argc, argv);
+		return !!run_command(cmd, argc, argv);
 	}
 
 	/* Parse the global options, stopping at the first non-option. */
@@ -381,7 +381,7 @@ int main(int argc, char *argv[], char *envp[])
 		argc -= optind;
 		argv += optind;
 		optind = 0;
-		return run_command(cmd, argc, argv);
+		return !!run_command(cmd, argc, argv);
 	}
 
 	/* Nope. We've no clue what we're being asked to do. */
diff --git a/tests/futility/run_test_scripts.sh b/tests/futility/run_test_scripts.sh
index 0f654d89..e8bbd232 100755
--- a/tests/futility/run_test_scripts.sh
+++ b/tests/futility/run_test_scripts.sh
@@ -51,6 +51,7 @@ ${SCRIPTDIR}/test_sign_firmware.sh
 ${SCRIPTDIR}/test_sign_fw_main.sh
 ${SCRIPTDIR}/test_sign_kernel.sh
 ${SCRIPTDIR}/test_sign_keyblocks.sh
+${SCRIPTDIR}/test_file_types.sh
 "
 
 # Get ready...
diff --git a/tests/futility/test_file_types.c b/tests/futility/test_file_types.c
index 47482e1e..a2d03bf7 100644
--- a/tests/futility/test_file_types.c
+++ b/tests/futility/test_file_types.c
@@ -29,6 +29,7 @@ static struct {
 	{FILE_TYPE_BIOS_IMAGE,      "tests/futility/data/bios_zgb_mp.bin"},
 	{FILE_TYPE_OLD_BIOS_IMAGE,  "tests/futility/data/bios_mario_mp.bin"},
 	{FILE_TYPE_KERN_PREAMBLE,   "tests/futility/data/kern_preamble.bin"},
+	/* We don't have a way to identify these (yet?) */
 	{FILE_TYPE_RAW_FIRMWARE,    },
 	{FILE_TYPE_RAW_KERNEL,	    },
 	{FILE_TYPE_CHROMIUMOS_DISK, },
diff --git a/tests/futility/test_file_types.sh b/tests/futility/test_file_types.sh
new file mode 100755
index 00000000..8acbcf72
--- /dev/null
+++ b/tests/futility/test_file_types.sh
@@ -0,0 +1,74 @@
+#!/bin/bash -eux
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+me=${0##*/}
+TMP="$me.tmp"
+
+# Work in scratch directory
+cd "$OUTDIR"
+
+
+# The first part of this is a script version of the compiled test by the same
+# name, to ensure we have working results.
+
+# Args are <expected_type>, <file_to_probe>
+test_case() {
+    local result
+    result=$(${FUTILITY} show -t "${SRCDIR}/$2" | awk '{print $NF}')
+    [ "$1" = "$result" ]
+}
+
+# Arg is <file_to_probe>
+fail_case() {
+    if ${FUTILITY} show -t "$1" ; then false; else true; fi
+}
+
+# Known types
+test_case "unknown"         "tests/futility/data/random_noise.bin"
+test_case "pubkey"          "tests/devkeys/root_key.vbpubk"
+test_case "keyblock"        "tests/devkeys/kernel.keyblock"
+test_case "fw_pre"          "tests/futility/data/fw_vblock.bin"
+test_case "gbb"	            "tests/futility/data/fw_gbb.bin"
+test_case "bios"            "tests/futility/data/bios_zgb_mp.bin"
+test_case "oldbios"         "tests/futility/data/bios_mario_mp.bin"
+test_case "kernel"          "tests/futility/data/kern_preamble.bin"
+# We don't have a way to identify these (yet?)
+# test_case "RAW_FIRMWARE"
+# test_case "RAW_KERNEL"
+# test_case "CHROMIUMOS_DISK"
+test_case "prikey"	    "tests/devkeys/root_key.vbprivk"
+test_case "pubkey21"        "tests/futility/data/sample.vbpubk2"
+test_case "prikey21"        "tests/futility/data/sample.vbprik2"
+test_case "pem"             "tests/testkeys/key_rsa2048.pem"
+
+# Expect failure here.
+fail_case "/Sir/Not/Appearing/In/This/Film"
+fail_case "${SRCDIR}"
+fail_case "/dev/zero"
+
+
+# Now test the show command when the file type is intentionally wrong. It
+# often won't work, but it certainly shouldn't core dump.
+
+# We'll ask futility to tell us what types it supports
+TYPES=$(${FUTILITY} show --type help | awk '/^ +/ {print $1}')
+
+# And we'll just reuse the same files above.
+FILES=$(awk '/^test_case / {print $NF}' $0 | tr -d '"')
+
+# futility should normally exit with either 0 or 1. Make sure that happens.
+# NOTE: /bin/bash returns values > 125 for special problems like signals.
+# I welcome patches to do this more portably.
+for type in $TYPES; do
+    for file in $FILES; do
+        ${FUTILITY} show --type ${type} "${SRCDIR}/${file}" && rc=$? || rc=$?
+        [ "$rc" -le 2 ]
+    done
+done
+
+
+# cleanup
+rm -rf ${TMP}*
+exit 0