Browse Source

implement GLWE key switching

gfhe-over-ring-trait
arnaucube 3 weeks ago
parent
commit
188bc7fa7f
8 changed files with 384 additions and 33 deletions
  1. +77
    -0
      arith/sage/ring.sage
  2. +6
    -0
      arith/src/ring.rs
  3. +39
    -0
      arith/src/ringq.rs
  4. +2
    -0
      arith/src/traits.rs
  5. +8
    -0
      arith/src/tuple_ring.rs
  6. +82
    -1
      arith/src/zq.rs
  7. +169
    -32
      gfhe/src/glwe.rs
  8. +1
    -0
      gfhe/src/lib.rs

+ 77
- 0
arith/sage/ring.sage

@ -0,0 +1,77 @@
def run_test(a, b):
print("\nnew test:")
print(a)
print(b)
c = a*b
print(c)
print(c.list())
n_iters = 100
Q= 65537
print(Q)
N=4
F = GF(Q)
R = QuotientRing(F[x], x^N + 1, names="X")
print(R)
a = R([4,2,1,0])
b = R([1,2,3,4])
run_test(a,b)
# print("Elements of the polynomial ring:")
# for e in R:
# print(e)
# Other:
# ======
#
# t = R.gen()
# a = 0 + t + 2*t^2 + 3*t^3 + 4*t^4 + 5*t^5
# b = 5 + 4*t + 3*t^2 + 2*t^3 + 1*t^4 + 0*t^5
# print("add", a+b)
# print("sub", a-b)
# print("mul", a*b)
# a = 0 + t + 2*t^2 + 3*t^3 + 4*t^4 + 5*t^5
# print("ring elem mul testvectors")
#
#
# def randvec(size=N):return [int(random()*(Q-1)) for t in range(size)]
#
# a_vecs = [None]*n_iters
# b_vecs = [None]*n_iters
# c_vecs = [None]*n_iters
#
# for i in range(n_iters):
# a_vec = randvec()
# b_vec = randvec()
# a_pol = R(a_vec)
# b_pol = R(b_vec)
#
# c_pol = a_pol*b_pol
#
# a_vecs[i] = a_pol.list()
# b_vecs[i] = b_pol.list()
# c_vecs[i] = c_pol.list()
#
# print("let a_vecs = vec!{};\n".format(a_vecs))
# print("let b_vecs = vec!{};\n".format(b_vecs))
# print("let c_vecs = vec!{};".format(c_vecs))
# # cyclotomic
#
# Q= 65537
# print(Q)
# N=4
# F = GF(Q)
# R = QuotientRing(F[x], x^N - 1, names="X")
# print(R)
#
# a = R([1,0,0,2])
# b = R([0,0,0,2])
# run_test(a, b)

+ 6
- 0
arith/src/ring.rs

@ -31,6 +31,12 @@ impl Ring for R {
// let coeffs: [C; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
// Self(coeffs)
}
// returns the decomposition of each polynomial coefficient
fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
unimplemented!();
// array::from_fn(|i| self.coeffs[i].decompose(beta, l))
}
}
impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> {

+ 39
- 0
arith/src/ringq.rs

@ -47,6 +47,18 @@ impl Ring for Rq {
evals: None,
}
}
// returns the decomposition of each polynomial coefficient, such
// decomposition will be a vecotor of length N, containint N vectors of Zq
fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
let elems: Vec<Vec<Zq<Q>>> = self.coeffs.iter().map(|r| r.decompose(beta, l)).collect();
// transpose it
let r: Vec<Vec<Zq<Q>>> = (0..elems[0].len())
.map(|i| (0..elems.len()).map(|j| elems[j][i]).collect())
.collect();
// convert it to Rq<Q,N>
r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
}
}
impl<const Q: u64, const N: usize> From<crate::ring::R<N>> for Rq<Q, N> {
@ -599,4 +611,31 @@ mod tests {
assert_eq!(c, expected_c);
Ok(())
}
#[test]
fn test_rq_decompose() -> Result<()> {
const Q: u64 = 16;
const N: usize = 4;
let beta = 4;
let l = 2;
let a = Rq::<Q, N>::from_vec_u64(vec![7u64, 14, 3, 6]);
let d = a.decompose(beta, l);
assert_eq!(
d[0],
vec![1u64, 3, 0, 1]
.iter()
.map(|e| Zq::<Q>::from_u64(*e))
.collect::<Vec<_>>()
);
assert_eq!(
d[1],
vec![3u64, 2, 3, 2]
.iter()
.map(|e| Zq::<Q>::from_u64(*e))
.collect::<Vec<_>>()
);
Ok(())
}
}

