bootm: Set working fdt address as part of the bootm flow

Set the fdt working address so "fdt FOO" commands can be used as part
of the bootm flow.  Also set an the environment variable "fdtaddr"
with the value.

Signed-off-by: Kumar Gala <galak@kernel.crashing.org>
diff --git a/common/cmd_bootm.c b/common/cmd_bootm.c
index 9a745f1..022cef8 100644
--- a/common/cmd_bootm.c
+++ b/common/cmd_bootm.c
@@ -44,6 +44,12 @@
 #include <hush.h>
 #endif
 
+#if defined(CONFIG_OF_LIBFDT)
+#include <fdt.h>
+#include <libfdt.h>
+#include <fdt_support.h>
+#endif
+
 DECLARE_GLOBAL_DATA_PTR;
 
 extern int gunzip (void *dst, int dstlen, unsigned char *src, unsigned long *lenp);
@@ -267,6 +273,8 @@
 			puts ("Could not find a valid device tree\n");
 			return 1;
 		}
+
+		set_working_fdt_addr(images.ft_addr);
 #endif
 	}
 
diff --git a/common/cmd_fdt.c b/common/cmd_fdt.c
index 94bca46..8bc900a 100644
--- a/common/cmd_fdt.c
+++ b/common/cmd_fdt.c
@@ -50,6 +50,16 @@
  */
 struct fdt_header *working_fdt;
 
+void set_working_fdt_addr(void *addr)
+{
+	char buf[17];
+
+	working_fdt = addr;
+
+	sprintf(buf, "%lx", (unsigned long)addr);
+	setenv("fdtaddr", buf);
+}
+
 /*
  * Flattened Device Tree command, see the help for parameter definitions.
  */
@@ -64,6 +74,7 @@
 	 * Set the address of the fdt
 	 ********************************************************************/
 	if (argv[1][0] == 'a') {
+		unsigned long addr;
 		/*
 		 * Set the address [and length] of the fdt.
 		 */
@@ -75,7 +86,8 @@
 			return 0;
 		}
 
-		working_fdt = (struct fdt_header *)simple_strtoul(argv[2], NULL, 16);
+		addr = simple_strtoul(argv[2], NULL, 16);
+		set_working_fdt_addr((void *)addr);
 
 		if (!fdt_valid()) {
 			return 1;
diff --git a/common/image.c b/common/image.c
index 18d667d..55c4cce 100644
--- a/common/image.c
+++ b/common/image.c
@@ -1245,6 +1245,7 @@
 		*of_size = of_len;
 	}
 
+	set_working_fdt_addr(*of_flat_tree);
 	return 0;
 
 error:
diff --git a/include/fdt_support.h b/include/fdt_support.h
index f2f2cd5..f14ab68 100644
--- a/include/fdt_support.h
+++ b/include/fdt_support.h
@@ -68,5 +68,7 @@
 void ft_pci_setup(void *blob, bd_t *bd);
 #endif
 
+void set_working_fdt_addr(void *addr);
+
 #endif /* ifdef CONFIG_OF_LIBFDT */
 #endif /* ifndef __FDT_SUPPORT_H */