From 2fe6e1b75468b750554cb69faf37d9aa818b64db Mon Sep 17 00:00:00 2001
From: Kyle Kotowick <kotowick@invictonlabs.com>
Date: Mon, 4 Nov 2024 22:16:27 -0500
Subject: [PATCH] Use constants instead of hardcoded values

---
 Cargo.toml   |  1 +
 src/ctap1.rs |  5 +++--
 src/ctap2.rs |  3 ++-
 src/sizes.rs | 13 +++++++++++++
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 629cb23..70480db 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ repository = "https://github.com/trussed-dev/ctap-types"
 arbitrary = { version = "1.3.2", features = ["derive"], optional = true }
 bitflags = "1.3"
 cbor-smol = { version = "0.5", features = ["heapless-bytes-v0-3"] }
+cfg-if = "1.0"
 cosey = "0.3.1"
 delog = "0.1"
 heapless = { version = "0.7", default-features = false, features = ["serde"] }
diff --git a/src/ctap1.rs b/src/ctap1.rs
index 0fc6153..aefe095 100644
--- a/src/ctap1.rs
+++ b/src/ctap1.rs
@@ -30,6 +30,7 @@ pub mod authenticate {
 
 pub mod register {
     use super::Bytes;
+    use crate::sizes::MAX_MESSAGE_LENGTH;
 
     #[derive(Clone, Debug, Eq, PartialEq)]
     pub struct Request<'a> {
@@ -42,7 +43,7 @@ pub mod register {
         pub header_byte: u8,
         pub public_key: Bytes<65>,
         pub key_handle: Bytes<255>,
-        pub attestation_certificate: Bytes<1024>,
+        pub attestation_certificate: Bytes<MAX_MESSAGE_LENGTH>,
         pub signature: Bytes<72>,
     }
 
@@ -52,7 +53,7 @@ pub mod register {
             public_key: &cosey::EcdhEsHkdf256PublicKey,
             key_handle: Bytes<255>,
             signature: Bytes<72>,
-            attestation_certificate: Bytes<1024>,
+            attestation_certificate: Bytes<MAX_MESSAGE_LENGTH>,
         ) -> Self {
             let mut public_key_bytes = Bytes::new();
             public_key_bytes.push(0x04).unwrap();
diff --git a/src/ctap2.rs b/src/ctap2.rs
index 8136be2..efc70d2 100644
--- a/src/ctap2.rs
+++ b/src/ctap2.rs
@@ -6,6 +6,7 @@ use bitflags::bitflags;
 use cbor_smol::cbor_deserialize;
 use serde::{Deserialize, Serialize};
 
+use crate::sizes::MAX_MESSAGE_LENGTH;
 use crate::{sizes::*, Bytes, TryFromStrError, Vec};
 
 pub use crate::operation::{Operation, VendorOperation};
@@ -301,7 +302,7 @@ pub struct PackedAttestationStatement {
     pub alg: i32,
     pub sig: Bytes<ASN1_SIGNATURE_LENGTH>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub x5c: Option<Vec<Bytes<1024>, 1>>,
+    pub x5c: Option<Vec<Bytes<MAX_MESSAGE_LENGTH>, 1>>,
 }
 
 #[derive(Clone, Debug, Default, Eq, PartialEq)]
diff --git a/src/sizes.rs b/src/sizes.rs
index 3f36595..43f51f4 100644
--- a/src/sizes.rs
+++ b/src/sizes.rs
@@ -30,3 +30,16 @@ pub const THEORETICAL_MAX_MESSAGE_SIZE: usize = PACKET_SIZE - 7 + 128 * (PACKET_
 pub const LARGE_BLOB_MAX_FRAGMENT_LENGTH: usize = 0;
 #[cfg(feature = "large-blobs")]
 pub const LARGE_BLOB_MAX_FRAGMENT_LENGTH: usize = 3008;
+
+// TODO: update these, and grab them from a common crate?
+cfg_if::cfg_if! {
+    if #[cfg(feature = "backend-dilithium5")] {
+        pub const MAX_MESSAGE_LENGTH: usize = 7523 + 57 + 30;
+    } else if #[cfg(feature = "backend-dilithium3")] {
+        pub const MAX_MESSAGE_LENGTH: usize = 6019 + 57 + 30;
+    } else if #[cfg(feature = "backend-dilithium2")] {
+        pub const MAX_MESSAGE_LENGTH: usize = 3907 + 57 + 30;
+    } else {
+        pub const MAX_MESSAGE_LENGTH: usize = 1024;
+    }
+}