diff options
| -rw-r--r-- | lib/efi_loader/efi_net.c | 61 | ||||
| -rw-r--r-- | net/wget.c | 18 |
2 files changed, 63 insertions, 16 deletions
diff --git a/lib/efi_loader/efi_net.c b/lib/efi_loader/efi_net.c index ce9272fa240..60aa076feaa 100644 --- a/lib/efi_loader/efi_net.c +++ b/lib/efi_loader/efi_net.c @@ -927,12 +927,15 @@ efi_status_t efi_net_register(void) &netobj->net); if (r != EFI_SUCCESS) goto failure_to_add_protocol; - if (!net_dp) - efi_net_set_dp("Net", NULL); - r = efi_add_protocol(&netobj->header, &efi_guid_device_path, - net_dp); + + if (net_dp) + r = efi_add_protocol(&netobj->header, &efi_guid_device_path, + net_dp); + else + r = efi_net_set_dp("Net", NULL); if (r != EFI_SUCCESS) goto failure_to_add_protocol; + r = efi_add_protocol(&netobj->header, &efi_pxe_base_code_protocol_guid, &netobj->pxe); if (r != EFI_SUCCESS) @@ -1057,18 +1060,58 @@ out_of_resources: */ efi_status_t efi_net_set_dp(const char *dev, const char *server) { - efi_free_pool(net_dp); + efi_status_t ret = EFI_SUCCESS; + struct efi_handler *phandler; + struct efi_device_path *old_net_dp, *new_net_dp; - net_dp = NULL; + old_net_dp = net_dp; + new_net_dp = NULL; if (!strcmp(dev, "Net")) - net_dp = efi_dp_from_eth(); + new_net_dp = efi_dp_from_eth(); else if (!strcmp(dev, "Http")) - net_dp = efi_dp_from_http(server); + new_net_dp = efi_dp_from_http(server); - if (!net_dp) + if (!new_net_dp) { return EFI_OUT_OF_RESOURCES; + } + + // If netobj is not started yet, end here. + if (!netobj) { + goto exit; + } + + phandler = NULL; + efi_search_protocol(&netobj->header, &efi_guid_device_path, &phandler); + + // If the device path protocol is not yet installed, install it + if (!phandler) + goto add; + + // If it is already installed, try to update it + ret = efi_reinstall_protocol_interface(&netobj->header, &efi_guid_device_path, + old_net_dp, new_net_dp); + if (ret != EFI_SUCCESS) + goto error; + + net_dp = new_net_dp; + efi_free_pool(old_net_dp); return EFI_SUCCESS; +add: + ret = efi_add_protocol(&netobj->header, &efi_guid_device_path, + new_net_dp); + if (ret != EFI_SUCCESS) + goto error; +exit: + net_dp = new_net_dp; + efi_free_pool(old_net_dp); + + return ret; +error: + // Failed, restore + efi_free_pool(new_net_dp); + + return ret; } /** diff --git a/net/wget.c b/net/wget.c index 0b082c61947..c73836cbc9d 100644 --- a/net/wget.c +++ b/net/wget.c @@ -53,6 +53,9 @@ static inline int store_block(uchar *src, unsigned int offset, unsigned int len) ulong store_addr = image_load_addr + offset; uchar *ptr; + // Avoid overflow + if (wget_info->buffer_size && wget_info->buffer_size < offset + len) + return -1; if (CONFIG_IS_ENABLED(LMB) && wget_info->set_bootdev) { if (store_addr < image_load_addr || lmb_read_check(store_addr, len)) { @@ -98,12 +101,6 @@ static void tcp_stream_on_closed(struct tcp_stream *tcp) net_set_state(wget_loop_state); if (wget_loop_state != NETLOOP_SUCCESS) { net_boot_file_size = 0; - if (wget_info->status_code == HTTP_STATUS_OK) { - wget_info->status_code = HTTP_STATUS_BAD; - wget_info->hdr_cont_len = 0; - if (wget_info->headers) - wget_info->headers[0] = 0; - } printf("\nwget: Transfer Fail, TCP status - %d\n", tcp->status); return; } @@ -212,6 +209,11 @@ static void tcp_stream_on_rcv_nxt_update(struct tcp_stream *tcp, u32 rx_bytes) "wget: Connected Len %lu\n", content_length); wget_info->hdr_cont_len = content_length; + if (wget_info->buffer_size && wget_info->buffer_size < wget_info->hdr_cont_len){ + tcp_stream_reset(tcp); + goto end; + } + } net_boot_file_size = rx_bytes - http_hdr_size; @@ -227,7 +229,9 @@ static int tcp_stream_rx(struct tcp_stream *tcp, u32 rx_offs, void *buf, int len if ((max_rx_pos == (u32)(-1)) || (max_rx_pos < rx_offs + len - 1)) max_rx_pos = rx_offs + len - 1; - store_block(buf, rx_offs - http_hdr_size, len); + // Avoid overflow + if (store_block(buf, rx_offs - http_hdr_size, len) < 0) + return -1; return len; } |