sandbox: Add facility to save/restore sandbox state

It is often useful to be able to save out the state from a sandbox test
run, for analysis or to restore it later to continue a test. Add generic
infrastructure for doing this using a device tree binary file. This is
a flexible tagged file format which is already supported by U-Boot, and
it supports hierarchy if needed.

Signed-off-by: Simon Glass <sjg@chromium.org>
Signed-off-by: Simon Glass <sjg@chromium.org>
Reviewed-by: Hung-ying Tyan <tyanh@chromium.org>
diff --git a/arch/sandbox/cpu/start.c b/arch/sandbox/cpu/start.c
index 452f2a9..1df21d4 100644
--- a/arch/sandbox/cpu/start.c
+++ b/arch/sandbox/cpu/start.c
@@ -126,19 +126,56 @@
 SANDBOX_CMDLINE_OPT_SHORT(memory, 'm', 1,
 			  "Read/write ram_buf memory contents from file");
 
+static int sandbox_cmdline_cb_state(struct sandbox_state *state,
+				    const char *arg)
+{
+	state->state_fname = arg;
+	return 0;
+}
+SANDBOX_CMDLINE_OPT_SHORT(state, 's', 1, "Specify the sandbox state FDT");
+
+static int sandbox_cmdline_cb_read(struct sandbox_state *state,
+				   const char *arg)
+{
+	state->read_state = true;
+	return 0;
+}
+SANDBOX_CMDLINE_OPT_SHORT(read, 'r', 0, "Read the state FDT on startup");
+
+static int sandbox_cmdline_cb_write(struct sandbox_state *state,
+				    const char *arg)
+{
+	state->write_state = true;
+	return 0;
+}
+SANDBOX_CMDLINE_OPT_SHORT(write, 'w', 0, "Write state FDT on exit");
+
+static int sandbox_cmdline_cb_ignore_missing(struct sandbox_state *state,
+					     const char *arg)
+{
+	state->ignore_missing_state_on_read = true;
+	return 0;
+}
+SANDBOX_CMDLINE_OPT_SHORT(ignore_missing, 'n', 0,
+			  "Ignore missing state on read");
+
 int main(int argc, char *argv[])
 {
 	struct sandbox_state *state;
-	int err;
+	int ret;
 
-	err = state_init();
-	if (err)
-		return err;
+	ret = state_init();
+	if (ret)
+		goto err;
 
 	state = state_get_current();
 	if (os_parse_args(state, argc, argv))
 		return 1;
 
+	ret = sandbox_read_state(state, state->state_fname);
+	if (ret)
+		goto err;
+
 	/* Do pre- and post-relocation init */
 	board_init_f(0);
 
@@ -146,4 +183,8 @@
 
 	/* NOTREACHED - board_init_r() does not return */
 	return 0;
+
+err:
+	printf("Error %d\n", ret);
+	return 1;
 }
diff --git a/arch/sandbox/cpu/state.c b/arch/sandbox/cpu/state.c
index 0380fe2..a145808 100644
--- a/arch/sandbox/cpu/state.c
+++ b/arch/sandbox/cpu/state.c
@@ -4,6 +4,8 @@
  */
 
 #include <common.h>
+#include <errno.h>
+#include <fdtdec.h>
 #include <os.h>
 #include <asm/state.h>
 
@@ -16,6 +18,324 @@
 	state->exit_type = exit_type;
 }
 
