From db10e49f14c3e57280db8fd1860bf29e49d9e89a Mon Sep 17 00:00:00 2001
From: Wu Jian Gang <wujiangang@espressif.com>
Date: Sun, 17 Jun 2018 23:39:56 +0800
Subject: [PATCH] feat(esptool): hack code to support esp8266 version 3 binary

---
 components/esptool_py/Kconfig.projbuild  |  10 +-
 components/esptool_py/Makefile.projbuild |   2 +-
 components/esptool_py/esptool/esptool.py | 211 ++++++++++++++++++++++-
 3 files changed, 215 insertions(+), 8 deletions(-)

diff --git a/components/esptool_py/Kconfig.projbuild b/components/esptool_py/Kconfig.projbuild
index 10b1aa1d..381136de 100644
--- a/components/esptool_py/Kconfig.projbuild
+++ b/components/esptool_py/Kconfig.projbuild
@@ -99,13 +99,13 @@ config ESPTOOLPY_FLASHFREQ
 
 choice ESPTOOLPY_FLASHSIZE
 	prompt "Flash size"
-	default ESPTOOLPY_FLASHSIZE_2MB_C1
+	default ESPTOOLPY_FLASHSIZE_2MB
 	help
 		SPI flash size, in megabytes
 
-config ESPTOOLPY_FLASHSIZE_2MB_C1
+config ESPTOOLPY_FLASHSIZE_2MB
 	bool "2 MB"
-config ESPTOOLPY_FLASHSIZE_4MB_C1
+config ESPTOOLPY_FLASHSIZE_4MB
 	bool "4 MB"
 config ESPTOOLPY_FLASHSIZE_8MB
 	bool "8 MB"
@@ -115,8 +115,8 @@ endchoice
 
 config ESPTOOLPY_FLASHSIZE
 	string
-	default "2MB-c1" if ESPTOOLPY_FLASHSIZE_2MB_C1
-	default "4MB-c1" if ESPTOOLPY_FLASHSIZE_4MB_C1
+	default "2MB" if ESPTOOLPY_FLASHSIZE_2MB
+	default "4MB" if ESPTOOLPY_FLASHSIZE_4MB
 	default "8MB" if ESPTOOLPY_FLASHSIZE_8MB
 	default "16MB" if ESPTOOLPY_FLASHSIZE_16MB
 
diff --git a/components/esptool_py/Makefile.projbuild b/components/esptool_py/Makefile.projbuild
index 1d8b170b..553c4149 100644
--- a/components/esptool_py/Makefile.projbuild
+++ b/components/esptool_py/Makefile.projbuild
@@ -30,7 +30,7 @@ ESPTOOL_WRITE_FLASH_OPTIONS := $(ESPTOOL_FLASH_OPTIONS)
 endif
 
 ifndef IS_BOOTLOADER_BUILD
-ESPTOOL_ELF2IMAGE_OPTIONS := --version=2
+ESPTOOL_ELF2IMAGE_OPTIONS := --version=3
 else
 ESPTOOL_ELF2IMAGE_OPTIONS :=
 endif
diff --git a/components/esptool_py/esptool/esptool.py b/components/esptool_py/esptool/esptool.py
index 17868ce6..6a0a5682 100755
--- a/components/esptool_py/esptool/esptool.py
+++ b/components/esptool_py/esptool/esptool.py
@@ -1460,6 +1460,211 @@ class ESP8266V2FirmwareImage(BaseFirmwareImage):
         with open(filename, 'ab') as f:
             f.write(struct.pack(b'<I', crc))
 
