Merge pull request #3 from PyryL/week4

Week4

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)

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`.

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

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

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)

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

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)
 

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)

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]) + ")"

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)

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:
     """

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