+static int state_ensure_space(int extra_size)
+{
+	void *blob = state->state_fdt;
+	int used, size, free;
+	void *buf;
+	int ret;
+
+	used = fdt_off_dt_strings(blob) + fdt_size_dt_strings(blob);
+	size = fdt_totalsize(blob);
+	free = size - used;
+	if (free > extra_size)
+		return 0;
+
+	size = used + extra_size;
+	buf = os_malloc(size);
+	if (!buf)
+		return -ENOMEM;
+
+	ret = fdt_open_into(blob, buf, size);
+	if (ret) {
+		os_free(buf);
+		return -EIO;
+	}
+
+	os_free(blob);
+	state->state_fdt = buf;
+	return 0;
+}
+
+static int state_read_file(struct sandbox_state *state, const char *fname)
+{
+	int size;
+	int ret;
+	int fd;
+
+	size = os_get_filesize(fname);
+	if (size < 0) {
+		printf("Cannot find sandbox state file '%s'\n", fname);
+		return -ENOENT;
+	}
+	state->state_fdt = os_malloc(size);
+	if (!state->state_fdt) {
+		puts("No memory to read sandbox state\n");
+		return -ENOMEM;
+	}
+	fd = os_open(fname, OS_O_RDONLY);
+	if (fd < 0) {
+		printf("Cannot open sandbox state file '%s'\n", fname);
+		ret = -EPERM;
+		goto err_open;
+	}
+	if (os_read(fd, state->state_fdt, size) != size) {
+		printf("Cannot read sandbox state file '%s'\n", fname);
+		ret = -EIO;
+		goto err_read;
+	}
+	os_close(fd);
+
+	return 0;
+err_read:
+	os_close(fd);
+err_open:
+	os_free(state->state_fdt);
+	state->state_fdt = NULL;
+
+	return ret;
+}
+
+/***
+ * sandbox_read_state_nodes() - Read state associated with a driver
+ *
+ * This looks through all compatible nodes and calls the read function on
+ * each one, to read in the state.
+ *
+ * If nothing is found, it still calls the read function once, to set up a
+ * single global state for that driver.
+ *
+ * @state: Sandbox state
+ * @io: Method to use for reading state
+ * @blob: FDT containing state
+ * @return 0 if OK, -EINVAL if the read function returned failure
+ */
+int sandbox_read_state_nodes(struct sandbox_state *state,
+			     struct sandbox_state_io *io, const void *blob)
+{
+	int count;
+	int node;
+	int ret;
+
+	debug("   - read %s\n", io->name);
+	if (!io->read)
+		return 0;
+
+	node = -1;
+	count = 0;
+	while (blob) {
+		node = fdt_node_offset_by_compatible(blob, node, io->compat);
+		if (node < 0)
+			return 0;	/* No more */
+		debug("   - read node '%s'\n", fdt_get_name(blob, node, NULL));
+		ret = io->read(blob, node);
+		if (ret) {
+			printf("Unable to read state for '%s'\n", io->compat);
+			return -EINVAL;
+		}
+		count++;
+	}
+
+	/*
+	 * If we got no saved state, call the read function once without a
+	 * node, to set up the global state.
+	 */
+	if (count == 0) {
+		debug("   - read global\n");
+		ret = io->read(NULL, -1);
+		if (ret) {
+			printf("Unable to read global state for '%s'\n",
+			       io->name);
+			return -EINVAL;
+		}
+	}
+
+	return 0;
+}
+
+int sandbox_read_state(struct sandbox_state *state, const char *fname)
+{
+	struct sandbox_state_io *io;
+	const void *blob;
+	bool got_err;
+	int ret;
+
+	if (state->read_state && fname) {
+		ret = state_read_file(state, fname);
+		if (ret == -ENOENT && state->ignore_missing_state_on_read)
+			ret = 0;
+		if (ret)
+			return ret;
+	}
+
+	/* Call all the state read funtcions */
+	got_err = false;
+	blob = state->state_fdt;
+	io = ll_entry_start(struct sandbox_state_io, state_io);
+	for (; io < ll_entry_end(struct sandbox_state_io, state_io); io++) {
+		ret = sandbox_read_state_nodes(state, io, blob);
+		if (ret < 0)
+			got_err = true;
+	}
+
+	if (state->read_state && fname) {
+		debug("Read sandbox state from '%s'%s\n", fname,
+		      got_err ? " (with errors)" : "");
+	}
+
+	return got_err ? -1 : 0;
+}
+
+/***
+ * sandbox_write_state_node() - Write state associated with a driver
+ *
+ * This calls the write function to write out global state for that driver.
+ *
+ * TODO(sjg@chromium.org): Support writing out state from multiple drivers
+ * of the same time. We don't need this yet,and it will be much easier to
+ * do when driver model is available.
+ *
+ * @state: Sandbox state
+ * @io: Method to use for writing state
+ * @return 0 if OK, -EIO if there is a fatal error (such as out of space
+ * for adding the data), -EINVAL if the write function failed.
+ */
+int sandbox_write_state_node(struct sandbox_state *state,
+			     struct sandbox_state_io *io)
+{
+	void *blob;
+	int node;
+	int ret;
+
+	if (!io->write)
+		return 0;
+
+	ret = state_ensure_space(SANDBOX_STATE_MIN_SPACE);
+	if (ret) {
+		printf("Failed to add more space for state\n");
+		return -EIO;
+	}
+
+	/* The blob location can change when the size increases */
+	blob = state->state_fdt;
+	node = fdt_node_offset_by_compatible(blob, -1, io->compat);
+	if (node == -FDT_ERR_NOTFOUND) {
+		node = fdt_add_subnode(blob, 0, io->name);
+		if (node < 0) {
+			printf("Cannot create node '%s': %s\n", io->name,
+			       fdt_strerror(node));
+			return -EIO;
+		}
+
+		if (fdt_setprop_string(blob, node, "compatible", io->compat)) {
+			puts("Cannot set compatible\n");
+			return -EIO;
+		}
+	} else if (node < 0) {
+		printf("Cannot access node '%s': %s\n", io->name,
+		       fdt_strerror(node));
+		return -EIO;
+	}
+	debug("Write state for '%s' to node %d\n", io->compat, node);
+	ret = io->write(blob, node);
+	if (ret) {
+		printf("Unable to write state for '%s'\n", io->compat);
+		return -EINVAL;
+	}
+
+	return 0;
+}
+
+int sandbox_write_state(struct sandbox_state *state, const char *fname)
+{
+	struct sandbox_state_io *io;
+	bool got_err;
+	int size;
+	int ret;
+	int fd;
+
+	/* Create a state FDT if we don't have one */
+	if (!state->state_fdt) {
+		size = 0x4000;
+		state->state_fdt = os_malloc(size);
+		if (!state->state_fdt) {
+			puts("No memory to create FDT\n");
+			return -ENOMEM;
+		}
+		ret = fdt_create_empty_tree(state->state_fdt, size);
+		if (ret < 0) {
+			printf("Cannot create empty state FDT: %s\n",
+			       fdt_strerror(ret));
+			ret = -EIO;
+			goto err_create;
+		}
+	}
+
+	/* Call all the state write funtcions */
+	got_err = false;
+	io = ll_entry_start(struct sandbox_state_io, state_io);
+	ret = 0;
+	for (; io < ll_entry_end(struct sandbox_state_io, state_io); io++) {
+		ret = sandbox_write_state_node(state, io);
+		if (ret == -EIO)
+			break;
+		else if (ret)
+			got_err = true;
+	}
+
+	if (ret == -EIO) {
+		printf("Could not write sandbox state\n");
+		goto err_create;
+	}
+
+	ret = fdt_pack(state->state_fdt);
+	if (ret < 0) {
+		printf("Cannot pack state FDT: %s\n", fdt_strerror(ret));
+		ret = -EINVAL;
+		goto err_create;
+	}
+	size = fdt_totalsize(state->state_fdt);
+	fd = os_open(fname, OS_O_WRONLY | OS_O_CREAT);
+	if (fd < 0) {
+		printf("Cannot open sandbox state file '%s'\n", fname);
+		ret = -EIO;
+		goto err_create;
+	}
+	if (os_write(fd, state->state_fdt, size) != size) {
+		printf("Cannot write sandbox state file '%s'\n", fname);
+		ret = -EIO;
+		goto err_write;
+	}
+	os_close(fd);
+
+	debug("Wrote sandbox state to '%s'%s\n", fname,
+	      got_err ? " (with errors)" : "");
+
+	return 0;
+err_write:
+	os_close(fd);
+err_create:
+	os_free(state->state_fdt);
+
+	return ret;
+}
+
+int state_setprop(int node, const char *prop_name, const void *data, int size)
+{
+	void *blob;
+	int len;
+	int ret;
+
+	fdt_getprop(state->state_fdt, node, prop_name, &len);
+
+	/* Add space for the new property, its name and some overhead */
+	ret = state_ensure_space(size - len + strlen(prop_name) + 32);
+	if (ret)
+		return ret;
+
+	/* This should succeed, barring a mutiny */
+	blob = state->state_fdt;
+	ret = fdt_setprop(blob, node, prop_name, data, size);
+	if (ret) {
+		printf("%s: Unable to set property '%s' in node '%s': %s\n",
+		       __func__, prop_name, fdt_get_name(blob, node, NULL),
+			fdt_strerror(ret));
+		return -ENOSPC;
+	}
+
+	return 0;
+}
+
 struct sandbox_state *state_get_current(void)
 {
 	assert(state);
@@ -53,5 +373,16 @@
 		}
 	}
 
