feat(coredump): add esp32s2 and esp32c3 support

components/espcoredump/corefile/loader.py

@@ -1,5 +1,5 @@
 #
-# Copyright 2021 Espressif Systems (Shanghai) PTE LTD
+# Copyright 2021 Espressif Systems (Shanghai) CO., LTD
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -25,12 +25,18 @@ import tempfile
 
 from construct import AlignedStruct, Bytes, GreedyRange, Int32ul, Padding, Struct, abs_, this
 
-from . import ESPCoreDumpLoaderError, _ArchMethodsBase, _TargetMethodsBase
+from . import ESPCoreDumpLoaderError
 from .elf import (TASK_STATUS_CORRECT, TASK_STATUS_TCB_CORRUPTED, ElfFile, ElfSegment, ESPCoreDumpElfFile,
                   EspTaskStatus, NoteSection)
-from .xtensa import _ArchMethodsXtensa, _TargetMethodsESP32
+from .riscv import Esp32c3Methods
+from .xtensa import Esp32Methods, Esp32S2Methods
 
-IDF_PATH = os.getenv('IDF_PATH')
+try:
+    from typing import Optional, Tuple
+except ImportError:
+    pass
+
+IDF_PATH = os.getenv('IDF_PATH', '')
 PARTTOOL_PY = os.path.join(IDF_PATH, 'components', 'partition_table', 'parttool.py')
 ESPTOOL_PY = os.path.join(IDF_PATH, 'components', 'esptool_py', 'esptool', 'esptool.py')
 
@@ -74,12 +80,14 @@ class EspCoreDumpVersion(object):
     # This class contains all version-dependent params
     ESP32 = 0
     ESP32S2 = 2
-
     XTENSA_CHIPS = [ESP32, ESP32S2]
 
-    ESP_COREDUMP_TARGETS = XTENSA_CHIPS
+    ESP32C3 = 5
+    RISCV_CHIPS = [ESP32C3]
 
-    def __init__(self, version=None):
+    COREDUMP_SUPPORTED_TARGETS = XTENSA_CHIPS + RISCV_CHIPS
+
+    def __init__(self, version=None):  # type: (int) -> None
         """Constructor for core dump version
         """
         super(EspCoreDumpVersion, self).__init__()
@@ -89,26 +97,26 @@ class EspCoreDumpVersion(object):
             self.set_version(version)
 
     @staticmethod
-    def make_dump_ver(major, minor):
+    def make_dump_ver(major, minor):  # type: (int, int) -> int
         return ((major & 0xFF) << 8) | ((minor & 0xFF) << 0)
 
-    def set_version(self, version):
+    def set_version(self, version):  # type: (int) -> None
         self.version = version
 
     @property
-    def chip_ver(self):
+    def chip_ver(self):  # type: () -> int
         return (self.version & 0xFFFF0000) >> 16
 
     @property
-    def dump_ver(self):
+    def dump_ver(self):  # type: () -> int
         return self.version & 0x0000FFFF
 
     @property
-    def major(self):
+    def major(self):  # type: () -> int
         return (self.version & 0x0000FF00) >> 8
 
     @property
-    def minor(self):
+    def minor(self):  # type: () -> int
         return self.version & 0x000000FF
 
 
@@ -119,42 +127,37 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
     ELF_CRC32 = EspCoreDumpVersion.make_dump_ver(1, 0)
     ELF_SHA256 = EspCoreDumpVersion.make_dump_ver(1, 1)
 
-    def __init__(self):
+    def __init__(self):  # type: () -> None
         super(EspCoreDumpLoader, self).__init__()
-        self.core_src_file = None
+        self.core_src_file = None  # type: Optional[str]
         self.core_src_struct = None
         self.core_src = None
 
-        self.core_elf_file = None
+        self.core_elf_file = None  # type: Optional[str]
 
         self.header = None
         self.header_struct = EspCoreDumpV1Header
         self.checksum_struct = CRC
 
-        # These two method classes will be assigned in ``reload_coredump``
-        self.target_method_cls = _TargetMethodsBase
-        self.arch_method_cls = _ArchMethodsBase
+        # target classes will be assigned in ``_reload_coredump``
+        self.target_methods = Esp32Methods()
 
