From df75c1796ab3cfc975828b54b849ebcdddbfe0a1 Mon Sep 17 00:00:00 2001
From: Andrew Tridgell <andrew@tridgell.net>
Date: Thu, 11 Aug 2022 18:03:25 +1000
Subject: [PATCH] waf: added --enable-check-firmware option

---
 Tools/ardupilotwaf/ardupilotwaf.py |  1 +
 Tools/ardupilotwaf/boards.py       | 10 ++++++++++
 Tools/ardupilotwaf/chibios.py      | 24 +++++++++---------------
 wscript                            |  4 ++++
 4 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/Tools/ardupilotwaf/ardupilotwaf.py b/Tools/ardupilotwaf/ardupilotwaf.py
index 1db0f43f9f..8495b8bff7 100644
--- a/Tools/ardupilotwaf/ardupilotwaf.py
+++ b/Tools/ardupilotwaf/ardupilotwaf.py
@@ -111,6 +111,7 @@ COMMON_VEHICLE_DEPENDENT_LIBRARIES = [
     'AP_CustomRotations',
     'AP_AIS',
     'AP_OpenDroneID',
+    'AP_CheckFirmware',
 ]
 
 def get_legacy_defines(sketch_name, bld):
diff --git a/Tools/ardupilotwaf/boards.py b/Tools/ardupilotwaf/boards.py
index 33133ee3af..a06feadacb 100644
--- a/Tools/ardupilotwaf/boards.py
+++ b/Tools/ardupilotwaf/boards.py
@@ -94,6 +94,16 @@ class Board:
                 ENABLE_ONVIF=0,
             )
 
+        # allow enable of firmware ID checking for any board
+        if cfg.options.enable_check_firmware:
+            env.CHECK_FIRMWARE_ENABLED = True
+            env.DEFINES.update(
+                AP_CHECK_FIRMWARE_ENABLED=1,
+            )
+            cfg.msg("Enabled firmware ID checking", 'yes')
+        else:
+            cfg.msg("Enabled firmware ID checking", 'no', color='YELLOW')
+
         d = env.get_merged_dict()
         # Always prepend so that arguments passed in the command line get
         # the priority.
diff --git a/Tools/ardupilotwaf/chibios.py b/Tools/ardupilotwaf/chibios.py
index 4c75d4f2a6..790e6c6096 100644
--- a/Tools/ardupilotwaf/chibios.py
+++ b/Tools/ardupilotwaf/chibios.py
@@ -236,17 +236,11 @@ class set_app_descriptor(Task.Task):
     def keyword(self):
         return "app_descriptor"
     def run(self):
-        if not 'APP_DESCRIPTOR' in self.env:
-            return
-        if self.env.APP_DESCRIPTOR == 'MissionPlanner':
-            descriptor = b'\x40\xa2\xe4\xf1\x64\x68\x91\x06'
-        else:
-            Logs.error("Bad APP_DESCRIPTOR %s" % self.env.APP_DESCRIPTOR)
-            return
+        descriptor = b'\x40\xa2\xe4\xf1\x64\x68\x91\x06'
         img = open(self.inputs[0].abspath(), 'rb').read()
         offset = img.find(descriptor)
         if offset == -1:
-            Logs.error("Failed to find %s APP_DESCRIPTOR" % self.env.APP_DESCRIPTOR)
+            Logs.info("No APP_DESCRIPTOR found")
             return
         offset += 8
         # next 8 bytes is 64 bit CRC. We set first 4 bytes to
@@ -263,7 +257,7 @@ class set_app_descriptor(Task.Task):
         githash = to_unsigned(int('0x' + os.environ.get('GIT_VERSION', self.generator.bld.git_head_hash(short=True)),16))
         desc = struct.pack('<IIII', crc1, crc2, len(img), githash)
         img = img[:offset] + desc + img[offset+desc_len:]
-        Logs.info("Applying %s APP_DESCRIPTOR %08x%08x" % (self.env.APP_DESCRIPTOR, crc1, crc2))
+        Logs.info("Applying APP_DESCRIPTOR %08x%08x" % (crc1, crc2))
         open(self.inputs[0].abspath(), 'wb').write(img)
 
 class generate_apj(Task.Task):
@@ -387,12 +381,12 @@ def chibios_firmware(self):
         default_params_task.set_run_after(self.link_task)
         generate_bin_task.set_run_after(default_params_task)
 
-    if self.env.APP_DESCRIPTOR:
-        app_descriptor_task = self.create_task('set_app_descriptor', src=bin_target)
-        app_descriptor_task.set_run_after(generate_bin_task)
-        generate_apj_task.set_run_after(app_descriptor_task)
-        if hex_task is not None:
-            hex_task.set_run_after(app_descriptor_task)
+    # we need to setup the app descriptor so the bootloader can validate the firmware
+    app_descriptor_task = self.create_task('set_app_descriptor', src=bin_target)
+    app_descriptor_task.set_run_after(generate_bin_task)
+    generate_apj_task.set_run_after(app_descriptor_task)
+    if hex_task is not None:
+        hex_task.set_run_after(app_descriptor_task)
 
     if self.bld.options.upload:
         _upload_task = self.create_task('upload_fw', src=apj_target)
diff --git a/wscript b/wscript
index 59dac5b0c1..b02ab41518 100644
--- a/wscript
+++ b/wscript
@@ -225,6 +225,10 @@ submodules at specific revisions.
                  default=False,
                  help="enable generation of scripting documentation")
 
+    g.add_option('--enable-check-firmware', action='store_true',
+                 default=False,
+                 help="Enables firmware ID checking on boot")
+    
     g = opt.ap_groups['linux']
 
     linux_options = ('--prefix', '--destdir', '--bindir', '--libdir')