1# SPDX-FileCopyrightText: 2014-2023 Fredrik Ahlberg, Angus Gratton,
2# Espressif Systems (Shanghai) CO LTD, other contributors as noted.
3#
4# SPDX-License-Identifier: GPL-2.0-or-later
5
6import os
7import struct
8from typing import Dict
9
10from .esp32 import ESP32ROM
11from ..loader import ESPLoader
12from ..reset import HardReset
13from ..util import FatalError, NotImplementedInROMError
14
15
16class ESP32S2ROM(ESP32ROM):
17    CHIP_NAME = "ESP32-S2"
18    IMAGE_CHIP_ID = 2
19
20    IROM_MAP_START = 0x40080000
21    IROM_MAP_END = 0x40B80000
22    DROM_MAP_START = 0x3F000000
23    DROM_MAP_END = 0x3F3F0000
24
25    CHIP_DETECT_MAGIC_VALUE = [0x000007C6]
26
27    SPI_REG_BASE = 0x3F402000
28    SPI_USR_OFFS = 0x18
29    SPI_USR1_OFFS = 0x1C
30    SPI_USR2_OFFS = 0x20
31    SPI_MOSI_DLEN_OFFS = 0x24
32    SPI_MISO_DLEN_OFFS = 0x28
33    SPI_W0_OFFS = 0x58
34
35    SPI_ADDR_REG_MSB = False
36
37    MAC_EFUSE_REG = 0x3F41A044  # ESP32-S2 has special block for MAC efuses
38
39    UART_CLKDIV_REG = 0x3F400014
40
41    SUPPORTS_ENCRYPTED_FLASH = True
42
43    FLASH_ENCRYPTED_WRITE_ALIGN = 16
44
45    # todo: use espefuse APIs to get this info
46    EFUSE_BASE = 0x3F41A000
47    EFUSE_RD_REG_BASE = EFUSE_BASE + 0x030  # BLOCK0 read base address
48    EFUSE_BLOCK1_ADDR = EFUSE_BASE + 0x044
49    EFUSE_BLOCK2_ADDR = EFUSE_BASE + 0x05C
50
51    EFUSE_PURPOSE_KEY0_REG = EFUSE_BASE + 0x34
52    EFUSE_PURPOSE_KEY0_SHIFT = 24
53    EFUSE_PURPOSE_KEY1_REG = EFUSE_BASE + 0x34
54    EFUSE_PURPOSE_KEY1_SHIFT = 28
55    EFUSE_PURPOSE_KEY2_REG = EFUSE_BASE + 0x38
56    EFUSE_PURPOSE_KEY2_SHIFT = 0
57    EFUSE_PURPOSE_KEY3_REG = EFUSE_BASE + 0x38
58    EFUSE_PURPOSE_KEY3_SHIFT = 4
59    EFUSE_PURPOSE_KEY4_REG = EFUSE_BASE + 0x38
60    EFUSE_PURPOSE_KEY4_SHIFT = 8
61    EFUSE_PURPOSE_KEY5_REG = EFUSE_BASE + 0x38
62    EFUSE_PURPOSE_KEY5_SHIFT = 12
63
64    EFUSE_DIS_DOWNLOAD_MANUAL_ENCRYPT_REG = EFUSE_RD_REG_BASE
65    EFUSE_DIS_DOWNLOAD_MANUAL_ENCRYPT = 1 << 19
66
67    EFUSE_SPI_BOOT_CRYPT_CNT_REG = EFUSE_BASE + 0x034
68    EFUSE_SPI_BOOT_CRYPT_CNT_MASK = 0x7 << 18
69
70    EFUSE_SECURE_BOOT_EN_REG = EFUSE_BASE + 0x038
71    EFUSE_SECURE_BOOT_EN_MASK = 1 << 20
72
73    EFUSE_RD_REPEAT_DATA3_REG = EFUSE_BASE + 0x3C
74    EFUSE_RD_REPEAT_DATA3_REG_FLASH_TYPE_MASK = 1 << 9
75
76    PURPOSE_VAL_XTS_AES256_KEY_1 = 2
77    PURPOSE_VAL_XTS_AES256_KEY_2 = 3
78    PURPOSE_VAL_XTS_AES128_KEY = 4
79
80    UARTDEV_BUF_NO = 0x3FFFFD14  # Variable in ROM .bss which indicates the port in use
81    UARTDEV_BUF_NO_USB_OTG = 2  # Value of the above indicating that USB-OTG is in use
82
83    USB_RAM_BLOCK = 0x800  # Max block size USB-OTG is used
84
85    GPIO_STRAP_REG = 0x3F404038
86    GPIO_STRAP_SPI_BOOT_MASK = 0x8  # Not download mode
87    GPIO_STRAP_VDDSPI_MASK = 1 << 4
88    RTC_CNTL_OPTION1_REG = 0x3F408128
89    RTC_CNTL_FORCE_DOWNLOAD_BOOT_MASK = 0x1  # Is download mode forced over USB?
90
91    MEMORY_MAP = [
92        [0x00000000, 0x00010000, "PADDING"],
93        [0x3F000000, 0x3FF80000, "DROM"],
94        [0x3F500000, 0x3FF80000, "EXTRAM_DATA"],
95        [0x3FF9E000, 0x3FFA0000, "RTC_DRAM"],
96        [0x3FF9E000, 0x40000000, "BYTE_ACCESSIBLE"],
97        [0x3FF9E000, 0x40072000, "MEM_INTERNAL"],
98        [0x3FFB0000, 0x40000000, "DRAM"],
99        [0x40000000, 0x4001A100, "IROM_MASK"],
100        [0x40020000, 0x40070000, "IRAM"],
101        [0x40070000, 0x40072000, "RTC_IRAM"],
102        [0x40080000, 0x40800000, "IROM"],
103        [0x50000000, 0x50002000, "RTC_DATA"],
104    ]
105
106    EFUSE_VDD_SPI_REG = EFUSE_BASE + 0x34
107    VDD_SPI_XPD = 1 << 4
108    VDD_SPI_TIEH = 1 << 5
109    VDD_SPI_FORCE = 1 << 6
110
111    UF2_FAMILY_ID = 0xBFDD4EEE
112
113    EFUSE_MAX_KEY = 5
114    KEY_PURPOSES: Dict[int, str] = {
115        0: "USER/EMPTY",
116        1: "RESERVED",
117        2: "XTS_AES_256_KEY_1",
118        3: "XTS_AES_256_KEY_2",
119        4: "XTS_AES_128_KEY",
120        5: "HMAC_DOWN_ALL",
121        6: "HMAC_DOWN_JTAG",
122        7: "HMAC_DOWN_DIGITAL_SIGNATURE",
123        8: "HMAC_UP",
124        9: "SECURE_BOOT_DIGEST0",
125        10: "SECURE_BOOT_DIGEST1",
126        11: "SECURE_BOOT_DIGEST2",
127    }
128
129    def get_pkg_version(self):
130        num_word = 4
131        return (self.read_reg(self.EFUSE_BLOCK1_ADDR + (4 * num_word)) >> 0) & 0x0F
132
133    def get_minor_chip_version(self):
134        hi_num_word = 3
135        hi = (self.read_reg(self.EFUSE_BLOCK1_ADDR + (4 * hi_num_word)) >> 20) & 0x01
136        low_num_word = 4
137        low = (self.read_reg(self.EFUSE_BLOCK1_ADDR + (4 * low_num_word)) >> 4) & 0x07
138        return (hi << 3) + low
139
140    def get_major_chip_version(self):
141        num_word = 3
142        return (self.read_reg(self.EFUSE_BLOCK1_ADDR + (4 * num_word)) >> 18) & 0x03
143
144    def get_flash_version(self):
145        num_word = 3
146        return (self.read_reg(self.EFUSE_BLOCK1_ADDR + (4 * num_word)) >> 21) & 0x0F
147
148    def get_flash_cap(self):
149        return self.get_flash_version()
150
151    def get_psram_version(self):
152        num_word = 3
153        return (self.read_reg(self.EFUSE_BLOCK1_ADDR + (4 * num_word)) >> 28) & 0x0F
154
155    def get_psram_cap(self):
156        return self.get_psram_version()
157
158    def get_block2_version(self):
159        # BLK_VERSION_MINOR
160        num_word = 4
161        return (self.read_reg(self.EFUSE_BLOCK2_ADDR + (4 * num_word)) >> 4) & 0x07
162
163    def get_chip_description(self):
164        chip_name = {
165            0: "ESP32-S2",
166            1: "ESP32-S2FH2",
167            2: "ESP32-S2FH4",
168            102: "ESP32-S2FNR2",
169            100: "ESP32-S2R2",
170        }.get(
171            self.get_flash_cap() + self.get_psram_cap() * 100,
172            "unknown ESP32-S2",
173        )
174        major_rev = self.get_major_chip_version()
175        minor_rev = self.get_minor_chip_version()
176        return f"{chip_name} (revision v{major_rev}.{minor_rev})"
177
178    def get_chip_features(self):
179        features = ["WiFi"]
180
181        if self.secure_download_mode:
182            features += ["Secure Download Mode Enabled"]
183
184        flash_version = {
185            0: "No Embedded Flash",
186            1: "Embedded Flash 2MB",
187            2: "Embedded Flash 4MB",
188        }.get(self.get_flash_cap(), "Unknown Embedded Flash")
189        features += [flash_version]
190
191        psram_version = {
192            0: "No Embedded PSRAM",
193            1: "Embedded PSRAM 2MB",
194            2: "Embedded PSRAM 4MB",
195        }.get(self.get_psram_cap(), "Unknown Embedded PSRAM")
196        features += [psram_version]
197
198        block2_version = {
199            0: "No calibration in BLK2 of efuse",
200            1: "ADC and temperature sensor calibration in BLK2 of efuse V1",
201            2: "ADC and temperature sensor calibration in BLK2 of efuse V2",
202        }.get(self.get_block2_version(), "Unknown Calibration in BLK2")
203        features += [block2_version]
204
205        return features
206
207    def get_crystal_freq(self):
208        # ESP32-S2 XTAL is fixed to 40MHz
209        return 40
210
211    def _get_rtc_cntl_flash_voltage(self):
212        return None  # not supported on ESP32-S2
213
214    def override_vddsdio(self, new_voltage):
215        raise NotImplementedInROMError(
216            "VDD_SDIO overrides are not supported for ESP32-S2"
217        )
218
219    def read_mac(self, mac_type="BASE_MAC"):
220        """Read MAC from EFUSE region"""
221        if mac_type != "BASE_MAC":
222            return None
223        mac0 = self.read_reg(self.MAC_EFUSE_REG)
224        mac1 = self.read_reg(self.MAC_EFUSE_REG + 4)  # only bottom 16 bits are MAC
225        bitstring = struct.pack(">II", mac1, mac0)[2:]
226        return tuple(bitstring)
227
228    def flash_type(self):
229        return (
230            1
231            if self.read_reg(self.EFUSE_RD_REPEAT_DATA3_REG)
232            & self.EFUSE_RD_REPEAT_DATA3_REG_FLASH_TYPE_MASK
233            else 0
234        )
235
236    def get_flash_crypt_config(self):
237        return None  # doesn't exist on ESP32-S2
238
239    def get_secure_boot_enabled(self):
240        return (
241            self.read_reg(self.EFUSE_SECURE_BOOT_EN_REG)
242            & self.EFUSE_SECURE_BOOT_EN_MASK
243        )
244
245    def get_key_block_purpose(self, key_block):
246        if key_block < 0 or key_block > self.EFUSE_MAX_KEY:
247            raise FatalError(
248                f"Valid key block numbers must be in range 0-{self.EFUSE_MAX_KEY}"
249            )
250
251        reg, shift = [
252            (self.EFUSE_PURPOSE_KEY0_REG, self.EFUSE_PURPOSE_KEY0_SHIFT),
253            (self.EFUSE_PURPOSE_KEY1_REG, self.EFUSE_PURPOSE_KEY1_SHIFT),
254            (self.EFUSE_PURPOSE_KEY2_REG, self.EFUSE_PURPOSE_KEY2_SHIFT),
255            (self.EFUSE_PURPOSE_KEY3_REG, self.EFUSE_PURPOSE_KEY3_SHIFT),
256            (self.EFUSE_PURPOSE_KEY4_REG, self.EFUSE_PURPOSE_KEY4_SHIFT),
257            (self.EFUSE_PURPOSE_KEY5_REG, self.EFUSE_PURPOSE_KEY5_SHIFT),
258        ][key_block]
259        return (self.read_reg(reg) >> shift) & 0xF
260
261    def is_flash_encryption_key_valid(self):
262        # Need to see either an AES-128 key or two AES-256 keys
263        purposes = [
264            self.get_key_block_purpose(b) for b in range(self.EFUSE_MAX_KEY + 1)
265        ]
266
267        if any(p == self.PURPOSE_VAL_XTS_AES128_KEY for p in purposes):
268            return True
269
270        return any(p == self.PURPOSE_VAL_XTS_AES256_KEY_1 for p in purposes) and any(
271            p == self.PURPOSE_VAL_XTS_AES256_KEY_2 for p in purposes
272        )
273
274    def uses_usb_otg(self):
275        """
276        Check the UARTDEV_BUF_NO register to see if USB-OTG console is being used
277        """
278        if self.secure_download_mode:
279            return False  # can't detect native USB in secure download mode
280        return self.get_uart_no() == self.UARTDEV_BUF_NO_USB_OTG
281
282    def _post_connect(self):
283        if self.uses_usb_otg():
284            self.ESP_RAM_BLOCK = self.USB_RAM_BLOCK
285
286    def _check_if_can_reset(self):
287        """
288        Check the strapping register to see if we can reset out of download mode.
289        """
290        if os.getenv("ESPTOOL_TESTING") is not None:
291            print("ESPTOOL_TESTING is set, ignoring strapping mode check")
292            # Esptool tests over USB-OTG run with GPIO0 strapped low,
293            # don't complain in this case.
294            return
295        strap_reg = self.read_reg(self.GPIO_STRAP_REG)
296        force_dl_reg = self.read_reg(self.RTC_CNTL_OPTION1_REG)
297        if (
298            strap_reg & self.GPIO_STRAP_SPI_BOOT_MASK == 0
299            and force_dl_reg & self.RTC_CNTL_FORCE_DOWNLOAD_BOOT_MASK == 0
300        ):
301            raise SystemExit(
302                f"Error: {self.get_chip_description()} chip was placed into download "
303                "mode using GPIO0.\nesptool.py can not exit the download mode over "
304                "USB. To run the app, reset the chip manually.\n"
305                "To suppress this note, set --after option to 'no_reset'."
306            )
307
308    def hard_reset(self):
309        uses_usb_otg = self.uses_usb_otg()
310        if uses_usb_otg:
311            self._check_if_can_reset()
312
313        print("Hard resetting via RTS pin...")
314        HardReset(self._port, uses_usb_otg)()
315
316    def change_baud(self, baud):
317        ESPLoader.change_baud(self, baud)
318
319    def check_spi_connection(self, spi_connection):
320        if not set(spi_connection).issubset(set(range(0, 22)) | set(range(26, 47))):
321            raise FatalError("SPI Pin numbers must be in the range 0-21, or 26-46.")
322        if any([v for v in spi_connection if v in [19, 20]]):
323            print(
324                "WARNING: GPIO pins 19 and 20 are used by USB-OTG, "
325                "consider using other pins for SPI flash connection."
326            )
327
328
329class ESP32S2StubLoader(ESP32S2ROM):
330    """Access class for ESP32-S2 stub loader, runs on top of ROM.
331
332    (Basically the same as ESP32StubLoader, but different base class.
333    Can possibly be made into a mixin.)
334    """
335
336    FLASH_WRITE_SIZE = 0x4000  # matches MAX_WRITE_BLOCK in stub_loader.c
337    STATUS_BYTES_LENGTH = 2  # same as ESP8266, different to ESP32 ROM
338    IS_STUB = True
339
340    def __init__(self, rom_loader):
341        self.secure_download_mode = rom_loader.secure_download_mode
342        self._port = rom_loader._port
343        self._trace_enabled = rom_loader._trace_enabled
344        self.cache = rom_loader.cache
345        self.flush_input()  # resets _slip_reader
346
347        if rom_loader.uses_usb_otg():
348            self.ESP_RAM_BLOCK = self.USB_RAM_BLOCK
349            self.FLASH_WRITE_SIZE = self.USB_RAM_BLOCK
350
351
352ESP32S2ROM.STUB_CLASS = ESP32S2StubLoader
353