helpers_codec/common/

pke.rs

// MIT License
//
// Copyright (c) 2026 Raja Lehtihet & Wael El Oraiby
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! Example RLWE-style public-key encryption API.
//!
//! This module is intentionally educational and not hardened for production use.

use rand_core::CryptoRng;

use super::codec::{
    CodecError, decode_message_scaled_bits_example, encode_message_scaled_bits_example,
};
use super::sampling::{SamplingError, sample_cbd_poly_example, sample_uniform_poly_example};
use nc_polynomial::{PolynomialError, RingContext, RingElem};

/// Public key for the example RLWE-style encryption API.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExamplePublicKey {
    /// Public random polynomial `a`.
    pub a: RingElem,
    /// Public polynomial `b = a*s + e`.
    pub b: RingElem,
}

/// Secret key for the example RLWE-style encryption API.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExampleSecretKey {
    /// Secret small-noise polynomial.
    pub s: RingElem,
}

/// Ciphertext for the example RLWE-style encryption API.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExampleCiphertext {
    /// First component `u = a*r + e1`.
    pub u: RingElem,
    /// Second component `v = b*r + e2 + m`.
    pub v: RingElem,
}

/// Errors returned by the example RLWE-style API.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExamplePkeError {
    /// Sampling helper failed.
    Sampling(SamplingError),
    /// Codec helper failed.
    Codec(CodecError),
    /// Ring arithmetic failed.
    Polynomial(PolynomialError),
    /// Input objects do not belong to the provided context.
    ParameterMismatch,
}

impl core::fmt::Display for ExamplePkeError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::Sampling(err) => write!(f, "sampling failed: {err}"),
            Self::Codec(err) => write!(f, "codec failed: {err}"),
            Self::Polynomial(err) => write!(f, "polynomial arithmetic failed: {err}"),
            Self::ParameterMismatch => write!(f, "parameters/context mismatch"),
        }
    }
}

impl std::error::Error for ExamplePkeError {}

impl From<SamplingError> for ExamplePkeError {
    fn from(value: SamplingError) -> Self {
        Self::Sampling(value)
    }
}

impl From<CodecError> for ExamplePkeError {
    fn from(value: CodecError) -> Self {
        Self::Codec(value)
    }
}

impl From<PolynomialError> for ExamplePkeError {
    fn from(value: PolynomialError) -> Self {
        Self::Polynomial(value)
    }
}

/// Generates an example keypair using centered-binomial secret/noise samples.
pub fn keygen_example<R: CryptoRng>(
    ctx: &RingContext,
    noise_eta: u8,
    rng: &mut R,
) -> Result<(ExamplePublicKey, ExampleSecretKey), ExamplePkeError> {
    // Public random polynomial.
    let a = sample_uniform_poly_example(ctx, rng)?;
    // Secret small-noise polynomial.
    let s = sample_cbd_poly_example(ctx, noise_eta, rng)?;
    // Public error polynomial.
    let e = sample_cbd_poly_example(ctx, noise_eta, rng)?;

    // RLWE public relation: b = a*s + e.
    let b = a.mul(&s)?.add(&e)?;

    Ok((ExamplePublicKey { a, b }, ExampleSecretKey { s }))
}

/// Encrypts a byte message using the example key.
///
/// The message is encoded using scaled bits from [`encode_message_scaled_bits_example`].
pub fn encrypt_example<R: CryptoRng>(
    ctx: &RingContext,
    public_key: &ExamplePublicKey,
    message: &[u8],
    noise_eta: u8,
    rng: &mut R,
) -> Result<ExampleCiphertext, ExamplePkeError> {
    // Ensure all key elements come from the same validated parameter set.
    ensure_element_matches_context(ctx, &public_key.a)?;
    ensure_element_matches_context(ctx, &public_key.b)?;

    // Encode message bits into ring element form.
    let m = encode_message_scaled_bits_example(ctx, message)?;
    // Fresh encryption randomness + masking errors.
    let r = sample_cbd_poly_example(ctx, noise_eta, rng)?;
    let e1 = sample_cbd_poly_example(ctx, noise_eta, rng)?;
    let e2 = sample_cbd_poly_example(ctx, noise_eta, rng)?;

    // RLWE encryption equations:
    //   u = a*r + e1
    //   v = b*r + e2 + m
    let u = public_key.a.mul(&r)?.add(&e1)?;
    let v = public_key.b.mul(&r)?.add(&e2)?.add(&m)?;

    Ok(ExampleCiphertext { u, v })
}

