N210: Additional checks on both the host and firmware sides of the firmware updater.

2 files changed, 30 insertions, 3 deletions

diff --git a/firmware/zpu/usrp2p/spi_flash.c b/firmware/zpu/usrp2p/spi_flash.c
index 25fc239be..2033b8035 100644
--- a/firmware/zpu/usrp2p/spi_flash.c
+++ b/firmware/zpu/usrp2p/spi_flash.c
@@ -51,6 +51,8 @@ void
 spi_flash_erase_sector_start(uint32_t flash_addr)
 {
   //uprintf(UART_DEBUG, "spi_flash_erase_sector_start: addr = 0x%x\n", flash_addr);
+  if(flash_addr > spi_flash_memory_size())
+    return;
 
   spi_flash_wait();
   spi_flash_write_enable();
@@ -65,6 +67,10 @@ spi_flash_page_program_start(uint32_t flash_addr, size_t nbytes, const void *buf
   if (nbytes == 0 || nbytes > SPI_FLASH_PAGE_SIZE)
     return false;
 
+  //please to not be writing past the end of the device
+  if ((flash_addr + nbytes) > spi_flash_memory_size())
+    return false;
+
   uint32_t local_buf[SPI_FLASH_PAGE_SIZE / sizeof(uint32_t)];
   memset(local_buf, 0xff, sizeof(local_buf));	// init to 0xff (nops when programming)
   memcpy(local_buf, buf, nbytes);
@@ -130,6 +136,8 @@ spi_flash_program(uint32_t flash_addr, size_t nbytes, const void *buf)
   const unsigned char *p = (const unsigned char *) buf;
   size_t n;
 
+  if ((nbytes + flash_addr) > spi_flash_memory_size())
+    return false;
   if (nbytes == 0)
     return true;
 
@@ -158,7 +166,7 @@ void
 spi_flash_async_erase_start(spi_flash_async_state_t *s,
 			    uint32_t flash_addr, size_t nbytes)
 {
-  if (nbytes == 0){
+  if ((nbytes == 0) || ((flash_addr + nbytes) > spi_flash_memory_size())){
     s->first = s->last = s->current = 0;
     return;
   }
diff --git a/host/utils/usrp_n2xx_net_burner.py b/host/utils/usrp_n2xx_net_burner.py
index 0b64f2008..6c2939cd4 100755
--- a/host/utils/usrp_n2xx_net_burner.py
+++ b/host/utils/usrp_n2xx_net_burner.py
@@ -139,16 +139,21 @@ class burner_socket(object):
 
         #  print "Incoming:\n\tVer: %i\n\tID: %c\n\tSeq: %i\n\tIP: %i\n" % (proto_ver, chr(pktid), rxseq, ip_addr)
 
+    memory_size_bytes = 0
+    sector_size_bytes = 0
     def get_flash_info(self):
+        if (self.memory_size_bytes != 0) and (self.sector_size_bytes != 0):
+            return (self.memory_size_bytes, self.sector_size_bytes)
+            
         out_pkt = pack_flash_args_fmt(USRP2_FW_PROTO_VERSION, update_id_t.USRP2_FW_UPDATE_ID_WATS_TEH_FLASH_INFO_LOL, seq(), 0, 0)
         in_pkt = self.send_and_recv(out_pkt)
 
-        (proto_ver, pktid, rxseq, sector_size_bytes, memory_size_bytes) = unpack_flash_info_fmt(in_pkt)
+        (proto_ver, pktid, rxseq, self.sector_size_bytes, self.memory_size_bytes) = unpack_flash_info_fmt(in_pkt)
 
         if pktid != update_id_t.USRP2_FW_UPDATE_ID_HERES_TEH_FLASH_INFO_OMG:
             raise Exception("Invalid reply %c from device." % (chr(pktid)))
 
-        return (memory_size_bytes, sector_size_bytes)
+        return (self.memory_size_bytes, self.sector_size_bytes)
 
     def burn_fw(self, fw, fpga, reset, safe):
         (flash_size, sector_size) = self.get_flash_info()
@@ -167,6 +172,9 @@ class burner_socket(object):
 
             if not is_valid_fpga_image(fpga_image):
                 raise Exception("Error: Invalid FPGA image file.")
+                
+            if (len(fpga_image) + image_location) > flash_size:
+                raise Exception("Error: Cannot write past end of device")
 
             print("Begin FPGA write: this should take about 1 minute...")
             start_time = time.time()
@@ -188,6 +196,9 @@ class burner_socket(object):
 
             if not is_valid_fw_image(fw_image):
                 raise Exception("Error: Invalid firmware image file.")
+                
+            if (len(fw_image) + image_location) > flash_size:
+                raise Exception("Error: Cannot write past end of device")
 
             print("Begin firmware write: this should take about 1 second...")
             start_time = time.time()
@@ -204,6 +215,10 @@ class burner_socket(object):
         self._status_cb("Writing")
         writedata = image
         #we split the image into smaller (256B) bits and send them down the wire
+        (mem_size, sector_size) = self.get_flash_info()
+        if (addr + len(writedata)) > mem_size:
+            raise Exception("Error: Cannot write past end of device")
+            
         while writedata:
             out_pkt = pack_flash_args_fmt(USRP2_FW_PROTO_VERSION, update_id_t.USRP2_FW_UPDATE_ID_WRITE_TEH_FLASHES_LOL, seq(), addr, FLASH_DATA_PACKET_SIZE, writedata[:FLASH_DATA_PACKET_SIZE])
             in_pkt = self.send_and_recv(out_pkt)
@@ -287,6 +302,10 @@ class burner_socket(object):
     def erase_image(self, addr, length):
         self._status_cb("Erasing")
         #get flash info first
+        (flash_size, sector_size) = self.get_flash_info()
+        if (addr + length) > flash_size:
+            raise Exception("Cannot erase past end of device")
+            
         out_pkt = pack_flash_args_fmt(USRP2_FW_PROTO_VERSION, update_id_t.USRP2_FW_UPDATE_ID_ERASE_TEH_FLASHES_LOL, seq(), addr, length)
         in_pkt = self.send_and_recv(out_pkt)