Implementation of lattice based NIKE cryptographic protocol (SWOOSH) not working

Question

I am making a rust library for Python using Pyo3 and maturin that implements different cryptographic protocols. I thought it would be nice to add one that could resist quantum attacks. Thats why I decided to implement the SWOOSH cryptographic protocol, which is based on lattice cryptography, specifically the Module Learning With Errors (M-LWE) problem.

In order to do this I tried to understand the paper which explains it and also the rust and jasmine implementation that they developed and I ended up with the following downsized implementation:

use pyo3::prelude::*;
use rand::{Rng, SeedableRng, thread_rng};
use rand_chacha::ChaCha20Rng;

const N: usize = 16; // Simplified from 8192
const Q: u64 = 12289; // A smaller prime modulus
const D: usize = 8; // Simplified from 256

type Poly = [i64; D];
type PolyVec = [Poly; N];

fn mod_q(x: i64) -> i64 {
    let r = x % Q as i64;
    if r < 0 {
        r + Q as i64
    } else {
        r
    }
}

fn poly_add(a: &Poly, b: &Poly) -> Poly {
    let mut res = [0i64; D];
    for i in 0..D {
        res[i] = mod_q(a[i] + b[i]);
    }
    res
}

fn poly_mul(a: &Poly, b: &Poly) -> Poly {
    let mut res = [0i64; D];
    for i in 0..D {
        for j in 0..D {
            let k = (i + j) % D;
            res[k] = mod_q(res[k] + a[i] * b[j]);
        }
    }
    res
}

fn sample_poly(rng: &mut ChaCha20Rng) -> Poly {
    let mut poly = [0i64; D];
    for i in 0..D {
        let r = rng.gen_range(0..3);
        poly[i] = if r == 0 { -1 } else if r == 1 { 0 } else { 1 };
    }
    poly
}

fn generate_matrix_a(rng: &mut ChaCha20Rng) -> [PolyVec; N] {
    let mut matrix = [[sample_poly(rng); N]; N];
    for i in 0..N {
        for j in 0..N {
            matrix[i][j] = sample_poly(rng);
        }
    }
    matrix
}

fn generate_keys() -> (PolyVec, PolyVec) {
    let seed: [u8; 32] = thread_rng().gen();
    let mut rng = ChaCha20Rng::from_seed(seed);
    let s = [sample_poly(&mut rng); N];
    let e = [sample_poly(&mut rng); N];
    let a = generate_matrix_a(&mut rng);
    let mut pk = e;
    for i in 0..N {
        let mut temp = [0i64; D];
        for j in 0..N {
            let prod = poly_mul(&a[i][j], &s[j]);
            temp = poly_add(&temp, &prod);
        }
        pk[i] = poly_add(&pk[i], &temp);
    }
    (s, pk)
}

fn derive_shared_key(s: &PolyVec, pk: &PolyVec) -> Vec<u8> {
    let mut k = [0i64; D];
    for i in 0..N {
        let prod = poly_mul(&s[i], &pk[i]);
        k = poly_add(&k, &prod);
    }
    reconcile(&k)
}

fn reconcile(k: &Poly) -> Vec<u8> {
    let mut key_bits = Vec::with_capacity(D);
    for i in 0..D {
        key_bits.push(if k[i] > Q as i64 / 2 { 1 } else { 0 });
    }
    key_bits
}

#[pyfunction]
pub fn swoosh_generate_keys() -> PyResult<(Vec<i64>, Vec<i64>)> {
    let (sk, pk) = generate_keys();
    Ok((sk.into_iter().flatten().collect(), pk.into_iter().flatten().collect()))
}

#[pyfunction]
pub fn swoosh_derive_shared_key(sk: Vec<i64>, pk: Vec<i64>) -> PyResult<Vec<u8>> {
    if sk.len() != N * D || pk.len() != N * D {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Invalid key length"));
    }
    
    let sk_polyvec: PolyVec = sk.chunks(D).map(|chunk| chunk.try_into().unwrap()).collect::<Vec<Poly>>().try_into().unwrap();
    let pk_polyvec: PolyVec = pk.chunks(D).map(|chunk| chunk.try_into().unwrap()).collect::<Vec<Poly>>().try_into().unwrap();
    
    Ok(derive_shared_key(&sk_polyvec, &pk_polyvec))
}

The problem is that the keys never match between each other, so it keeps on failing basically and I cannot understand why it is happening despite my debugging, and because it is a complex implementation I don´t know if the values I get before are correct. So how can I fix it and is my approach even remotely correct?

Here is the testing code:

def test_swoosh():
    print("\nTesting SWOOSH key exchange...")
    try:
        print("Generating first key pair...")
        private_key1, public_key1 = shadow_crypt.swoosh_generate_keys()
        print(f"SWOOSH Private Key 1: {private_key1}")
        print(f"SWOOSH Public Key 1: {public_key1}")

        print("Generating second key pair...")
        private_key2, public_key2 = shadow_crypt.swoosh_generate_keys()
        print(f"SWOOSH Private Key 2: {private_key2}")
        print(f"SWOOSH Public Key 2: {public_key2}")

        print("Deriving shared keys...")
        shared_key1 = shadow_crypt.swoosh_derive_shared_key(private_key1, public_key2)
        shared_key2 = shadow_crypt.swoosh_derive_shared_key(private_key2, public_key1)

        print(f"SWOOSH Shared Key 1: {shared_key1}")
        print(f"SWOOSH Shared Key 2: {shared_key2}")


        if shared_key1 == shared_key2:
            print("SWOOSH shared keys match!")
        else:
            print("WARNING: SWOOSH shared keys do not match!")

        print(f"SWOOSH Private key length: {len(private_key1)} integers")
        print(f"SWOOSH Public key length: {len(public_key1)} integers")
        print(f"SWOOSH Shared key length: {len(shared_key1)} integers")

        return shared_key1 == shared_key2

    except Exception as e:
        print(f"An error occurred: {e}")
        import traceback
        traceback.print_exc()
        return False

And here is the output:

Testing SWOOSH key exchange...
Generating first key pair...
SWOOSH Private Key 1: [1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1, 1, 0, -1, 0, -1, 1, -1, 1]
SWOOSH Public Key 1: [2, 3, 12281, 12287, 12288, 10, 12278, 8, 8, 3, 11, 12271, 12286, 5, 12285, 12288, 1, 12, 12283, 12288, 2, 12286, 12282, 3, 12279, 5, 12281, 20, 12278, 12288, 1, 5, 12285, 12287, 11, 10, 12281, 2, 12284, 12286, 12282, 5, 0, 12287, 13, 12275, 12287, 8, 15, 12282, 12287, 12279, 12, 12285, 5, 12281, 12282, 2, 10, 1, 3, 12279, 0, 2, 0, 12283, 12280, 12288, 12, 6, 12284, 4, 12287, 8, 12286, 12287, 2, 10, 12274, 3, 12285, 3, 12286, 8, 12286, 8, 0, 12281, 12285, 17, 12267, 3, 12275, 6, 4, 11, 12287, 12281, 12278, 12284, 4, 12288, 13, 11, 9, 12283, 6, 12286, 1, 12283, 2, 12287, 0, 12285, 12279, 8, 0, 12287, 4, 5, 12278, 7, 10, 3, 12275, 3, 12280, 12]
Generating second key pair...
SWOOSH Private Key 2: [-1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1, -1, 1, 0, -1, 0, 0, 0, 1]
SWOOSH Public Key 2: [12266, 11, 7, 12275, 12, 12285, 12286, 14, 12279, 0, 12288, 9, 12285, 5, 12284, 6, 3, 8, 12281, 0, 12286, 3, 12288, 12287, 1, 2, 4, 12281, 12284, 2, 12284, 9, 12284, 7, 12283, 5, 3, 12280, 1, 4, 6, 12287, 12282, 7, 3, 12286, 1, 12284, 1, 7, 12281, 2, 6, 12284, 0, 12286, 12285, 15, 12285, 12281, 15, 12277, 6, 12281, 12280, 6, 2, 1, 12287, 7, 12281, 3, 12287, 12286, 7, 12279, 2, 12282, 12287, 15, 12279, 14, 12272, 8, 3, 12284, 10, 12286, 12279, 9, 6, 12279, 10, 12276, 4, 4, 12274, 16, 12285, 12283, 12287, 2, 1, 8, 12287, 13, 12271, 1, 2, 12282, 14, 12286, 12284, 7, 12288, 12287, 2, 12278, 18, 12281, 12285, 12288, 3, 12285, 12285, 6, 4, 0]
Deriving shared keys...
SWOOSH Shared Key 1: ([1, 0, 1, 1, 0, 1, 0, 1], [12253, 151, 12179, 12244, 116, 12074, 270, 12158])
SWOOSH Shared Key 2: ([0, 1, 0, 1, 1, 1, 0, 1], [89, 12236, 25, 12257, 12280, 12280, 79, 12199])
WARNING: SWOOSH shared keys do not match!
SWOOSH Private key length: 128 integers
SWOOSH Public key length: 128 integers
SWOOSH Shared key length: 2 integers

Answer 1

Thanks for your interest in our paper — I'm one of the authors.

I've taken a look at your code, and while I haven't gone through every detail, I can offer a few observations that might help.

Parameter Selection

First off, the parameters you've selected for N and D are much too small, which means there is no security in your current implementation (ignoring any side-channels, etc). I'm assuming that this is for demonstration or educational purposes rather than practical use.

Correctness Error

More critically, the parameters you've chosen lead to a significant correctness error. As mentioned in Lemma 4 of our paper, the probability that the keys generated by the protocol do not match is given by:

\frac{4 * \beta^2 * D^2 * N}{q}

For the parameters you've selected:

\frac{4 * 1^2 * 8^2 * 16}{12289} \approx 0.33

This means there's roughly a 33% chance that your keys won't match, even if everything else is implemented correctly.

Sampling from the Correct Distribution

Another issue I noticed is the way you're sampling coefficients. Currently, you're using a uniform distribution where each coefficient takes the value -1, 0, or 1 with equal probability. However, what you should be using is a centred binomial distribution where p(-1) = p(1) = 0.25 and p(0) = 0.5.