/// Decrypts an example ciphertext and decodes the requested number of output bytes.
pub fn decrypt_example(
    ctx: &RingContext,
    secret_key: &ExampleSecretKey,
    ciphertext: &ExampleCiphertext,
    output_len_bytes: usize,
) -> Result<Vec<u8>, ExamplePkeError> {
    // Context checks prevent cross-parameter misuse.
    ensure_element_matches_context(ctx, &secret_key.s)?;
    ensure_element_matches_context(ctx, &ciphertext.u)?;
    ensure_element_matches_context(ctx, &ciphertext.v)?;

    // Recover encoded message:
    //   v - s*u = (a*s + e)r + e2 + m - s(a*r + e1)
    //           = m + (e*r + e2 - s*e1)
    // Small noise lets threshold decoder recover bits.
    let recovered = ciphertext.v.sub(&secret_key.s.mul(&ciphertext.u)?)?;
    Ok(decode_message_scaled_bits_example(
        &recovered,
        output_len_bytes,
    )?)
}

fn ensure_element_matches_context(
    ctx: &RingContext,
    element: &RingElem,
) -> Result<(), ExamplePkeError> {
    // Structural parameter equality check.
    if element.params() != ctx.params() {
        return Err(ExamplePkeError::ParameterMismatch);
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{
        ExampleCiphertext, ExamplePkeError, ExamplePublicKey, decrypt_example, encrypt_example,
        keygen_example,
    };
    use core::convert::Infallible;
    use nc_polynomial::RingContext;
    use rand_core::{Rng, TryCryptoRng, TryRng};

    #[derive(Debug, Clone)]
    struct TestRng {
        state: u64,
    }

    impl TestRng {
        fn new(seed: u64) -> Self {
            Self { state: seed }
        }

        fn step(&mut self) -> u64 {
            let mut x = self.state;
            x ^= x >> 12;
            x ^= x << 25;
            x ^= x >> 27;
            self.state = x;
            x.wrapping_mul(0x2545_F491_4F6C_DD1D)
        }
    }

    impl TryRng for TestRng {
        type Error = Infallible;

        fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
            Ok(self.step() as u32)
        }

        fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
            Ok(self.step())
        }

        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
            let mut offset = 0;
            while offset < dest.len() {
                let word = self.step().to_le_bytes();
                let chunk = (dest.len() - offset).min(8);
                dest[offset..offset + chunk].copy_from_slice(&word[..chunk]);
                offset += chunk;
            }
            Ok(())
        }
    }

    impl TryCryptoRng for TestRng {}

    fn ctx(max_degree: usize) -> RingContext {
        let mut modulus_poly = vec![0_u64; max_degree + 1];
        modulus_poly[0] = 1;
        modulus_poly[max_degree] = 1;

        RingContext::from_parts(max_degree, 998_244_353, &modulus_poly, 3)
            .expect("context should build")
    }

    #[test]
    fn example_pke_known_answer_vector() {
        // Deterministic seeds make vectors stable and reproducible.
        let ctx = ctx(16);

        let mut key_rng = TestRng::new(0x1234_5678_90AB_CDEF);
        let (pk, sk) = keygen_example(&ctx, 2, &mut key_rng).expect("keygen should work");

        let message = [0xA5_u8, 0x5A];
        let mut enc_rng = TestRng::new(0x0FED_CBA9_8765_4321);
        let ct =
            encrypt_example(&ctx, &pk, &message, 2, &mut enc_rng).expect("encrypt should work");

        assert_eq!(
            sk.s.trimmed_coefficients(),
            vec![
                998244352, 2, 0, 998244352, 0, 2, 998244352, 1, 1, 0, 998244352, 998244352, 2,
                998244352, 1,
            ]
        );
        assert_eq!(
            pk.a.trimmed_coefficients(),
            vec![
                981506489, 745670232, 465336855, 4228489, 817594600, 449616339, 214586882,
                863302342, 412683612, 739770493, 484384582, 958972188, 465458099, 901548526,
                136656154, 185069671
            ]
        );
        assert_eq!(
            pk.b.trimmed_coefficients(),
            vec![
                523016218, 151678120, 905400193, 336053334, 886938703, 345066676, 653029107,
                919882161, 719988960, 738092331, 49398785, 531018501, 299386522, 26450153,
                202169127, 149348579
            ]
        );
        assert_eq!(
            ct.u.trimmed_coefficients(),
            vec![
                185825251, 992594928, 222735870, 419735116, 492318984, 871565634, 299285265,
                894961651, 551797064, 802414206, 845696604, 713894635, 456686131, 629006526,
                242614853, 521523065
            ]
        );
        assert_eq!(
            ct.v.trimmed_coefficients(),
            vec![
                342654269, 292514575, 651371554, 799848083, 605114155, 22609066, 790897921,
                525265555, 989917947, 531922498, 781997577, 952513402, 961072975, 455278710,
                776307615, 91714186
            ]
        );

        let decrypted =
            decrypt_example(&ctx, &sk, &ct, message.len()).expect("decrypt should work");
        // Reference vector is only valid if decrypt round-trip is correct.
        assert_eq!(decrypted, message);
    }

    #[test]
    fn example_pke_encrypt_decrypt_round_trip_many_messages() {
        let ctx = ctx(32);

        // Vary seeds and message lengths to cover different bit layouts.
        for seed in 0_u64..64 {
            let mut key_rng = TestRng::new(0xAA55_AA55_0000_0000 ^ seed);
            let (pk, sk) = keygen_example(&ctx, 2, &mut key_rng).expect("keygen should work");

            let msg_len = (seed as usize) % 5;
            let mut msg_rng = TestRng::new(0xC0DE_CAFE_0000_0000 ^ seed);
            let mut message = vec![0_u8; msg_len];
            msg_rng.fill_bytes(&mut message);

            let mut enc_rng = TestRng::new(0xDEAD_BEEF_0000_0000 ^ seed);
            let ct =
                encrypt_example(&ctx, &pk, &message, 2, &mut enc_rng).expect("encrypt should work");
            let recovered =
                decrypt_example(&ctx, &sk, &ct, message.len()).expect("decrypt should work");

            assert_eq!(recovered, message, "seed {seed}");
        }
    }

    #[test]
    fn example_pke_rejects_parameter_mismatch_inputs() {
        let ctx_a = ctx(16);
        let ctx_b = ctx(32);

        let mut rng = TestRng::new(101);
        let (pk_a, sk_a) = keygen_example(&ctx_a, 2, &mut rng).expect("keygen should work");

        let bad_pk = ExamplePublicKey {
            // Mismatched context in `a`.
            a: ctx_b.one_element(),
            b: pk_a.b.clone(),
        };
        let mut enc_rng = TestRng::new(102);
        let err = encrypt_example(&ctx_a, &bad_pk, b"a", 2, &mut enc_rng)
            .expect_err("expected mismatch error");
        assert_eq!(err, ExamplePkeError::ParameterMismatch);

        let mut good_enc_rng = TestRng::new(103);
        let good_ct =
            encrypt_example(&ctx_a, &pk_a, b"a", 2, &mut good_enc_rng).expect("encrypt works");
        let bad_ct = ExampleCiphertext {
            // Mismatched context in `u`.
            u: ctx_b.one_element(),
            v: good_ct.v,
        };
        let err = decrypt_example(&ctx_a, &sk_a, &bad_ct, 1).expect_err("expected mismatch");
        assert_eq!(err, ExamplePkeError::ParameterMismatch);
    }

    #[test]
    fn example_pke_property_like_round_trips_many_seeds() {
        let ctx = ctx(32);

        // Property-like deterministic sweep:
        // different eta values, message lengths, and seed mixes.
        for seed in 0_u64..128 {
            let mut key_rng = TestRng::new(0x1111_2222_3333_4444 ^ seed.rotate_left(5));
            let (pk, sk) = keygen_example(&ctx, 1 + (seed % 2) as u8, &mut key_rng)
                .expect("keygen should work");

            let msg_len = ((seed * 7) as usize) % 5;
            let mut msg_rng = TestRng::new(0x5555_6666_7777_8888 ^ seed.rotate_left(11));
            let mut message = vec![0_u8; msg_len];
            msg_rng.fill_bytes(&mut message);

            let mut enc_rng = TestRng::new(0x9999_AAAA_BBBB_CCCC ^ seed.rotate_left(17));
            let ciphertext =
                encrypt_example(&ctx, &pk, &message, 1 + (seed % 2) as u8, &mut enc_rng)
                    .expect("encrypt should work");
            let recovered = decrypt_example(&ctx, &sk, &ciphertext, message.len())
                .expect("decrypt should work");
            assert_eq!(recovered, message, "seed {seed}");
        }
    }
}