From 586651cbbb67afe33ea0c2386a9a1524215f5c5b Mon Sep 17 00:00:00 2001
From: supinie <86788874+supinie@users.noreply.github.com>
Date: Wed, 10 Apr 2024 16:35:38 +0100
Subject: [PATCH] adds public key packing

---
 src/errors.rs      |  8 ++++++
 src/indcpa.rs      | 63 +++++++++++++++++++++++++++++++++++-----------
 src/matrix.rs      |  8 ++++--
 src/polynomials.rs |  2 +-
 src/vectors.rs     |  4 +--
 5 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/src/errors.rs b/src/errors.rs
index 6cb7953..d7ac12a 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -1,5 +1,6 @@
 use crate::params::{SecurityLevel, K};
 use core::{
+    array::TryFromSliceError,
     fmt::{Display, Formatter},
     num::TryFromIntError,
 };
@@ -28,6 +29,7 @@ pub enum PackingError {
     Crystals(CrystalsError),
     TryFromPrimitive(TryFromPrimitiveError<K>),
     TryFromInt(TryFromIntError),
+    TryFromSlice(TryFromSliceError),
 }
 
 impl From<CrystalsError> for PackingError {
@@ -47,3 +49,9 @@ impl From<TryFromIntError> for PackingError {
         Self::TryFromInt(error)
     }
 }
+
+impl From<TryFromSliceError> for PackingError {
+    fn from(error: TryFromSliceError) -> Self {
+        Self::TryFromSlice(error)
+    }
+}
diff --git a/src/indcpa.rs b/src/indcpa.rs
index aa14bb5..28072a3 100644
--- a/src/indcpa.rs
+++ b/src/indcpa.rs
@@ -1,11 +1,11 @@
 use core::num::TryFromIntError;
 
 use crate::{
+    errors::{CrystalsError, PackingError},
     matrix::Matrix,
-    params::{POLYBYTES, SYMBYTES},
-    polynomials::{Poly, State, Normalised, Unnormalised},
-    vectors::{PolyVec, unpack_to_polyvec},
-    errors::PackingError,
+    params::{SecurityLevel, K, POLYBYTES, SYMBYTES},
+    polynomials::{Normalised, Poly, State, Unnormalised},
+    vectors::{unpack_to_polyvec, PolyVec},
 };
 use sha3::{Digest, Sha3_512};
 
@@ -15,28 +15,61 @@ pub struct PrivateKey {
 }
 
 #[derive(Default, PartialEq, Debug, Eq)]
-pub struct PublicKey<S: State> {
+pub struct PublicKey {
     rho: [u8; SYMBYTES],
-    noise: PolyVec<S>,
-    a_t: Matrix<S>,
+    noise: PolyVec<Normalised>,
+    a_t: Matrix<Unnormalised>,
 }
 
 impl PrivateKey {
-    pub fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
+    fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
         self.secret.pack(buf)
     }
 }
 
 fn unpack_to_private_key(buf: &[u8]) -> Result<PrivateKey, PackingError> {
     let secret = unpack_to_polyvec(buf)?.normalise();
-    Ok(PrivateKey {
-        secret
-    })
+    Ok(PrivateKey { secret })
+}
+
+impl PublicKey {
+    fn sec_level(&self) -> Result<SecurityLevel, CrystalsError> {
+        if self.noise.sec_level() == self.a_t.sec_level() {
+            Ok(self.noise.sec_level())
+        } else {
+            Err(CrystalsError::MismatchedSecurityLevels(
+                self.noise.sec_level(),
+                self.a_t.sec_level(),
+            ))
+        }
+    }
+
+    fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
+        let k: usize = self.sec_level()?.k().into();
+
+        let break_point: usize = POLYBYTES * k;
+        if buf[break_point..].len() == SYMBYTES {
+            self.noise.pack(&mut buf[..break_point])?;
+            buf[break_point..].copy_from_slice(&self.rho[..]);
+            Ok(())
+        } else {
+            Err(CrystalsError::IncorrectBufferLength(buf.len(), break_point + SYMBYTES).into())
+        }
+    }
+}
+
+fn unpack_to_public_key(buf: &[u8]) -> Result<PublicKey, PackingError> {
+    let k = K::try_from((buf.len() - SYMBYTES) / POLYBYTES)?;
+    let k_value: usize = k.into();
+    let break_point: usize = POLYBYTES * k_value;
+
+    let noise = unpack_to_polyvec(&buf[..break_point])?.normalise();
+    let rho: [u8; SYMBYTES] = buf[break_point..].try_into()?;
+
+    let a_t = Matrix::derive(&rho, true, k)?;
+
+    Ok(PublicKey { rho, noise, a_t })
 }
