From 78768b2d57c51272e714d5a794130a81989e9941 Mon Sep 17 00:00:00 2001
From: Carl Lundin <clundin@google.com>
Date: Tue, 14 Jan 2025 13:20:39 -0800
Subject: [PATCH] Add vendor and owner signature injection to image generator
 tool

---
 builder/bin/image_gen.rs | 44 +++++++++++++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/builder/bin/image_gen.rs b/builder/bin/image_gen.rs
index 6f93a466aa..817bb990c7 100644
--- a/builder/bin/image_gen.rs
+++ b/builder/bin/image_gen.rs
@@ -5,12 +5,14 @@ use caliptra_builder::version;
 use caliptra_builder::ImageOptions;
 use caliptra_image_types::ImageHeader;
 use caliptra_image_types::ImageManifest;
+use caliptra_image_types::ImageSignatures;
 use clap::{arg, value_parser, Command};
 use memoffset::{offset_of, span_of};
 use serde_json::{json, to_string_pretty};
 use sha2::{Digest, Sha384};
 use std::collections::HashSet;
 use std::path::PathBuf;
+use zerocopy::FromBytes;
 
 fn main() {
     let args = Command::new("image-gen")
@@ -42,6 +44,8 @@ fn main() {
             arg!(--"hashes" [FILE] "File path for output JSON file containing image bundle header hashes for external signing tools")
                 .value_parser(value_parser!(PathBuf)),
         )
+        .arg(arg!(--"owner-sig-override" [FILE] "Manually overwrite the owner_sigs of the FW bundle image with the contents of binary [FILE]. The signature should be an ECC signature concatenated with an LMS signature").value_parser(value_parser!(PathBuf)))
+        .arg(arg!(--"vendor-sig-override" [FILE] "Manually overwrite the vendor_sigs of the FW bundle image with the contents of binary [FILE]. The signature should be an ECC signature concatenated with an LMS signature").value_parser(value_parser!(PathBuf)))
         .get_matches();
 
     if let Some(path) = args.get_one::<PathBuf>("rom-no-log") {
@@ -72,18 +76,34 @@ fn main() {
 
     if let Some(path) = args.get_one::<PathBuf>("fw") {
         // Generate Image Bundle
-        let image = caliptra_builder::build_and_sign_image(
-            &firmware::FMC_WITH_UART,
-            &firmware::APP_WITH_UART,
-            ImageOptions {
-                fmc_version: version::get_fmc_version(),
-                app_version: version::get_runtime_version(),
-                fmc_svn,
-                app_svn,
-                ..Default::default()
-            },
-        )
-        .unwrap();
+        let image = {
+            let mut image = caliptra_builder::build_and_sign_image(
+                &firmware::FMC_WITH_UART,
+                &firmware::APP_WITH_UART,
+                ImageOptions {
+                    fmc_version: version::get_fmc_version(),
+                    app_version: version::get_runtime_version(),
+                    fmc_svn,
+                    app_svn,
+                    ..Default::default()
+                },
+            )
+            .unwrap();
+
+            if let Some(path) = args.get_one::<PathBuf>("owner-sig-override") {
+                let sig_override = std::fs::read(path).unwrap();
+                image.manifest.preamble.owner_sigs =
+                    ImageSignatures::read_from_bytes(&sig_override).unwrap();
+            }
+
+            if let Some(path) = args.get_one::<PathBuf>("vendor-sig-override") {
+                let sig_override = std::fs::read(path).unwrap();
+                image.manifest.preamble.vendor_sigs =
+                    ImageSignatures::read_from_bytes(&sig_override).unwrap();
+            }
+
+            image
+        };
 
         let contents = image.to_bytes().unwrap();
         std::fs::write(path, contents.clone()).unwrap();