-        self._temp_files = []
+        self.temp_files = []  # type: list[str]
 
-    def __del__(self):
-        if self.core_src_file:
-            self.core_src_file.close()
-        if self.core_elf_file:
-            self.core_elf_file.close()
-        for f in self._temp_files:
-            try:
-                os.remove(f)
-            except OSError:
-                pass
-
-    def _create_temp_file(self):
+    def _create_temp_file(self):  # type: () -> str
         t = tempfile.NamedTemporaryFile('wb', delete=False)
-        self._temp_files.append(t.name)
-        return t
+        # Here we close this at first to make sure the read/write is wrapped in context manager
+        # Otherwise the result will be wrong if you read while open in another session
+        t.close()
+        self.temp_files.append(t.name)
+        return t.name
 
-    def _reload_coredump(self):
-        with open(self.core_src_file.name, 'rb') as fr:
+    def _load_core_src(self):  # type: () -> str
+        """
+        Write core elf into ``self.core_src``,
+        Return the target str by reading core elf
+        """
+        with open(self.core_src_file, 'rb') as fr:  # type: ignore
             coredump_bytes = fr.read()
 
         _header = EspCoreDumpV1Header.parse(coredump_bytes)  # first we use V1 format to get version
@@ -179,23 +182,28 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
             'data' / Bytes(this.header.tot_len - self.header_struct.sizeof() - self.checksum_struct.sizeof()),
             'checksum' / self.checksum_struct,
         )
-        self.core_src = self.core_src_struct.parse(coredump_bytes)
+        self.core_src = self.core_src_struct.parse(coredump_bytes)  # type: ignore
 
         # Reload header if header struct changes after parsing
         if self.header_struct != EspCoreDumpV1Header:
             self.header = EspCoreDumpV2Header.parse(coredump_bytes)
 
-        if self.chip_ver in self.ESP_COREDUMP_TARGETS:
+        if self.chip_ver in self.COREDUMP_SUPPORTED_TARGETS:
             if self.chip_ver == self.ESP32:
-                self.target_method_cls = _TargetMethodsESP32
-
-            if self.chip_ver in self.XTENSA_CHIPS:
-                self.arch_method_cls = _ArchMethodsXtensa
+                self.target_methods = Esp32Methods()  # type: ignore
+            elif self.chip_ver == self.ESP32S2:
+                self.target_methods = Esp32S2Methods()  # type: ignore
+            elif self.chip_ver == self.ESP32C3:
+                self.target_methods = Esp32c3Methods()  # type: ignore
+            else:
+                raise NotImplementedError
         else:
             raise ESPCoreDumpLoaderError('Core dump chip "0x%x" is not supported!' % self.chip_ver)
 
-    def _validate_dump_file(self):
-        if self.chip_ver not in self.ESP_COREDUMP_TARGETS:
+        return self.target_methods.TARGET  # type: ignore
+
+    def _validate_dump_file(self):  # type: () -> None
+        if self.chip_ver not in self.COREDUMP_SUPPORTED_TARGETS:
             raise ESPCoreDumpLoaderError('Invalid core dump chip version: "{}", should be <= "0x{:X}"'
                                          .format(self.chip_ver, self.ESP32S2))
 
@@ -204,20 +212,24 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
         elif self.checksum_struct == SHA256:
             self._sha256_validate()
 
-    def _crc_validate(self):
-        data_crc = binascii.crc32(EspCoreDumpV2Header.build(self.core_src.header) + self.core_src.data) & 0xffffffff
-        if data_crc != self.core_src.checksum:
-            raise ESPCoreDumpLoaderError('Invalid core dump CRC %x, should be %x' % (data_crc, self.core_src.crc))
+    def _crc_validate(self):  # type: () -> None
+        data_crc = binascii.crc32(
+            EspCoreDumpV2Header.build(self.core_src.header) + self.core_src.data) & 0xffffffff  # type: ignore
+        if data_crc != self.core_src.checksum:  # type: ignore
+            raise ESPCoreDumpLoaderError(
+                'Invalid core dump CRC %x, should be %x' % (data_crc, self.core_src.crc))  # type: ignore
 
