Upload template command

Also includes a fix for downloading templates

crosec/src/commands/fp_download.rs

@@ -41,7 +41,7 @@ pub fn fp_download(
         DownloadType::SimpleImage => (fp_info.get_simple_image_size(), FP_FRAME_INDEX_RAW_IMAGE),
         DownloadType::RawImage => (fp_info.frame_size as usize, FP_FRAME_INDEX_RAW_IMAGE),
         DownloadType::Template(template_index) => {
-            (fp_info.template_size as usize, *template_index as u32)
+            (fp_info.template_size as usize, *template_index as u32 + 1)
         }
     };
     // The template may be (and probably is) bigger than the max output size, so we need to download it in chunks
@@ -81,18 +81,23 @@ pub fn fp_download(
 
 /// A safe wrapper around the actual template so you don't try to upload arbitrary data
 pub struct FpTemplate {
-    buffer: Vec<u8>,
+    vec: Vec<u8>,
 }
 
 impl Into<Vec<u8>> for FpTemplate {
     fn into(self) -> Vec<u8> {
-        self.buffer
+        self.vec
     }
 }
 
 impl FpTemplate {
     pub fn buffer(&self) -> &Vec<u8> {
-        &self.buffer
+        &self.vec
+    }
+
+    /// Make sure your buffer is actually a compatible fp template
+    pub unsafe fn from_vec_unchecked(vec: Vec<u8>) -> Self {
+        FpTemplate { vec }
     }
 }
 
@@ -103,6 +108,6 @@ pub fn fp_download_template(
     index: usize,
 ) -> FpTemplate {
     FpTemplate {
-        buffer: fp_download(file, fp_info, protocol_info, &DownloadType::Template(index)),
+        vec: fp_download(file, fp_info, protocol_info, &DownloadType::Template(index)),
     }
 }

crosec/src/commands/fp_info.rs

@@ -16,16 +16,23 @@ pub struct EcResponseFpInfo {
     pub product_id: u32,
     pub model_id: u32,
     pub version: u32,
+    /// The size of the PGM image, in bytes
     pub frame_size: u32,
     pub pixel_format: u32,
     pub width: u16,
     pub height: u16,
     pub bpp: u16,
     pub errors: u16,
+    /// The template size, in bytes
     pub template_size: u32,
+    /// The maximum number of templates the FP can store and match at once
     pub template_max: u16,
+    /// The number of templates loaded into the FP
     pub template_valid: u16,
+    /// The first bit (the rightmost) represents template 0, the 2nd bit form the right represents template 1, etc.
+    /// If the bit is 1, that means that the template has been updated by the FP and the updated version has not been downloaded yet.
     pub template_dirty: u32,
+    /// This version could increase after an update to the FP firmware
     pub template_version: u32,
 }
 impl EcResponseFpInfo {

crosec/src/commands/fp_upload_template.rs

@@ -0,0 +1,69 @@
+use std::{fs::File, mem::offset_of, os::fd::AsRawFd};
+
+use bytemuck::{bytes_of, Pod, Zeroable};
+
+use crate::{ec_command::ec_command_with_dynamic_output_size, EcCmdResult};
+
+use super::{
+    fp_download::FpTemplate, fp_info::EcResponseFpInfo,
+    get_protocol_info::EcResponseGetProtocolInfo, CrosEcCmd,
+};
+
+#[derive(Pod, Zeroable, Clone, Copy)]
+#[repr(C)]
+struct EcParamsFpTemplateWithoutData {
+    offset: u32,
+    size: u32,
+    data: [u8; 0],
+}
+
+/// Flag in the 'size' field indicating that the full template has been sent
+const FP_TEMPLATE_COMMIT: u32 = 0x80000000;
+
+pub fn fp_upload_template(
+    file: &mut File,
+    protocol_info: &EcResponseGetProtocolInfo,
+    fp_info: &EcResponseFpInfo,
+    template: &FpTemplate,
+) -> EcCmdResult<()> {
+    assert_eq!(
+        template.buffer().len(),
+        fp_info.template_size as usize,
+        "The given template must match the fp sensor's template size"
+    );
+    // TODO(b/78544921): removing 32 bits is a workaround for the MCU bug
+    // Idk what this bug is, but the ChromiumOS ectool removes 4 bytes, so we should too
+    let max_chunk_size =
+        protocol_info.max_ec_output_size() - offset_of!(EcParamsFpTemplateWithoutData, data) - 4;
+    let number_of_chunks = template
+        .buffer()
+        .len()
+        .div_ceil(protocol_info.max_ec_input_size());
+    for chunk_index in 0..number_of_chunks {
+        ec_command_with_dynamic_output_size(
+            CrosEcCmd::FpTemplate,
+            0,
+            &{
+                let bytes_uploaded = chunk_index * max_chunk_size;
+                let size = (template.buffer().len() - bytes_uploaded).min(max_chunk_size);
+                let mut vec = bytes_of(&EcParamsFpTemplateWithoutData {
+                    offset: bytes_uploaded as u32,
+                    size: {
+                        let mut size = size as u32;
+                        if chunk_index == number_of_chunks - 1 {
+                            size |= FP_TEMPLATE_COMMIT;
+                        }
+                        size
+                    },
+                    data: [],
+                })
+                .to_vec();
+                vec.extend_from_slice(&template.buffer()[bytes_uploaded..bytes_uploaded + size]);
+                vec
+            },
+            0,
+            file.as_raw_fd(),
+        )?;
+    }
+    Ok(())
+}

crosec/src/commands/get_protocol_info.rs

@@ -4,7 +4,7 @@ use bytemuck::{Pod, Zeroable};
 
 use crate::{commands::CrosEcCmd, ec_command::ec_command_bytemuck, EcCmdResult};
 
-#[derive(Pod, Zeroable, Clone, Copy)]
+#[derive(Pod, Zeroable, Clone, Copy, Debug)]
 #[repr(C)]
 pub struct EcResponseGetProtocolInfo {
     protocol_versions: u32,

crosec/src/commands/mod.rs

@@ -32,6 +32,7 @@ pub mod fp_info;
 pub mod fp_mode;
 pub mod fp_set_seed;
 pub mod fp_stats;
+pub mod fp_upload_template;
 pub mod get_chip_info;
 pub mod get_cmd_versions;
 pub mod get_features;

ectool/src/fp_upload_template_command.rs

@@ -0,0 +1,26 @@
+use std::{
+    fs::File,
+    io::{stdin, Read},
+};
+
+use color_eyre::eyre::Result;
+use crosec::{
+    commands::{
+        fp_download::FpTemplate, fp_info::fp_info, fp_upload_template::fp_upload_template,
+        get_protocol_info::get_protocol_info,
+    },
+    CROS_FP_PATH,
+};
+
+pub fn fp_upload_template_command() -> Result<()> {
+    let mut buf = Default::default();
+    println!("Reading from stdin. If this command is taking a long time, it's probably because there is no EOF inputted from stdin.");
+    stdin().read_to_end(&mut buf)?;
+    let template = unsafe { FpTemplate::from_vec_unchecked(buf) };
+    let mut file = File::open(CROS_FP_PATH)?;
+    let protocol_info = get_protocol_info(&mut file)?;
+    let fp_info = fp_info(&mut file)?;
+    fp_upload_template(&mut file, &protocol_info, &fp_info, &template)?;
+    println!("Uploaded template");
+    Ok(())
+}

ectool/src/main.rs

@@ -11,6 +11,7 @@ use crosec::commands::fp_stats::fp_stats;
 use crosec::commands::get_protocol_info::get_protocol_info;
 use crosec::wait_event::{event::EcMkbpEventType, wait_event};
 use fp_download_subcommand::{fp_download_subcommand, FpDownloadSubcommand};
+use fp_upload_template_command::fp_upload_template_command;
 use num_traits::cast::FromPrimitive;
 
 use crosec::battery::battery;
@@ -31,6 +32,7 @@ use crosec::{
 
 mod charge_control_subcommand;
 mod fp_download_subcommand;
+mod fp_upload_template_command;
 
 #[derive(Parser)]
 struct Cli {
@@ -112,6 +114,8 @@ enum Commands {
         #[command(subcommand)]
         command: FpDownloadSubcommand,
     },
+    /// Uploads template from stdin
+    FpUploadTemplate,
 }
 
 fn main() -> Result<()> {
@@ -270,6 +274,7 @@ fn main() -> Result<()> {
             println!("{result:#?}");
         }
         Commands::FpDownload { command } => fp_download_subcommand(command)?,
+        Commands::FpUploadTemplate => fp_upload_template_command()?,
     }
 
     Ok(())