From 13f6a2afb5f913b155f3244c4b6906d30c57e5e6 Mon Sep 17 00:00:00 2001
From: Josh Pieper <jjp@pobox.com>
Date: Wed, 13 Jan 2021 17:10:27 -0500
Subject: [PATCH] Prepare for more advanced firmware upgrades

* The "firmware_version" is now stored in a symbol in the ELF file for
  easier extraction.
* moteus_tool extracts the firmware from the ELF file, and requires
  that it be a version it knows about in order to flash.
* moteus_tool also stops telemetry and flushes before trying to flash.
  This will make flashing work even if tview was shut down uncleanly.

Currently, no advanced upgrade rules exist.  This just makes
moteus_tool properly know if it is able to handle a particular
version.
---
 fw/moteus.cc                     |  6 ++-
 lib/python/moteus/moteus.py      | 12 +++++
 lib/python/moteus/moteus_tool.py | 76 ++++++++++++++++++++++++++++----
 3 files changed, 85 insertions(+), 9 deletions(-)

diff --git a/fw/moteus.cc b/fw/moteus.cc
index a114fc24..49249d5b 100644
--- a/fw/moteus.cc
+++ b/fw/moteus.cc
@@ -43,6 +43,10 @@
 #error "Unknown target"
 #endif
 
+extern "C" {
+  uint32_t kMoteusFirmwareVersion = MOTEUS_FIRMWARE_VERSION;
+}
+
 auto* const MyDWT = DWT;
 auto* const MyFLASH = FLASH;
 
@@ -265,7 +269,7 @@ int main(void) {
 
   SystemInfo system_info(pool, telemetry_manager);
   FirmwareInfo firmware_info(pool, telemetry_manager,
-                             MOTEUS_FIRMWARE_VERSION, MOTEUS_MODEL_NUMBER);
+                             kMoteusFirmwareVersion, MOTEUS_MODEL_NUMBER);
   ClockManager clock(&timer, persistent_config, command_manager);
 
   MoteusController moteus_controller(
diff --git a/lib/python/moteus/moteus.py b/lib/python/moteus/moteus.py
index df06e802..f8e2f5c5 100644
--- a/lib/python/moteus/moteus.py
+++ b/lib/python/moteus/moteus.py
@@ -655,6 +655,18 @@ async def read(self, size, block=True):
         to_return, self._read_data = self._read_data[0:size], self._read_data[size:]
         return to_return
 
+    async def flush_read(self, timeout=0.01):
+        self._read_data = b''
+
+        try:
+            await asyncio.wait_for(self.read(65536), timeout)
+            raise RuntimeError("More data to flush than expected")
+        except asyncio.TimeoutError:
+            # This is the expected path.
+            pass
+
+        self._read_data = b''
+
     async def _read_maybe_empty_line(self):
         while b'\n' not in self._read_data and b'\r' not in self._read_data:
             async with self.lock:
diff --git a/lib/python/moteus/moteus_tool.py b/lib/python/moteus/moteus_tool.py
index 539e0d87..abf03e5d 100644
--- a/lib/python/moteus/moteus_tool.py
+++ b/lib/python/moteus/moteus_tool.py
@@ -25,6 +25,7 @@
 import io
 import math
 import os
+import struct
 import sys
 import tempfile
 
@@ -36,6 +37,25 @@
 MAX_FLASH_BLOCK_SIZE = 32
 
 
+class FirmwareUpgrade:
+    '''This encodes "magic" rules about upgrading firmware, largely about
+    how to munge configuration options so as to not cause behavior
+    change upon firmware changes.
+    '''
+
+    def __init__(self, old, new):
+        self.old = old
+        self.new = new
+
+        if new > 0x0100:
+            raise RuntimeError("Firmware to be flashed has a newer version than we support")
+
+    def fix_config(self, old_config):
+        lines = old_config.split(b'\n')
+
+        return b'\n'.join(lines)
+
+
 def _get_log_directory():
     moteus_cal_dir = os.environ.get("MOTEUS_CAL_DIR", None)
     if moteus_cal_dir:
@@ -110,6 +130,11 @@ def logical_to_physical(self, address, size):
         raise RuntimeError(f"no mapping found for {address:x}")
 
 
+class ElfData:
+    sections = []
+    firmware_version = None
+
+
 def _read_elf(filename, sections):
     fp = elftools.elf.elffile.ELFFile(open(filename, "rb"))
 
@@ -131,9 +156,33 @@ def _read_elf(filename, sections):
             section['sh_addr'], section['sh_size'])
         sections[physical_address] = data
 
-    result = []
+    result = ElfData()
+    result.sections = []
+
     for key in sorted(sections.keys()):
-        result.append((key, sections[key]))
+        result.sections.append((key, sections[key]))
+
+    # Try to find the firmware version.
+    symtab = fp.get_section_by_name('.symtab')
+    if symtab:
+        maybe_version = symtab.get_symbol_by_name('kMoteusFirmwareVersion')
+        if maybe_version:
+            version = maybe_version[0]
+            if version['st_size'] != 4:
+                raise RuntimeError('Version in firmware image corrupt')
+
+            symbol_address = version['st_value']
+            symbol_section = version['st_shndx']
+
+            sec = fp.get_section(symbol_section)
+            sec_start = symbol_address - sec['sh_addr']
+            ver, = struct.unpack('<i', sec.data()[sec_start:sec_start+4])
+
+            result.firmware_version = ver
+
+    if result.firmware_version is None:
+        print("WARNING: Could not find firmware version in elf file, assuming 0x00000100")
+        result.firmware_version = 0x00000100
 
     return result
 
@@ -286,11 +335,22 @@ async def write_config_stream(self, fp):
             print()
 
     async def do_flash(self, elffile):
-        elfs = _read_elf(elffile, [".text", ".ARM.extab", ".ARM.exidx",
-                                   ".data", ".ccmram", ".isr_vector"])
-        count_bytes = sum([len(section) for address, section in elfs])
+        elf = _read_elf(elffile, [".text", ".ARM.extab", ".ARM.exidx",
+                                  ".data", ".ccmram", ".isr_vector"])
+        count_bytes = sum([len(section) for address, section in elf.sections])
+
+        fw = '0x{:08x}'.format(elf.firmware_version)
+        print(f"Read ELF file version {fw}: {count_bytes} bytes")
+
+        await self.write_message("tel stop")
+
+        # Discard anything that might have been en route.
+        await self.stream.flush_read()
+
+        # Get the current firmware version.
+        old_firmware = await self.read_data("firmware")
 
-        print(f"Read ELF file: {count_bytes} bytes")
+        upgrade = FirmwareUpgrade(old_firmware.version, elf.firmware_version)
 
         if not self.args.no_restore_config:
             # Read our old config.
@@ -311,7 +371,7 @@ async def do_flash(self, elffile):
         await self.stream.readline()
 
         await self.command("unlock")
-        await self.write_flash(elfs)
+        await self.write_flash(elf.sections)
         await self.command("lock")
         # This will reset the controller, so we don't expect a response.
         await self.write_message("reset")
@@ -319,7 +379,7 @@ async def do_flash(self, elffile):
         await asyncio.sleep(1.0)
 
         if not self.args.no_restore_config:
-            await self.restore_config(old_config)
+            await self.restore_config(upgrade.fix_config(old_config))
 
     def _emit_flash_progress(self, ctx, type):
         if self.args.verbose: