1 files changed, 5 insertions, 42 deletions

diff --git a/net/wget.c b/net/wget.c
index 4a168641c65..b4251e0f293 100644
--- a/net/wget.c
+++ b/net/wget.c
@@ -8,6 +8,7 @@
 #include <command.h>
 #include <display_options.h>
 #include <env.h>
+#include <efi_loader.h>
 #include <image.h>
 #include <lmb.h>
 #include <mapmem.h>
@@ -64,26 +65,6 @@ static unsigned int retry_tcp_ack_num;	/* TCP retry acknowledge number*/
 static unsigned int retry_tcp_seq_num;	/* TCP retry sequence number */
 static int retry_len;			/* TCP retry length */
 
-static ulong wget_load_size;
-
-/**
- * wget_init_max_size() - initialize maximum load size
- *
- * Return:	0 if success, -1 if fails
- */
-static int wget_init_load_size(void)
-{
-	phys_size_t max_size;
-
-	max_size = lmb_get_free_size(image_load_addr);
-	if (!max_size)
-		return -1;
-
-	wget_load_size = max_size;
-
-	return 0;
-}
-
 /**
  * store_block() - store block in memory
  * @src: source of data
@@ -97,13 +78,8 @@ static inline int store_block(uchar *src, unsigned int offset, unsigned int len)
 	uchar *ptr;
 
 	if (CONFIG_IS_ENABLED(LMB)) {
-		ulong end_addr = image_load_addr + wget_load_size;
-
-		if (!end_addr)
-			end_addr = ULONG_MAX;
-
 		if (store_addr < image_load_addr ||
-		    store_addr + len > end_addr) {
+		    lmb_read_check(store_addr, len)) {
 			printf("\nwget error: ");
 			printf("trying to overwrite reserved memory...\n");
 			return -1;
@@ -196,13 +172,6 @@ void wget_fail(char *error_message, unsigned int tcp_seq_num,
 	wget_send(action, tcp_seq_num, tcp_ack_num, 0);
 }
 
-void wget_success(u8 action, unsigned int tcp_seq_num,
-		  unsigned int tcp_ack_num, int len, int packets)
-{
-	printf("Packets received %d, Transfer Successful\n", packets);
-	wget_send(action, tcp_seq_num, tcp_ack_num, len);
-}
-
 /*
  * Interfaces of U-BOOT
  */
@@ -432,6 +401,9 @@ static void wget_handler(uchar *pkt, u16 dport,
 	case WGET_TRANSFERRED:
 		printf("Packets received %d, Transfer Successful\n", packets);
 		net_set_state(wget_loop_state);
+		efi_set_bootdev("Net", "", image_url,
+				map_sysmem(image_load_addr, 0),
+				net_boot_file_size);
 		break;
 	}
 }
@@ -493,15 +465,6 @@ void wget_start(void)
 	debug_cond(DEBUG_WGET,
 		   "\nwget:Load address: 0x%lx\nLoading: *\b", image_load_addr);
 
-	if (CONFIG_IS_ENABLED(LMB)) {
-		if (wget_init_load_size()) {
-			printf("\nwget error: ");
-			printf("trying to overwrite reserved memory...\n");
-			net_set_state(NETLOOP_FAIL);
-			return;
-		}
-	}
-
 	net_set_timeout_handler(wget_timeout, wget_timeout_handler);
 	tcp_set_tcp_handler(wget_handler);