From 0300fc5bb2bd6c595127f58a023e0cbb33a7d423 Mon Sep 17 00:00:00 2001
From: Enrico Marconi <enrico.marconi@hotmail.it>
Date: Tue, 15 Oct 2024 11:41:10 +0200
Subject: [PATCH] review comments

---
 src/builder.rs                | 20 ++++++++++++++------
 src/decoder.rs                |  6 +++++-
 src/key_binding_jwt_claims.rs | 18 ++++++++++++++++++
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/builder.rs b/src/builder.rs
index 77435ee..ac7ba18 100644
--- a/src/builder.rs
+++ b/src/builder.rs
@@ -108,16 +108,21 @@ impl<H: Hasher> SdJwtBuilder<H> {
   }
 
   /// Adds a new claim to the underlying object.
-  pub fn insert_claim<'a, K, V>(mut self, key: K, value: V) -> Self
+  pub fn insert_claim<'a, K, V>(mut self, key: K, value: V) -> Result<Self>
   where
     K: Into<Cow<'a, str>>,
     V: Serialize,
   {
     let key = key.into().into_owned();
-    let value = serde_json::to_value(value).unwrap();
-    self.encoder.object.as_object_mut().unwrap().insert(key, value);
-
+    let value = serde_json::to_value(value).map_err(|e| Error::DeserializationError(e.to_string()))?;
     self
+      .encoder
+      .object
+      .as_object_mut()
+      .expect("encoder::object is a JSON Object")
+      .insert(key, value);
+
+    Ok(self)
   }
 
   /// Adds a decoy digest to the specified path.
@@ -157,7 +162,10 @@ impl<H: Hasher> SdJwtBuilder<H> {
     // Add key binding requirement as `cnf`.
     if let Some(key_bind) = key_bind {
       let key_bind = serde_json::to_value(key_bind).map_err(|e| Error::DeserializationError(e.to_string()))?;
-      object.as_object_mut().unwrap().insert("cnf".to_string(), key_bind);
+      object
+        .as_object_mut()
+        .expect("encoder::object is a JSON Object")
+        .insert("cnf".to_string(), key_bind);
     }
 
     // Check mandatory header properties or insert them.
@@ -173,7 +181,7 @@ impl<H: Hasher> SdJwtBuilder<H> {
     header.insert("alg".to_string(), Value::String(alg.to_string()));
 
     let jws = signer
-      .sign(&header, object.as_object().unwrap())
+      .sign(&header, object.as_object().expect("encoder::object is a JSON Object"))
       .await
       .map_err(|e| anyhow::anyhow!("jws failed: {e}"))
       .and_then(|jws_bytes| String::from_utf8(jws_bytes).context("invalid JWS"))
diff --git a/src/decoder.rs b/src/decoder.rs
index 5c95332..50826ee 100644
--- a/src/decoder.rs
+++ b/src/decoder.rs
@@ -88,7 +88,11 @@ impl SdObjectDecoder {
               output.insert(claim_name, recursively_decoded);
             }
           }
-          if output.get(DIGESTS_KEY).unwrap().is_array() {
+          if output
+            .get(DIGESTS_KEY)
+            .expect("output has a `DIGEST_KEY` property")
+            .is_array()
+          {
             output.remove(DIGESTS_KEY);
           }
         }
diff --git a/src/key_binding_jwt_claims.rs b/src/key_binding_jwt_claims.rs
index 5475753..3bdb40a 100644
--- a/src/key_binding_jwt_claims.rs
+++ b/src/key_binding_jwt_claims.rs
@@ -51,14 +51,17 @@ impl FromStr for KeyBindingJwt {
 }
 
 impl KeyBindingJwt {
+  /// Returns a [`KeyBindingJwtBuilder`] that allows the creation of a [`KeyBindingJwt`].
   pub fn builder() -> KeyBindingJwtBuilder {
     KeyBindingJwtBuilder::default()
   }
+  /// Returns a reference to this [`KeyBindingJwt`] claim set.
   pub fn claims(&self) -> &KeyBindingJwtClaims {
     &self.0.claims
   }
 }
 
+/// Builder-style struct to ease the creation of an [`KeyBindingJwt`].
 #[derive(Debug, Default, Clone)]
 pub struct KeyBindingJwtBuilder {
   header: JsonObject,
@@ -66,23 +69,32 @@ pub struct KeyBindingJwtBuilder {
 }
 
 impl KeyBindingJwtBuilder {
+  /// Creates a new [`KeyBindingJwtBuilder`].
   pub fn new() -> Self {
     Self::default()
   }
+
+  /// Creates a new [`KeyBindingJwtBuilder`] using `object` as its payload.
   pub fn from_object(object: JsonObject) -> Self {
     Self {
       header: JsonObject::default(),
       payload: object,
     }
   }
+
+  /// Sets the JWT's header.
   pub fn header(mut self, header: JsonObject) -> Self {
     self.header = header;
     self
   }
+
+  /// Sets the [iat](https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.6) property.
   pub fn iat(mut self, iat: i64) -> Self {
     self.payload.insert("iat".to_string(), iat.into());
     self
   }
+
+  /// Sets the [aud](https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.3) property.
   pub fn aud<'a, S>(mut self, aud: S) -> Self
   where
     S: Into<Cow<'a, str>>,
@@ -90,6 +102,8 @@ impl KeyBindingJwtBuilder {
     self.payload.insert("aud".to_string(), aud.into().into_owned().into());
     self
   }
+
+  /// Sets the `nonce` property.
   pub fn nonce<'a, S>(mut self, nonce: S) -> Self
   where
     S: Into<Cow<'a, str>>,
@@ -99,10 +113,14 @@ impl KeyBindingJwtBuilder {
       .insert("nonce".to_string(), nonce.into().into_owned().into());
     self
   }
+
+  /// Inserts a given property with key `name` and value `value` in the payload.
   pub fn insert_property(mut self, name: &str, value: Value) -> Self {
     self.payload.insert(name.to_string(), value);
     self
   }
+
+  /// Builds an [`KeyBindingJwt`] from the data provided to builder.
   pub async fn finish<S>(
     self,
     sd_jwt: &SdJwt,