From 9ff988ed74bdf9eebec9f30e8093994e7dddef2e Mon Sep 17 00:00:00 2001
From: mintsuki <mintsuki@protonmail.com>
Date: Fri, 28 Apr 2023 04:59:08 +0200
Subject: [PATCH] pxe: Add and use constant PXE ACK packet size and fix BIOS
 PXE ACK struct

---
 common/protos/multiboot2.c |  8 ++++----
 common/pxe/pxe.h           |  8 +++++---
 common/pxe/tftp.s2.c       | 16 ++++++++++------
 3 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/common/protos/multiboot2.c b/common/protos/multiboot2.c
index 816aa887..65f678c0 100644
--- a/common/protos/multiboot2.c
+++ b/common/protos/multiboot2.c
@@ -57,7 +57,7 @@ static size_t get_multiboot2_info_size(
                 ALIGN_UP(sizeof(struct multiboot_tag_efi64_ih), MULTIBOOT_TAG_ALIGN) +                                  // EFI image handle 64
             #endif
         #endif
-        ALIGN_UP(sizeof(struct multiboot_tag_network) + cached_dhcp_packet_len, MULTIBOOT_TAG_ALIGN) +                  // network info
+        ALIGN_UP(sizeof(struct multiboot_tag_network) + DHCP_ACK_PACKET_LEN, MULTIBOOT_TAG_ALIGN) +                  // network info
         ALIGN_UP(sizeof(struct multiboot_tag), MULTIBOOT_TAG_ALIGN);                                                    // end
 }
 
@@ -788,14 +788,14 @@ skip_modeset:;
     // Create network info tag
     //////////////////////////////////////////////
     {
-        if (cached_dhcp_packet_len) {
+        if (cached_dhcp_ack_valid) {
             struct multiboot_tag_network *tag = (struct multiboot_tag_network *)(mb2_info + info_idx);
 
             tag->type = MULTIBOOT_TAG_TYPE_NETWORK;
-            tag->size = sizeof(struct multiboot_tag_network) + cached_dhcp_packet_len;
+            tag->size = sizeof(struct multiboot_tag_network) + DHCP_ACK_PACKET_LEN;
 
             // Copy over the DHCP packet.
-            memcpy(tag->dhcpack, cached_dhcp_packet, cached_dhcp_packet_len);
+            memcpy(tag->dhcpack, cached_dhcp_packet, DHCP_ACK_PACKET_LEN);
             append_tag(info_idx, tag);
         }
     }
diff --git a/common/pxe/pxe.h b/common/pxe/pxe.h
index 20c51119..49d094b6 100644
--- a/common/pxe/pxe.h
+++ b/common/pxe/pxe.h
@@ -4,8 +4,10 @@
 #include <stdint.h>
 #include <lib/part.h>
 
-extern uint8_t cached_dhcp_packet[1472];
-extern int cached_dhcp_packet_len;
+#define DHCP_ACK_PACKET_LEN 296
+
+extern uint8_t cached_dhcp_packet[DHCP_ACK_PACKET_LEN];
+extern bool cached_dhcp_ack_valid;
 
 #if defined (BIOS)
 
@@ -36,7 +38,7 @@ struct bootph {
         struct bootph_vendor_v {
             uint8_t magic[4];
             uint32_t flags;
-            uint8_t pad[56];
+            uint8_t pad[52];
         } v;
     } vendor;
 };
diff --git a/common/pxe/tftp.s2.c b/common/pxe/tftp.s2.c
index ff7c8078..e871270a 100644
--- a/common/pxe/tftp.s2.c
+++ b/common/pxe/tftp.s2.c
@@ -11,8 +11,8 @@
 #include <lib/misc.h>
 
 // cache the dhcp packet
-uint8_t cached_dhcp_packet[1472] = { 0 };
-int cached_dhcp_packet_len = 0;
+uint8_t cached_dhcp_packet[DHCP_ACK_PACKET_LEN] = { 0 };
+bool cached_dhcp_ack_valid = false;
 
 #if defined (BIOS)
 
@@ -21,8 +21,10 @@ static uint32_t get_boot_server_info(void) {
     cachedinfo.packet_type = PXENV_PACKET_TYPE_CACHED_REPLY;
     pxe_call(PXENV_GET_CACHED_INFO, ((uint16_t)rm_seg(&cachedinfo)), (uint16_t)rm_off(&cachedinfo));
     struct bootph *ph = (struct bootph*)(void *) (((((uint32_t)cachedinfo.buffer) >> 16) << 4) + (((uint32_t)cachedinfo.buffer) & 0xFFFF));
-    memcpy(&cached_dhcp_packet, ph, sizeof(struct bootph));
-    cached_dhcp_packet_len = sizeof(struct bootph);
+    if (!cached_dhcp_ack_valid) {
+        memcpy(cached_dhcp_packet, ph, DHCP_ACK_PACKET_LEN);
+        cached_dhcp_ack_valid = true;
+    }
     return ph->sip;
 }
 
@@ -147,8 +149,10 @@ static EFI_IP_ADDRESS *parse_ip_addr(struct volume *part, const char *server_add
         else if (part->pxe_base_code->Mode->ProxyOfferReceived) packet = &part->pxe_base_code->Mode->ProxyOffer;
         else packet = &part->pxe_base_code->Mode->DhcpAck;
         memcpy(out.Addr, packet->Dhcpv4.BootpSiAddr, 4);
-        memcpy(cached_dhcp_packet, packet, sizeof(EFI_PXE_BASE_CODE_PACKET));
-        cached_dhcp_packet_len = sizeof(EFI_PXE_BASE_CODE_PACKET);
+        if (!cached_dhcp_ack_valid) {
+            memcpy(cached_dhcp_packet, packet, DHCP_ACK_PACKET_LEN);
+            cached_dhcp_ack_valid = true;
+        }
     } else {
         if (inet_pton(server_addr, &out.Addr)) {
             panic(true, "tftp: Invalid IPv4 address: \"%s\"", server_addr);