From b8539602f0d9f667b0da760b20acf632856b5472 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rodolphe=20Br=C3=A9ard?= <v.code.ngyvz7vf@what.tf>
Date: Sun, 17 Mar 2024 14:31:01 +0100
Subject: [PATCH] Move the context objects to a dedicated module

---
 src/cipher_box.rs | 21 ++------------------
 src/context.rs    | 50 +++++++++++++++++++++++++++++++++++++++++++++++
 src/kdf.rs        | 33 +------------------------------
 src/lib.rs        |  8 +++++---
 4 files changed, 58 insertions(+), 54 deletions(-)
 create mode 100644 src/context.rs

diff --git a/src/cipher_box.rs b/src/cipher_box.rs
index faed85f..7e4b29f 100644
--- a/src/cipher_box.rs
+++ b/src/cipher_box.rs
@@ -1,6 +1,7 @@
 use crate::canonicalization::{canonicalize, join_canonicalized_str};
+use crate::context::{DataContext, KeyContext};
 use crate::error::Result;
-use crate::kdf::{derive_key, KeyContext};
+use crate::kdf::derive_key;
 use crate::{storage, IkmId, InputKeyMaterialList};
 use std::time::{SystemTime, UNIX_EPOCH};
 
@@ -8,24 +9,6 @@ pub(crate) type DecryptionFunction = dyn Fn(&[u8], &EncryptedData, &str) -> Resu
 pub(crate) type EncryptionFunction = dyn Fn(&[u8], &[u8], &[u8], &str) -> Result<EncryptedData>;
 pub(crate) type GenNonceFunction = dyn Fn() -> Result<Vec<u8>>;
 
-pub struct DataContext {
-	ctx: Vec<String>,
-}
-
-impl DataContext {
-	pub(crate) fn get_ctx_elems(&self) -> &[String] {
-		self.ctx.as_ref()
-	}
-}
-
-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(),
-		}
-	}
-}
-
 #[derive(Debug)]
 pub(crate) struct EncryptedData {
 	pub(crate) nonce: Vec<u8>,
diff --git a/src/context.rs b/src/context.rs
new file mode 100644
index 0000000..115c143
--- /dev/null
+++ b/src/context.rs
@@ -0,0 +1,50 @@
+use std::num::NonZeroU64;
+
+pub struct DataContext {
+	ctx: Vec<String>,
+}
+
+impl DataContext {
+	pub(crate) fn get_ctx_elems(&self) -> &[String] {
+		self.ctx.as_ref()
+	}
+}
+
+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(),
+		}
+	}
+}
+
+pub struct KeyContext {
+	pub(crate) ctx: Vec<String>,
+	pub(crate) periodicity: Option<u64>,
+}
+
+impl KeyContext {
+	pub fn set_static(&mut self) {
+		self.periodicity = None;
+	}
+
+	pub fn set_periodicity(&mut self, periodicity: NonZeroU64) {
+		self.periodicity = Some(periodicity.get());
+	}
+
+	pub(crate) fn get_ctx_elems(&self, time_period: Option<u64>) -> Vec<Vec<u8>> {
+		let mut ret: Vec<Vec<u8>> = self.ctx.iter().map(|s| s.as_bytes().to_vec()).collect();
+		if let Some(tp) = time_period {
+			ret.push(tp.to_le_bytes().to_vec());
+		}
+		ret
+	}
+
+	pub(crate) fn get_time_period(&self, timestamp: u64) -> Option<u64> {
+		self.periodicity.map(|p| timestamp / p)
+	}
+
+	pub(crate) fn is_periodic(&self) -> bool {
+		self.periodicity.is_some()
+	}
+}
diff --git a/src/kdf.rs b/src/kdf.rs
index ad842bd..d416deb 100644
--- a/src/kdf.rs
+++ b/src/kdf.rs
@@ -1,40 +1,9 @@
 use crate::canonicalization::canonicalize;
+use crate::context::KeyContext;
 use crate::ikm::InputKeyMaterial;
-use std::num::NonZeroU64;
 
 pub(crate) type KdfFunction = dyn Fn(&str, &[u8]) -> Vec<u8>;
 
-pub struct KeyContext {
-	ctx: Vec<String>,
-	periodicity: Option<u64>,
-}
-
-impl KeyContext {
-	pub fn set_static(&mut self) {
-		self.periodicity = None;
-	}
-
-	pub fn set_periodicity(&mut self, periodicity: NonZeroU64) {
-		self.periodicity = Some(periodicity.get());
-	}
-
-	pub(crate) fn get_ctx_elems(&self, time_period: Option<u64>) -> Vec<Vec<u8>> {
-		let mut ret: Vec<Vec<u8>> = self.ctx.iter().map(|s| s.as_bytes().to_vec()).collect();
-		if let Some(tp) = time_period {
-			ret.push(tp.to_le_bytes().to_vec());
-		}
-		ret
-	}
-
-	pub(crate) fn get_time_period(&self, timestamp: u64) -> Option<u64> {
-		self.periodicity.map(|p| timestamp / p)
-	}
-
-	pub(crate) fn is_periodic(&self) -> bool {
-		self.periodicity.is_some()
-	}
-}
-
 impl<const N: usize> From<[&str; N]> for KeyContext {
 	fn from(ctx: [&str; N]) -> Self {
 		Self {
diff --git a/src/lib.rs b/src/lib.rs
index 1c07dec..c6acce9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,8 @@
 mod canonicalization;
 #[cfg(feature = "encryption")]
 mod cipher_box;
+#[cfg(feature = "encryption")]
+mod context;
 #[cfg(any(feature = "encryption", feature = "ikm-management"))]
 mod error;
 #[cfg(any(feature = "encryption", feature = "ikm-management"))]
@@ -14,13 +16,13 @@ mod scheme;
 mod storage;
 
 #[cfg(feature = "encryption")]
-pub use cipher_box::{CipherBox, DataContext};
+pub use cipher_box::CipherBox;
+#[cfg(feature = "encryption")]
+pub use context::{DataContext, KeyContext};
 #[cfg(any(feature = "encryption", feature = "ikm-management"))]
 pub use error::Error;
 #[cfg(any(feature = "encryption", feature = "ikm-management"))]
 pub use ikm::{IkmId, InputKeyMaterial, InputKeyMaterialList};
-#[cfg(feature = "encryption")]
-pub use kdf::KeyContext;
 #[cfg(any(feature = "encryption", feature = "ikm-management"))]
 pub use scheme::Scheme;