+	if (state->write_state) {
+		if (sandbox_write_state(state, state->state_fname)) {
+			printf("Failed to write sandbox state\n");
+			return -1;
+		}
+	}
+
+	if (state->state_fdt)
+		os_free(state->state_fdt);
+	memset(state, '\0', sizeof(*state));
+
 	return 0;
 }
diff --git a/arch/sandbox/include/asm/state.h b/arch/sandbox/include/asm/state.h
index e4b4c72..e8e4fea 100644
--- a/arch/sandbox/include/asm/state.h
+++ b/arch/sandbox/include/asm/state.h
@@ -8,6 +8,7 @@
 
 #include <config.h>
 #include <stdbool.h>
+#include <linux/stringify.h>
 
 /* How we exited U-Boot */
 enum exit_type_id {
@@ -34,12 +35,82 @@
 	unsigned int ram_size;		/* Size of RAM buffer */
 	const char *ram_buf_fname;	/* Filename to use for RAM buffer */
 	bool write_ram_buf;		/* Write RAM buffer on exit */
+	const char *state_fname;	/* File containing sandbox state */
+	void *state_fdt;		/* Holds saved state for sandbox */
+	bool read_state;		/* Read sandbox state on startup */
+	bool write_state;		/* Write sandbox state on exit */
+	bool ignore_missing_state_on_read;	/* No error if state missing */
 
 	/* Pointer to information for each SPI bus/cs */
 	struct sandbox_spi_info spi[CONFIG_SANDBOX_SPI_MAX_BUS]
 					[CONFIG_SANDBOX_SPI_MAX_CS];
 };
 
