mlkem1024.go

     1  // Code generated by generate1024.go. DO NOT EDIT.
     2  
     3  package mlkem
     4  
     5  import (
     6  	"bytes"
     7  	"crypto/internal/fips140"
     8  	"crypto/internal/fips140/drbg"
     9  	"crypto/internal/fips140/sha3"
    10  	"crypto/internal/fips140/subtle"
    11  	"errors"
    12  )
    13  
    14  // A DecapsulationKey1024 is the secret key used to decapsulate a shared key from a
    15  // ciphertext. It includes various precomputed values.
    16  type DecapsulationKey1024 struct {
    17  	d [32]byte // decapsulation key seed
    18  	z [32]byte // implicit rejection sampling seed
    19  
    20  	ρ [32]byte // sampleNTT seed for A, stored for the encapsulation key
    21  	h [32]byte // H(ek), stored for ML-KEM.Decaps_internal
    22  
    23  	encryptionKey1024
    24  	decryptionKey1024
    25  }
    26  
    27  // Bytes returns the decapsulation key as a 64-byte seed in the "d || z" form.
    28  //
    29  // The decapsulation key must be kept secret.
    30  func (dk *DecapsulationKey1024) Bytes() []byte {
    31  	var b [SeedSize]byte
    32  	copy(b[:], dk.d[:])
    33  	copy(b[32:], dk.z[:])
    34  	return b[:]
    35  }
    36  
    37  // TestingOnlyExpandedBytes1024 returns the decapsulation key as a byte slice
    38  // using the full expanded NIST encoding.
    39  //
    40  // This should only be used for ACVP testing. For all other purposes prefer
    41  // the Bytes method that returns the (much smaller) seed.
    42  func TestingOnlyExpandedBytes1024(dk *DecapsulationKey1024) []byte {
    43  	b := make([]byte, 0, decapsulationKeySize1024)
    44  
    45  	// ByteEncode₁₂(s)
    46  	for i := range dk.s {
    47  		b = polyByteEncode(b, dk.s[i])
    48  	}
    49  
    50  	// ByteEncode₁₂(t) || ρ
    51  	for i := range dk.t {
    52  		b = polyByteEncode(b, dk.t[i])
    53  	}
    54  	b = append(b, dk.ρ[:]...)
    55  
    56  	// H(ek) || z
    57  	b = append(b, dk.h[:]...)
    58  	b = append(b, dk.z[:]...)
    59  
    60  	return b
    61  }
    62  
    63  // EncapsulationKey returns the public encapsulation key necessary to produce
    64  // ciphertexts.
    65  func (dk *DecapsulationKey1024) EncapsulationKey() *EncapsulationKey1024 {
    66  	return &EncapsulationKey1024{
    67  		ρ:                 dk.ρ,
    68  		h:                 dk.h,
    69  		encryptionKey1024: dk.encryptionKey1024,
    70  	}
    71  }
    72  
    73  // An EncapsulationKey1024 is the public key used to produce ciphertexts to be
    74  // decapsulated by the corresponding [DecapsulationKey1024].
    75  type EncapsulationKey1024 struct {
    76  	ρ [32]byte // sampleNTT seed for A
    77  	h [32]byte // H(ek)
    78  	encryptionKey1024
    79  }
    80  
    81  // Bytes returns the encapsulation key as a byte slice.
    82  func (ek *EncapsulationKey1024) Bytes() []byte {
    83  	// The actual logic is in a separate function to outline this allocation.
    84  	b := make([]byte, 0, EncapsulationKeySize1024)
    85  	return ek.bytes(b)
    86  }
    87  
    88  func (ek *EncapsulationKey1024) bytes(b []byte) []byte {
    89  	for i := range ek.t {
    90  		b = polyByteEncode(b, ek.t[i])
    91  	}
    92  	b = append(b, ek.ρ[:]...)
    93  	return b
    94  }
    95  
    96  // encryptionKey1024 is the parsed and expanded form of a PKE encryption key.
    97  type encryptionKey1024 struct {
    98  	t [k1024]nttElement         // ByteDecode₁₂(ek[:384k])
    99  	a [k1024 * k1024]nttElement // A[i*k+j] = sampleNTT(ρ, j, i)
   100  }
   101  
   102  // decryptionKey1024 is the parsed and expanded form of a PKE decryption key.
   103  type decryptionKey1024 struct {
   104  	s [k1024]nttElement // ByteDecode₁₂(dk[:decryptionKey1024Size])
   105  }
   106  
   107  // GenerateKey1024 generates a new decapsulation key, drawing random bytes from
   108  // a DRBG. The decapsulation key must be kept secret.
   109  func GenerateKey1024() (*DecapsulationKey1024, error) {
   110  	// The actual logic is in a separate function to outline this allocation.
   111  	dk := &DecapsulationKey1024{}
   112  	return generateKey1024(dk)
   113  }
   114  
   115  func generateKey1024(dk *DecapsulationKey1024) (*DecapsulationKey1024, error) {
   116  	fipsSelfTest()
   117  	var d [32]byte
   118  	drbg.Read(d[:])
   119  	var z [32]byte
   120  	drbg.Read(z[:])
   121  	kemKeyGen1024(dk, &d, &z)
   122  	fips140.PCT("ML-KEM PCT", func() error { return kemPCT1024(dk) })
   123  	fips140.RecordApproved()
   124  	return dk, nil
   125  }
   126  
   127  // GenerateKeyInternal1024 is a derandomized version of GenerateKey1024,
   128  // exclusively for use in tests.
   129  func GenerateKeyInternal1024(d, z *[32]byte) *DecapsulationKey1024 {
   130  	fipsSelfTest()
   131  	dk := &DecapsulationKey1024{}
   132  	kemKeyGen1024(dk, d, z)
   133  	return dk
   134  }
   135  
   136  // NewDecapsulationKey1024 parses a decapsulation key from a 64-byte
   137  // seed in the "d || z" form. The seed must be uniformly random.
   138  func NewDecapsulationKey1024(seed []byte) (*DecapsulationKey1024, error) {
   139  	// The actual logic is in a separate function to outline this allocation.
   140  	dk := &DecapsulationKey1024{}
   141  	return newKeyFromSeed1024(dk, seed)
   142  }
   143  
   144  func newKeyFromSeed1024(dk *DecapsulationKey1024, seed []byte) (*DecapsulationKey1024, error) {
   145  	fipsSelfTest()
   146  	if len(seed) != SeedSize {
   147  		return nil, errors.New("mlkem: invalid seed length")
   148  	}
   149  	d := (*[32]byte)(seed[:32])
   150  	z := (*[32]byte)(seed[32:])
   151  	kemKeyGen1024(dk, d, z)
   152  	fips140.RecordApproved()
   153  	return dk, nil
   154  }
   155  
   156  // TestingOnlyNewDecapsulationKey1024 parses a decapsulation key from its expanded NIST format.
   157  //
   158  // Bytes() must not be called on the returned key, as it will not produce the
   159  // original seed.
   160  //
   161  // This function should only be used for ACVP testing. Prefer NewDecapsulationKey1024 for all
   162  // other purposes.
   163  func TestingOnlyNewDecapsulationKey1024(b []byte) (*DecapsulationKey1024, error) {
   164  	if len(b) != decapsulationKeySize1024 {
   165  		return nil, errors.New("mlkem: invalid NIST decapsulation key length")
   166  	}
   167  
   168  	dk := &DecapsulationKey1024{}
   169  	for i := range dk.s {
   170  		var err error
   171  		dk.s[i], err = polyByteDecode[nttElement](b[:encodingSize12])
   172  		if err != nil {
   173  			return nil, errors.New("mlkem: invalid secret key encoding")
   174  		}
   175  		b = b[encodingSize12:]
   176  	}
   177  
   178  	ek, err := NewEncapsulationKey1024(b[:EncapsulationKeySize1024])
   179  	if err != nil {
   180  		return nil, err
   181  	}
   182  	dk.ρ = ek.ρ
   183  	dk.h = ek.h
   184  	dk.encryptionKey1024 = ek.encryptionKey1024
   185  	b = b[EncapsulationKeySize1024:]
   186  
   187  	if !bytes.Equal(dk.h[:], b[:32]) {
   188  		return nil, errors.New("mlkem: inconsistent H(ek) in encoded bytes")
   189  	}
   190  	b = b[32:]
   191  
   192  	copy(dk.z[:], b)
   193  
   194  	// Generate a random d value for use in Bytes(). This is a safety mechanism
   195  	// that avoids returning a broken key vs a random key if this function is
   196  	// called in contravention of the TestingOnlyNewDecapsulationKey1024 function
   197  	// comment advising against it.
   198  	drbg.Read(dk.d[:])
   199  
   200  	return dk, nil
   201  }
   202  
   203  // kemKeyGen1024 generates a decapsulation key.
   204  //
   205  // It implements ML-KEM.KeyGen_internal according to FIPS 203, Algorithm 16, and
   206  // K-PKE.KeyGen according to FIPS 203, Algorithm 13. The two are merged to save
   207  // copies and allocations.
   208  func kemKeyGen1024(dk *DecapsulationKey1024, d, z *[32]byte) {
   209  	dk.d = *d
   210  	dk.z = *z
   211  
   212  	g := sha3.New512()
   213  	g.Write(d[:])
   214  	g.Write([]byte{k1024}) // Module dimension as a domain separator.
   215  	G := g.Sum(make([]byte, 0, 64))
   216  	ρ, σ := G[:32], G[32:]
   217  	dk.ρ = [32]byte(ρ)
   218  
   219  	A := &dk.a
   220  	for i := byte(0); i < k1024; i++ {
   221  		for j := byte(0); j < k1024; j++ {
   222  			A[i*k1024+j] = sampleNTT(ρ, j, i)
   223  		}
   224  	}
   225  
   226  	var N byte
   227  	s := &dk.s
   228  	for i := range s {
   229  		s[i] = ntt(samplePolyCBD(σ, N))
   230  		N++
   231  	}
   232  	e := make([]nttElement, k1024)
   233  	for i := range e {
   234  		e[i] = ntt(samplePolyCBD(σ, N))
   235  		N++
   236  	}
   237  
   238  	t := &dk.t
   239  	for i := range t { // t = A ◦ s + e
   240  		t[i] = e[i]
   241  		for j := range s {
   242  			t[i] = polyAdd(t[i], nttMul(A[i*k1024+j], s[j]))
   243  		}
   244  	}
   245  
   246  	H := sha3.New256()
   247  	ek := dk.EncapsulationKey().Bytes()
   248  	H.Write(ek)
   249  	H.Sum(dk.h[:0])
   250  }
   251  
   252  // kemPCT1024 performs a Pairwise Consistency Test per FIPS 140-3 IG 10.3.A
   253  // Additional Comment 1: "For key pairs generated for use with approved KEMs in
   254  // FIPS 203, the PCT shall consist of applying the encapsulation key ek to
   255  // encapsulate a shared secret K leading to ciphertext c, and then applying
   256  // decapsulation key dk to retrieve the same shared secret K. The PCT passes if
   257  // the two shared secret K values are equal. The PCT shall be performed either
   258  // when keys are generated/imported, prior to the first exportation, or prior to
   259  // the first operational use (if not exported before the first use)."
   260  func kemPCT1024(dk *DecapsulationKey1024) error {
   261  	ek := dk.EncapsulationKey()
   262  	K, c := ek.Encapsulate()
   263  	K1, err := dk.Decapsulate(c)
   264  	if err != nil {
   265  		return err
   266  	}
   267  	if subtle.ConstantTimeCompare(K, K1) != 1 {
   268  		return errors.New("mlkem: PCT failed")
   269  	}
   270  	return nil
   271  }
   272  
   273  // Encapsulate generates a shared key and an associated ciphertext from an
   274  // encapsulation key, drawing random bytes from a DRBG.
   275  //
   276  // The shared key must be kept secret.
   277  func (ek *EncapsulationKey1024) Encapsulate() (sharedKey, ciphertext []byte) {
   278  	// The actual logic is in a separate function to outline this allocation.
   279  	var cc [CiphertextSize1024]byte
   280  	return ek.encapsulate(&cc)
   281  }
   282  
   283  func (ek *EncapsulationKey1024) encapsulate(cc *[CiphertextSize1024]byte) (sharedKey, ciphertext []byte) {
   284  	fipsSelfTest()
   285  	var m [messageSize]byte
   286  	drbg.Read(m[:])
   287  	// Note that the modulus check (step 2 of the encapsulation key check from
   288  	// FIPS 203, Section 7.2) is performed by polyByteDecode in parseEK1024.
   289  	fips140.RecordApproved()
   290  	return kemEncaps1024(cc, ek, &m)
   291  }
   292  
   293  // EncapsulateInternal is a derandomized version of Encapsulate, exclusively for
   294  // use in tests.
   295  func (ek *EncapsulationKey1024) EncapsulateInternal(m *[32]byte) (sharedKey, ciphertext []byte) {
   296  	fipsSelfTest()
   297  	cc := &[CiphertextSize1024]byte{}
   298  	return kemEncaps1024(cc, ek, m)
   299  }
   300  
   301  // kemEncaps1024 generates a shared key and an associated ciphertext.
   302  //
   303  // It implements ML-KEM.Encaps_internal according to FIPS 203, Algorithm 17.
   304  func kemEncaps1024(cc *[CiphertextSize1024]byte, ek *EncapsulationKey1024, m *[messageSize]byte) (K, c []byte) {
   305  	g := sha3.New512()
   306  	g.Write(m[:])
   307  	g.Write(ek.h[:])
   308  	G := g.Sum(nil)
   309  	K, r := G[:SharedKeySize], G[SharedKeySize:]
   310  	c = pkeEncrypt1024(cc, &ek.encryptionKey1024, m, r)
   311  	return K, c
   312  }
   313  
   314  // NewEncapsulationKey1024 parses an encapsulation key from its encoded form.
   315  // If the encapsulation key is not valid, NewEncapsulationKey1024 returns an error.
   316  func NewEncapsulationKey1024(encapsulationKey []byte) (*EncapsulationKey1024, error) {
   317  	// The actual logic is in a separate function to outline this allocation.
   318  	ek := &EncapsulationKey1024{}
   319  	return parseEK1024(ek, encapsulationKey)
   320  }
   321  
   322  // parseEK1024 parses an encryption key from its encoded form.
   323  //
   324  // It implements the initial stages of K-PKE.Encrypt according to FIPS 203,
   325  // Algorithm 14.
   326  func parseEK1024(ek *EncapsulationKey1024, ekPKE []byte) (*EncapsulationKey1024, error) {
   327  	if len(ekPKE) != EncapsulationKeySize1024 {
   328  		return nil, errors.New("mlkem: invalid encapsulation key length")
   329  	}
   330  
   331  	h := sha3.New256()
   332  	h.Write(ekPKE)
   333  	h.Sum(ek.h[:0])
   334  
   335  	for i := range ek.t {
   336  		var err error
   337  		ek.t[i], err = polyByteDecode[nttElement](ekPKE[:encodingSize12])
   338  		if err != nil {
   339  			return nil, err
   340  		}
   341  		ekPKE = ekPKE[encodingSize12:]
   342  	}
   343  	copy(ek.ρ[:], ekPKE)
   344  
   345  	for i := byte(0); i < k1024; i++ {
   346  		for j := byte(0); j < k1024; j++ {
   347  			ek.a[i*k1024+j] = sampleNTT(ek.ρ[:], j, i)
   348  		}
   349  	}
   350  
   351  	return ek, nil
   352  }
   353  
   354  // pkeEncrypt1024 encrypt a plaintext message.
   355  //
   356  // It implements K-PKE.Encrypt according to FIPS 203, Algorithm 14, although the
   357  // computation of t and AT is done in parseEK1024.
   358  func pkeEncrypt1024(cc *[CiphertextSize1024]byte, ex *encryptionKey1024, m *[messageSize]byte, rnd []byte) []byte {
   359  	var N byte
   360  	r, e1 := make([]nttElement, k1024), make([]ringElement, k1024)
   361  	for i := range r {
   362  		r[i] = ntt(samplePolyCBD(rnd, N))
   363  		N++
   364  	}
   365  	for i := range e1 {
   366  		e1[i] = samplePolyCBD(rnd, N)
   367  		N++
   368  	}
   369  	e2 := samplePolyCBD(rnd, N)
   370  
   371  	u := make([]ringElement, k1024) // NTT⁻¹(AT ◦ r) + e1
   372  	for i := range u {
   373  		var uHat nttElement
   374  		for j := range r {
   375  			// Note that i and j are inverted, as we need the transposed of A.
   376  			uHat = polyAdd(uHat, nttMul(ex.a[j*k1024+i], r[j]))
   377  		}
   378  		u[i] = polyAdd(e1[i], inverseNTT(uHat))
   379  	}
   380  
   381  	μ := ringDecodeAndDecompress1(m)
   382  
   383  	var vNTT nttElement // t⊺ ◦ r
   384  	for i := range ex.t {
   385  		vNTT = polyAdd(vNTT, nttMul(ex.t[i], r[i]))
   386  	}
   387  	v := polyAdd(polyAdd(inverseNTT(vNTT), e2), μ)
   388  
   389  	c := cc[:0]
   390  	for _, f := range u {
   391  		c = ringCompressAndEncode11(c, f)
   392  	}
   393  	c = ringCompressAndEncode5(c, v)
   394  
   395  	return c
   396  }
   397  
   398  // Decapsulate generates a shared key from a ciphertext and a decapsulation key.
   399  // If the ciphertext is not valid, Decapsulate returns an error.
   400  //
   401  // The shared key must be kept secret.
   402  func (dk *DecapsulationKey1024) Decapsulate(ciphertext []byte) (sharedKey []byte, err error) {
   403  	fipsSelfTest()
   404  	if len(ciphertext) != CiphertextSize1024 {
   405  		return nil, errors.New("mlkem: invalid ciphertext length")
   406  	}
   407  	c := (*[CiphertextSize1024]byte)(ciphertext)
   408  	// Note that the hash check (step 3 of the decapsulation input check from
   409  	// FIPS 203, Section 7.3) is foregone as a DecapsulationKey is always
   410  	// validly generated by ML-KEM.KeyGen_internal.
   411  	return kemDecaps1024(dk, c), nil
   412  }
   413  
   414  // kemDecaps1024 produces a shared key from a ciphertext.
   415  //
   416  // It implements ML-KEM.Decaps_internal according to FIPS 203, Algorithm 18.
   417  func kemDecaps1024(dk *DecapsulationKey1024, c *[CiphertextSize1024]byte) (K []byte) {
   418  	fips140.RecordApproved()
   419  	m := pkeDecrypt1024(&dk.decryptionKey1024, c)
   420  	g := sha3.New512()
   421  	g.Write(m[:])
   422  	g.Write(dk.h[:])
   423  	G := g.Sum(make([]byte, 0, 64))
   424  	Kprime, r := G[:SharedKeySize], G[SharedKeySize:]
   425  	J := sha3.NewShake256()
   426  	J.Write(dk.z[:])
   427  	J.Write(c[:])
   428  	Kout := make([]byte, SharedKeySize)
   429  	J.Read(Kout)
   430  	var cc [CiphertextSize1024]byte
   431  	c1 := pkeEncrypt1024(&cc, &dk.encryptionKey1024, (*[32]byte)(m), r)
   432  
   433  	subtle.ConstantTimeCopy(subtle.ConstantTimeCompare(c[:], c1), Kout, Kprime)
   434  	return Kout
   435  }
   436  
   437  // pkeDecrypt1024 decrypts a ciphertext.
   438  //
   439  // It implements K-PKE.Decrypt according to FIPS 203, Algorithm 15,
   440  // although s is retained from kemKeyGen1024.
   441  func pkeDecrypt1024(dx *decryptionKey1024, c *[CiphertextSize1024]byte) []byte {
   442  	u := make([]ringElement, k1024)
   443  	for i := range u {
   444  		b := (*[encodingSize11]byte)(c[encodingSize11*i : encodingSize11*(i+1)])
   445  		u[i] = ringDecodeAndDecompress11(b)
   446  	}
   447  
   448  	b := (*[encodingSize5]byte)(c[encodingSize11*k1024:])
   449  	v := ringDecodeAndDecompress5(b)
   450  
   451  	var mask nttElement // s⊺ ◦ NTT(u)
   452  	for i := range dx.s {
   453  		mask = polyAdd(mask, nttMul(dx.s[i], ntt(u[i])))
   454  	}
   455  	w := polySub(v, inverseNTT(mask))
   456  
   457  	return ringCompressAndEncode1(nil, w)
   458  }
   459
View as plain text