From 5d873444b2cf801ba73f4a457993260df3a412b8 Mon Sep 17 00:00:00 2001
From: Adeeb Shihadeh <adeebshihadeh@gmail.com>
Date: Sun, 16 Jul 2023 19:33:18 -0700
Subject: [PATCH] SPI: connect by VERSION command (#1495)

* SPI: connect by VERSION command

* shorter timeout

* add exception

* simple test

* fallback

* bootstub check

* update comments

---------

Co-authored-by: Comma Device <device@comma.ai>
---
 __init__.py         |  1 +
 python/__init__.py  | 32 ++++++++++++++++++++++++++------
 python/spi.py       |  8 +++++++-
 tests/hitl/8_spi.py | 16 ++++++++++++++--
 4 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/__init__.py b/__init__.py
index 21e35b9e98..19cbee3e4f 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,4 +1,5 @@
 from .python.constants import McuType, BASEDIR, FW_PATH  # noqa: F401
+from .python.spi import PandaSpiException, PandaProtocolMismatch  # noqa: F401
 from .python.serial import PandaSerial  # noqa: F401
 from .python import (Panda, PandaDFU, # noqa: F401
                      pack_can_buffer, unpack_can_buffer, calculate_checksum, unpack_log,
diff --git a/python/__init__.py b/python/__init__.py
index 26a70ff260..0349d161b7 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -17,7 +17,7 @@
 from .constants import FW_PATH, McuType
 from .dfu import PandaDFU
 from .isotp import isotp_send, isotp_recv
-from .spi import PandaSpiHandle, PandaSpiException
+from .spi import PandaSpiHandle, PandaSpiException, PandaProtocolMismatch
 from .usb import PandaUsbHandle
 
 __version__ = '0.0.10'
@@ -326,16 +326,30 @@ def connect(self, claim=True, wait=False):
       self.set_power_save(0)
 
   @staticmethod
-  def spi_connect(serial):
+  def spi_connect(serial, ignore_version=False):
     # get UID to confirm slave is present and up
     handle = None
     spi_serial = None
     bootstub = None
+    spi_version = None
     try:
       handle = PandaSpiHandle()
-      dat = handle.controlRead(Panda.REQUEST_IN, 0xc3, 0, 0, 12, timeout=100)
-      spi_serial = binascii.hexlify(dat).decode()
-      bootstub = Panda.flasher_present(handle)
+
+      # connect by protcol version
+      try:
+        dat = handle.get_protocol_version()
+        spi_serial = binascii.hexlify(dat[:12]).decode()
+        pid = dat[13]
+        if pid not in (0xcc, 0xee):
+          raise PandaSpiException("invalid bootstub status")
+        bootstub = pid == 0xee
+        spi_version = dat[14]
+      except PandaSpiException:
+        # fallback, we'll raise a protocol mismatch below
+        dat = handle.controlRead(Panda.REQUEST_IN, 0xc3, 0, 0, 12, timeout=100)
+        spi_serial = binascii.hexlify(dat).decode()
+        bootstub = Panda.flasher_present(handle)
+        spi_version = 0
     except PandaSpiException:
       pass
 
@@ -345,6 +359,12 @@ def spi_connect(serial):
       spi_serial = None
       bootstub = False
 
+    # ensure our protocol version matches the panda
+    if handle is not None and not ignore_version:
+      if spi_version != handle.PROTOCOL_VERSION:
+        err = f"panda protocol mismatch: expected {handle.PROTOCOL_VERSION}, got {spi_version}. reflash panda"
+        raise PandaProtocolMismatch(err)
+
     return handle, spi_serial, bootstub, None
 
   @staticmethod
@@ -416,7 +436,7 @@ def usb_list():
 
   @staticmethod
   def spi_list():
-    _, serial, _, _ = Panda.spi_connect(None)
+    _, serial, _, _ = Panda.spi_connect(None, ignore_version=True)
     if serial is not None:
       return [serial, ]
     return []
diff --git a/python/spi.py b/python/spi.py
index b73aea1dba..55bfa6d2e4 100644
--- a/python/spi.py
+++ b/python/spi.py
@@ -40,6 +40,9 @@
 class PandaSpiException(Exception):
   pass
 
+class PandaProtocolMismatch(PandaSpiException):
+  pass
+
 class PandaSpiUnavailable(PandaSpiException):
   pass
 
@@ -109,6 +112,9 @@ class PandaSpiHandle(BaseHandle):
   """
   A class that mimics a libusb1 handle for panda SPI communications.
   """
+
+  PROTOCOL_VERSION = 1
+
   def __init__(self):
     self.dev = SpiDevice()
 
@@ -225,7 +231,7 @@ def _get_version(spi) -> bytes:
         version_bytes = spi.readbytes(len(vers_str) + 2)
         if bytes(version_bytes).startswith(vers_str):
           break
-        if (time.monotonic() - start) > 0.5:
+        if (time.monotonic() - start) > 0.01:
           raise PandaSpiMissingAck
 
       rlen = struct.unpack("<H", bytes(version_bytes[-2:]))[0]
diff --git a/tests/hitl/8_spi.py b/tests/hitl/8_spi.py
index 04d674f2a6..b3ef487ae2 100644
--- a/tests/hitl/8_spi.py
+++ b/tests/hitl/8_spi.py
@@ -4,7 +4,7 @@
 from unittest.mock import patch
 
 from panda import Panda
-from panda.python.spi import PandaSpiNackResponse
+from panda.python.spi import PandaProtocolMismatch, PandaSpiNackResponse
 
 pytestmark = [
   pytest.mark.test_panda_types((Panda.HW_TYPE_TRES, ))
@@ -19,7 +19,19 @@ def _ping(self, mocker, panda):
     mocker.stop(spy)
 
   @pytest.mark.expected_logs(2)
-  def test_protocol_version(self, p):
+  def test_protocol_version_check(self, p):
+    for bootstub in (False, True):
+      p.reset(enter_bootstub=bootstub)
+      with patch('panda.python.spi.PandaSpiHandle.PROTOCOL_VERSION', return_value="abc"):
+        # list should still work with wrong version
+        assert p._serial in Panda.list()
+
+        # connect but raise protocol error
+        with pytest.raises(PandaProtocolMismatch):
+          Panda(p._serial)
+
+  @pytest.mark.expected_logs(2)
+  def test_protocol_version_data(self, p):
     for bootstub in (False, True):
       p.reset(enter_bootstub=bootstub)
       v = p._handle.get_protocol_version()