-    def _sha256_validate(self):
-        data_sha256 = hashlib.sha256(EspCoreDumpV2Header.build(self.core_src.header) + self.core_src.data)
+    def _sha256_validate(self):  # type: () -> None
+        data_sha256 = hashlib.sha256(
+            EspCoreDumpV2Header.build(self.core_src.header) + self.core_src.data)  # type: ignore
         data_sha256_str = data_sha256.hexdigest()
-        sha256_str = binascii.hexlify(self.core_src.checksum).decode('ascii')
+        sha256_str = binascii.hexlify(self.core_src.checksum).decode('ascii')  # type: ignore
         if data_sha256_str != sha256_str:
             raise ESPCoreDumpLoaderError('Invalid core dump SHA256 "{}", should be "{}"'
                                          .format(data_sha256_str, sha256_str))
 
-    def create_corefile(self, exe_name=None):  # type: (str) -> None
+    def create_corefile(self, exe_name=None, e_machine=ESPCoreDumpElfFile.EM_XTENSA):
+        # type: (Optional[str], int) -> None
         """
         Creates core dump ELF file
         """
@@ -226,22 +238,21 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
 
         if self.dump_ver in [self.ELF_CRC32,
                              self.ELF_SHA256]:
-            self._extract_elf_corefile(exe_name)
+            self._extract_elf_corefile(exe_name, e_machine)
         elif self.dump_ver in [self.BIN_V1,
                                self.BIN_V2]:
-            self._extract_bin_corefile()
+            self._extract_bin_corefile(e_machine)
         else:
             raise NotImplementedError
 
-    def _extract_elf_corefile(self, exe_name=None):
+    def _extract_elf_corefile(self, exe_name=None, e_machine=ESPCoreDumpElfFile.EM_XTENSA):  # type: (str, int) -> None
         """
         Reads the ELF formatted core dump image and parse it
         """
-        self.core_elf_file.write(self.core_src.data)
-        # Need to be closed before read. Otherwise the result will be wrong
-        self.core_elf_file.close()
+        with open(self.core_elf_file, 'wb') as fw:  # type: ignore
+            fw.write(self.core_src.data)  # type: ignore
 
-        core_elf = ESPCoreDumpElfFile(self.core_elf_file.name)
+        core_elf = ESPCoreDumpElfFile(self.core_elf_file, e_machine=e_machine)  # type: ignore
 
         # Read note segments from core file which are belong to tasks (TCB or stack)
         for seg in core_elf.note_segments:
@@ -259,7 +270,7 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
                     coredump_sha256 = coredump_sha256_struct.parse(note_sec.desc[:coredump_sha256_struct.sizeof()])
                     if coredump_sha256.sha256 != app_sha256:
                         raise ESPCoreDumpLoaderError(
-                            'Invalid application image for coredump: coredump SHA256({}) != app SHA256({}).'
+                            'Invalid application image for coredump: coredump SHA256({!r}) != app SHA256({!r}).'
                             .format(coredump_sha256, app_sha256))
                     if coredump_sha256.ver != self.version:
                         raise ESPCoreDumpLoaderError(
@@ -267,46 +278,43 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
                             .format(coredump_sha256.ver, self.version))
 
     @staticmethod
