mlkem1024.go

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