CRYSTALS - Kyber core

In this article we will go through how CRYSTALS-Kyber is implemented and refer to a (pedagogical) reference implementation you can find from my github. This article relies on mathematical machinery introduced in the previous article.

We will first implement the basic PKE-scheme. Full KEM is implemented in the next article.

The toolset

For reminder, here are the key definitions from the previous article. We are working in the polynomial ring \[R_q = \mathbb{Z}_q[x]/(x^n+1).\] We define \( R_q^k \) to be a (column) vector of polynomials from the ring \[\mathbb{Z}_q[x]/(x^n+1).\] Let \(f(x)\) be a polynomial. Its infinity-norm \[ \|f(x)\|_{\infty}=\max_i \|a_i\|_{\infty}.\] We use \(S_{\eta}\) to denote the set of small polynomials \[ S_{\eta} = \{ f\in R_q\, |\,\|f\|_{\infty}\leq\eta\}.\]

Encrypt and decrypt

We are ready to go through the basic encryption and decryption. Fix positive integers \(q,n,k,\eta_1,\eta_2\). Let Alice construct public and private keys. First Alice chooses randomly matrix \(A\in R^{k\times k}_q\) and vectors \(s\in S_{\eta_1}^k \) and \(e\in S_{\eta_2}^k\).

Here "random" has to be taken carefully. For matrix \(A\) we sample the values of coefficients from uniform distribution in \([0,q-1]\), this is straightforward. For vector coefficients we use binomial distribution \(\text{Binomial}(\eta,1/2)\) to have \(\eta\) many samples, each either \(0\) or \(1\) and sum them up and store the result in \(X\). Then we repeat and store the result to \(Y\). Now result is \(X-Y\). The resulting distribution is Centered Binomial Distribution and it is usually denoted as \(\text{CBD}(\eta)\). The used distributions have been carefully chosen for the best balance between efficiency and security.

The sampling is, again, implemented in the Zq class; our beloved workhorse.

    @classmethod
    def random_uniform(cls, q):
        # For public matrix A
        return cls(q, random.randrange(q))

    @classmethod
    def random_binomial(cls, q, eta):
        # For secret/error terms
        bits1 = [random.randint(0, 1) for _ in range(eta)]
        bits2 = [random.randint(0, 1) for _ in range(eta)]
        value = sum(bits1) - sum(bits2)
        result = cls(q,value)
        result = result.to_symmetric()
        return result

Next Alice computes \[t=As+e\] and \(s\) is now her private key/secret. Her public key is the pair \((A,t)\). Notice that computing \(s\) from \((A,t)\) is an instance of MLWE-problem which is believed to hard.

The preparations are within the constructor of the class Person.

        self.A = PolyMatrix(self.kyber_params.k, self.kyber_params.k, self.kyber_params.q, self.kyber_params.n)
        self.A.random_uniform()

        self.s = PolyMatrix(self.kyber_params.k, 1, self.kyber_params.q, self.kyber_params.n)
        self.s = self.s.random_binomial(self.kyber_params.eta1)

        self.e = PolyMatrix(self.kyber_params.k, 1, self.kyber_params.q, self.kyber_params.n)
        self.e.random_binomial(self.kyber_params.eta2)

        self.t = self.A @ self.s + self.e

Now that she has generated the public and private keys, we need to define how Bob could use \((A,t)\) to encrypt a message. Let \(m\) be the message. The length of \(m\) is \(n\) (remember we are only required to exchange a message of fixed size, namely the symmetric encryption key). Now \(m\) is a string of bits, which we will present as a list of zeroes and ones in our Python code. In the setup, \(m\) is encoded in coefficients of a polynomial of degree \(n-1\).

To encode the message, Bob proceed as follows. He chooses random vectors \[r\in S_{\eta_1}^k\] \[e_1\in S_{\eta_2}^k\] \[e_2\in S_{\eta_2}^k.\] Next he computes \[ u = A^Tr + e_1.\] Notice that \(u\) is now a polynomial vector of length \(k\). Next he computes \[v = t^Tr + e_2 + \lceil q/2 \rfloor m\] which is a polynomial. Here \(\lceil x\rfloor\) denotes the integer closest to \(x\) ties rounded up.