-    def _get_aligned_size(size, align_with=4):
+    def _get_aligned_size(size, align_with=4):  # type: (int, int) -> int
         if size % align_with:
             return align_with * (size // align_with + 1)
         return size
 
     @staticmethod
-    def _build_note_section(name, sec_type, desc):
-        name = bytearray(name, encoding='ascii') + b'\0'
-        return NoteSection.build({
-            'namesz': len(name),
+    def _build_note_section(name, sec_type, desc):  # type: (str, int, str) -> bytes
+        b_name = bytearray(name, encoding='ascii') + b'\0'
+        return NoteSection.build({  # type: ignore
+            'namesz': len(b_name),
             'descsz': len(desc),
             'type': sec_type,
-            'name': name,
+            'name': b_name,
             'desc': desc,
         })
 
-    def _extract_bin_corefile(self):
+    def _extract_bin_corefile(self, e_machine=ESPCoreDumpElfFile.EM_XTENSA):  # type: (int) -> None
         """
         Creates core dump ELF file
         """
-        tcbsz_aligned = self._get_aligned_size(self.header.tcbsz)
-
         coredump_data_struct = Struct(
             'tasks' / GreedyRange(
                 AlignedStruct(
                     4,
                     'task_header' / TaskHeader,
-                    'tcb' / Bytes(self.header.tcbsz),
-                    'stack' / Bytes(abs_(this.task_header.stack_top - this.task_header.stack_end)),
+                    'tcb' / Bytes(self.header.tcbsz),  # type: ignore
+                    'stack' / Bytes(abs_(this.task_header.stack_top - this.task_header.stack_end)),  # type: ignore
                 )
             ),
-            'mem_seg_headers' / MemSegmentHeader[self.core_src.header.segs_num]
+            'mem_seg_headers' / MemSegmentHeader[self.core_src.header.segs_num]  # type: ignore
         )
-
-        core_elf = ESPCoreDumpElfFile()
+        core_elf = ESPCoreDumpElfFile(e_machine=e_machine)
         notes = b''
         core_dump_info_notes = b''
         task_info_notes = b''
 
-        coredump_data = coredump_data_struct.parse(self.core_src.data)
+        coredump_data = coredump_data_struct.parse(self.core_src.data)  # type: ignore
         for i, task in enumerate(coredump_data.tasks):
             stack_len_aligned = self._get_aligned_size(abs(task.task_header.stack_top - task.task_header.stack_end))
             task_status_kwargs = {
@@ -314,32 +322,34 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
                 'task_flags': TASK_STATUS_CORRECT,
                 'task_tcb_addr': task.task_header.tcb_addr,
                 'task_stack_start': min(task.task_header.stack_top, task.task_header.stack_end),
+                'task_stack_end': max(task.task_header.stack_top, task.task_header.stack_end),
                 'task_stack_len': stack_len_aligned,
                 'task_name': Padding(16).build({})  # currently we don't have task_name, keep it as padding
             }
 
             # Write TCB
             try:
-                if self.target_method_cls.tcb_is_sane(task.task_header.tcb_addr, tcbsz_aligned):
+                if self.target_methods.tcb_is_sane(task.task_header.tcb_addr, self.header.tcbsz):  # type: ignore
                     core_elf.add_segment(task.task_header.tcb_addr,
                                          task.tcb,
                                          ElfFile.PT_LOAD,
                                          ElfSegment.PF_R | ElfSegment.PF_W)
-                elif task.task_header.tcb_addr and self.target_method_cls.addr_is_fake(task.task_header.tcb_addr):
+                elif task.task_header.tcb_addr and self.target_methods.addr_is_fake(task.task_header.tcb_addr):
                     task_status_kwargs['task_flags'] |= TASK_STATUS_TCB_CORRUPTED
             except ESPCoreDumpLoaderError as e:
                 logging.warning('Skip TCB {} bytes @ 0x{:x}. (Reason: {})'
-                                .format(tcbsz_aligned, task.task_header.tcb_addr, e))
+                                .format(self.header.tcbsz, task.task_header.tcb_addr, e))  # type: ignore
 
             # Write stack
             try:
-                if self.target_method_cls.stack_is_sane(task_status_kwargs['task_stack_start']):
+                if self.target_methods.stack_is_sane(task_status_kwargs['task_stack_start'],
+                                                     task_status_kwargs['task_stack_end']):
                     core_elf.add_segment(task_status_kwargs['task_stack_start'],
                                          task.stack,
                                          ElfFile.PT_LOAD,
                                          ElfSegment.PF_R | ElfSegment.PF_W)
-                elif task_status_kwargs['task_stack_start'] \
-                        and self.target_method_cls.addr_is_fake(task_status_kwargs['task_stack_start']):
+                elif (task_status_kwargs['task_stack_start']
+                      and self.target_methods.addr_is_fake(task_status_kwargs['task_stack_start'])):
                     task_status_kwargs['task_flags'] |= TASK_STATUS_TCB_CORRUPTED
                     core_elf.add_segment(task_status_kwargs['task_stack_start'],
                                          task.stack,
@@ -355,7 +365,7 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
             try:
                 logging.debug('Stack start_end: 0x{:x} @ 0x{:x}'
                               .format(task.task_header.stack_top, task.task_header.stack_end))
-                task_regs, extra_regs = self.arch_method_cls.get_registers_from_stack(
+                task_regs, extra_regs = self.target_methods.get_registers_from_stack(
                     task.stack,
                     task.task_header.stack_end > task.task_header.stack_top
                 )
@@ -367,23 +377,24 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
                                                         EspTaskStatus.build(task_status_kwargs))
             notes += self._build_note_section('CORE',
                                               ElfFile.PT_LOAD,
-                                              self.arch_method_cls.build_prstatus_data(task.task_header.tcb_addr,
-                                                                                       task_regs))
+                                              self.target_methods.build_prstatus_data(task.task_header.tcb_addr,
+                                                                                      task_regs))
 
-            if extra_regs and len(core_dump_info_notes) == 0:
-                # actually there will be only one such note - for crashed task
+            if len(core_dump_info_notes) == 0:  # the first task is the crashed task
                 core_dump_info_notes += self._build_note_section('ESP_CORE_DUMP_INFO',
                                                                  ESPCoreDumpElfFile.PT_INFO,
-                                                                 Int32ul.build(self.header.ver))
+                                                                 Int32ul.build(self.header.ver))  # type: ignore
+                _regs = [task.task_header.tcb_addr]
+
+                # For xtensa, we need to put the exception registers into the extra info as well
+                if e_machine == ESPCoreDumpElfFile.EM_XTENSA and extra_regs:
+                    for reg_id in extra_regs:
+                        _regs.extend([reg_id, extra_regs[reg_id]])
 
-                exc_regs = []
-                for reg_id in extra_regs:
-                    exc_regs.extend([reg_id, extra_regs[reg_id]])
-                _regs = [task.task_header.tcb_addr] + exc_regs
                 core_dump_info_notes += self._build_note_section(
                     'EXTRA_INFO',
                     ESPCoreDumpElfFile.PT_EXTRA_INFO,
-                    Int32ul[1 + len(exc_regs)].build(_regs)
+                    Int32ul[len(_regs)].build(_regs)
                 )
 
         if self.dump_ver == self.BIN_V2:
@@ -409,30 +420,29 @@ class EspCoreDumpLoader(EspCoreDumpVersion):
                             .format(len(task_info_notes), 0, e))
         # dump core ELF
         core_elf.e_type = ElfFile.ET_CORE