+class ESP8266V3FirmwareImage(BaseFirmwareImage):
+    """ ESP32 firmware image is very similar to V1 ESP8266 image,
+    except with an additional 16 byte reserved header at top of image,
+    and because of new flash mapping capabilities the flash-mapped regions
+    can be placed in the normal image (just @ 64kB padded offsets).
+    """
+
+    ROM_LOADER = ESP32ROM
+
+    # ROM bootloader will read the wp_pin field if SPI flash
+    # pins are remapped via flash. IDF actually enables QIO only
+    # from software bootloader, so this can be ignored. But needs
+    # to be set to this value so ROM bootloader will skip it.
+    WP_PIN_DISABLED = 0xEE
+
+    EXTENDED_HEADER_STRUCT_FMT = "B" * 16
+
+    def __init__(self, load_file=None):
+        super(ESP8266V3FirmwareImage, self).__init__()
+        self.flash_mode = 0
+        self.flash_size_freq = 0
+        self.version = 1
+        self.wp_pin = self.WP_PIN_DISABLED
+        # SPI pin drive levels
+        self.clk_drv = 0
+        self.q_drv = 0
+        self.d_drv = 0
+        self.cs_drv = 0
+        self.hd_drv = 0
+        self.wp_drv = 0
+
+        self.append_digest = True
+
+        if load_file is not None:
+            start = load_file.tell()
+
+            segments = self.load_common_header(load_file, ESPLoader.ESP_IMAGE_MAGIC)
+            self.load_extended_header(load_file)
+
+            for _ in range(segments):
+                self.load_segment(load_file)
+            self.checksum = self.read_checksum(load_file)
+
+            if self.append_digest:
+                end = load_file.tell()
+                self.stored_digest = load_file.read(32)
+                load_file.seek(start)
+                calc_digest = hashlib.sha256()
+                calc_digest.update(load_file.read(end - start))
+                self.calc_digest = calc_digest.digest()  # TODO: decide what to do here?
+
+    def is_flash_addr(self, addr):
+    	#print(' %x' % addr)
+    	#print(' %x' % (0x40200000 <= addr))
+    	#print(' %x' % ESP32ROM.DROM_MAP_START)
+    	return (0x40200000 <= addr)
+        #return (ESP32ROM.IROM_MAP_START <= addr < ESP32ROM.IROM_MAP_END) \
+        #    or (ESP32ROM.DROM_MAP_START <= addr < ESP32ROM.DROM_MAP_END)
+
+    def default_output_name(self, input_file):
+        """ Derive a default output name from the ELF name. """
+        return "%s.bin" % (os.path.splitext(input_file)[0])
+
+    def warn_if_unusual_segment(self, offset, size, is_irom_segment):
+        pass  # TODO: add warnings for ESP32 segment offset/size combinations that are wrong
+
+    def save(self, filename):
+        total_segments = 0
+        with io.BytesIO() as f:  # write file to memory first
+            self.write_common_header(f, self.segments)
+
+            # first 4 bytes of header are read by ROM bootloader for SPI
+            # config, but currently unused
+            #self.save_extended_header(f)
+
+            checksum = ESPLoader.ESP_CHECKSUM_MAGIC
+
+            # split segments into flash-mapped vs ram-loaded, and take copies so we can mutate them
+            flash_segments = [copy.deepcopy(s) for s in sorted(self.segments, key=lambda s:s.addr) if self.is_flash_addr(s.addr)]
+            ram_segments = [copy.deepcopy(s) for s in sorted(self.segments, key=lambda s:s.addr) if not self.is_flash_addr(s.addr)]
+
+            IROM_ALIGN = 65536
+
+            # check for multiple ELF sections that are mapped in the same flash mapping region.
+            # this is usually a sign of a broken linker script, but if you have a legitimate
+            # use case then let us know (we can merge segments here, but as a rule you probably
+            # want to merge them in your linker script.)
+            a = len(flash_segments)
+            
+            if len(flash_segments) > 0:
+                last_addr = flash_segments[0].addr
+                #print('%x' % last_addr)
+                for segment in flash_segments[1:]:
+                    if segment.addr // IROM_ALIGN == last_addr // IROM_ALIGN:
+                        raise FatalError(("Segment loaded at 0x%08x lands in same 64KB flash mapping as segment loaded at 0x%08x. " +
+                                          "Can't generate binary. Suggest changing linker script or ELF to merge sections.") %
+                                         (segment.addr, last_addr))
+                    last_addr = segment.addr
+                    print('%x' % last_addr)
+
+            def get_alignment_data_needed(segment):
+                # Actual alignment (in data bytes) required for a segment header: positioned so that
+                # after we write the next 8 byte header, file_offs % IROM_ALIGN == segment.addr % IROM_ALIGN
+                #
+                # (this is because the segment's vaddr may not be IROM_ALIGNed, more likely is aligned
+                # IROM_ALIGN+0x18 to account for the binary file header
+                align_past = (segment.addr % IROM_ALIGN) - self.SEG_HEADER_LEN
+                pad_len = (IROM_ALIGN - (f.tell() % IROM_ALIGN)) + align_past
+                if pad_len == 0 or pad_len == IROM_ALIGN:
+                    return 0  # already aligned
+
+                # subtract SEG_HEADER_LEN a second time, as the padding block has a header as well
+                pad_len -= self.SEG_HEADER_LEN
+                if pad_len < 0:
+                    pad_len += IROM_ALIGN
+                return pad_len
+
+            # try to fit each flash segment on a 64kB aligned boundary
+            # by padding with parts of the non-flash segments...
+            while len(flash_segments) > 0:
+                segment = flash_segments[0]
+                print('yyy %x' % segment.addr)
+                pad_len = 0
+                #get_alignment_data_needed(segment)
+                print('pad len %x' % pad_len)
+                if pad_len > 0:  # need to pad
+                    #if len(ram_segments) > 0 and pad_len > self.SEG_HEADER_LEN:
+                    #    pad_segment = ram_segments[0].split_image(pad_len)
+                    #    if len(ram_segments[0].data) == 0:
+                    #        ram_segments.pop(0)
+                    #else:
+                    pad_segment = ImageSegment(0, b'\x00' * pad_len, f.tell())
+                    print('qqq %x' % pad_segment.addr)
+                    checksum = self.save_segment(f, pad_segment, checksum)
+                    total_segments += 1
+                else:
+                    # write the flash segment
+                    #assert (f.tell() + 8) % IROM_ALIGN == segment.addr % IROM_ALIGN
+                    checksum = self.save_segment(f, segment, checksum)
+                    flash_segments.pop(0)
+                    total_segments += 1
+
+            # flash segments all written, so write any remaining RAM segments
+            for segment in ram_segments:
+                checksum = self.save_segment(f, segment, checksum)
+                total_segments += 1
+
+            # done writing segments
+            self.append_checksum(f, checksum)
+            # kinda hacky: go back to the initial header and write the new segment count
+            # that includes padding segments. This header is not checksummed
+            image_length = f.tell()
+            f.seek(1)
+            try:
+                f.write(chr(total_segments))
+            except TypeError:  # Python 3
+                f.write(bytes([total_segments]))
+
+            if self.append_digest:
+                # calculate the SHA256 of the whole file and append it
+                f.seek(0)
+                digest = hashlib.sha256()
+                digest.update(f.read(image_length))
+                f.write(digest.digest())
+
+            with open(filename, 'wb') as real_file:
+                real_file.write(f.getvalue())
+
+    def load_extended_header(self, load_file):
+        def split_byte(n):
+            return (n & 0x0F, (n >> 4) & 0x0F)
+
+        fields = list(struct.unpack(self.EXTENDED_HEADER_STRUCT_FMT, load_file.read(16)))
+
+        self.wp_pin = fields[0]
+
+        # SPI pin drive stengths are two per byte
+        self.clk_drv, self.q_drv = split_byte(fields[1])
+        self.d_drv, self.cs_drv = split_byte(fields[2])
+        self.hd_drv, self.wp_drv = split_byte(fields[3])
+
+        if fields[15] in [0, 1]:
+            self.append_digest = (fields[15] == 1)
+        else:
+            raise RuntimeError("Invalid value for append_digest field (0x%02x). Should be 0 or 1.", fields[15])
+
+        # remaining fields in the middle should all be zero
+        if any(f for f in fields[4:15] if f != 0):
+            print("Warning: some reserved header fields have non-zero values. This image may be from a newer esptool.py?")
+
+    def save_extended_header(self, save_file):
+        def join_byte(ln,hn):
+            return (ln & 0x0F) + ((hn & 0x0F) << 4)
+
+        append_digest = 1 if self.append_digest else 0
+
+        fields = [self.wp_pin,
+                  join_byte(self.clk_drv, self.q_drv),
+                  join_byte(self.d_drv, self.cs_drv),
+                  join_byte(self.hd_drv, self.wp_drv)]
+        fields += [0] * 11
+        fields += [append_digest]
+
+        packed = struct.pack(self.EXTENDED_HEADER_STRUCT_FMT, *fields)
+        save_file.write(packed)
 
 # Backwards compatibility for previous API, remove in esptool.py V3
 ESPFirmwareImage = ESP8266ROMFirmwareImage
@@ -2140,8 +2345,10 @@ def elf2image(args):
         image = ESP32FirmwareImage()
     elif args.version == '1':  # ESP8266
         image = ESP8266ROMFirmwareImage()
+    elif args.version == '2':  # ESP8266
+    	    image = ESP8266V2FirmwareImage()
     else:
-        image = ESP8266V2FirmwareImage()
+        image = ESP8266V3FirmwareImage()
     image.entrypoint = e.entrypoint
     image.segments = e.sections  # ELFSection is a subclass of ImageSegment
     image.flash_mode = {'qio':0, 'qout':1, 'dio':2, 'dout': 3}[args.flash_mode]
@@ -2406,7 +2613,7 @@ def main():
         help='Create an application image from ELF file')
     parser_elf2image.add_argument('input', help='Input ELF file')
     parser_elf2image.add_argument('--output', '-o', help='Output filename prefix (for version 1 image), or filename (for version 2 single image)', type=str)
-    parser_elf2image.add_argument('--version', '-e', help='Output image version', choices=['1','2'], default='1')
+    parser_elf2image.add_argument('--version', '-e', help='Output image version', choices=['1','2','3'], default='1')
 
     add_spi_flash_subparsers(parser_elf2image, is_elf2image=True)