diff --git a/README.md b/README.md index 2974723..ee8758a 100644 --- a/README.md +++ b/README.md @@ -18,3 +18,4 @@ Implementation of [CRYSTALS-Kyber](https://pq-crystals.org/kyber/index.shtml) en * [Week 1](docs/week-1.md) * [Week 2](docs/week-2.md) * [Week 3](docs/week-3.md) +* [Week 4](docs/week-4.md) diff --git a/docs/tests.md b/docs/tests.md index 4adb789..da178cc 100644 --- a/docs/tests.md +++ b/docs/tests.md @@ -53,3 +53,9 @@ TOTAL 331 0 92 0 100% ``` For more detailed report, run `poetry run invoke coverage-report` and then open `htmlcov/index.html`. + +## Performance tests + +The asymmetric encryption part of Kyber (CPAPKE in the specification document) only works with fixed-lengthed input, but we can split larger payload into 32-byte chunks and encrypt them separately. Ciphertexts can be concatenated and the whole process can be reversed during decryption. Using this method, encryption is tested with about 10 kibibytes of random payload. + +Performance tests can be run with `poetry run invoke performance`. diff --git a/docs/usage.md b/docs/usage.md index 052ef2e..a827f18 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -50,6 +50,14 @@ poetry run invoke coverage-report after which the report will appear at `htmlcov/index.html`. +### Performance tests + +Run performance tests with + +``` +poetry run invoke performance +``` + ### Lint Run static style cheking with diff --git a/docs/week-4.md b/docs/week-4.md new file mode 100644 index 0000000..d0eda2b --- /dev/null +++ b/docs/week-4.md @@ -0,0 +1,7 @@ +# Week 4 + +_20. – 26.11.2023_ + +This week I started implementing own data structures. First, I added `PolynomialRing` which, as the name tells, represents polynomial ring structure and handles its core calculations. My original plan was also to implement matrix structure, but it turned out to be much slower than Numpy array, and therefore will not be added. I also added more unit tests and started adding performance tests. + +Total working time: 5 hours diff --git a/kyber/ccakem.py b/kyber/ccakem.py index 2752f01..1e7b84c 100644 --- a/kyber/ccakem.py +++ b/kyber/ccakem.py @@ -21,7 +21,7 @@ def ccakem_generate_keys() -> tuple[bytes, bytes]: pk # public key ) -def ccakem_encrypt(public_key: bytes) -> tuple[bytes, bytes]: +def ccakem_encrypt(public_key: bytes, shared_secret_length: int = 32) -> tuple[bytes, bytes]: """ Takes public key as input and returns (ciphertext, shered_secret) as a tuple. Shared secret is 32 bytes in length. @@ -33,14 +33,14 @@ def ccakem_encrypt(public_key: bytes) -> tuple[bytes, bytes]: Kr = G(m + H(public_key)) K, r = Kr[:32], Kr[32:] c = Encrypt(public_key, m, r).encrypt() - K = kdf(K + H(c), 32) + K = kdf(K + H(c), shared_secret_length) return ( c, # ciphertext K # shared secret ) -def ccakem_decrypt(ciphertext: bytes, private_key: bytes) -> bytes: +def ccakem_decrypt(ciphertext: bytes, private_key: bytes, shared_secret_length: int = 32) -> bytes: """ Decrypts the given ciphertext with the private key. :returns Decrypted 32-byte shared secret. @@ -62,5 +62,5 @@ def ccakem_decrypt(ciphertext: bytes, private_key: bytes) -> bytes: c = Encrypt(pk, m, r).encrypt() if c == ciphertext: - return kdf(K + H(c), 32) - return kdf(z + H(c), 32) + return kdf(K + H(c), shared_secret_length) + return kdf(z + H(c), shared_secret_length) diff --git a/kyber/encryption/decrypt.py b/kyber/encryption/decrypt.py index cdca206..25b19f3 100644 --- a/kyber/encryption/decrypt.py +++ b/kyber/encryption/decrypt.py @@ -1,9 +1,8 @@ import numpy as np -from numpy.polynomial.polynomial import Polynomial from kyber.utils.compression import compress, decompress from kyber.utils.encoding import encode, decode -from kyber.utils.modulo import polmod from kyber.constants import n, k, du, dv +from kyber.entities.polring import PolynomialRing class Decrypt: def __init__(self, private_key, ciphertext) -> None: @@ -30,8 +29,7 @@ def decrypt(self) -> bytes: u = np.array([decompress(pol, du) for pol in u]) v = decompress(v, dv) - m: Polynomial = v - np.matmul(s.T, u) - m = polmod(m) + m: PolynomialRing = v - np.matmul(s.T, u) m: bytes = encode(compress([m], 1), 1) assert len(m) == 32 diff --git a/kyber/encryption/encrypt.py b/kyber/encryption/encrypt.py index fbcb215..90ade12 100644 --- a/kyber/encryption/encrypt.py +++ b/kyber/encryption/encrypt.py @@ -1,15 +1,14 @@ from random import randbytes import numpy as np -from numpy.polynomial.polynomial import Polynomial from kyber.utils.cbd import cbd from kyber.utils.pseudo_random import prf -from kyber.utils.modulo import matmod, polmod from kyber.utils.compression import compress, decompress from kyber.utils.encoding import encode, decode from kyber.constants import k, eta1, eta2, n, du, dv from kyber.utils.byte_conversion import int_to_bytes from kyber.utils.parse import parse from kyber.utils.pseudo_random import xof +from kyber.entities.polring import PolynomialRing class Encrypt: def __init__(self, public_key: bytes, m: bytes = None, r: bytes = None) -> None: @@ -38,33 +37,27 @@ def encrypt(self): t, rho = self._pk[:-32], self._pk[-32:] t = np.array(decode(t, 12)) - A = np.empty((k, k), Polynomial) + A = np.empty((k, k), PolynomialRing) for i in range(k): for j in range(k): A[i][j] = parse(xof(rho, int_to_bytes(i), int_to_bytes(j))) N = 0 - r = np.empty((k, ), Polynomial) + r = np.empty((k, ), PolynomialRing) for i in range(k): r[i] = cbd(prf(rb, int_to_bytes(N)), eta1) - r[i] = polmod(r[i]) N += 1 - e1 = np.empty((k, ), Polynomial) + e1 = np.empty((k, ), PolynomialRing) for i in range(k): e1[i] = cbd(prf(rb, int_to_bytes(N)), eta2) - e1[i] = polmod(e1[i]) N += 1 e2 = cbd(prf(rb, int_to_bytes(N)), eta2) - e2 = polmod(e2) u = np.matmul(A.T, r) + e1 v = np.matmul(t.T, r) + e2 + decompress(decode(m, 1)[0], 1) - u = matmod(u) - v = polmod(v) - u = compress(u, du) v = compress([v], dv) diff --git a/kyber/encryption/keygen.py b/kyber/encryption/keygen.py index 3e3c99b..1951d23 100644 --- a/kyber/encryption/keygen.py +++ b/kyber/encryption/keygen.py @@ -1,13 +1,12 @@ from random import randbytes import numpy as np -from numpy.polynomial.polynomial import Polynomial from kyber.constants import k, eta1 from kyber.utils.pseudo_random import prf, G, xof from kyber.utils.cbd import cbd -from kyber.utils.modulo import matmod, polmod from kyber.utils.byte_conversion import int_to_bytes from kyber.utils.encoding import encode from kyber.utils.parse import parse +from kyber.entities.polring import PolynomialRing def generate_keys() -> tuple: """ @@ -18,26 +17,23 @@ def generate_keys() -> tuple: d = randbytes(32) rho, sigma = G(d)[:32], G(d)[32:] - A = np.empty((k, k), Polynomial) + A = np.empty((k, k), PolynomialRing) for i in range(k): for j in range(k): A[i][j] = parse(xof(rho, int_to_bytes(i), int_to_bytes(j))) N = 0 - s = np.empty((k, ), Polynomial) + s = np.empty((k, ), PolynomialRing) for i in range(k): s[i] = cbd(prf(sigma, int_to_bytes(N)), eta1) - s[i] = polmod(s[i]) N += 1 - e = np.empty((k, ), Polynomial) + e = np.empty((k, ), PolynomialRing) for i in range(k): e[i] = cbd(prf(sigma, int_to_bytes(N)), eta1) - e[i] = polmod(e[i]) N += 1 t = np.matmul(A, s) + e # t is a polynomial matrix with shape (k, ) - t = matmod(t) s: bytes = encode(s, 12) t: bytes = encode(t, 12) diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py new file mode 100644 index 0000000..a1c69f2 --- /dev/null +++ b/kyber/entities/polring.py @@ -0,0 +1,72 @@ +from kyber.constants import q, n + +class PolynomialRing: + def __init__(self, coefs: list[int], check_limits: bool = True) -> None: + """ + :param coefs: Coefficients of the polynomial. E.g. `[1, 2, 3]` represents `1+2x+3x^2`. + :param check_limits: Set to `False` if coefs is already taken to modulo. + """ + self._coefs = [int(c) for c in coefs] + if check_limits: + self._apply_limits() + + @property + def coefs(self) -> list[int]: + """Coefficients of the polynomial. E.g. `[1, 2, 3]` represents `1+2x+3x^2`.""" + return self._coefs + + def _apply_limits(self) -> None: + """Take this polynomial to modulo `x^n+1` and coefs to modulo `q`.""" + # apply degree limit by dividing self by x^n+1 + self._apply_polynomial_modulo_limit() + + # apply coef limit + self._coefs = [c % q for c in self._coefs] + + # remove trailing zero coefficients + while len(self._coefs) > 0 and self._coefs[-1] == 0: + self._coefs.pop() + + def _apply_polynomial_modulo_limit(self) -> None: + """Replaces `self._coefs` with the remainder of division `self._coefs / (x^n+1)`.""" + # this is an optimal version of polynomial long division + coef_count = len(self._coefs) + while coef_count >= n+1: + self._coefs[-n-1] -= self._coefs[-1] + self._coefs[-1] = 0 + while self._coefs[-1] == 0: + self._coefs.pop() + coef_count -= 1 + + def __add__(self, other: "PolynomialRing") -> "PolynomialRing": + result = [] + self_length = len(self._coefs) + other_length = len(other.coefs) + for i in range(max(self_length, other_length)): + self_coef = self.coefs[i] if i < self_length else 0 + other_coef = other.coefs[i] if i < other_length else 0 + result.append(self_coef + other_coef) + return PolynomialRing(result) + + def __sub__(self, other: "PolynomialRing") -> "PolynomialRing": + result = [] + self_length = len(self._coefs) + other_length = len(other.coefs) + for i in range(max(self_length, other_length)): + self_coef = self.coefs[i] if i < self_length else 0 + other_coef = other.coefs[i] if i < other_length else 0 + result.append(self_coef - other_coef) + return PolynomialRing(result) + + def __mul__(self, other: "PolynomialRing") -> "PolynomialRing": + result = [0 for _ in range(len(self.coefs) + len(other.coefs) - 1)] + for a in range(len(self.coefs)): + for b in range(len(other.coefs)): + result[a+b] += self.coefs[a] * other.coefs[b] + return PolynomialRing(result) + + def __eq__(self, other: "PolynomialRing") -> bool: + return self.coefs == other.coefs + + def __repr__(self) -> str: + return "PolRing(" + ", ".join([str(c) for c in self.coefs]) + ")" diff --git a/kyber/utils/cbd.py b/kyber/utils/cbd.py index 60abbee..d87878b 100644 --- a/kyber/utils/cbd.py +++ b/kyber/utils/cbd.py @@ -1,7 +1,7 @@ -from numpy.polynomial.polynomial import Polynomial from kyber.utils.byte_conversion import bytes_to_bits +from kyber.entities.polring import PolynomialRing -def cbd(b: bytes, eta: int) -> Polynomial: +def cbd(b: bytes, eta: int) -> PolynomialRing: """ Deterministically creates and returns a polynomial (degree 255) from the given byte array (length 64*eta). @@ -20,4 +20,4 @@ def cbd(b: bytes, eta: int) -> Polynomial: b += bits[2 * i * eta + eta + j] f.append(a-b) assert len(f) == 256 - return Polynomial(f) + return PolynomialRing(f) diff --git a/kyber/utils/compression.py b/kyber/utils/compression.py index 6dc47c3..5f290d3 100644 --- a/kyber/utils/compression.py +++ b/kyber/utils/compression.py @@ -1,25 +1,25 @@ from math import log2, ceil -from numpy.polynomial.polynomial import Polynomial from kyber.constants import q from kyber.utils.round import normal_round +from kyber.entities.polring import PolynomialRing -def compress(pols: list[Polynomial], d: int) -> list[Polynomial]: +def compress(pols: list[PolynomialRing], d: int) -> list[PolynomialRing]: """ Reduces every coefficient of every polynomial in the given list to range `0...2**d-1` (inclusive). """ result = [] for pol in pols: - f = [compress_int(c, d) for c in pol.coef] - result.append(Polynomial(f)) + f = [compress_int(c, d) for c in pol.coefs] + result.append(PolynomialRing(f)) return result -def decompress(pol: Polynomial, d: int) -> Polynomial: +def decompress(pol: PolynomialRing, d: int) -> PolynomialRing: """ Multiplies each coefficient of the given polynomial by `q/(2**d)`. Each coefficient of the given polynomial must be in range `0...2^d-1` (inclusive). """ - return Polynomial([decompress_int(c, d) for c in pol.coef]) + return PolynomialRing([decompress_int(c, d) for c in pol.coefs]) def compress_int(x: int, d: int) -> int: """ diff --git a/kyber/utils/encoding.py b/kyber/utils/encoding.py index 2aaa533..cbd2d05 100644 --- a/kyber/utils/encoding.py +++ b/kyber/utils/encoding.py @@ -1,8 +1,8 @@ import numpy as np -from numpy.polynomial.polynomial import Polynomial from kyber.utils.byte_conversion import bytes_to_bits, bits_to_bytes +from kyber.entities.polring import PolynomialRing -def encode(pols: list[Polynomial], l: int) -> bytes: +def encode(pols: list[PolynomialRing], l: int) -> bytes: """ Converts the given polynomial (degree 255, each coefficient in range `0...2**l-1` inclusive) into a byte array of lenght `32*l`. @@ -12,9 +12,7 @@ def encode(pols: list[Polynomial], l: int) -> bytes: result = bytearray() for pol in pols: - if len(pol.coef) > 256: - raise ValueError("too high polynomial degree") - f = list(pol.coef) + [0 for _ in range(256-len(pol.coef))] + f = list(pol.coefs) + [0 for _ in range(256-len(pol.coefs))] bits = np.empty((256*l, )) for i in range(256): f_item = f[i] @@ -32,7 +30,7 @@ def encode(pols: list[Polynomial], l: int) -> bytes: assert len(result) == 32*l*len(pols) return bytes(result) -def decode(b: bytes, l: int) -> list[Polynomial]: +def decode(b: bytes, l: int) -> list[PolynomialRing]: """ Converts the given byte array (length `32*l*x` for some integer x) into a list of polynomials (length x, each degree 255) @@ -48,5 +46,5 @@ def decode(b: bytes, l: int) -> list[Polynomial]: for i in range(256): f[i] = sum(bits[i*l+j]*2**j for j in range(l)) # accesses each bit exactly once assert 0 <= f[i] and f[i] <= 2**l-1 - result.append(Polynomial(f)) + result.append(PolynomialRing(f)) return result diff --git a/kyber/utils/modulo.py b/kyber/utils/modulo.py deleted file mode 100644 index 8bd5624..0000000 --- a/kyber/utils/modulo.py +++ /dev/null @@ -1,44 +0,0 @@ -import numpy as np -from numpy.polynomial.polynomial import Polynomial -from kyber.constants import n, q - -def matmod( - matrix: np.ndarray, - pol_divisor: Polynomial = None, - int_divisor: int = None - ) -> np.ndarray: - """ - Applies `polmod` to each element of the given matrix of polynomials. - :param matrix A matrix of polynomials - :param pol_divisor A polynomial divisor. Only for testing purposes. (default: x^n+1) - :param int_divisor An integer divisor. Only for testing purposes. (default: constant q) - :returns A new matrix of polynomials. - """ - - result = matrix.copy() - for index, item in np.ndenumerate(matrix): - result[index] = polmod(item, pol_divisor, int_divisor) - return result - -def polmod(pol: Polynomial, pol_divisor: Polynomial = None, int_divisor: int = None) -> Polynomial: - """ - Finds given polynomial modulo another polynomial and coefficients modulo an integer. - :param pol The polynomial to be modded - :param pol_divisor A polynomial divisor. Only for testing purposes. (default: x^n+1) - :param int_divisor An integer divisor. Only for testing purposes. (default: constant q) - :returns The remainder of polynomial division whose coefficients are taken modulo int_divisor. - """ - - # divide the given polynomial by another polynomial and take the remainder - if pol_divisor is not None: - divisor = pol_divisor - else: - divisor = Polynomial([1] + [0 for _ in range(n-1)] + [1]) # x^n + 1 - pol = pol % divisor - - # divide each coefficient of the polynomial by an integer and take the remainder - divisor = int_divisor if int_divisor is not None else q - for i in range(len(pol.coef)): - pol.coef[i] %= divisor - - return pol diff --git a/kyber/utils/parse.py b/kyber/utils/parse.py index 1f320bb..2e96e59 100644 --- a/kyber/utils/parse.py +++ b/kyber/utils/parse.py @@ -1,14 +1,14 @@ from math import floor from typing import Generator -from numpy.polynomial.polynomial import Polynomial import numpy as np from kyber.constants import n, q +from kyber.entities.polring import PolynomialRing def byte_to_int(b: bytes) -> int: """Returns the unsigned integer that the given big-endian byte array represents.""" return int.from_bytes(b) -def parse(stream: Generator[bytes, None, None]) -> Polynomial: +def parse(stream: Generator[bytes, None, None]) -> PolynomialRing: """ Deterministically creates a polynomial (degree n-1, each coefficient in range `0...4095` inclusive) from the given bytestream. @@ -27,4 +27,4 @@ def parse(stream: Generator[bytes, None, None]) -> Polynomial: a[j] = d2 j += 1 i += 3 - return Polynomial(a) + return PolynomialRing(a) diff --git a/perf_tests/__main__.py b/perf_tests/__main__.py new file mode 100644 index 0000000..a6b752e --- /dev/null +++ b/perf_tests/__main__.py @@ -0,0 +1,3 @@ +from perf_tests.test_encryption import runner as encryption_test_runner + +encryption_test_runner() diff --git a/perf_tests/test_encryption.py b/perf_tests/test_encryption.py new file mode 100644 index 0000000..0f7e48b --- /dev/null +++ b/perf_tests/test_encryption.py @@ -0,0 +1,41 @@ +from random import randbytes +from time import time +from kyber.encryption import Encrypt, Decrypt, generate_keys +from kyber.constants import k, n, du, dv + +def run(payload: bytes) -> tuple[float, float, float]: + t0 = time() + + private_key, public_key = generate_keys() + + t1 = time() + + ciphertext = bytearray() + for i in range(0, len(payload), 32): + ciphertext += Encrypt(public_key, payload[i:i+32]).encrypt() + + t2 = time() + + ciphertext_chunk_size = du*k*n//8 + dv*n//8 + restored_payload = bytearray() + for i in range(0, len(ciphertext), ciphertext_chunk_size): + restored_payload += Decrypt(private_key, ciphertext[i:i+ciphertext_chunk_size]).decrypt() + + t3 = time() + + assert payload == restored_payload + + return (t1-t0, t2-t1, t3-t2) + +def runner(): + payload = randbytes(10_048) # about 10 kiB, multipla of 32 + print("Starting encryption performance test (about 1 min)") + durations = run(payload) + print("Results:") + print(f"Keypair generation: {durations[0]:.2f} sec") + print(f"Encryption: {durations[1]:.2f} sec") + print(f"Decryption: {durations[2]:.2f} sec") + print(f"Total: {sum(durations):.2f} sec") + +if __name__ == "__main__": + runner() diff --git a/tasks.py b/tasks.py index acc40fc..29e232e 100644 --- a/tasks.py +++ b/tasks.py @@ -1,9 +1,14 @@ +import os from invoke import task @task def test(ctx): ctx.run("pytest tests/", pty=True) +@task +def performance(ctx): + ctx.run("python perf_tests/", pty=True) + @task def coverage(ctx): ctx.run("coverage run -m pytest --ignore tests/test_integration.py tests/", pty=True) diff --git a/tests/test_cbd.py b/tests/test_cbd.py index c946f2f..7006af3 100644 --- a/tests/test_cbd.py +++ b/tests/test_cbd.py @@ -26,7 +26,14 @@ def test_cbd_throws_with_incorrect_argument_length(self): cbd(argument, eta) def test_cbd_result_polynomial_degree(self): + # test that CBD always outputs a polynomial with degree <= 255 + # also make sure that the degree is not always less than 255 eta = 5 - argument = randbytes(320) - result = cbd(argument, eta) - self.assertEqual(len(result.coef), 256) + some_had_degree_255 = False + for _ in range(100): + argument = randbytes(320) + result = cbd(argument, eta) + self.assertLessEqual(len(result.coefs), 256) + if len(result.coefs) == 256: + some_had_degree_255 = True + self.assertTrue(some_had_degree_255) diff --git a/tests/test_ccakem.py b/tests/test_ccakem.py index e8c0a5d..eaf2a3c 100644 --- a/tests/test_ccakem.py +++ b/tests/test_ccakem.py @@ -39,3 +39,11 @@ def test_ccakem_integration(self): ciphertext, shared_secret1 = ccakem_encrypt(public_key) shared_secret2 = ccakem_decrypt(ciphertext, private_key) self.assertEqual(shared_secret1, shared_secret2) + + def test_ccakem_symmetry_with_long_shared_secret(self): + # generate 1000-byte shared secret + private_key, public_key = ccakem_generate_keys() + ciphertext, shared_secret1 = ccakem_encrypt(public_key, 1000) + shared_secret2 = ccakem_decrypt(ciphertext, private_key, 1000) + self.assertEqual(len(shared_secret1), 1000) + self.assertEqual(shared_secret1, shared_secret2) diff --git a/tests/test_compression.py b/tests/test_compression.py index cac6d40..6c8d33f 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -1,43 +1,43 @@ import unittest -from numpy.polynomial.polynomial import Polynomial from random import seed, randint from kyber.utils.compression import compress, decompress +from kyber.entities.polring import PolynomialRing class TestCompression(unittest.TestCase): def test_compression_symmetry(self): seed(42) - polynomial = Polynomial([randint(0, 2047) for _ in range(256)]) + polynomial = PolynomialRing([randint(0, 2047) for _ in range(256)]) decompressed = decompress(polynomial, 11) compressed = compress([decompressed], 11)[0] - self.assertListEqual(list(polynomial.coef), list(compressed.coef)) + self.assertListEqual(list(polynomial.coefs), list(compressed.coefs)) def test_compression(self): - polynomial = Polynomial([416, 2913, 0, 1248]) - expected_result = Polynomial([1, 7, 0, 3]) + polynomial = PolynomialRing([416, 2913, 0, 1248]) + expected_result = PolynomialRing([1, 7, 0, 3]) self.assertEqual(compress([polynomial], 3)[0], expected_result) def test_compression_result_coefficients_in_range(self): # each coefficient in the result should be in range 0...2**d-1 (inclusive) d = 11 - polynomial = Polynomial([4000, -700, 0, 32]) + polynomial = PolynomialRing([4000, -700, 0, 32]) result = compress([polynomial], d)[0] - for c in result.coef: + for c in result.coefs: self.assertTrue(0 <= c and c <= 2**d-1) def test_decompression(self): - polynomial = Polynomial([1, 7, 0, 3]) - expected_result = Polynomial([416, 2913, 0, 1248]) + polynomial = PolynomialRing([1, 7, 0, 3]) + expected_result = PolynomialRing([416, 2913, 0, 1248]) self.assertEqual(decompress(polynomial, 3), expected_result) def test_decompression_raises_with_negative_coefficient(self): # coefficient should not be negative - polynomial = Polynomial([2, -1, 3]) + polynomial = PolynomialRing([2, -1, 3]) with self.assertRaises(ValueError): decompress(polynomial, 3) def test_decompression_raises_with_too_large_coefficient(self): # coefficient should not be greather than 2**d-1 = 7 - polynomial = Polynomial([2, 8, 3]) + polynomial = PolynomialRing([2, 8, 3]) with self.assertRaises(ValueError): decompress(polynomial, 3) diff --git a/tests/test_encoding.py b/tests/test_encoding.py index 0199d5a..051870c 100644 --- a/tests/test_encoding.py +++ b/tests/test_encoding.py @@ -1,15 +1,15 @@ import unittest -from numpy.polynomial.polynomial import Polynomial from random import seed, randbytes, randint from kyber.utils.encoding import encode, decode +from kyber.entities.polring import PolynomialRing class TestEncoding(unittest.TestCase): def setUp(self): seed(42) self.l = 5 self.data = randbytes(32 * self.l) - self.polynomial = Polynomial([randint(0, 1) for _ in range(256)]) - self.polynomial2 = Polynomial([randint(0, 1) for _ in range(256)]) + self.polynomial = PolynomialRing([randint(0, 1) for _ in range(256)]) + self.polynomial2 = PolynomialRing([randint(0, 1) for _ in range(256)]) def test_encoding_symmetry(self): polynomials = decode(self.data, self.l) @@ -18,12 +18,12 @@ def test_encoding_symmetry(self): def test_decode_coefficients(self): polynomial = decode(self.data, self.l)[0] - for c in polynomial.coef: + for c in polynomial.coefs: self.assertTrue(0 <= int(c) or int(c) <= 2**self.l-1) def test_decode_degree(self): polynomial = decode(self.data, self.l)[0] - self.assertEqual(len(polynomial.coef), 256) + self.assertEqual(len(polynomial.coefs), 256) def test_decode_raises_with_invalid_argument_length(self): data = self.data + bytes([42]) @@ -42,13 +42,6 @@ def test_encode_raises_with_too_large_coefficient(self): # each coefficient should be in range 0...2**l-1 (inclusive) coefs = [i % (2**self.l-1) for i in range(1, 257)] coefs[4] = 2**self.l - polynomial = Polynomial(coefs) - with self.assertRaises(ValueError): - encode([polynomial], self.l) - - def test_encode_raises_with_too_high_degree(self): - # encode should be given a polynomial with degree of 255, not 256 - coefs = [1 for _ in range(257)] - polynomial = Polynomial(coefs) + polynomial = PolynomialRing(coefs) with self.assertRaises(ValueError): encode([polynomial], self.l) diff --git a/tests/test_modulo.py b/tests/test_modulo.py deleted file mode 100644 index b23a780..0000000 --- a/tests/test_modulo.py +++ /dev/null @@ -1,26 +0,0 @@ -import unittest -import numpy as np -from numpy.polynomial.polynomial import Polynomial -from kyber.utils.modulo import matmod, polmod - -class TestModulo(unittest.TestCase): - def setUp(self): - self.polynomial = Polynomial([72, -2, 0, 52, -3, 17]) - self.pol_divisor = Polynomial([1, 0, 0, 1]) - self.int_divisor = 7 - - def test_polmod(self): - # remainder of division polynomial/pol_divisor is -17x^2+x+20 - # then each coefficient is taken modulo int_divisor to get the expected_result - result = polmod(self.polynomial, self.pol_divisor, self.int_divisor) - expected_result = Polynomial([6, 1, 4]) - self.assertEqual(result, expected_result) - - def test_matmod(self): - # basically the same test than `test_polmod` but for a matrix of two identical polynomials - matrix = np.array([self.polynomial, self.polynomial]) - expected_result = Polynomial([6, 1, 4]) - result = matmod(matrix, self.pol_divisor, self.int_divisor) - self.assertEqual(result.shape, (2, )) - self.assertEqual(result[0], expected_result) - self.assertEqual(result[1], expected_result) diff --git a/tests/test_parse.py b/tests/test_parse.py index 631a593..23a9288 100644 --- a/tests/test_parse.py +++ b/tests/test_parse.py @@ -12,8 +12,8 @@ def sample_generator(seed: int = 0) -> Generator[bytes, None, None]: class TestParse(unittest.TestCase): def test_parse_outputted_polynomial_characteristics(self): pol = parse(sample_generator()) - self.assertEqual(len(pol.coef), 256) - for c in pol.coef: + self.assertEqual(len(pol.coefs), 256) + for c in pol.coefs: self.assertTrue(0 <= c and c <= 4095) def test_parse_outputs_same_polynomial_with_same_input(self): diff --git a/tests/test_polring.py b/tests/test_polring.py new file mode 100644 index 0000000..06052d2 --- /dev/null +++ b/tests/test_polring.py @@ -0,0 +1,51 @@ +import unittest +from random import seed, randint +from numpy.polynomial.polynomial import Polynomial +from kyber.entities.polring import PolynomialRing +from kyber.constants import q, n + +# q = 3329 + +class TestPolynomialRing(unittest.TestCase): + def test_initialization(self): + pol = PolynomialRing([71, -5, 0, 1, 3329, 3330, 3328, 15]) + self.assertListEqual(pol.coefs, [71, 3324, 0, 1, 0, 1, 3328, 15]) + + def test_init_with_random_inputs(self): + # test with 1000 samples that randomily initialized polring matches expected + seed(42) + for _ in range(100): + coef_count = randint(1, 500) + coefs = [ + (1 if randint(1, 2) == 1 else -1) * randint(1, 10000) + for _ in range(coef_count) + ] + pol = PolynomialRing(coefs) + self.assertEqual(len(pol.coefs), min(coef_count, 256)) + expected_result = Polynomial(coefs) % Polynomial([1] + [0 for _ in range(n-1)] + [1]) + for i in range(len(pol.coefs)): + self.assertEqual(pol.coefs[i], expected_result.coef[i] % q) + + def test_sum(self): + pol1 = PolynomialRing([581, -50, 100, 31, -4500, 4567, 11, 12]) + pol2 = PolynomialRing([1986, -150, -99, 34, 500, 0]) + # sum before mod [2567, -200, 1, 65, -4000, 4567, 11, 12] + self.assertListEqual((pol1 + pol2).coefs, [2567, 3129, 1, 65, 2658, 1238, 11, 12]) + + def test_multiplication(self): + pol1 = PolynomialRing([15, -13, 0, 7, 472, -88, 112, 5]) + pol2 = PolynomialRing([2, 17, 8, 590, -11, -101, 91]) + # product before mod [30, 229, -101, 8760, -6772, 6532, 9312, 278430, -56838, 20053, 53558, -19375, 9687, 455] + self.assertListEqual((pol1 * pol2).coefs, [30, 229, 3228, 2102, 3215, 3203, 2654, 2123, 3084, 79, 294, 599, 3029, 455]) + + def test_multiplication_with_random_inputs(self): + seed(42) + for _ in range(100): + pol1_degree, pol2_degree = randint(1, 500), randint(1, 500) + coefs_1 = [randint(-4000, 4000) for _ in range(pol1_degree)] + coefs_2 = [randint(-4000, 4000) for _ in range(pol2_degree)] + result = PolynomialRing(coefs_1) * PolynomialRing(coefs_2) + expected = (Polynomial(coefs_1) * Polynomial(coefs_2)) % Polynomial([1] + [0 for _ in range(n-1)] + [1]) + self.assertEqual(len(result.coefs), len(expected.coef)) + for i in range(len(result.coefs)): + self.assertEqual(result.coefs[i], expected.coef[i] % q)