-        core_elf.e_machine = ESPCoreDumpElfFile.EM_XTENSA
-        core_elf.dump(self.core_elf_file.name)
+        core_elf.dump(self.core_elf_file)  # type: ignore
 
 
 class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
     ESP_COREDUMP_PART_TABLE_OFF = 0x8000
 
-    def __init__(self, offset, target='esp32', port=None, baud=None):
+    def __init__(self, offset, target=None, port=None, baud=None):
+        # type: (int, Optional[str], Optional[str], Optional[int]) -> None
         super(ESPCoreDumpFlashLoader, self).__init__()
         self.port = port
         self.baud = baud
-        self.target = target
 
-        self._get_coredump(offset)
-        self._reload_coredump()
+        self._get_core_src(offset, target)
+        self.target = self._load_core_src()
 
-    def _get_coredump(self, off):
+    def _get_core_src(self, off, target=None):  # type: (int, Optional[str]) -> None
         """
         Loads core dump from flash using parttool or elftool (if offset is set)
         """
         try:
             if off:
                 logging.info('Invoke esptool to read image.')
-                self._invoke_esptool(off=off)
+                self._invoke_esptool(off=off, target=target)
             else:
                 logging.info('Invoke parttool to read image.')
                 self._invoke_parttool()
@@ -440,15 +450,14 @@ class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
             if e.output:
                 logging.info(e.output)
             logging.error('Error during the subprocess execution')
