Browse Source

add TGSW encryption & decryption

main
arnaucube 1 week ago
parent
commit
f053e9a904
4 changed files with 136 additions and 48 deletions
  1. +1
    -0
      tfhe/src/lib.rs
  2. +74
    -0
      tfhe/src/tgsw.rs
  3. +19
    -10
      tfhe/src/tlev.rs
  4. +42
    -38
      tfhe/src/tlwe.rs

+ 1
- 0
tfhe/src/lib.rs

@ -5,5 +5,6 @@
#![allow(clippy::upper_case_acronyms)]
#![allow(dead_code)] // TMP
pub mod tgsw;
pub mod tlev;
pub mod tlwe;

+ 74
- 0
tfhe/src/tgsw.rs

@ -0,0 +1,74 @@
use anyhow::Result;
use itertools::zip_eq;
use rand::Rng;
use std::array;
use std::ops::{Add, Mul};
use arith::{Ring, Rq, Tn, T64, TR};
use crate::tlev::TLev;
use crate::tlwe::{PublicKey, SecretKey, TLWE};
use gfhe::glwe::GLWE;
/// vector of length K+1 = [K], [1]
#[derive(Clone, Debug)]
pub struct TGSW<const K: usize>(pub(crate) Vec<TLev<K>>, TLev<K>);
impl<const K: usize> TGSW<K> {
pub fn encrypt_s(
mut rng: impl Rng,
beta: u32,
l: u32,
sk: &SecretKey<K>,
m: &T64,
) -> Result<Self> {
let a: Vec<TLev<K>> = (0..K)
.map(|i| TLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0[i] * *m)))
.collect::<Result<Vec<_>>>()?;
let b: TLev<K> = TLev::encrypt_s(&mut rng, beta, l, sk, m)?;
Ok(Self(a, b))
}
pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> T64 {
self.1.decrypt(sk, beta)
}
pub fn from_tlwe(_tlwe: TLWE<K>) -> Self {
todo!()
}
}
#[cfg(test)]
mod tests {
use anyhow::Result;
use rand::distributions::Uniform;
use super::*;
#[test]
fn test_encrypt_decrypt() -> Result<()> {
const T: u64 = 2; // plaintext modulus
const K: usize = 16;
type S = TGSW<K>;
let beta: u32 = 2;
let l: u32 = 16;
let mut rng = rand::thread_rng();
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..50 {
let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
let p: T64 = TLev::<K>::encode::<T>(&m); // plaintext
let c = S::encrypt_s(&mut rng, beta, l, &sk, &p)?;
let p_recovered = c.decrypt(&sk, beta);
let m_recovered = TLev::<K>::decode::<T>(&p_recovered);
assert_eq!(m, m_recovered);
}
Ok(())
}
}

+ 19
- 10
tfhe/src/tlev.rs