+ 2
- 0
arith/src/traits.rs

@ -27,4 +27,6 @@ pub trait Ring:
fn zero() -> Self;
// note/wip/warning: dist (0,q) with f64, will output more '0=q' elements than other values
fn rand(rng: impl Rng, dist: impl Distribution<f64>) -> Self;
fn decompose(&self, beta: u32, l: u32) -> Vec<Self>;
}

+ 8
- 0
arith/src/tuple_ring.rs

@ -16,6 +16,9 @@ use crate::Ring;
pub struct TR<R: Ring, const K: usize>(pub Vec<R>);
impl<R: Ring, const K: usize> TR<R, K> {
pub fn zero() -> Self {
Self((0..K).into_iter().map(|_| R::zero()).collect())
}
pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
Self(
(0..K)
@ -24,6 +27,11 @@ impl TR {
.collect(),
)
}
// returns the decomposition of each polynomial element
pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
unimplemented!()
// self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>>
}
}
impl<R: Ring, const K: usize> ops::Add<TR<R, K>> for TR<R, K> {

+ 82
- 1
arith/src/zq.rs

@ -3,7 +3,7 @@ use rand::{distributions::Distribution, Rng};
use std::fmt;
use std::ops;
// Z_q, integers modulus q, not necessarily prime
/// Z_q, integers modulus q, not necessarily prime
#[derive(Clone, Copy, PartialEq)]
pub struct Zq<const Q: u64>(pub u64);
@ -130,6 +130,28 @@ impl Zq {
pub fn mod_switch<const Q2: u64>(&self) -> Zq<Q2> {
Zq::<Q2>::from_u64(((self.0 as f64 * Q2 as f64) / Q as f64).round() as u64)
}
// TODO more efficient method for when decomposing with base 2 (beta=2)
pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
let mut rem: u64 = self.0;
// next if is for cases in which beta does not divide Q (concretely
// beta^l!=Q). round to the nearest multiple of q/beta^l
if rem >= beta.pow(l) as u64 {
// rem = Q - 1 - (Q / beta as u64); // floor
return vec![Zq(beta as u64 - 1); l as usize];
}
let mut x: Vec<Self> = vec![];
for i in 1..l + 1 {
let den = Q / beta.pow(i) as u64;
let x_i = rem / den; // division between u64 already does floor
x.push(Self::from_u64(x_i));
if x_i != 0 {
rem = rem % den;
}
}
x
}
}
impl<const Q: u64> Zq<Q> {
@ -243,6 +265,7 @@ impl fmt::Debug for Zq {
#[cfg(test)]
mod tests {
use super::*;
use rand::distributions::Uniform;
#[test]
fn exp() {
@ -262,4 +285,62 @@ mod tests {
let b = Zq::<Q>::from_f64(-1.0);
assert_eq!(-a, a * b);
}
fn recompose<const Q: u64>(beta: u32, l: u32, d: Vec<Zq<Q>>) -> Zq<Q> {
let mut x = 0u64;
for i in 0..l {
x += d[i as usize].0 * Q / beta.pow(i + 1) as u64;
}
Zq::from_u64(x)
}
#[test]
fn test_decompose() {
const Q1: u64 = 16;
let beta: u32 = 2;
let l: u32 = 4;
let x = Zq::<Q1>::from_u64(9);
let d = x.decompose(beta, l);
assert_eq!(recompose::<Q1>(beta, l, d), x);
const Q: u64 = 5u64.pow(3);
let beta: u32 = 5;
let l: u32 = 3;
let dist = Uniform::new(0_u64, Q);
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let x = Zq::<Q>::from_u64(dist.sample(&mut rng));
let d = x.decompose(beta, l);
assert_eq!(recompose::<Q>(beta, l, d), x);
}
}
#[test]
fn test_decompose_approx() {
const Q: u64 = 2u64.pow(4) + 1;
let beta: u32 = 2;
let l: u32 = 4;
let x = Zq::<Q>::from_u64(16); // in q, but bigger than beta^l
let d = x.decompose(beta, l);
assert_eq!(recompose::<Q>(beta, l, d), Zq(15));
const Q2: u64 = 5u64.pow(3) + 1;
let beta: u32 = 5;
let l: u32 = 3;
let x = Zq::<Q2>::from_u64(125); // in q, but bigger than beta^l
let d = x.decompose(beta, l);
assert_eq!(recompose::<Q2>(beta, l, d), Zq(124));
const Q3: u64 = 2u64.pow(16) + 1;
let beta: u32 = 2;
let l: u32 = 16;
let x = Zq::<Q3>::from_u64(Q3 - 1); // in q, but bigger than beta^l
let d = x.decompose(beta, l);
assert_eq!(recompose::<Q3>(beta, l, d), Zq(beta.pow(l) as u64 - 1));
}
}

+ 169
- 32
gfhe/src/glwe.rs

@ -1,12 +1,17 @@
use anyhow::Result;
use itertools::zip_eq;
use rand::Rng;
use rand_distr::{Normal, Uniform};
use std::ops::{Add, Mul};
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, Sub};
use arith::{Ring, Rq, TR};
use arith::{Ring, Rq, Zq, TR};
use crate::glev::GLev;
const ERR_SIGMA: f64 = 3.2;
#[derive(Clone, Debug)]
pub struct GLWE<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>, Rq<Q, N>);
#[derive(Clone, Debug)]
@ -14,7 +19,15 @@ pub struct SecretKey(TR,
#[derive(Clone, Debug)]
pub struct PublicKey<const Q: u64, const N: usize, const K: usize>(Rq<Q, N>, TR<Rq<Q, N>, K>);
// K GLevs, each KSK_i=l GLWEs
#[derive(Clone, Debug)]
pub struct KSK<const Q: u64, const N: usize, const K: usize>(Vec<GLev<Q, N, K>>);
impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
pub fn zero() -> Self {
Self(TR::zero(), Rq::zero())
}
pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<Q, N, K>, PublicKey<Q, N, K>)> {
let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
@ -27,32 +40,69 @@ impl GLWE {
Ok((SecretKey(s), pk))
}
// TODO delta not as input
pub fn encrypt_s<const T: u64>(
pub fn new_ksk(
mut rng: impl Rng,
beta: u32,
l: u32,
sk: &SecretKey<Q, N, K>,
new_sk: &SecretKey<Q, N, K>,
) -> Result<KSK<Q, N, K>> {
let r: Vec<GLev<Q, N, K>> = (0..K)
.into_iter()
.map(|i|
// treat sk_i as the msg being encrypted
GLev::<Q, N, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
.collect::<Result<Vec<_>>>()?;
Ok(KSK(r))
}
pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<Q, N, K>) -> Self {
let (a, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1);
let lhs: GLWE<Q, N, K> = GLWE(TR::zero(), b);
// K iterations, ksk.0 contains K times GLev
let rhs: GLWE<Q, N, K> = zip_eq(a.0, ksk.0.clone())
.map(|(a_i, ksk_i)| Self::dot_prod(a_i.decompose(beta, l), ksk_i))
.sum();
lhs - rhs
}
// note: a_decomp is of length N
fn dot_prod(a_decomp: Vec<Rq<Q, N>>, ksk_i: GLev<Q, N, K>) -> GLWE<Q, N, K> {
// l times GLWES
let glwes: Vec<GLWE<Q, N, K>> = ksk_i.0;
// l iterations
let r: GLWE<Q, N, K> = zip_eq(a_decomp, glwes)
.map(|(a_d_i, glwe_i)| glwe_i * a_d_i)
.sum();
r
}
// encrypts with the given SecretKey (instead of PublicKey)
pub fn encrypt_s(
mut rng: impl Rng,
sk: &SecretKey<Q, N, K>,
m: &Rq<T, N>,
m: &Rq<Q, N>,
// TODO delta not as input
delta: u64,
) -> Result<Self> {
let m: Rq<Q, N> = m.remodule::<Q>();
let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key);
let e = Rq::<Q, N>::rand(&mut rng, Xi_err);
let b: Rq<Q, N> = (&a * &sk.0) + m * delta + e;
let b: Rq<Q, N> = (&a * &sk.0) + *m * delta + e;
Ok(Self(a, b))
}
pub fn encrypt<const T: u64>(
pub fn encrypt(
mut rng: impl Rng,
pk: &PublicKey<Q, N, K>,
m: &Rq<T, N>,
m: &Rq<Q, N>,
delta: u64,
) -> Result<Self> {
let m: Rq<Q, N> = m.remodule::<Q>();
let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
@ -61,23 +111,16 @@ impl GLWE {
let e0 = Rq::<Q, N>::rand(&mut rng, Xi_err);
let e1 = TR::<Rq<Q, N>, K>::rand(&mut rng, Xi_err);
let b: Rq<Q, N> = pk.0 * u + m * delta + e0;
let b: Rq<Q, N> = pk.0 * u + *m * delta + e0;
let d: TR<Rq<Q, N>, K> = &pk.1 * &u + e1;
Ok(Self(d, b))
}
pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<T, N> {
pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<Q, N> {
let (d, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1);
let r: Rq<Q, N> = b - &d * &sk.0;
let r = r.mul_div_round(T, Q);
// let r_scaled: Vec<f64> = r
// .coeffs()
// .iter()
// // .map(|e| (e.0 as f64 / delta as f64).round())
// .map(|e| e.mul_div_round(T, Q))
// .collect();
// let r = Rq::<Q, N>::from_vec_f64(r_scaled);
r.remodule::<T>()
r
}
pub fn mod_switch<const P: u64>(&self) -> GLWE<P, N, K> {
@ -104,6 +147,36 @@ impl Add> for GLWE
Self(a, b)
}
}
impl<const Q: u64, const N: usize, const K: usize> AddAssign for GLWE<Q, N, K> {
fn add_assign(&mut self, rhs: Self) {
for i in 0..K {
self.0 .0[i] = self.0 .0[i] + rhs.0 .0[i];
}
self.1 = self.1 + rhs.1;
}
}
impl<const Q: u64, const N: usize, const K: usize> Sum<GLWE<Q, N, K>> for GLWE<Q, N, K> {
fn sum<I>(iter: I) -> Self
where
I: Iterator<Item = Self>,
{
let mut acc = GLWE::<Q, N, K>::zero();
for e in iter {
acc += e;
}
acc
}
}
impl<const Q: u64, const N: usize, const K: usize> Sub<GLWE<Q, N, K>> for GLWE<Q, N, K> {
type Output = Self;
fn sub(self, other: Self) -> Self {
let a: TR<Rq<Q, N>, K> = self.0 - other.0;
let b: Rq<Q, N> = self.1 - other.1;
Self(a, b)
}
}
impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Q, N, K> {
type Output = Self;
fn mul(self, plaintext: Rq<Q, N>) -> Self {
@ -118,6 +191,15 @@ impl Mul> for GLWE
}
}
impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
type Output = Self;
fn mul(self, e: Zq<Q>) -> Self {
let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
let b: Rq<Q, N> = self.1 * e;
Self(a, b)
}
}
#[cfg(test)]
mod tests {
use anyhow::Result;
@ -141,11 +223,18 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m: Rq<Q, N> = m.remodule::<Q>();
let c = S::encrypt(&mut rng, &pk, &m, delta)?;
let m_recovered = c.decrypt(&sk, delta);
let m_recovered = c.decrypt::<T>(&sk, delta);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
// same but using encrypt_s (with sk instead of pk))
let c = S::encrypt_s(&mut rng, &sk, &m, delta)?;
let m_recovered = c.decrypt::<T>(&sk, delta);
assert_eq!(m, m_recovered);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
}
Ok(())
@ -168,15 +257,17 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m1: Rq<Q, N> = m1.remodule::<Q>();
let m2: Rq<Q, N> = m2.remodule::<Q>();
let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?;
let c2 = S::encrypt(&mut rng, &pk, &m2, delta)?;
let c3 = c1 + c2;
let m3_recovered = c3.decrypt(&sk, delta);
let m3_recovered = c3.decrypt::<T>(&sk, delta);
assert_eq!(m1 + m2, m3_recovered);
assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
}
Ok(())
@ -199,15 +290,17 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2_scaled: Rq<Q, N> = m2.remodule::<Q>() * delta;
let m1: Rq<Q, N> = m1.remodule::<Q>();
let m2: Rq<Q, N> = m2.remodule::<Q>();
let m2_scaled: Rq<Q, N> = m2 * delta;
let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?;
let c3 = c1 + m2_scaled;
let m3_recovered = c3.decrypt(&sk, delta);
let m3_recovered = c3.decrypt::<T>(&sk, delta);
assert_eq!(m1 + m2, m3_recovered);
assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
}
Ok(())
@ -230,12 +323,14 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m1: Rq<Q, N> = m1.remodule::<Q>();
let m2: Rq<Q, N> = m2.remodule::<Q>();
let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?;
let c3 = c1 * m2;
let m3_recovered: Rq<T, N> = c3.decrypt(&sk, delta);
let m3_recovered: Rq<Q, N> = c3.decrypt::<T>(&sk, delta);
let m3_recovered: Rq<T, N> = m3_recovered.remodule::<T>();
assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
}
@ -265,19 +360,61 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m: Rq<Q, N> = m.remodule::<Q>();
let c = S::encrypt(&mut rng, &pk, &m, delta)?;
// let c = S::encrypt_s(&mut rng, &sk, &m, delta)?;
let c2 = c.mod_switch::<P>();
let sk2: SecretKey<P, N, K> =
SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect()));
let delta2: u64 = ((P as f64 * delta as f64) / Q as f64).round() as u64;
let m_recovered = c2.decrypt(&sk2, delta2);
let m_recovered = c2.decrypt::<T>(&sk2, delta2);
assert_eq!(m, m_recovered);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
}
Ok(())
}
#[test]
fn test_key_switch() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 128;
const T: u64 = 2; // plaintext modulus
const K: usize = 16;
type S = GLWE<Q, N, K>;
let beta: u32 = 2;
let l: u32 = 16;
let delta: u64 = Q / T; // floored
let mut rng = rand::thread_rng();
let (sk, pk) = S::new_key(&mut rng)?;
let (sk2, _) = S::new_key(&mut rng)?;
// ksk to switch from sk to sk2
let ksk = S::new_ksk(&mut rng, beta, l, &sk, &sk2)?;
let msg_dist = Uniform::new(0_u64, T);
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m: Rq<Q, N> = m.remodule::<Q>();
let c = S::encrypt_s(&mut rng, &sk, &m, delta)?;
let c2 = c.key_switch(beta, l, &ksk);
// decrypt with the 2nd secret key
let m_recovered = c2.decrypt::<T>(&sk2, delta);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
// do the same but now encrypting with pk
// let c = S::encrypt(&mut rng, &pk, &m, delta)?;
// let c2 = c.key_switch(beta, l, &ksk);
// let m_recovered = c2.decrypt::<T>(&sk2, delta);
// assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
Ok(())
}
}

+ 1
- 0
gfhe/src/lib.rs

@ -5,4 +5,5 @@
#![allow(clippy::upper_case_acronyms)]
#![allow(dead_code)] // TMP
pub mod glev;
pub mod glwe;

Loading…
Cancel
Save