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 <xynnn@linux.alibaba.com>
  • Loading branch information
Xynnn007 committed Dec 10, 2024
1 parent ee031a4 commit 08421a0
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 12 deletions.
3 changes: 2 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ license = "Apache-2.0"
[features]
default = [ "std" ]
alloc = [ "serde/alloc", "serde_json/alloc" ]
std = [ "serde/std", "serde_json/std" ]
std = ["serde/std", "serde_json/std", "thiserror/std"]
tee-sev = [ "sev" ]
tee-snp = [ "sev" ]

Expand All @@ -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 = { version = "2.0.3", default-features = false }

[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> = core::result::Result<T, KbsTypesError>;

#[derive(Error, Debug)]
pub enum KbsTypesError {
#[error("Serialize/Deserialize error")]
Serde,
}
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, vec::Vec};
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).map_err(|_| KbsTypesError::Serde)?;
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>
) -> core::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,
) -> core::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) -> core::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) -> core::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,
) -> core::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,
) -> core::result::Result<Option<Vec<u8>>, D::Error>
where
D: serde::Deserializer<'de>,
{
Expand Down Expand Up @@ -273,7 +282,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 08421a0

Please sign in to comment.