4 files changed, 62 insertions, 12 deletions
diff --git a/net/lwip/Makefile b/net/lwip/Makefile
index f2558f8763a..79dd6b3fb50 100644
--- a/net/lwip/Makefile
+++ b/net/lwip/Makefile
@@ -5,4 +5,4 @@ obj-$(CONFIG_CMD_DHCP) += dhcp.o
 obj-$(CONFIG_CMD_DNS) += dns.o
 obj-$(CONFIG_CMD_PING) += ping.o
 obj-$(CONFIG_CMD_TFTPBOOT) += tftp.o
-obj-$(CONFIG_CMD_WGET) += wget.o
+obj-$(CONFIG_WGET) += wget.o
diff --git a/net/lwip/dhcp.c b/net/lwip/dhcp.c
index 9b882cf5b87..e7d9147455c 100644
--- a/net/lwip/dhcp.c
+++ b/net/lwip/dhcp.c
@@ -3,6 +3,7 @@
 
 #include <command.h>
 #include <console.h>
+#include <log.h>
 #include <dm/device.h>
 #include <linux/delay.h>
 #include <linux/errno.h>
@@ -112,10 +113,17 @@ static int dhcp_loop(struct udevice *udev)
 int do_dhcp(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
 	int ret;
+	struct udevice *dev;
 
 	eth_set_current();
 
-	ret = dhcp_loop(eth_get_dev());
+	dev = eth_get_dev();
+	if (!dev) {
+		log_err("No network device\n");
+		return CMD_RET_FAILURE;
+	}
+
+	ret = dhcp_loop(dev);
 	if (ret)
 		return ret;
 
diff --git a/net/lwip/dns.c b/net/lwip/dns.c
index 4b937feaee1..1de63c9998b 100644
--- a/net/lwip/dns.c
+++ b/net/lwip/dns.c
@@ -56,7 +56,7 @@ static int dns_loop(struct udevice *udev, const char *name, const char *var)
 
 	netif = net_lwip_new_netif(udev);
 	if (!netif)
-		return -1;
+		return CMD_RET_FAILURE;
 
 	dns_init();
 
diff --git a/net/lwip/wget.c b/net/lwip/wget.c
index 5501ffdd004..b76f6c0f1d9 100644
--- a/net/lwip/wget.c
+++ b/net/lwip/wget.c
@@ -27,6 +27,8 @@ enum done_state {
 };
 
 struct wget_ctx {
+	char server_name[SERVER_NAME_SIZE];
+	u16 port;
 	char *path;
 	ulong daddr;
 	ulong saved_daddr;
@@ -36,6 +38,23 @@ struct wget_ctx {
 	enum done_state done;
 };
 
+static void wget_lwip_fill_info(struct pbuf *hdr, u16_t hdr_len, u32_t hdr_cont_len)
+{
+	if (wget_info->headers) {
+		if (hdr_len < MAX_HTTP_HEADERS_SIZE)
+			pbuf_copy_partial(hdr, (void *)wget_info->headers, hdr_len, 0);
+		else
+			hdr_len = 0;
+		wget_info->headers[hdr_len] = 0;
+	}
+	wget_info->hdr_cont_len = (u32)hdr_cont_len;
+}
+
+static void wget_lwip_set_file_size(u32_t rx_content_len)
+{
+	wget_info->file_size = (ulong)rx_content_len;
+}
+
 bool wget_validate_uri(char *uri);
 
 int mbedtls_hardware_poll(void *data, unsigned char *output, size_t len,
@@ -215,6 +234,13 @@ static void httpc_result_cb(void *arg, httpc_result_t httpc_result,
 	struct wget_ctx *ctx = arg;
 	ulong elapsed;
 
+	wget_info->status_code = (u32)srv_res;
+
+	if (err == ERR_BUF) {
+		ctx->done = FAILURE;
+		return;
+	}
+
 	if (httpc_result != HTTPC_RESULT_OK) {
 		log_err("\nHTTP client error %d\n", httpc_result);
 		ctx->done = FAILURE;
@@ -234,8 +260,10 @@ static void httpc_result_cb(void *arg, httpc_result_t httpc_result,
 	printf("%u bytes transferred in %lu ms (", rx_content_len, elapsed);
 	print_size(rx_content_len / elapsed * 1000, "/s)\n");
 	printf("Bytes transferred = %lu (%lx hex)\n", ctx->size, ctx->size);
-	efi_set_bootdev("Net", "", ctx->path, map_sysmem(ctx->saved_daddr, 0),
-			rx_content_len);
+	if (wget_info->set_bootdev)
+		efi_set_bootdev("Http", ctx->server_name, ctx->path, map_sysmem(ctx->saved_daddr, 0),
+				rx_content_len);
+	wget_lwip_set_file_size(rx_content_len);
 	if (env_set_hex("filesize", rx_content_len) ||
 	    env_set_hex("fileaddr", ctx->saved_daddr)) {
 		log_err("Could not set filesize or fileaddr\n");
@@ -246,9 +274,19 @@ static void httpc_result_cb(void *arg, httpc_result_t httpc_result,
 	ctx->done = SUCCESS;
 }
 
+static err_t httpc_headers_done_cb(httpc_state_t *connection, void *arg, struct pbuf *hdr,
+				   u16_t hdr_len, u32_t content_len)
+{
+	wget_lwip_fill_info(hdr, hdr_len, content_len);
+
+	if (wget_info->check_buffer_size && (ulong)content_len > wget_info->buffer_size)
+		return ERR_BUF;
+
+	return ERR_OK;
+}
+
 static int wget_loop(struct udevice *udev, ulong dst_addr, char *uri)
 {
-	char server_name[SERVER_NAME_SIZE];
 #if defined CONFIG_WGET_HTTPS
 	altcp_allocator_t tls_allocator;
 #endif
@@ -257,7 +295,6 @@ static int wget_loop(struct udevice *udev, ulong dst_addr, char *uri)
 	struct netif *netif;
 	struct wget_ctx ctx;
 	char *path;
-	u16 port;
 	bool is_https;
 
 	ctx.daddr = dst_addr;
@@ -267,7 +304,7 @@ static int wget_loop(struct udevice *udev, ulong dst_addr, char *uri)
 	ctx.prevsize = 0;
 	ctx.start_time = 0;
 
-	if (parse_url(uri, server_name, &port, &path, &is_https))
+	if (parse_url(uri, ctx.server_name, &ctx.port, &path, &is_https))
 		return CMD_RET_USAGE;
 
 	netif = net_lwip_new_netif(udev);
@@ -279,7 +316,7 @@ static int wget_loop(struct udevice *udev, ulong dst_addr, char *uri)
 	if (is_https) {
 		tls_allocator.alloc = &altcp_tls_alloc;
 		tls_allocator.arg =
-			altcp_tls_create_config_client(NULL, 0, server_name);
+			altcp_tls_create_config_client(NULL, 0, ctx.server_name);
 
 		if (!tls_allocator.arg) {
 			log_err("error: Cannot create a TLS connection\n");
@@ -292,8 +329,9 @@ static int wget_loop(struct udevice *udev, ulong dst_addr, char *uri)
 #endif
 
 	conn.result_fn = httpc_result_cb;
+	conn.headers_done_fn = httpc_headers_done_cb;
 	ctx.path = path;
-	if (httpc_get_file_dns(server_name, port, path, &conn, httpc_recv_cb,
+	if (httpc_get_file_dns(ctx.server_name, ctx.port, path, &conn, httpc_recv_cb,
 			       &ctx, &state)) {
 		net_lwip_remove_netif(netif);
 		return CMD_RET_FAILURE;
@@ -314,10 +352,13 @@ static int wget_loop(struct udevice *udev, ulong dst_addr, char *uri)
 	return -1;
 }
 
-int wget_with_dns(ulong dst_addr, char *uri)
+int wget_do_request(ulong dst_addr, char *uri)
 {
 	eth_set_current();
 
+	if (!wget_info)
+		wget_info = &default_wget_info;
+
 	return wget_loop(eth_get_dev(), dst_addr, uri);
 }
 
@@ -344,7 +385,8 @@ int do_wget(struct cmd_tbl *cmdtp, int flag, int argc, char * const argv[])
 	if (parse_legacy_arg(url, nurl, sizeof(nurl)))
 		return CMD_RET_FAILURE;
 
-	if (wget_with_dns(dst_addr, nurl))
+	wget_info = &default_wget_info;
+	if (wget_do_request(dst_addr, nurl))
 		return CMD_RET_FAILURE;
 
 	return CMD_RET_SUCCESS;