From c6f2680f9d0288f1c577651a1ecf8e722f14e54b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sosth=C3=A8ne=20Gu=C3=A9don?= <sosthene@nitrokey.com>
Date: Mon, 17 Oct 2022 14:45:58 +0200
Subject: [PATCH] Fix sign_count size

It should always be 3 bytes, and not depend on usize
---
 src/command/private_key_template.rs | 13 ++-----------
 src/state.rs                        | 18 ++++++------------
 2 files changed, 8 insertions(+), 23 deletions(-)

diff --git a/src/command/private_key_template.rs b/src/command/private_key_template.rs
index c6f5593..8f13b8c 100644
--- a/src/command/private_key_template.rs
+++ b/src/command/private_key_template.rs
@@ -15,7 +15,7 @@ const CONCATENATION_KEY_DATA_DO: u16 = 0x5F48;
 
 // § 4.4.3.12
 pub fn put_private_key_template<const R: usize, T: trussed::Client>(
-    mut ctx: LoadedContext<'_, R, T>,
+    ctx: LoadedContext<'_, R, T>,
 ) -> Result<(), Status> {
     let data = get_do(&[PRIVATE_KEY_TEMPLATE_DO], ctx.data).ok_or_else(|| {
         warn!("Got put private key template without 4D DO");
@@ -26,16 +26,7 @@ pub fn put_private_key_template<const R: usize, T: trussed::Client>(
     debug!("Importing {key_type:?} key");
 
     match key_type {
-        KeyType::Sign => {
-            put_sign(ctx.lend())?;
-            ctx.state
-                .internal
-                .set_sign_count(0, ctx.backend.client_mut())
-                .map_err(|_err| {
-                    warn!("Failed to save sign count: {_err}");
-                    Status::UnspecifiedNonpersistentExecutionError
-                })?;
-        }
+        KeyType::Sign => put_sign(ctx)?,
         KeyType::Dec => put_dec(ctx)?,
         KeyType::Aut => put_aut(ctx)?,
     }
diff --git a/src/state.rs b/src/state.rs
index 0f85f3b..30827b8 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -286,7 +286,7 @@ pub struct Internal {
     cardholder_name: Bytes<39>,
     cardholder_sex: Sex,
     language_preferences: Bytes<8>,
-    sign_count: usize,
+    sign_count: u32,
     uif_sign: Uif,
     uif_dec: Uif,
     uif_aut: Uif,
@@ -615,19 +615,10 @@ impl Internal {
         self.save(client)
     }
 
-    pub fn sign_count(&self) -> usize {
+    pub fn sign_count(&self) -> u32 {
         self.sign_count
     }
 
-    pub fn set_sign_count(
-        &mut self,
-        count: usize,
-        client: &mut impl trussed::Client,
-    ) -> Result<(), Error> {
-        self.sign_count = count;
-        self.save(client)
-    }
-
     pub fn key_id(&self, ty: KeyType) -> Option<KeyId> {
         match ty {
             KeyType::Sign => self.signing_key,
@@ -654,7 +645,10 @@ impl Internal {
         client: &mut impl trussed::Client,
     ) -> Result<Option<(KeyId, KeyOrigin)>, Error> {
         match ty {
-            KeyType::Sign => swap(&mut self.signing_key, &mut new),
+            KeyType::Sign => {
+                self.sign_count = 0;
+                swap(&mut self.signing_key, &mut new)
+            }
             KeyType::Dec => swap(&mut self.confidentiality_key, &mut new),
             KeyType::Aut => swap(&mut self.aut_key, &mut new),
         }