From ae19a16531980823ac910739693fe99fa5e39f46 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rodolphe=20Br=C3=A9ard?= <v.code.ngyvz7vf@what.tf>
Date: Sun, 24 Mar 2024 09:47:36 +0100
Subject: [PATCH] Improve the context and IKM list APIs

---
 benches/data.rs   |  4 ++--
 src/cipher_box.rs |  4 ++--
 src/context.rs    | 51 ++++++++++++++++++++++++++++++++++++++++++++---
 src/ikm.rs        | 31 ++++++++++++++++++++++++++++
 src/kdf.rs        |  9 ---------
 5 files changed, 83 insertions(+), 16 deletions(-)

diff --git a/benches/data.rs b/benches/data.rs
index 01f88cb..4201d52 100644
--- a/benches/data.rs
+++ b/benches/data.rs
@@ -1,6 +1,6 @@
 pub const MEASUREMENT_TIME: u64 = 30;
-pub const KEY_CTX: [&str; 3] = ["database_name", "table_name", "column_name"];
-pub const DATA_CTX: [&str; 2] = [
+pub const KEY_CTX: &[&str] = &["database_name", "table_name", "column_name"];
+pub const DATA_CTX: &[&str] = &[
 	"b3b21eb1-70d7-4dc6-9a2a-439e17d8491d",
 	"8dfa06bc-de19-455a-8e43-2f5d8019442f",
 ];
diff --git a/src/cipher_box.rs b/src/cipher_box.rs
index ab6e8e2..f17f743 100644
--- a/src/cipher_box.rs
+++ b/src/cipher_box.rs
@@ -79,8 +79,8 @@ mod tests {
 
 	const TEST_CIPHERTEXT: &str = "AQAAAA:W-nzcGkPU6eWj_JjjqLpQk6WSe_CIUPF:we_HR8yD3XnQ9aaJlZFvqPitnDlQHexw4QPaYaOTzpHSWNW86QQrLRRZOg:NgAAAAAAAAA";
 	const TEST_DATA: &[u8] = b"Lorem ipsum dolor sit amet.";
-	const TEST_KEY_CTX: [&str; 3] = ["db_name", "table_name", "column_name"];
-	const TEST_DATA_CTX: [&str; 1] = ["018db876-3d9d-79af-9460-55d17da991d8"];
+	const TEST_KEY_CTX: &[&str] = &["db_name", "table_name", "column_name"];
+	const TEST_DATA_CTX: &[&str] = &["018db876-3d9d-79af-9460-55d17da991d8"];
 
 	fn get_static_key_ctx() -> KeyContext {
 		let mut ctx: KeyContext = TEST_KEY_CTX.into();
diff --git a/src/context.rs b/src/context.rs
index 115c143..cadfbcf 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -1,5 +1,13 @@
 use std::num::NonZeroU64;
 
+macro_rules! data_ctx_from_iter {
+	($self: ident, $ctx: ident) => {
+		$self {
+			ctx: $ctx.iter().map(|s| s.to_string()).collect(),
+		}
+	};
+}
+
 pub struct DataContext {
 	ctx: Vec<String>,
 }
@@ -12,12 +20,31 @@ impl DataContext {
 
 impl<const N: usize> From<[&str; N]> for DataContext {
 	fn from(ctx: [&str; N]) -> Self {
-		Self {
-			ctx: ctx.iter().map(|s| s.to_string()).collect(),
-		}
+		data_ctx_from_iter!(Self, ctx)
 	}
 }
 
+impl<const N: usize> From<&[&str; N]> for DataContext {
+	fn from(ctx: &[&str; N]) -> Self {
+		data_ctx_from_iter!(Self, ctx)
+	}
+}
+
+impl From<&[&str]> for DataContext {
+	fn from(ctx: &[&str]) -> Self {
+		data_ctx_from_iter!(Self, ctx)
+	}
+}
+
+macro_rules! key_ctx_from_iter {
+	($self: ident, $ctx: ident) => {
+		$self {
+			ctx: $ctx.iter().map(|s| s.to_string()).collect(),
+			periodicity: Some(crate::DEFAULT_KEY_CTX_PERIODICITY),
+		}
+	};
+}
+
 pub struct KeyContext {
 	pub(crate) ctx: Vec<String>,
 	pub(crate) periodicity: Option<u64>,
@@ -48,3 +75,21 @@ impl KeyContext {
 		self.periodicity.is_some()
 	}
 }
+
+impl<const N: usize> From<[&str; N]> for KeyContext {
+	fn from(ctx: [&str; N]) -> Self {
+		key_ctx_from_iter!(Self, ctx)
+	}
+}
+
+impl<const N: usize> From<&[&str; N]> for KeyContext {
+	fn from(ctx: &[&str; N]) -> Self {
+		key_ctx_from_iter!(Self, ctx)
+	}
+}
+
+impl From<&[&str]> for KeyContext {
+	fn from(ctx: &[&str]) -> Self {
+		key_ctx_from_iter!(Self, ctx)
+	}
+}
diff --git a/src/ikm.rs b/src/ikm.rs
index 79c6ca0..3ae2a11 100644
--- a/src/ikm.rs
+++ b/src/ikm.rs
@@ -154,6 +154,14 @@ impl InputKeyMaterialList {
 	}
 }
 
+impl std::str::FromStr for InputKeyMaterialList {
+	type Err = Error;
+
+	fn from_str(s: &str) -> Result<Self, Self::Err> {
+		Self::import(s)
+	}
+}
+
 #[cfg(feature = "ikm-management")]
 impl std::ops::Deref for InputKeyMaterialList {
 	type Target = Vec<InputKeyMaterial>;
@@ -166,6 +174,7 @@ impl std::ops::Deref for InputKeyMaterialList {
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use std::str::FromStr;
 
 	#[test]
 	fn import() {
@@ -188,6 +197,28 @@ mod tests {
 		);
 		assert_eq!(ikm.is_revoked, false);
 	}
+
+	#[test]
+	fn from_str() {
+		let s =
+			"AQAAAA:AQAAAAEAAAC_vYEw1ujVG5i-CtoPYSzik_6xaAq59odjPm5ij01-e6zz4mUAAAAALJGBiwAAAAAA";
+		let res = InputKeyMaterialList::from_str(s);
+		assert!(res.is_ok(), "res: {res:?}");
+		let lst = res.unwrap();
+		assert_eq!(lst.id_counter, 1);
+		assert_eq!(lst.ikm_lst.len(), 1);
+		let ikm = lst.ikm_lst.first().unwrap();
+		assert_eq!(ikm.id, 1);
+		assert_eq!(ikm.scheme, Scheme::XChaCha20Poly1305WithBlake3);
+		assert_eq!(
+			ikm.content,
+			[
+				191, 189, 129, 48, 214, 232, 213, 27, 152, 190, 10, 218, 15, 97, 44, 226, 147, 254,
+				177, 104, 10, 185, 246, 135, 99, 62, 110, 98, 143, 77, 126, 123
+			]
+		);
+		assert_eq!(ikm.is_revoked, false);
+	}
 }
 
 #[cfg(all(test, feature = "ikm-management"))]
diff --git a/src/kdf.rs b/src/kdf.rs
index d416deb..7cdb1c4 100644
--- a/src/kdf.rs
+++ b/src/kdf.rs
@@ -4,15 +4,6 @@ use crate::ikm::InputKeyMaterial;
 
 pub(crate) type KdfFunction = dyn Fn(&str, &[u8]) -> Vec<u8>;
 
-impl<const N: usize> From<[&str; N]> for KeyContext {
-	fn from(ctx: [&str; N]) -> Self {
-		Self {
-			ctx: ctx.iter().map(|s| s.to_string()).collect(),
-			periodicity: Some(crate::DEFAULT_KEY_CTX_PERIODICITY),
-		}
-	}
-}
-
 pub(crate) fn derive_key(
 	ikm: &InputKeyMaterial,
 	ctx: &KeyContext,