+/* Minimum space we guarantee in the state FDT when calling read/write*/
+#define SANDBOX_STATE_MIN_SPACE		0x1000
+
+/**
+ * struct sandbox_state_io - methods to saved/restore sandbox state
+ * @name: Name of of the device tree node, also the name of the variable
+ *	holding this data so it should be an identifier (use underscore
+ *	instead of minus)
+ * @compat: Compatible string for the node containing this state
+ *
+ * @read: Function to read state from FDT
+ *	If data is available, then blob and node will provide access to it. If
+ *	not (blob == NULL and node == -1) this function should set up an empty
+ *	data set for start-of-day.
+ *	@param blob: Pointer to device tree blob, or NULL if no data to read
+ *	@param node: Node offset to read from
+ *	@return 0 if OK, -ve on error
+ *
+ * @write: Function to write state to FDT
+ *	The caller will ensure that there is a node ready for the state. The
+ *	node may already contain the old state, in which case it should be
+ *	overridden. There is guaranteed to be SANDBOX_STATE_MIN_SPACE bytes
+ *	of free space, so error checking is not required for fdt_setprop...()
+ *	calls which add up to less than this much space.
+ *
+ *	For adding larger properties, use state_setprop().
+ *
+ * @param blob: Device tree blob holding state
+ * @param node: Node to write our state into
+ *
+ * Note that it is possible to save data as large blobs or as individual
+ * hierarchical properties. However, unless you intend to keep state files
+ * around for a long time and be able to run an old state file on a new
+ * sandbox, it might not be worth using individual properties for everything.
+ * This is certainly supported, it is just a matter of the effort you wish
+ * to put into the state read/write feature.
+ */
+struct sandbox_state_io {
+	const char *name;
+	const char *compat;
+	int (*write)(void *blob, int node);
+	int (*read)(const void *blob, int node);
+};
+
+/**
+ * SANDBOX_STATE_IO - Declare sandbox state to read/write
+ *
+ * Sandbox permits saving state from one run and restoring it in another. This
+ * allows the test system to retain state between runs and thus better
+ * emulate a real system. Examples of state that might be useful to save are
+ * the emulated GPIOs pin settings, flash memory contents and TPM private
+ * data. U-Boot memory contents is dealth with separately since it is large
+ * and it is not normally useful to save it (since a normal system does not
+ * preserve DRAM between runs). See the '-m' option for this.
+ *
+ * See struct sandbox_state_io above for member documentation.
+ */
+#define SANDBOX_STATE_IO(_name, _compat, _read, _write) \
+	ll_entry_declare(struct sandbox_state_io, _name, state_io) = { \
+		.name = __stringify(_name), \
+		.read = _read, \
+		.write = _write, \
+		.compat = _compat, \
+	}
+
 /**
  * Record the exit type to be reported by the test program.
  *
@@ -55,6 +126,49 @@
 struct sandbox_state *state_get_current(void);
 
 /**
+ * Read the sandbox state from the supplied device tree file
+ *
+ * This calls all registered state handlers to read in the sandbox state
+ * from a previous test run.
+ *
+ * @param state		Sandbox state to update
+ * @param fname		Filename of device tree file to read from
+ * @return 0 if OK, -ve on error
+ */
+int sandbox_read_state(struct sandbox_state *state, const char *fname);
+
+/**
+ * Write the sandbox state to the supplied device tree file
+ *
+ * This calls all registered state handlers to write out the sandbox state
+ * so that it can be preserved for a future test run.
+ *
+ * If the file exists it is overwritten.
+ *
+ * @param state		Sandbox state to update
+ * @param fname		Filename of device tree file to write to
+ * @return 0 if OK, -ve on error
+ */
+int sandbox_write_state(struct sandbox_state *state, const char *fname);
+
+/**
+ * Add a property to a sandbox state node
+ *
+ * This is equivalent to fdt_setprop except that it automatically enlarges
+ * the device tree if necessary. That means it is safe to write any amount
+ * of data here.
+ *
+ * This function can only be called from within struct sandbox_state_io's
+ * ->write method, i.e. within state I/O drivers.
+ *
+ * @param node		Device tree node to write to
+ * @param prop_name	Property to write
+ * @param data		Data to write into property
+ * @param size		Size of data to write into property
+ */
+int state_setprop(int node, const char *prop_name, const void *data, int size);
+
+/**
  * Initialize the test system state
  */
 int state_init(void);