Bob is now left with pair \[ c=(u,v)\,\in R_q^k\times R_q\] which is his cipher, i.e. encrypted message to Alice.

The encryption is implemented in the method encrypt in class Person.

    def encrypt(self, message, others_a, others_t):
        """ Encrypt the message. Message is a list of zeroes and ones of length n."""
        print(f"{self.name} is encrypting a message")
        self.r = PolyMatrix(self.kyber_params.k, 1, self.kyber_params.q, self.kyber_params.n)
        self.r.random_binomial(self.kyber_params.eta1)

        self.e1 = PolyMatrix(self.kyber_params.k, 1, self.kyber_params.q, self.kyber_params.n)
        self.e1.random_binomial(self.kyber_params.eta2)

        self.e2 = PolyMatrix(1, 1, self.kyber_params.q, self.kyber_params.n)
        self.e2.random_binomial(self.kyber_params.eta2)

        m_new = [int((self.kyber_params.q / 2) * x + 0.5) for x in message]
        m_matrix = PolyMatrix(1, 1, self.kyber_params.q, self.kyber_params.n)
        m_matrix[(0, 0)] = ZqPolynomial(self.kyber_params.q, m_new)
        u = others_a.T @ self.r + self.e1
        v = others_t.T @ self.r + self.e2 + m_matrix
        return (u,v)

Alice can decrypt the message simply by \[m' = \text{Round}_q(v-s^Tu)\] where \(s\) is the private key. Rounding is simply rounding the coefficients to \(0\) or \(1\) defined as \[\text{Round}_q(x) = \begin{cases} 0 & \text{if \(-q/4 < x < q/4\)}\\ 1 & \text{otherwise.}\end{cases}\]

The decrypt-method in class Person is as simple as it can get.

        decoded = v - self.s.T @ u
        decoded = decoded.round()
        decodedmessage = decoded.extract_message()
        return decodedmessage

The full process can be tested with the Person class.

if __name__ == "__main__":

    KYBER512 = KyberParams(k=2, n=256, q=3329, eta1=3, eta2=2, du=10, dv=4)

    alice = Person("Alice", KYBER512)
    bob = Person("Bob", KYBER512)

    # The message to be transferred
    message = [random.randint(0,1) for _ in range(KYBER512.n)]
    u, v = bob.encrypt(message, alice.A, alice.t)

    decoded_message = alice.decrypt(u,v)

    difference = sum(a != b for a, b in zip(message, decoded_message))
    print(f"The total amount of bits that differ is {difference}")

Understanding the full process

Why this works? Let us rewind to see what happens in decrypt. We have \(c=(u,v)\) where \[u = A^Tr + e_1\] and \[v = t^Tr + e_2 + \lceil q/2 \rfloor m.\] If we substitute these to the equation \(v-s^Tu\) we get \[v-s^Tu=(t^Tr+e_2+\lceil q/2 \rfloor m)-s^T(A^Tr + e_1).\] Substitute \(t=As+e\) to get \[v-s^Tu=((As+e)^Tr+e_2+\lceil q/2 \rfloor m)-s^T(A^Tr+e_1)\]\[=(s^TA^Tr+e^Tr+e_2+\lceil q/2 \rfloor m)-(s^TA^Tr+s^Te_1)\]\[=e^Tr+e_2-s^Te_1+\lceil q/2 \rfloor m\] and we can we can see that the error term \[e^Tr+e_2-s^Te_1\] is small (assuming proper parametrisation).

Clearly now \(\lceil q/2 \rfloor m\) dominates the expression and the rounding \(\text{Round}_q(v-s^Tu)\) provides us the original message. Sure as rock, isn't it?

Well, actually no. It happens that there are cases where decrypt doesn't recover the original message. What the probability is depends heavily on how the process is parametrized. There are three parameter sets that are suggested by the crypto-community at the time being.

Here \(d\) is \(\delta\), being the probability that decryption of cipher fails. The NIST lvl refers to NIST's classification of the algorithm based on how difficult it is to break. It compares it to different AES/SHA configurations.

As we notice, the probability that decrypt fails is astronomically small and hence not really an issue.

Future work

The full implementation of Kyber includes optimizations and FO-transform to implement full KEM. You can find KEM and compression from the next article. NTT will be handled later on separate article.