From 25f86bb6681be3c740203b1e1114af8684290e46 Mon Sep 17 00:00:00 2001
From: Yush G <aayushgupta05@gmail.com>
Date: Wed, 15 May 2024 23:07:53 +0900
Subject: [PATCH] Made more fields public + added get_canonicalized_header
 function to authenticated_message (#12)

* made authenticated message members pub

* made more things public + added print

* more pub(crate)s -> pub

* fix print on headers midway thru verify dkim

* fixed print again

* debug signature

* print canonicalized header?

* finished debugging and adding get_canonicalized_header function

* reduced pub variables, cleared prints, cleaned code

* fixed comments on function
---
 src/common/headers.rs    |  6 +++---
 src/dkim/canonicalize.rs |  2 +-
 src/dkim/mod.rs          |  4 ++--
 src/dkim/verify.rs       | 30 ++++++++++++++++++++++++++++++
 src/lib.rs               | 24 ++++++++++++------------
 5 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/src/common/headers.rs b/src/common/headers.rs
index 2cf0265..4bec251 100644
--- a/src/common/headers.rs
+++ b/src/common/headers.rs
@@ -60,9 +60,9 @@ pub(crate) enum AuthenticatedHeader<'x> {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Header<'x, T> {
-    pub(crate) name: &'x [u8],
-    pub(crate) value: &'x [u8],
-    pub(crate) header: T,
+    pub name: &'x [u8],
+    pub value: &'x [u8],
+    pub header: T,
 }
 
 impl<'x> HeaderParser<'x> {
diff --git a/src/dkim/canonicalize.rs b/src/dkim/canonicalize.rs
index d27ac40..7a21b6c 100644
--- a/src/dkim/canonicalize.rs
+++ b/src/dkim/canonicalize.rs
@@ -158,7 +158,7 @@ impl Canonicalization {
 }
 
 impl Signature {
-    pub(crate) fn canonicalize<'x>(
+    pub fn canonicalize<'x>(
         &self,
         mut message: impl HeaderStream<'x>,
     ) -> (usize, CanonicalHeaders<'x>, Vec<String>, CanonicalBody<'x>) {
diff --git a/src/dkim/mod.rs b/src/dkim/mod.rs
index ee4b110..d30306e 100644
--- a/src/dkim/mod.rs
+++ b/src/dkim/mod.rs
@@ -36,8 +36,8 @@ pub enum Canonicalization {
 #[derive(Debug, PartialEq, Eq, Clone, Default)]
 pub struct DkimSigner<T: SigningKey, State = NeedDomain> {
     _state: std::marker::PhantomData<State>,
-    pub(crate) key: T,
-    pub(crate) template: Signature,
+    pub key: T,
+    pub template: Signature,
 }
 
 pub struct NeedDomain;
diff --git a/src/dkim/verify.rs b/src/dkim/verify.rs
index a590339..31690f9 100644
--- a/src/dkim/verify.rs
+++ b/src/dkim/verify.rs
@@ -240,6 +240,36 @@ impl Resolver {
 }
 
 impl<'x> AuthenticatedMessage<'x> {
+    pub async fn get_canonicalized_header(&self) -> Result<Vec<u8>, Error> {
+        // Based on verify_dkim_ function
+        // Iterate through possible DKIM headers
+        let mut data = Vec::with_capacity(256);
+        for header in &self.dkim_headers {
+            // Ensure signature is not obviously invalid
+            let signature = match &header.header {
+                Ok(signature) => {
+                    if signature.x == 0 || (signature.x > signature.t) {
+                        signature
+                    } else {
+                        continue;
+                    }
+                }
+                Err(_err) => {
+                    continue;
+                }
+            };
+
+            // Get pre-hashed but canonically ordered headers, who's hash is signed
+            let dkim_hdr_value = header.value.strip_signature();
+            let headers = self.signed_headers(&signature.h, header.name, &dkim_hdr_value);
+            signature.ch.canonicalize_headers(headers, &mut data);
+
+            return Ok(data);
+        }
+        // Return not ok
+        Err(Error::FailedBodyHashMatch)
+    }
+
     pub fn signed_headers<'z: 'x>(
         &'z self,
         headers: &'x [String],
diff --git a/src/lib.rs b/src/lib.rs
index 67d3c91..7957986 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -335,18 +335,18 @@ pub struct MX {
 
 #[derive(Debug, Clone)]
 pub struct AuthenticatedMessage<'x> {
-    pub(crate) headers: Vec<(&'x [u8], &'x [u8])>,
-    pub(crate) from: Vec<String>,
-    pub(crate) raw_message: &'x [u8],
-    pub(crate) body_offset: usize,
-    pub(crate) body_hashes: Vec<(Canonicalization, HashAlgorithm, u64, Vec<u8>)>,
-    pub(crate) dkim_headers: Vec<Header<'x, crate::Result<dkim::Signature>>>,
-    pub(crate) ams_headers: Vec<Header<'x, crate::Result<arc::Signature>>>,
-    pub(crate) as_headers: Vec<Header<'x, crate::Result<arc::Seal>>>,
-    pub(crate) aar_headers: Vec<Header<'x, crate::Result<arc::Results>>>,
-    pub(crate) received_headers_count: usize,
-    pub(crate) date_header_present: bool,
-    pub(crate) message_id_header_present: bool,
+    pub headers: Vec<(&'x [u8], &'x [u8])>,
+    pub from: Vec<String>,
+    pub raw_message: &'x [u8],
+    pub body_offset: usize,
+    pub body_hashes: Vec<(Canonicalization, HashAlgorithm, u64, Vec<u8>)>,
+    pub dkim_headers: Vec<Header<'x, crate::Result<dkim::Signature>>>,
+    pub ams_headers: Vec<Header<'x, crate::Result<arc::Signature>>>,
+    pub as_headers: Vec<Header<'x, crate::Result<arc::Seal>>>,
+    pub aar_headers: Vec<Header<'x, crate::Result<arc::Results>>>,
+    pub received_headers_count: usize,
+    pub date_header_present: bool,
+    pub message_id_header_present: bool,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]