arnaucube/fhe-study
1
1
Fork 0
mirror of https://github.com/arnaucube/fhe-study.git synced 2026-01-24 04:33:52 +01:00
Code Issues Packages Projects Releases Wiki Activity

implement GLWE key switching

Browse Source
This commit is contained in:
arnaucube
2025-07-16 21:02:07 +02:00
parent 1e2ea824fd
commit 188bc7fa7f
8 changed files with 384 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

201
gfhe/src/glwe.rs
View File

@@ -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<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>,
#[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<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
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<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
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<const Q: u64, const N: usize, const K: usize> Add<Rq<Q, N>> for GLWE<Q, N,
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<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Q, N,
}
}
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, m_recovered);
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.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(())
}
}