Skip to content

Commit

Permalink
Add Error handling for Serialization
Browse files Browse the repository at this point in the history
This commit adds error handling for the crate using thiserror crate.
Also, based on the error handling, making serialization of `AAD`
eleganttly be handled with result mechanism.

Signed-off-by: Xynnn007 <[email protected]>
  • Loading branch information
Xynnn007 committed Dec 9, 2024
1 parent a31e162 commit b8d56d9
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 11 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ base64 = "0.22.1"
serde = { version = "1.0", default-features = false, features = ["derive"] }
serde_json = { version = "1.0", default-features = false }
sev = { version = "3.2.0", features = ["openssl"], optional = true }
thiserror = "2.0.3"

[dev-dependencies]
codicon = "3.0.0"
Expand Down
9 changes: 9 additions & 0 deletions src/error.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
use thiserror::Error;

pub type Result<T> = std::result::Result<T, KbsTypesError>;

#[derive(Error, Debug)]
pub enum KbsTypesError {
#[error("Serialize/Deserialize error")]
Serde(#[from] serde_json::Error),
}
31 changes: 20 additions & 11 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@
#[cfg(feature = "alloc")]
extern crate alloc;

mod error;
pub use error::{KbsTypesError, Result};

#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::{collections::btree_map::BTreeMap, string::String};
use base64::{prelude::BASE64_URL_SAFE_NO_PAD, Engine};
Expand Down Expand Up @@ -100,18 +103,17 @@ pub struct ProtectedHeader {

impl ProtectedHeader {
/// The generation of AAD for JWE follows [A.3.5 RFC7516](https://www.rfc-editor.org/rfc/rfc7516#appendix-A.3.5)
pub fn generate_aad(&self) -> Vec<u8> {
let protected_utf8 =
serde_json::to_string(&self).expect("unexpected OOM when serializing ProtectedHeader");
pub fn generate_aad(&self) -> Result<Vec<u8>> {
let protected_utf8 = serde_json::to_string(&self)?;
let aad = BASE64_URL_SAFE_NO_PAD.encode(protected_utf8);
aad.into_bytes()
Ok(aad.into_bytes())
}
}

fn serialize_base64_protected_header<S>(
sub: &ProtectedHeader,
serializer: S,
) -> Result<S::Ok, S::Error>
) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
Expand All @@ -120,7 +122,9 @@ where
serializer.serialize_str(&encoded)
}

fn deserialize_base64_protected_header<'de, D>(deserializer: D) -> Result<ProtectedHeader, D::Error>
fn deserialize_base64_protected_header<'de, D>(
deserializer: D,
) -> std::result::Result<ProtectedHeader, D::Error>
where
D: serde::Deserializer<'de>,
{
Expand All @@ -133,15 +137,15 @@ where
Ok(protected_header)
}

fn serialize_base64<S>(sub: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
fn serialize_base64<S>(sub: &Vec<u8>, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let encoded = BASE64_URL_SAFE_NO_PAD.encode(sub);
serializer.serialize_str(&encoded)
}

fn deserialize_base64<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
fn deserialize_base64<'de, D>(deserializer: D) -> std::result::Result<Vec<u8>, D::Error>
where
D: serde::Deserializer<'de>,
{
Expand All @@ -153,7 +157,10 @@ where
Ok(decoded)
}

fn serialize_base64_option<S>(sub: &Option<Vec<u8>>, serializer: S) -> Result<S::Ok, S::Error>
fn serialize_base64_option<S>(
sub: &Option<Vec<u8>>,
serializer: S,
) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
Expand All @@ -166,7 +173,9 @@ where
}
}

fn deserialize_base64_option<'de, D>(deserializer: D) -> Result<Option<Vec<u8>>, D::Error>
fn deserialize_base64_option<'de, D>(
deserializer: D,
) -> std::result::Result<Option<Vec<u8>>, D::Error>
where
D: serde::Deserializer<'de>,
{
Expand Down Expand Up @@ -270,7 +279,7 @@ mod tests {
other_fields: BTreeMap::new(),
};

let aad = protected_header.generate_aad();
let aad = protected_header.generate_aad().unwrap();

assert_eq!(
aad,
Expand Down

0 comments on commit b8d56d9

Please sign in to comment.