From 85c482678286cecad178cfef17000576f977cd26 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rodolphe=20Br=C3=A9ard?= <rodolphe@what.tf>
Date: Sat, 22 Jul 2023 18:02:05 +0200
Subject: [PATCH] Implement code generation

---
 Cargo.toml     |  2 ++
 src/address.rs | 19 +++++++++++++++++--
 src/code.rs    | 32 ++++++++++++++++++++++++++++++++
 src/main.rs    |  1 +
 src/service.rs |  6 ++----
 start_test.py  |  6 +++++-
 6 files changed, 59 insertions(+), 7 deletions(-)
 create mode 100644 src/code.rs

diff --git a/Cargo.toml b/Cargo.toml
index 1cff1a6..6196ee8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,3 +11,5 @@ publish = false
 anyhow = { version = "1.0.71", default-features = false, features = ["std"] }
 clap = { version = "4.3.11", default-features = false, features = ["derive", "std"] }
 data-encoding = { version = "2.4.0", default-features = false, features = ["std"] }
+hmac = { version = "0.12.1", default-features = false }
+sha2 = { version = "0.10.7", default-features = false, features = ["std", "asm"] }
diff --git a/src/address.rs b/src/address.rs
index 9ea840f..85d47f7 100644
--- a/src/address.rs
+++ b/src/address.rs
@@ -1,3 +1,4 @@
+use crate::code::generate_code;
 use anyhow::{ensure, Error, Result};
 use data_encoding::{BASE32_NOPAD, BASE64};
 use std::hash::{Hash, Hasher};
@@ -9,6 +10,7 @@ pub struct CodedAddress {
 	sub_addr: Option<String>,
 	code: Vec<u8>,
 	domain: Option<String>,
+	separator: char,
 }
 
 impl CodedAddress {
@@ -38,6 +40,7 @@ impl CodedAddress {
 			sub_addr,
 			code,
 			domain,
+			separator,
 		})
 	}
 }
@@ -51,8 +54,20 @@ pub struct KeyedAddress {
 
 impl KeyedAddress {
 	pub fn check_code(&self, addr: &CodedAddress) -> bool {
-		// TODO
-		false
+		if addr.local_part.is_empty() || self.key.is_empty() {
+			return false;
+		}
+		match &addr.sub_addr {
+			Some(sub_addr) => {
+				if !sub_addr.is_empty() {
+					addr.code
+						== generate_code(&addr.local_part, addr.separator, sub_addr, &self.key)
+				} else {
+					false
+				}
+			}
+			None => false,
+		}
 	}
 }
 
diff --git a/src/code.rs b/src/code.rs
new file mode 100644
index 0000000..705ba63
--- /dev/null
+++ b/src/code.rs
@@ -0,0 +1,32 @@
+use hmac::{Hmac, Mac};
+use sha2::Sha256;
+
+const NB_BYTES: usize = 5;
+
+pub fn generate_code(local_part: &str, separator: char, sub_addr: &str, key: &[u8]) -> Vec<u8> {
+	// Compute the HMAC-SHA-256
+	let mut hmac = Hmac::<Sha256>::new_from_slice(key).unwrap();
+	hmac.update(local_part.as_bytes());
+	hmac.update(separator.to_string().as_bytes());
+	hmac.update(sub_addr.as_bytes());
+	let result = hmac.finalize().into_bytes();
+
+	// Reduce the result to NB_BYTES using a dynamic offset truncation
+	let offset = (result[result.len() - 1] & 0xf) as usize;
+	result[offset..offset + NB_BYTES].to_vec()
+}
+
+#[cfg(test)]
+mod tests {
+	use super::generate_code;
+
+	#[test]
+	fn code_generation() {
+		let key: &[u8] = &[
+			0xd7, 0x5b, 0xe8, 0x89, 0xe7, 0xca, 0xe4, 0xf8, 0x02, 0x5f, 0x91, 0x75, 0x4d, 0x37,
+			0x2e, 0xa1,
+		];
+		let code = generate_code("a", '+', "test", key);
+		assert_eq!(code, vec![0x7d, 0xd8, 0xd7, 0x1c, 0x8e]);
+	}
+}
diff --git a/src/main.rs b/src/main.rs
index f8b1fa8..a8b13f0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 use std::process::ExitCode;
 
 mod address;
+mod code;
 mod config;
 mod input;
 mod service;
diff --git a/src/service.rs b/src/service.rs
index 2f906b5..3e02f49 100644
--- a/src/service.rs
+++ b/src/service.rs
@@ -95,15 +95,13 @@ mod tests {
 	}
 
 	#[test]
-	#[ignore]
 	fn test_valid_code_domain() {
-		assert!(run_test_with_addr("a+test+TODO@example.org"));
+		assert!(run_test_with_addr("a+test+pxmnoheo@example.org"));
 	}
 
 	#[test]
-	#[ignore]
 	fn test_valid_code_no_domain() {
-		assert!(run_test_with_addr("b+test+TODO@example.org"));
+		assert!(run_test_with_addr("b+test+uivojtoa@example.org"));
 	}
 
 	#[test]
diff --git a/start_test.py b/start_test.py
index 4f4bbdc..7a9eace 100755
--- a/start_test.py
+++ b/start_test.py
@@ -122,8 +122,12 @@ def get_maildir():
 def start_tests(test_dir, smtp_port):
     to_addrs = [
         ("test@example.org", True),
-        ("test@nope.example.org", False),
         ("a@example.com", True),
+        ("a+anything@example.com", True),
+        ("a+anything+test@example.com", True),
+        ("a+test+pxmnoheo@example.org", True),
+        ("b+test+uivojtoa@example.org", True),
+        ("test@nope.example.org", False),
         ("a@example.org", False),
         ("a+invalid@example.org", False),
         ("a+invalid+input@example.org", False),