@ -1,4 +1,5 @@
use anyhow::Result;
use itertools::zip_eq;
use rand::Rng;
use std::array;
use std::ops::{Add, Mul};
@ -11,11 +12,11 @@ use crate::tlwe::{PublicKey, SecretKey, TLWE};
pub struct TLev<const K: usize>(pub(crate) Vec<TLWE<K>>);
impl<const K: usize> TLev<K> {
pub fn encode<const T: u64>(m: &Rq<T, 1>) -> Tn<1> {
pub fn encode<const T: u64>(m: &Rq<T, 1>) -> T64 {
let coeffs = m.coeffs();
Tn(array::from_fn(|i| T64(coeffs[i].0)))
T64(coeffs[0].0) // N=1, so take the only coeff
}
pub fn decode<const T: u64>(p: &Tn<1>) -> Rq<T, 1> {
pub fn decode<const T: u64>(p: &T64) -> Rq<T, 1> {
Rq::<T, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
}
pub fn encrypt(
@ -23,7 +24,7 @@ impl TLev {
beta: u32,
l: u32,
pk: &PublicKey<K>,
m: &Tn<1>,
m: &T64,
) -> Result<Self> {
let tlev: Vec<TLWE<K>> = (1..l + 1)
.map(|i| {
@ -35,25 +36,33 @@ impl TLev {
}
pub fn encrypt_s(
mut rng: impl Rng,
beta: u32,
_beta: u32, // TODO rm, and make beta=2 always
l: u32,
sk: &SecretKey<K>,
m: &Tn<1>,
m: &T64,
) -> Result<Self> {
let tlev: Vec<TLWE<K>> = (1..l + 1)
let tlev: Vec<TLWE<K>> = (1..l as u64 + 1)
.map(|i| {
TLWE::<K>::encrypt_s(&mut rng, sk, &(*m * (u64::MAX / beta.pow(i as u32) as u64)))
let aux = if i < 64 {
*m * (u64::MAX / (1u64 << i))
} else {
// 1<<64 would overflow, and anyways we're dividing u64::MAX
// by it, which would be equal to 1
*m
};
TLWE::<K>::encrypt_s(&mut rng, sk, &aux)
})
.collect::<Result<Vec<_>>>()?;
Ok(Self(tlev))
}
pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> Tn<1> {
pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> T64 {
let pt = self.0[0].decrypt(sk);
pt.mul_div_round(beta as u64, u64::MAX)
}
}
// TODO review u64::MAX, since is -1 of the value we actually want
#[cfg(test)]
mod tests {
@ -78,7 +87,7 @@ mod tests {
let (sk, pk) = TLWE::<K>::new_key(&mut rng)?;
let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
let p: Tn<1> = S::encode::<T>(&m); // plaintext
let p: T64 = S::encode::<T>(&m); // plaintext
let c = S::encrypt(&mut rng, beta, l, &pk, &p)?;
let p_recovered = c.decrypt(&sk, beta);

+ 42
- 38
tfhe/src/tlwe.rs

@ -10,47 +10,50 @@ use std::ops::{Add, AddAssign, Mul, Sub};
use arith::{Ring, Rq, Tn, T64, TR};
use gfhe::{glwe, GLWE};
const ERR_SIGMA: f64 = 3.2;
// #[derive(Clone, Debug)]
// pub struct SecretKey<const K: usize>(glwe::SecretKey<T64, K>);
pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>;
#[derive(Clone, Debug)]
pub struct SecretKey<const K: usize>(glwe::SecretKey<Tn<1>, K>);
#[derive(Clone, Debug)]
pub struct PublicKey<const K: usize>(glwe::PublicKey<Tn<1>, K>);
// #[derive(Clone, Debug)]
// pub struct PublicKey<const K: usize>(glwe::PublicKey<T64, K>);
pub type PublicKey<const K: usize> = glwe::PublicKey<T64, K>;
#[derive(Clone, Debug)]
pub struct TLWE<const K: usize>(pub GLWE<Tn<1>, K>);
pub struct TLWE<const K: usize>(pub GLWE<T64, K>);
impl<const K: usize> TLWE<K> {
pub fn zero() -> Self {
Self(GLWE::<Tn<1>, K>::zero())
Self(GLWE::<T64, K>::zero())
}
pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
let (sk, pk) = GLWE::new_key(rng)?;
Ok((SecretKey(sk), PublicKey(pk)))
// Ok((SecretKey(sk), PublicKey(pk)))
Ok((sk, pk))
}
pub fn encode<const P: u64>(m: &Rq<P, 1>) -> Tn<1> {
pub fn encode<const P: u64>(m: &Rq<P, 1>) -> T64 {
let delta = u64::MAX / P; // floored
let coeffs = m.coeffs();
Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
// Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
T64(coeffs[0].0 * delta)
}
pub fn decode<const P: u64>(p: &Tn<1>) -> Rq<P, 1> {
pub fn decode<const P: u64>(p: &T64) -> Rq<P, 1> {
let p = p.mul_div_round(P, u64::MAX);
Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
}
// encrypts with the given SecretKey (instead of PublicKey)
pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &Tn<1>) -> Result<Self> {
let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &T64) -> Result<Self> {
let glwe = GLWE::encrypt_s(rng, &sk, p)?;
Ok(Self(glwe))
}
pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &Tn<1>) -> Result<Self> {
let glwe = GLWE::encrypt(rng, &pk.0, p)?;
pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &T64) -> Result<Self> {
let glwe = GLWE::encrypt(rng, &pk, p)?;
Ok(Self(glwe))
}
pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> {
self.0.decrypt(&sk.0)
pub fn decrypt(&self, sk: &SecretKey<K>) -> T64 {
self.0.decrypt(&sk)
}
}
@ -86,29 +89,29 @@ impl Sub> for TLWE {
}
// plaintext addition
impl<const K: usize> Add<Tn<1>> for TLWE<K> {
impl<const K: usize> Add<T64> for TLWE<K> {
type Output = Self;
fn add(self, plaintext: Tn<1>) -> Self {
let a: TR<Tn<1>, K> = self.0 .0;
let b: Tn<1> = self.0 .1 + plaintext;
fn add(self, plaintext: T64) -> Self {
let a: TR<T64, K> = self.0 .0;
let b: T64 = self.0 .1 + plaintext;
Self(GLWE(a, b))
}
}
// plaintext substraction
impl<const K: usize> Sub<Tn<1>> for TLWE<K> {
impl<const K: usize> Sub<T64> for TLWE<K> {
type Output = Self;
fn sub(self, plaintext: Tn<1>) -> Self {
let a: TR<Tn<1>, K> = self.0 .0;
let b: Tn<1> = self.0 .1 - plaintext;
fn sub(self, plaintext: T64) -> Self {
let a: TR<T64, K> = self.0 .0;
let b: T64 = self.0 .1 - plaintext;
Self(GLWE(a, b))
}
}
// plaintext multiplication
impl<const K: usize> Mul<Tn<1>> for TLWE<K> {
impl<const K: usize> Mul<T64> for TLWE<K> {
type Output = Self;
fn mul(self, plaintext: Tn<1>) -> Self {
let a: TR<Tn<1>, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect());
let b: Tn<1> = self.0 .1 * plaintext;
fn mul(self, plaintext: T64) -> Self {
let a: TR<T64, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect());
let b: T64 = self.0 .1 * plaintext;
Self(GLWE(a, b))
}
}
@ -134,7 +137,7 @@ mod tests {
let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
dbg!(&m);
let p: Tn<1> = S::encode::<T>(&m);
let p: T64 = S::encode::<T>(&m);
dbg!(&p);
let c = S::encrypt(&mut rng, &pk, &p)?;
@ -168,8 +171,8 @@ mod tests {
let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
let p1: Tn<1> = S::encode::<T>(&m1); // plaintext
let p2: Tn<1> = S::encode::<T>(&m2); // plaintext
let p1: T64 = S::encode::<T>(&m1); // plaintext
let p2: T64 = S::encode::<T>(&m2); // plaintext
let c1 = S::encrypt(&mut rng, &pk, &p1)?;
let c2 = S::encrypt(&mut rng, &pk, &p2)?;
@ -199,8 +202,8 @@ mod tests {
let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
let p1: Tn<1> = S::encode::<T>(&m1); // plaintext
let p2: Tn<1> = S::encode::<T>(&m2); // plaintext
let p1: T64 = S::encode::<T>(&m1); // plaintext
let p2: T64 = S::encode::<T>(&m2); // plaintext
let c1 = S::encrypt(&mut rng, &pk, &p1)?;
@ -209,7 +212,7 @@ mod tests {
let p3_recovered = c3.decrypt(&sk);
let m3_recovered = S::decode::<T>(&p3_recovered);
assert_eq!((m1 + m2).remodule::<T>(), m3_recovered);
assert_eq!(m1 + m2, m3_recovered);
}
Ok(())
@ -229,15 +232,16 @@ mod tests {
let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
let p1: Tn<1> = S::encode::<T>(&m1);
let p1: T64 = S::encode::<T>(&m1);
// don't scale up p2, set it directly from m2
let p2: Tn<1> = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
// let p2: T64 = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
let p2: T64 = T64(m2.coeffs()[0].0);
let c1 = S::encrypt(&mut rng, &pk, &p1)?;
let c3 = c1 * p2;
let p3_recovered: Tn<1> = c3.decrypt(&sk);
let p3_recovered: T64 = c3.decrypt(&sk);
let m3_recovered = S::decode::<T>(&p3_recovered);
assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
}

Loading…
Cancel
Save