mlkem768.go

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package mlkem implements the quantum-resistant key encapsulation method
     6  // ML-KEM (formerly known as Kyber), as specified in [NIST FIPS 203].
     7  //
     8  // [NIST FIPS 203]: https://doi.org/10.6028/NIST.FIPS.203
     9  package mlkem
    10  
    11  // This package targets security, correctness, simplicity, readability, and
    12  // reviewability as its primary goals. All critical operations are performed in
    13  // constant time.
    14  //
    15  // Variable and function names, as well as code layout, are selected to
    16  // facilitate reviewing the implementation against the NIST FIPS 203 document.
    17  //
    18  // Reviewers unfamiliar with polynomials or linear algebra might find the
    19  // background at https://words.filippo.io/kyber-math/ useful.
    20  //
    21  // This file implements the recommended parameter set ML-KEM-768. The ML-KEM-1024
    22  // parameter set implementation is auto-generated from this file.
    23  //
    24  //go:generate go run generate1024.go -input mlkem768.go -output mlkem1024.go
    25  
    26  import (
    27  	"bytes"
    28  	"crypto/internal/fips140"
    29  	"crypto/internal/fips140/drbg"
    30  	"crypto/internal/fips140/sha3"
    31  	"crypto/internal/fips140/subtle"
    32  	"errors"
    33  )
    34  
    35  const (
    36  	// ML-KEM global constants.
    37  	n = 256
    38  	q = 3329
    39  
    40  	// encodingSizeX is the byte size of a ringElement or nttElement encoded
    41  	// by ByteEncode_X (FIPS 203, Algorithm 5).
    42  	encodingSize12 = n * 12 / 8
    43  	encodingSize11 = n * 11 / 8
    44  	encodingSize10 = n * 10 / 8
    45  	encodingSize5  = n * 5 / 8
    46  	encodingSize4  = n * 4 / 8
    47  	encodingSize1  = n * 1 / 8
    48  
    49  	messageSize = encodingSize1
    50  
    51  	SharedKeySize = 32
    52  	SeedSize      = 32 + 32
    53  )
    54  
    55  // ML-KEM-768 parameters.
    56  const (
    57  	k = 3
    58  
    59  	CiphertextSize768       = k*encodingSize10 + encodingSize4
    60  	EncapsulationKeySize768 = k*encodingSize12 + 32
    61  	decapsulationKeySize768 = k*encodingSize12 + EncapsulationKeySize768 + 32 + 32
    62  )
    63  
    64  // ML-KEM-1024 parameters.
    65  const (
    66  	k1024 = 4
    67  
    68  	CiphertextSize1024       = k1024*encodingSize11 + encodingSize5
    69  	EncapsulationKeySize1024 = k1024*encodingSize12 + 32
    70  	decapsulationKeySize1024 = k1024*encodingSize12 + EncapsulationKeySize1024 + 32 + 32
    71  )
    72  
    73  // A DecapsulationKey768 is the secret key used to decapsulate a shared key from a
    74  // ciphertext. It includes various precomputed values.
    75  type DecapsulationKey768 struct {
    76  	d [32]byte // decapsulation key seed
    77  	z [32]byte // implicit rejection sampling seed
    78  
    79  	ρ [32]byte // sampleNTT seed for A, stored for the encapsulation key
    80  	h [32]byte // H(ek), stored for ML-KEM.Decaps_internal
    81  
    82  	encryptionKey
    83  	decryptionKey
    84  }
    85  
    86  // Bytes returns the decapsulation key as a 64-byte seed in the "d || z" form.
    87  //
    88  // The decapsulation key must be kept secret.
    89  func (dk *DecapsulationKey768) Bytes() []byte {
    90  	var b [SeedSize]byte
    91  	copy(b[:], dk.d[:])
    92  	copy(b[32:], dk.z[:])
    93  	return b[:]
    94  }
    95  
    96  // TestingOnlyExpandedBytes768 returns the decapsulation key as a byte slice
    97  // using the full expanded NIST encoding.
    98  //
    99  // This should only be used for ACVP testing. For all other purposes prefer
   100  // the Bytes method that returns the (much smaller) seed.
   101  func TestingOnlyExpandedBytes768(dk *DecapsulationKey768) []byte {
   102  	b := make([]byte, 0, decapsulationKeySize768)
   103  
   104  	// ByteEncode₁₂(s)
   105  	for i := range dk.s {
   106  		b = polyByteEncode(b, dk.s[i])
   107  	}
   108  
   109  	// ByteEncode₁₂(t) || ρ
   110  	for i := range dk.t {
   111  		b = polyByteEncode(b, dk.t[i])
   112  	}
   113  	b = append(b, dk.ρ[:]...)
   114  
   115  	// H(ek) || z
   116  	b = append(b, dk.h[:]...)
   117  	b = append(b, dk.z[:]...)
   118  
   119  	return b
   120  }
   121  
   122  // EncapsulationKey returns the public encapsulation key necessary to produce
   123  // ciphertexts.
   124  func (dk *DecapsulationKey768) EncapsulationKey() *EncapsulationKey768 {
   125  	return &EncapsulationKey768{
   126  		ρ:             dk.ρ,
   127  		h:             dk.h,
   128  		encryptionKey: dk.encryptionKey,
   129  	}
   130  }
   131  
   132  // An EncapsulationKey768 is the public key used to produce ciphertexts to be
   133  // decapsulated by the corresponding [DecapsulationKey768].
   134  type EncapsulationKey768 struct {
   135  	ρ [32]byte // sampleNTT seed for A
   136  	h [32]byte // H(ek)
   137  	encryptionKey
   138  }
   139  
   140  // Bytes returns the encapsulation key as a byte slice.
   141  func (ek *EncapsulationKey768) Bytes() []byte {
   142  	// The actual logic is in a separate function to outline this allocation.
   143  	b := make([]byte, 0, EncapsulationKeySize768)
   144  	return ek.bytes(b)
   145  }
   146  
   147  func (ek *EncapsulationKey768) bytes(b []byte) []byte {
   148  	for i := range ek.t {
   149  		b = polyByteEncode(b, ek.t[i])
   150  	}
   151  	b = append(b, ek.ρ[:]...)
   152  	return b
   153  }
   154  
   155  // encryptionKey is the parsed and expanded form of a PKE encryption key.
   156  type encryptionKey struct {
   157  	t [k]nttElement     // ByteDecode₁₂(ek[:384k])
   158  	a [k * k]nttElement // A[i*k+j] = sampleNTT(ρ, j, i)
   159  }
   160  
   161  // decryptionKey is the parsed and expanded form of a PKE decryption key.
   162  type decryptionKey struct {
   163  	s [k]nttElement // ByteDecode₁₂(dk[:decryptionKeySize])
   164  }
   165  
   166  // GenerateKey768 generates a new decapsulation key, drawing random bytes from
   167  // a DRBG. The decapsulation key must be kept secret.
   168  func GenerateKey768() (*DecapsulationKey768, error) {
   169  	// The actual logic is in a separate function to outline this allocation.
   170  	dk := &DecapsulationKey768{}
   171  	return generateKey(dk)
   172  }
   173  
   174  func generateKey(dk *DecapsulationKey768) (*DecapsulationKey768, error) {
   175  	fipsSelfTest()
   176  	var d [32]byte
   177  	drbg.Read(d[:])
   178  	var z [32]byte
   179  	drbg.Read(z[:])
   180  	kemKeyGen(dk, &d, &z)
   181  	fips140.PCT("ML-KEM PCT", func() error { return kemPCT(dk) })
   182  	fips140.RecordApproved()
   183  	return dk, nil
   184  }
   185  
   186  // GenerateKeyInternal768 is a derandomized version of GenerateKey768,
   187  // exclusively for use in tests.
   188  func GenerateKeyInternal768(d, z *[32]byte) *DecapsulationKey768 {
   189  	fipsSelfTest()
   190  	dk := &DecapsulationKey768{}
   191  	kemKeyGen(dk, d, z)
   192  	return dk
   193  }
   194  
   195  // NewDecapsulationKey768 parses a decapsulation key from a 64-byte
   196  // seed in the "d || z" form. The seed must be uniformly random.
   197  func NewDecapsulationKey768(seed []byte) (*DecapsulationKey768, error) {
   198  	// The actual logic is in a separate function to outline this allocation.
   199  	dk := &DecapsulationKey768{}
   200  	return newKeyFromSeed(dk, seed)
   201  }
   202  
   203  func newKeyFromSeed(dk *DecapsulationKey768, seed []byte) (*DecapsulationKey768, error) {
   204  	fipsSelfTest()
   205  	if len(seed) != SeedSize {
   206  		return nil, errors.New("mlkem: invalid seed length")
   207  	}
   208  	d := (*[32]byte)(seed[:32])
   209  	z := (*[32]byte)(seed[32:])
   210  	kemKeyGen(dk, d, z)
   211  	fips140.RecordApproved()
   212  	return dk, nil
   213  }
   214  
   215  // TestingOnlyNewDecapsulationKey768 parses a decapsulation key from its expanded NIST format.
   216  //
   217  // Bytes() must not be called on the returned key, as it will not produce the
   218  // original seed.
   219  //
   220  // This function should only be used for ACVP testing. Prefer NewDecapsulationKey768 for all
   221  // other purposes.
   222  func TestingOnlyNewDecapsulationKey768(b []byte) (*DecapsulationKey768, error) {
   223  	if len(b) != decapsulationKeySize768 {
   224  		return nil, errors.New("mlkem: invalid NIST decapsulation key length")
   225  	}
   226  
   227  	dk := &DecapsulationKey768{}
   228  	for i := range dk.s {
   229  		var err error
   230  		dk.s[i], err = polyByteDecode[nttElement](b[:encodingSize12])
   231  		if err != nil {
   232  			return nil, errors.New("mlkem: invalid secret key encoding")
   233  		}
   234  		b = b[encodingSize12:]
   235  	}
   236  
   237  	ek, err := NewEncapsulationKey768(b[:EncapsulationKeySize768])
   238  	if err != nil {
   239  		return nil, err
   240  	}
   241  	dk.ρ = ek.ρ
   242  	dk.h = ek.h
   243  	dk.encryptionKey = ek.encryptionKey
   244  	b = b[EncapsulationKeySize768:]
   245  
   246  	if !bytes.Equal(dk.h[:], b[:32]) {
   247  		return nil, errors.New("mlkem: inconsistent H(ek) in encoded bytes")
   248  	}
   249  	b = b[32:]
   250  
   251  	copy(dk.z[:], b)
   252  
   253  	// Generate a random d value for use in Bytes(). This is a safety mechanism
   254  	// that avoids returning a broken key vs a random key if this function is
   255  	// called in contravention of the TestingOnlyNewDecapsulationKey768 function
   256  	// comment advising against it.
   257  	drbg.Read(dk.d[:])
   258  
   259  	return dk, nil
   260  }
   261  
   262  // kemKeyGen generates a decapsulation key.
   263  //
   264  // It implements ML-KEM.KeyGen_internal according to FIPS 203, Algorithm 16, and
   265  // K-PKE.KeyGen according to FIPS 203, Algorithm 13. The two are merged to save
   266  // copies and allocations.
   267  func kemKeyGen(dk *DecapsulationKey768, d, z *[32]byte) {
   268  	dk.d = *d
   269  	dk.z = *z
   270  
   271  	g := sha3.New512()
   272  	g.Write(d[:])
   273  	g.Write([]byte{k}) // Module dimension as a domain separator.
   274  	G := g.Sum(make([]byte, 0, 64))
   275  	ρ, σ := G[:32], G[32:]
   276  	dk.ρ = [32]byte(ρ)
   277  
   278  	A := &dk.a
   279  	for i := byte(0); i < k; i++ {
   280  		for j := byte(0); j < k; j++ {
   281  			A[i*k+j] = sampleNTT(ρ, j, i)
   282  		}
   283  	}
   284  
   285  	var N byte
   286  	s := &dk.s
   287  	for i := range s {
   288  		s[i] = ntt(samplePolyCBD(σ, N))
   289  		N++
   290  	}
   291  	e := make([]nttElement, k)
   292  	for i := range e {
   293  		e[i] = ntt(samplePolyCBD(σ, N))
   294  		N++
   295  	}
   296  
   297  	t := &dk.t
   298  	for i := range t { // t = A ◦ s + e
   299  		t[i] = e[i]
   300  		for j := range s {
   301  			t[i] = polyAdd(t[i], nttMul(A[i*k+j], s[j]))
   302  		}
   303  	}
   304  
   305  	H := sha3.New256()
   306  	ek := dk.EncapsulationKey().Bytes()
   307  	H.Write(ek)
   308  	H.Sum(dk.h[:0])
   309  }
   310  
   311  // kemPCT performs a Pairwise Consistency Test per FIPS 140-3 IG 10.3.A
   312  // Additional Comment 1: "For key pairs generated for use with approved KEMs in
   313  // FIPS 203, the PCT shall consist of applying the encapsulation key ek to
   314  // encapsulate a shared secret K leading to ciphertext c, and then applying
   315  // decapsulation key dk to retrieve the same shared secret K. The PCT passes if
   316  // the two shared secret K values are equal. The PCT shall be performed either
   317  // when keys are generated/imported, prior to the first exportation, or prior to
   318  // the first operational use (if not exported before the first use)."
   319  func kemPCT(dk *DecapsulationKey768) error {
   320  	ek := dk.EncapsulationKey()
   321  	K, c := ek.Encapsulate()
   322  	K1, err := dk.Decapsulate(c)
   323  	if err != nil {
   324  		return err
   325  	}
   326  	if subtle.ConstantTimeCompare(K, K1) != 1 {
   327  		return errors.New("mlkem: PCT failed")
   328  	}
   329  	return nil
   330  }
   331  
   332  // Encapsulate generates a shared key and an associated ciphertext from an
   333  // encapsulation key, drawing random bytes from a DRBG.
   334  //
   335  // The shared key must be kept secret.
   336  func (ek *EncapsulationKey768) Encapsulate() (sharedKey, ciphertext []byte) {
   337  	// The actual logic is in a separate function to outline this allocation.
   338  	var cc [CiphertextSize768]byte
   339  	return ek.encapsulate(&cc)
   340  }
   341  
   342  func (ek *EncapsulationKey768) encapsulate(cc *[CiphertextSize768]byte) (sharedKey, ciphertext []byte) {
   343  	fipsSelfTest()
   344  	var m [messageSize]byte
   345  	drbg.Read(m[:])
   346  	// Note that the modulus check (step 2 of the encapsulation key check from
   347  	// FIPS 203, Section 7.2) is performed by polyByteDecode in parseEK.
   348  	fips140.RecordApproved()
   349  	return kemEncaps(cc, ek, &m)
   350  }
   351  
   352  // EncapsulateInternal is a derandomized version of Encapsulate, exclusively for
   353  // use in tests.
   354  func (ek *EncapsulationKey768) EncapsulateInternal(m *[32]byte) (sharedKey, ciphertext []byte) {
   355  	fipsSelfTest()
   356  	cc := &[CiphertextSize768]byte{}
   357  	return kemEncaps(cc, ek, m)
   358  }
   359  
   360  // kemEncaps generates a shared key and an associated ciphertext.
   361  //
   362  // It implements ML-KEM.Encaps_internal according to FIPS 203, Algorithm 17.
   363  func kemEncaps(cc *[CiphertextSize768]byte, ek *EncapsulationKey768, m *[messageSize]byte) (K, c []byte) {
   364  	g := sha3.New512()
   365  	g.Write(m[:])
   366  	g.Write(ek.h[:])
   367  	G := g.Sum(nil)
   368  	K, r := G[:SharedKeySize], G[SharedKeySize:]
   369  	c = pkeEncrypt(cc, &ek.encryptionKey, m, r)
   370  	return K, c
   371  }
   372  
   373  // NewEncapsulationKey768 parses an encapsulation key from its encoded form.
   374  // If the encapsulation key is not valid, NewEncapsulationKey768 returns an error.
   375  func NewEncapsulationKey768(encapsulationKey []byte) (*EncapsulationKey768, error) {
   376  	// The actual logic is in a separate function to outline this allocation.
   377  	ek := &EncapsulationKey768{}
   378  	return parseEK(ek, encapsulationKey)
   379  }
   380  
   381  // parseEK parses an encryption key from its encoded form.
   382  //
   383  // It implements the initial stages of K-PKE.Encrypt according to FIPS 203,
   384  // Algorithm 14.
   385  func parseEK(ek *EncapsulationKey768, ekPKE []byte) (*EncapsulationKey768, error) {
   386  	if len(ekPKE) != EncapsulationKeySize768 {
   387  		return nil, errors.New("mlkem: invalid encapsulation key length")
   388  	}
   389  
   390  	h := sha3.New256()
   391  	h.Write(ekPKE)
   392  	h.Sum(ek.h[:0])
   393  
   394  	for i := range ek.t {
   395  		var err error
   396  		ek.t[i], err = polyByteDecode[nttElement](ekPKE[:encodingSize12])
   397  		if err != nil {
   398  			return nil, err
   399  		}
   400  		ekPKE = ekPKE[encodingSize12:]
   401  	}
   402  	copy(ek.ρ[:], ekPKE)
   403  
   404  	for i := byte(0); i < k; i++ {
   405  		for j := byte(0); j < k; j++ {
   406  			ek.a[i*k+j] = sampleNTT(ek.ρ[:], j, i)
   407  		}
   408  	}
   409  
   410  	return ek, nil
   411  }
   412  
   413  // pkeEncrypt encrypt a plaintext message.
   414  //
   415  // It implements K-PKE.Encrypt according to FIPS 203, Algorithm 14, although the
   416  // computation of t and AT is done in parseEK.
   417  func pkeEncrypt(cc *[CiphertextSize768]byte, ex *encryptionKey, m *[messageSize]byte, rnd []byte) []byte {
   418  	var N byte
   419  	r, e1 := make([]nttElement, k), make([]ringElement, k)
   420  	for i := range r {
   421  		r[i] = ntt(samplePolyCBD(rnd, N))
   422  		N++
   423  	}
   424  	for i := range e1 {
   425  		e1[i] = samplePolyCBD(rnd, N)
   426  		N++
   427  	}
   428  	e2 := samplePolyCBD(rnd, N)
   429  
   430  	u := make([]ringElement, k) // NTT⁻¹(AT ◦ r) + e1
   431  	for i := range u {
   432  		var uHat nttElement
   433  		for j := range r {
   434  			// Note that i and j are inverted, as we need the transposed of A.
   435  			uHat = polyAdd(uHat, nttMul(ex.a[j*k+i], r[j]))
   436  		}
   437  		u[i] = polyAdd(e1[i], inverseNTT(uHat))
   438  	}
   439  
   440  	μ := ringDecodeAndDecompress1(m)
   441  
   442  	var vNTT nttElement // t⊺ ◦ r
   443  	for i := range ex.t {
   444  		vNTT = polyAdd(vNTT, nttMul(ex.t[i], r[i]))
   445  	}
   446  	v := polyAdd(polyAdd(inverseNTT(vNTT), e2), μ)
   447  
   448  	c := cc[:0]
   449  	for _, f := range u {
   450  		c = ringCompressAndEncode10(c, f)
   451  	}
   452  	c = ringCompressAndEncode4(c, v)
   453  
   454  	return c
   455  }
   456  
   457  // Decapsulate generates a shared key from a ciphertext and a decapsulation key.
   458  // If the ciphertext is not valid, Decapsulate returns an error.
   459  //
   460  // The shared key must be kept secret.
   461  func (dk *DecapsulationKey768) Decapsulate(ciphertext []byte) (sharedKey []byte, err error) {
   462  	fipsSelfTest()
   463  	if len(ciphertext) != CiphertextSize768 {
   464  		return nil, errors.New("mlkem: invalid ciphertext length")
   465  	}
   466  	c := (*[CiphertextSize768]byte)(ciphertext)
   467  	// Note that the hash check (step 3 of the decapsulation input check from
   468  	// FIPS 203, Section 7.3) is foregone as a DecapsulationKey is always
   469  	// validly generated by ML-KEM.KeyGen_internal.
   470  	return kemDecaps(dk, c), nil
   471  }
   472  
   473  var testingOnlyRejectionOutcome func(int)
   474  
   475  // kemDecaps produces a shared key from a ciphertext.
   476  //
   477  // It implements ML-KEM.Decaps_internal according to FIPS 203, Algorithm 18.
   478  func kemDecaps(dk *DecapsulationKey768, c *[CiphertextSize768]byte) (K []byte) {
   479  	fips140.RecordApproved()
   480  	m := pkeDecrypt(&dk.decryptionKey, c)
   481  	g := sha3.New512()
   482  	g.Write(m[:])
   483  	g.Write(dk.h[:])
   484  	G := g.Sum(make([]byte, 0, 64))
   485  	Kprime, r := G[:SharedKeySize], G[SharedKeySize:]
   486  	J := sha3.NewShake256()
   487  	J.Write(dk.z[:])
   488  	J.Write(c[:])
   489  	Kout := make([]byte, SharedKeySize)
   490  	J.Read(Kout)
   491  	var cc [CiphertextSize768]byte
   492  	c1 := pkeEncrypt(&cc, &dk.encryptionKey, (*[32]byte)(m), r)
   493  
   494  	if testingOnlyRejectionOutcome != nil {
   495  		testingOnlyRejectionOutcome(subtle.ConstantTimeCompare(c[:], c1))
   496  	}
   497  
   498  	subtle.ConstantTimeCopy(subtle.ConstantTimeCompare(c[:], c1), Kout, Kprime)
   499  	return Kout
   500  }
   501  
   502  // pkeDecrypt decrypts a ciphertext.
   503  //
   504  // It implements K-PKE.Decrypt according to FIPS 203, Algorithm 15,
   505  // although s is retained from kemKeyGen.
   506  func pkeDecrypt(dx *decryptionKey, c *[CiphertextSize768]byte) []byte {
   507  	u := make([]ringElement, k)
   508  	for i := range u {
   509  		b := (*[encodingSize10]byte)(c[encodingSize10*i : encodingSize10*(i+1)])
   510  		u[i] = ringDecodeAndDecompress10(b)
   511  	}
   512  
   513  	b := (*[encodingSize4]byte)(c[encodingSize10*k:])
   514  	v := ringDecodeAndDecompress4(b)
   515  
   516  	var mask nttElement // s⊺ ◦ NTT(u)
   517  	for i := range dx.s {
   518  		mask = polyAdd(mask, nttMul(dx.s[i], ntt(u[i])))
   519  	}
   520  	w := polySub(v, inverseNTT(mask))
   521  
   522  	return ringCompressAndEncode1(nil, w)
   523  }
   524
View as plain text