From 79f7bb195f364ca7f29bc077b6f97c6e215ec1e0 Mon Sep 17 00:00:00 2001
From: Justus K <justus.k@protonmail.com>
Date: Sat, 4 Nov 2023 15:48:25 +0100
Subject: [PATCH] fix: verify key length in `compatible_with` method

---
 src/jwa/pbes2.rs          |  2 +-
 src/jwk.rs                | 42 +++++++++++++++++++++++++++++++++------
 src/jwk/symmetric.rs      | 12 +++++++++++
 src/jwk/symmetric/hmac.rs |  5 ++++-
 4 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/src/jwa/pbes2.rs b/src/jwa/pbes2.rs
index 3a7f922..76bbfd9 100644
--- a/src/jwa/pbes2.rs
+++ b/src/jwa/pbes2.rs
@@ -4,7 +4,7 @@
 /// [section 4.8 of RFC 7518]: <https://datatracker.ietf.org/doc/html/rfc7518#section-4.8>
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum Pbes2 {
-    /// PBES2 with HMAC SHA56 and "A128KW" wrapping
+    /// PBES2 with HMAC SHA-256 and "A128KW" wrapping
     Hs256Aes128,
     /// PBES2 with HMAC SHA-384 and "A192KW" wrapping
     Hs384Aes192,
diff --git a/src/jwk.rs b/src/jwk.rs
index 7e60a9c..828f7b7 100644
--- a/src/jwk.rs
+++ b/src/jwk.rs
@@ -10,13 +10,17 @@ use hashbrown::HashSet;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use crate::{
-    jwa::{EcDSA, JsonWebAlgorithm, JsonWebEncryptionAlgorithm, JsonWebSigningAlgorithm},
+    jwa::{
+        AesGcm, AesKw, EcDSA, Hmac, JsonWebAlgorithm, JsonWebEncryptionAlgorithm,
+        JsonWebSigningAlgorithm, Pbes2,
+    },
     jwk::{
         ec::{EcPrivate, EcPublic},
         okp::{
             curve25519::{Curve25519Private, Curve25519Public},
             OkpPrivate, OkpPublic,
         },
+        symmetric::hmac::{HmacVariant, Hs256, Hs384, Hs512},
     },
     policy::{Checkable, Checked, CryptographicOperation, Policy},
     sealed::Sealed,
@@ -727,13 +731,39 @@ impl JsonWebKeyType {
         // it is unreadable with the matches! macro and there's no benefit
         #[allow(clippy::match_like_matches_macro)]
         match (self, alg) {
+            (
+                Symmetric(SymmetricJsonWebKey::OctetSequence(key)),
+                Signing(JsonWebSigningAlgorithm::Hmac(hmac)),
+            ) => match (hmac, key.len()) {
+                (Hmac::Hs256, Hs256::OUTPUT_SIZE..)
+                | (Hmac::Hs384, Hs384::OUTPUT_SIZE..)
+                | (Hmac::Hs512, Hs512::OUTPUT_SIZE..) => true,
+                _ => false,
+            },
+            (
+                Symmetric(SymmetricJsonWebKey::OctetSequence(key)),
+                Encryption(JsonWebEncryptionAlgorithm::AesKw(aes)),
+            ) => matches!(
+                (aes, key.len()),
+                (AesKw::Aes128, 16) | (AesKw::Aes192, 24) | (AesKw::Aes256, 32)
+            ),
+            (
+                Symmetric(SymmetricJsonWebKey::OctetSequence(key)),
+                Encryption(JsonWebEncryptionAlgorithm::AesGcmKw(aes)),
+            ) => matches!(
+                (aes, key.len()),
+                (AesGcm::Aes128, 16) | (AesGcm::Aes192, 24) | (AesGcm::Aes256, 32)
+            ),
+            (
+                Symmetric(SymmetricJsonWebKey::OctetSequence(key)),
+                Encryption(JsonWebEncryptionAlgorithm::Pbes2(pbes2)),
+            ) => matches!(
+                (pbes2, key.len()),
+                (Pbes2::Hs256Aes128, 16) | (Pbes2::Hs384Aes192, 24) | (Pbes2::Hs512Aes256, 32)
+            ),
             (
                 Symmetric(SymmetricJsonWebKey::OctetSequence(..)),
-                Signing(JsonWebSigningAlgorithm::Hmac(..))
-                | Encryption(JsonWebEncryptionAlgorithm::Direct)
-                | Encryption(JsonWebEncryptionAlgorithm::AesKw(..))
-                | Encryption(JsonWebEncryptionAlgorithm::AesGcmKw(..))
-                | Encryption(JsonWebEncryptionAlgorithm::Pbes2(..)),
+                Encryption(JsonWebEncryptionAlgorithm::Direct),
             ) => true,
             (Asymmetric(key), alg) => match (&**key, alg) {
                 (
diff --git a/src/jwk/symmetric.rs b/src/jwk/symmetric.rs
index 7238b5e..11a556b 100644
--- a/src/jwk/symmetric.rs
+++ b/src/jwk/symmetric.rs
@@ -36,6 +36,18 @@ impl OctetSequence {
     pub(crate) fn new(x: impl Into<Vec<u8>>) -> Self {
         Self(Base64UrlBytes(x.into()))
     }
+
+    /// Returns the number of bytes that are in this octet sequence.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.0 .0.len()
+    }
+
+    /// Returns `true` if this octet sequence has a length of zero.
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
 }
 
 impl crate::sealed::Sealed for OctetSequence {}
diff --git a/src/jwk/symmetric/hmac.rs b/src/jwk/symmetric/hmac.rs
index 6e509c8..5b8d069 100644
--- a/src/jwk/symmetric/hmac.rs
+++ b/src/jwk/symmetric/hmac.rs
@@ -22,6 +22,9 @@ pub trait HmacVariant: Sealed {
     /// The JWA algorithm for this variant.
     const ALGORITHM: JsonWebSigningAlgorithm;
 
+    /// The output size of this HMAC variant in bytes.
+    const OUTPUT_SIZE: usize = <<Self::HmacType as OutputSizeUser>::OutputSize as Unsigned>::USIZE;
+
     /// The [`hmac::Hmac`] type for this variant.
     type HmacType: Clone + KeyInit + FixedOutput + FixedOutputReset + Update + fmt::Debug;
 }
@@ -186,7 +189,7 @@ impl<H: HmacVariant> FromKey<&OctetSequence> for HmacKey<H> {
                 // This check is not required for normal Hmac implementations based on RFC 2104
                 // but RFC 7518 section 3.2 requires this check and forbids keys with a length <
                 // output
-                if key.len() < <<H::HmacType as OutputSizeUser>::OutputSize as Unsigned>::USIZE {
+                if key.len() < <H as HmacVariant>::OUTPUT_SIZE {
                     return Err(digest::InvalidLength.into());
                 }