-//     pub fn unpack(&mut self, buf: &[u8]) {
-//         self.secret.unpack(buf);
-//         self.secret.normalise();
-//     }
 
 // impl<PV: PolyVecOperations, M: MatOperations + GetSecLevel + LinkSecLevel<PV>> PublicKey<PV, M> {
 //     pub fn pack(&self, buf: &mut [u8]) {
diff --git a/src/matrix.rs b/src/matrix.rs
index a255dc8..8357f2c 100644
--- a/src/matrix.rs
+++ b/src/matrix.rs
@@ -13,7 +13,7 @@ pub struct Matrix<S: State> {
 }
 
 impl<S: State + Copy> Matrix<S> {
-    const fn sec_level(&self) -> SecurityLevel {
+    pub(crate) const fn sec_level(&self) -> SecurityLevel {
         SecurityLevel::new(self.sec_level)
     }
 
@@ -67,7 +67,11 @@ impl Matrix<Normalised> {
 }
 
 impl Matrix<Unnormalised> {
-    fn derive(seed: &[u8], transpose: bool, sec_level: K) -> Result<Self, CrystalsError> {
+    pub(crate) fn derive(
+        seed: &[u8],
+        transpose: bool,
+        sec_level: K,
+    ) -> Result<Self, CrystalsError> {
         let mut polyvecs = ArrayVec::<[PolyVec<Unnormalised>; 4]>::new();
         if transpose {
             for i in 0..sec_level.into() {
diff --git a/src/polynomials.rs b/src/polynomials.rs
index 98251a6..2bb32f1 100644
--- a/src/polynomials.rs
+++ b/src/polynomials.rs
@@ -19,7 +19,7 @@ pub struct Poly<S: State> {
 // Normalised coefficients lie within {0..q-1}
 #[derive(Default, Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Normalised;
-#[derive(Default, Copy, Clone)]
+#[derive(Default, Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Unnormalised;
 #[derive(Default, Copy, Clone)]
 pub struct Noise;
diff --git a/src/vectors.rs b/src/vectors.rs
index c2c5bb2..18cfd8e 100644
--- a/src/vectors.rs
+++ b/src/vectors.rs
@@ -16,7 +16,7 @@ pub struct PolyVec<S: State> {
 
 impl<S: State> PolyVec<S> {
     // Gets the security level of the given polyvec.
-    const fn sec_level(&self) -> SecurityLevel {
+    pub(crate) const fn sec_level(&self) -> SecurityLevel {
         SecurityLevel::new(self.sec_level)
     }
 
@@ -193,7 +193,7 @@ impl PolyVec<Noise> {
 // The buffer should be of length k * POLYBYTES.
 // If the length of the buffer is incorrect, the operation can still succeed provided it is a valid
 // multiple of POLYBYTES, and will result in a polyvec of incorrect security level.
-pub(crate) fn unpack_to_polyvec(buf: &[u8]) -> Result<PolyVec<Unnormalised>, PackingError> {
+pub fn unpack_to_polyvec(buf: &[u8]) -> Result<PolyVec<Unnormalised>, PackingError> {
     let sec_level = K::try_from(buf.len() / POLYBYTES)?; // If this fails then we know the
                                                          // buffer is not of the right size and
                                                          // so no further checks are needed.