-        else:
-            # Need to be closed before read. Otherwise the result will be wrong
-            self.core_src_file.close()
 
-    def _invoke_esptool(self, off=None):
+    def _invoke_esptool(self, off=None, target=None):  # type: (Optional[int], Optional[str]) -> None
         """
         Loads core dump from flash using elftool
         """
-        tool_args = [sys.executable, ESPTOOL_PY, '-c', self.target]
+        if target is None:
+            target = 'auto'
+        tool_args = [sys.executable, ESPTOOL_PY, '-c', target]
         if self.port:
             tool_args.extend(['-p', self.port])
         if self.baud:
@@ -466,14 +475,14 @@ class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
 
             # Here we use V1 format to locate the size
             tool_args.extend(['read_flash', str(off), str(EspCoreDumpV1Header.sizeof())])
-            tool_args.append(self.core_src_file.name)
+            tool_args.append(self.core_src_file)  # type: ignore
 
             # read core dump length
             et_out = subprocess.check_output(tool_args)
             if et_out:
                 logging.info(et_out.decode('utf-8'))
 
-            header = EspCoreDumpV1Header.parse(open(self.core_src_file.name, 'rb').read())
+            header = EspCoreDumpV1Header.parse(open(self.core_src_file, 'rb').read())  # type: ignore
             if not header or not 0 < header.tot_len <= part_size:
                 logging.error('Incorrect size of core dump image: {}, use partition size instead: {}'
                               .format(header.tot_len, part_size))
@@ -492,7 +501,7 @@ class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
             logging.debug(e.output)
             raise e
 
-    def _invoke_parttool(self):
+    def _invoke_parttool(self):  # type: () -> None
         """
         Loads core dump from flash using parttool
         """
@@ -503,7 +512,7 @@ class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
 
         self.core_src_file = self._create_temp_file()
         try:
-            tool_args.append(self.core_src_file.name)
+            tool_args.append(self.core_src_file)  # type: ignore
             # read core dump partition
             et_out = subprocess.check_output(tool_args)
             if et_out:
@@ -515,7 +524,7 @@ class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
             logging.debug(e.output)
             raise e
 
-    def _get_core_dump_partition_info(self, part_off=None):
+    def _get_core_dump_partition_info(self, part_off=None):  # type: (Optional[int]) -> Tuple[int, int]
         """
         Get core dump partition info using parttool
         """
@@ -545,28 +554,27 @@ class ESPCoreDumpFlashLoader(EspCoreDumpLoader):
 
 
 class ESPCoreDumpFileLoader(EspCoreDumpLoader):
-    def __init__(self, path, is_b64=False):
+    def __init__(self, path, is_b64=False):  # type: (str, bool) -> None
         super(ESPCoreDumpFileLoader, self).__init__()
         self.is_b64 = is_b64
 
-        self._get_coredump(path)
-        self._reload_coredump()
+        self._get_core_src(path)
+        self.target = self._load_core_src()
 
-    def _get_coredump(self, path):
+    def _get_core_src(self, path):  # type: (str) -> None
         """
         Loads core dump from (raw binary or base64-encoded) file
         """
         logging.debug('Load core dump from "%s", %s format', path, 'b64' if self.is_b64 else 'raw')
         if not self.is_b64:
-            self.core_src_file = open(path, mode='rb')
+            self.core_src_file = path
         else:
             self.core_src_file = self._create_temp_file()
-            with open(path, 'rb') as fb64:
-                while True:
-                    line = fb64.readline()
-                    if len(line) == 0:
-                        break
-                    data = base64.standard_b64decode(line.rstrip(b'\r\n'))
-                    self.core_src_file.write(data)
-                self.core_src_file.flush()
-                self.core_src_file.seek(0)
+            with open(self.core_src_file, 'wb') as fw:
+                with open(path, 'rb') as fb64:
+                    while True:
+                        line = fb64.readline()
+                        if len(line) == 0:
+                            break
+                        data = base64.standard_b64decode(line.rstrip(b'\r\n'))
+                        fw.write(data)  # type: ignore