From 18ed803bf60befe26156d5aeda4fdf24ab76572c Mon Sep 17 00:00:00 2001
From: Thor Kampefner <thorck@pm.me>
Date: Tue, 6 Aug 2024 12:41:04 -0700
Subject: [PATCH] refactor: main -> lib. mv Witness -> witness mod

---
 src/{main.rs => lib.rs} | 60 +----------------------------------------
 src/proof.rs            |  2 +-
 src/utils.rs            | 45 +++++++++++++++++++++++++++++++
 src/witness.rs          | 12 +++++++--
 4 files changed, 57 insertions(+), 62 deletions(-)
 rename src/{main.rs => lib.rs} (53%)
 create mode 100644 src/utils.rs

diff --git a/src/main.rs b/src/lib.rs
similarity index 53%
rename from src/main.rs
rename to src/lib.rs
index 4aa5d21..166ce55 100644
--- a/src/main.rs
+++ b/src/lib.rs
@@ -15,6 +15,7 @@ use ark_circom::CircomBuilder;
 use ark_ec::pairing::Pairing;
 
 mod proof;
+mod utils;
 mod witness;
 
 /// Circom compilation artifacts
@@ -27,65 +28,6 @@ const AES_256_CRT_R1CS: &str = "./build/aes_256_ctr_test.r1cs";
 pub type AAD = [u8; 5];
 pub type Nonce = [u8; 12];
 
-/// Witness bytes generated by this binary
-pub struct Witness {
-    pub key: Vec<u8>,
-    pub iv:  Vec<u8>,
-    pub ct:  Vec<u8>,
-    pub pt:  Vec<u8>,
-}
-
-// TODO(TK 2024-08-06): move these to tests
-#[tokio::main]
-async fn main() -> io::Result<()> {
-    // aes_gcm_siv_test().await?;
-
-    // aes_256ctr_test().await?;
-    Ok(())
-}
-
-// TODO(TK 2024-08-06): refactor; move util methods to utils
-// TODO(TK 2024-08-06): test with test vectors at bottom of rfc 8452
-//
-/// construct the nonce from the `iv` and `seq` as specified in RFC 8452
-/// https://www.rfc-editor.org/rfc/rfc8452
-pub fn make_nonce(iv: [u8; 12], seq: u64) -> Nonce {
-    let mut nonce = [0u8; 12];
-    nonce[4..].copy_from_slice(&seq.to_be_bytes());
-
-    nonce.iter_mut().zip(iv).for_each(|(nonce, iv)| {
-        *nonce ^= iv;
-    });
-
-    nonce
-}
-
-fn make_tls13_aad(len: usize) -> AAD {
-    [
-        0x17, // ContentType::ApplicationData
-        0x3,  // ProtocolVersion (major)
-        0x3,  // ProtocolVersion (minor)
-        (len >> 8) as u8,
-        len as u8,
-    ]
-}
-
-// TODO(TK 2024-08-06): @devloper, document
-fn push_bytes_as_bits<T: Pairing>(
-    mut builder: CircomBuilder<T>,
-    field: &str,
-    bytes: &[u8],
-) -> CircomBuilder<T> {
-    for byte in bytes {
-        for i in 0..8 {
-            let bit = (byte >> (7 - i)) & 1;
-            builder.push_input(field, bit as u64);
-        }
-    }
-
-    builder
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/proof.rs b/src/proof.rs
index 3fc6231..2e0452a 100644
--- a/src/proof.rs
+++ b/src/proof.rs
@@ -9,7 +9,7 @@ use ark_std::rand::thread_rng;
 type GrothBn = Groth16<Bn254>;
 use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
 
-use crate::{push_bytes_as_bits, Witness};
+use crate::{utils::push_bytes_as_bits, witness::Witness};
 
 // TODO(TK 2024-08-06): refactor kludge
 // load up the circom
diff --git a/src/utils.rs b/src/utils.rs
new file mode 100644
index 0000000..f35fab6
--- /dev/null
+++ b/src/utils.rs
@@ -0,0 +1,45 @@
+use ark_circom::CircomBuilder;
+use ark_ec::pairing::Pairing;
+
+use crate::{Nonce, AAD};
+
+// TODO(TK 2024-08-06): test with test vectors at bottom of rfc 8452
+//
+/// construct the nonce from the `iv` and `seq` as specified in RFC 8452
+/// https://www.rfc-editor.org/rfc/rfc8452
+pub(crate) fn make_nonce(iv: [u8; 12], seq: u64) -> Nonce {
+    let mut nonce = [0u8; 12];
+    nonce[4..].copy_from_slice(&seq.to_be_bytes());
+
+    nonce.iter_mut().zip(iv).for_each(|(nonce, iv)| {
+        *nonce ^= iv;
+    });
+
+    nonce
+}
+
+pub(crate) fn make_tls13_aad(len: usize) -> AAD {
+    [
+        0x17, // ContentType::ApplicationData
+        0x3,  // ProtocolVersion (major)
+        0x3,  // ProtocolVersion (minor)
+        (len >> 8) as u8,
+        len as u8,
+    ]
+}
+
+// TODO(TK 2024-08-06): @devloper, document
+pub(crate) fn push_bytes_as_bits<T: Pairing>(
+    mut builder: CircomBuilder<T>,
+    field: &str,
+    bytes: &[u8],
+) -> CircomBuilder<T> {
+    for byte in bytes {
+        for i in 0..8 {
+            let bit = (byte >> (7 - i)) & 1;
+            builder.push_input(field, bit as u64);
+        }
+    }
+
+    builder
+}
diff --git a/src/witness.rs b/src/witness.rs
index 5e90c84..c49b56b 100644
--- a/src/witness.rs
+++ b/src/witness.rs
@@ -8,11 +8,21 @@ use aes_gcm::{
 };
 use cipher::consts::U16;
 
+use crate::utils::{make_nonce, make_tls13_aad};
+
 type Ctr32BE<Aes128> = ctr::CtrCore<Aes128, ctr::flavors::Ctr32BE>;
 type Aes256Ctr32BE = ctr::Ctr32BE<Aes256>;
 type Block = GenericArray<u8, U16>;
 type Aes128Ctr32BE = ctr::Ctr32BE<aes::Aes128>; // Note: Ctr32BE is used in AES GCM
 
+/// Witness bytes generated by this binary
+pub struct Witness {
+    pub key: Vec<u8>,
+    pub iv:  Vec<u8>,
+    pub ct:  Vec<u8>,
+    pub pt:  Vec<u8>,
+}
+
 pub enum CipherMode {
     Vanilla, // no IV Here
     Ctr256,
@@ -21,8 +31,6 @@ pub enum CipherMode {
     Ctr128,
 }
 
-use crate::{make_nonce, make_tls13_aad, Witness};
-
 const KEY_ASCII: &str = "1111111111111111"; // 16 bytes
 const IV_ASCII: &str = "111111111111"; // 12 bytes
 const MESSAGE: &str = "test000000000000";