pq.go

     1  // Copyright 2025 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 hpke
     6  
     7  import (
     8  	"bytes"
     9  	"crypto"
    10  	"crypto/ecdh"
    11  	"crypto/fips140"
    12  	"crypto/internal/fips140/drbg"
    13  	"crypto/internal/rand"
    14  	"crypto/mlkem"
    15  	"crypto/sha3"
    16  	"errors"
    17  	"internal/byteorder"
    18  )
    19  
    20  var mlkem768X25519 = &hybridKEM{
    21  	id: 0x647a,
    22  	label: /**/ `\./` +
    23  		/*   */ `/^\`,
    24  	curve: ecdh.X25519(),
    25  
    26  	curveSeedSize:    32,
    27  	curvePointSize:   32,
    28  	pqEncapsKeySize:  mlkem.EncapsulationKeySize768,
    29  	pqCiphertextSize: mlkem.CiphertextSize768,
    30  
    31  	pqNewPublicKey: func(data []byte) (crypto.Encapsulator, error) {
    32  		return mlkem.NewEncapsulationKey768(data)
    33  	},
    34  	pqNewPrivateKey: func(data []byte) (crypto.Decapsulator, error) {
    35  		return mlkem.NewDecapsulationKey768(data)
    36  	},
    37  }
    38  
    39  // MLKEM768X25519 returns a KEM implementing MLKEM768-X25519 (a.k.a. X-Wing)
    40  // from draft-ietf-hpke-pq.
    41  func MLKEM768X25519() KEM {
    42  	return mlkem768X25519
    43  }
    44  
    45  var mlkem768P256 = &hybridKEM{
    46  	id:    0x0050,
    47  	label: "MLKEM768-P256",
    48  	curve: ecdh.P256(),
    49  
    50  	curveSeedSize:    32,
    51  	curvePointSize:   65,
    52  	pqEncapsKeySize:  mlkem.EncapsulationKeySize768,
    53  	pqCiphertextSize: mlkem.CiphertextSize768,
    54  
    55  	pqNewPublicKey: func(data []byte) (crypto.Encapsulator, error) {
    56  		return mlkem.NewEncapsulationKey768(data)
    57  	},
    58  	pqNewPrivateKey: func(data []byte) (crypto.Decapsulator, error) {
    59  		return mlkem.NewDecapsulationKey768(data)
    60  	},
    61  }
    62  
    63  // MLKEM768P256 returns a KEM implementing MLKEM768-P256 from draft-ietf-hpke-pq.
    64  func MLKEM768P256() KEM {
    65  	return mlkem768P256
    66  }
    67  
    68  var mlkem1024P384 = &hybridKEM{
    69  	id:    0x0051,
    70  	label: "MLKEM1024-P384",
    71  	curve: ecdh.P384(),
    72  
    73  	curveSeedSize:    48,
    74  	curvePointSize:   97,
    75  	pqEncapsKeySize:  mlkem.EncapsulationKeySize1024,
    76  	pqCiphertextSize: mlkem.CiphertextSize1024,
    77  
    78  	pqNewPublicKey: func(data []byte) (crypto.Encapsulator, error) {
    79  		return mlkem.NewEncapsulationKey1024(data)
    80  	},
    81  	pqNewPrivateKey: func(data []byte) (crypto.Decapsulator, error) {
    82  		return mlkem.NewDecapsulationKey1024(data)
    83  	},
    84  }
    85  
    86  // MLKEM1024P384 returns a KEM implementing MLKEM1024-P384 from draft-ietf-hpke-pq.
    87  func MLKEM1024P384() KEM {
    88  	return mlkem1024P384
    89  }
    90  
    91  type hybridKEM struct {
    92  	id    uint16
    93  	label string
    94  	curve ecdh.Curve
    95  
    96  	curveSeedSize    int
    97  	curvePointSize   int
    98  	pqEncapsKeySize  int
    99  	pqCiphertextSize int
   100  
   101  	pqNewPublicKey  func(data []byte) (crypto.Encapsulator, error)
   102  	pqNewPrivateKey func(data []byte) (crypto.Decapsulator, error)
   103  }
   104  
   105  func (kem *hybridKEM) ID() uint16 {
   106  	return kem.id
   107  }
   108  
   109  func (kem *hybridKEM) encSize() int {
   110  	return kem.pqCiphertextSize + kem.curvePointSize
   111  }
   112  
   113  func (kem *hybridKEM) sharedSecret(ssPQ, ssT, ctT, ekT []byte) []byte {
   114  	h := sha3.New256()
   115  	h.Write(ssPQ)
   116  	h.Write(ssT)
   117  	h.Write(ctT)
   118  	h.Write(ekT)
   119  	h.Write([]byte(kem.label))
   120  	return h.Sum(nil)
   121  }
   122  
   123  type hybridPublicKey struct {
   124  	kem *hybridKEM
   125  	t   *ecdh.PublicKey
   126  	pq  crypto.Encapsulator
   127  }
   128  
   129  // NewHybridPublicKey returns a PublicKey implementing one of
   130  //
   131  //   - MLKEM768-X25519 (a.k.a. X-Wing)
   132  //   - MLKEM768-P256
   133  //   - MLKEM1024-P384
   134  //
   135  // from draft-ietf-hpke-pq, depending on the underlying curve of t
   136  // ([ecdh.X25519], [ecdh.P256], or [ecdh.P384]) and the type of pq (either
   137  // *[mlkem.EncapsulationKey768] or *[mlkem.EncapsulationKey1024]).
   138  //
   139  // This function is meant for applications that already have instantiated
   140  // crypto/ecdh and crypto/mlkem public keys. Otherwise, applications should use
   141  // the [KEM.NewPublicKey] method of e.g. [MLKEM768X25519].
   142  func NewHybridPublicKey(pq crypto.Encapsulator, t *ecdh.PublicKey) (PublicKey, error) {
   143  	switch t.Curve() {
   144  	case ecdh.X25519():
   145  		if _, ok := pq.(*mlkem.EncapsulationKey768); !ok {
   146  			return nil, errors.New("invalid PQ KEM for X25519 hybrid")
   147  		}
   148  		return &hybridPublicKey{mlkem768X25519, t, pq}, nil
   149  	case ecdh.P256():
   150  		if _, ok := pq.(*mlkem.EncapsulationKey768); !ok {
   151  			return nil, errors.New("invalid PQ KEM for P-256 hybrid")
   152  		}
   153  		return &hybridPublicKey{mlkem768P256, t, pq}, nil
   154  	case ecdh.P384():
   155  		if _, ok := pq.(*mlkem.EncapsulationKey1024); !ok {
   156  			return nil, errors.New("invalid PQ KEM for P-384 hybrid")
   157  		}
   158  		return &hybridPublicKey{mlkem1024P384, t, pq}, nil
   159  	default:
   160  		return nil, errors.New("unsupported curve")
   161  	}
   162  }
   163  
   164  func (kem *hybridKEM) NewPublicKey(data []byte) (PublicKey, error) {
   165  	if len(data) != kem.pqEncapsKeySize+kem.curvePointSize {
   166  		return nil, errors.New("invalid public key size")
   167  	}
   168  	pq, err := kem.pqNewPublicKey(data[:kem.pqEncapsKeySize])
   169  	if err != nil {
   170  		return nil, err
   171  	}
   172  	var k *ecdh.PublicKey
   173  	fips140.WithoutEnforcement(func() { // Hybrid of ML-KEM, which is Approved.
   174  		k, err = kem.curve.NewPublicKey(data[kem.pqEncapsKeySize:])
   175  	})
   176  	if err != nil {
   177  		return nil, err
   178  	}
   179  	return NewHybridPublicKey(pq, k)
   180  }
   181  
   182  func (pk *hybridPublicKey) KEM() KEM {
   183  	return pk.kem
   184  }
   185  
   186  func (pk *hybridPublicKey) Bytes() []byte {
   187  	return append(pk.pq.Bytes(), pk.t.Bytes()...)
   188  }
   189  
   190  var testingOnlyEncapsulate func() (ss, ct []byte)
   191  
   192  func (pk *hybridPublicKey) encap() (sharedSecret []byte, encapPub []byte, err error) {
   193  	var skE *ecdh.PrivateKey
   194  	fips140.WithoutEnforcement(func() { // Hybrid of ML-KEM, which is Approved.
   195  		skE, err = pk.t.Curve().GenerateKey(rand.Reader)
   196  	})
   197  	if err != nil {
   198  		return nil, nil, err
   199  	}
   200  	if testingOnlyGenerateKey != nil {
   201  		skE = testingOnlyGenerateKey()
   202  	}
   203  	var ssT []byte
   204  	fips140.WithoutEnforcement(func() {
   205  		ssT, err = skE.ECDH(pk.t)
   206  	})
   207  	if err != nil {
   208  		return nil, nil, err
   209  	}
   210  	ctT := skE.PublicKey().Bytes()
   211  
   212  	ssPQ, ctPQ := pk.pq.Encapsulate()
   213  	if testingOnlyEncapsulate != nil {
   214  		ssPQ, ctPQ = testingOnlyEncapsulate()
   215  	}
   216  
   217  	ss := pk.kem.sharedSecret(ssPQ, ssT, ctT, pk.t.Bytes())
   218  	ct := append(ctPQ, ctT...)
   219  	return ss, ct, nil
   220  }
   221  
   222  type hybridPrivateKey struct {
   223  	kem  *hybridKEM
   224  	seed []byte // can be nil
   225  	t    ecdh.KeyExchanger
   226  	pq   crypto.Decapsulator
   227  }
   228  
   229  // NewHybridPrivateKey returns a PrivateKey implementing
   230  //
   231  //   - MLKEM768-X25519 (a.k.a. X-Wing)
   232  //   - MLKEM768-P256
   233  //   - MLKEM1024-P384
   234  //
   235  // from draft-ietf-hpke-pq, depending on the underlying curve of t
   236  // ([ecdh.X25519], [ecdh.P256], or [ecdh.P384]) and the type of pq.Encapsulator()
   237  // (either *[mlkem.EncapsulationKey768] or *[mlkem.EncapsulationKey1024]).
   238  //
   239  // This function is meant for applications that already have instantiated
   240  // crypto/ecdh and crypto/mlkem private keys, or another implementation of a
   241  // [ecdh.KeyExchanger] and [crypto.Decapsulator] (e.g. a hardware key).
   242  // Otherwise, applications should use the [KEM.NewPrivateKey] method of e.g.
   243  // [MLKEM768X25519].
   244  func NewHybridPrivateKey(pq crypto.Decapsulator, t ecdh.KeyExchanger) (PrivateKey, error) {
   245  	return newHybridPrivateKey(pq, t, nil)
   246  }
   247  
   248  func (kem *hybridKEM) GenerateKey() (PrivateKey, error) {
   249  	seed := make([]byte, 32)
   250  	drbg.Read(seed)
   251  	return kem.NewPrivateKey(seed)
   252  }
   253  
   254  func (kem *hybridKEM) NewPrivateKey(priv []byte) (PrivateKey, error) {
   255  	if len(priv) != 32 {
   256  		return nil, errors.New("hpke: invalid hybrid KEM secret length")
   257  	}
   258  
   259  	s := sha3.NewSHAKE256()
   260  	s.Write(priv)
   261  
   262  	seedPQ := make([]byte, mlkem.SeedSize)
   263  	s.Read(seedPQ)
   264  	pq, err := kem.pqNewPrivateKey(seedPQ)
   265  	if err != nil {
   266  		return nil, err
   267  	}
   268  
   269  	seedT := make([]byte, kem.curveSeedSize)
   270  	for {
   271  		s.Read(seedT)
   272  		var k ecdh.KeyExchanger
   273  		fips140.WithoutEnforcement(func() { // Hybrid of ML-KEM, which is Approved.
   274  			k, err = kem.curve.NewPrivateKey(seedT)
   275  		})
   276  		if err != nil {
   277  			continue
   278  		}
   279  		return newHybridPrivateKey(pq, k, priv)
   280  	}
   281  }
   282  
   283  func newHybridPrivateKey(pq crypto.Decapsulator, t ecdh.KeyExchanger, seed []byte) (PrivateKey, error) {
   284  	switch t.Curve() {
   285  	case ecdh.X25519():
   286  		if _, ok := pq.Encapsulator().(*mlkem.EncapsulationKey768); !ok {
   287  			return nil, errors.New("invalid PQ KEM for X25519 hybrid")
   288  		}
   289  		return &hybridPrivateKey{mlkem768X25519, bytes.Clone(seed), t, pq}, nil
   290  	case ecdh.P256():
   291  		if _, ok := pq.Encapsulator().(*mlkem.EncapsulationKey768); !ok {
   292  			return nil, errors.New("invalid PQ KEM for P-256 hybrid")
   293  		}
   294  		return &hybridPrivateKey{mlkem768P256, bytes.Clone(seed), t, pq}, nil
   295  	case ecdh.P384():
   296  		if _, ok := pq.Encapsulator().(*mlkem.EncapsulationKey1024); !ok {
   297  			return nil, errors.New("invalid PQ KEM for P-384 hybrid")
   298  		}
   299  		return &hybridPrivateKey{mlkem1024P384, bytes.Clone(seed), t, pq}, nil
   300  	default:
   301  		return nil, errors.New("unsupported curve")
   302  	}
   303  }
   304  
   305  func (kem *hybridKEM) DeriveKeyPair(ikm []byte) (PrivateKey, error) {
   306  	suiteID := byteorder.BEAppendUint16([]byte("KEM"), kem.id)
   307  	dk, err := SHAKE256().labeledDerive(suiteID, ikm, "DeriveKeyPair", nil, 32)
   308  	if err != nil {
   309  		return nil, err
   310  	}
   311  	return kem.NewPrivateKey(dk)
   312  }
   313  
   314  func (k *hybridPrivateKey) KEM() KEM {
   315  	return k.kem
   316  }
   317  
   318  func (k *hybridPrivateKey) Bytes() ([]byte, error) {
   319  	if k.seed == nil {
   320  		return nil, errors.New("private key seed not available")
   321  	}
   322  	return k.seed, nil
   323  }
   324  
   325  func (k *hybridPrivateKey) PublicKey() PublicKey {
   326  	return &hybridPublicKey{
   327  		kem: k.kem,
   328  		t:   k.t.PublicKey(),
   329  		pq:  k.pq.Encapsulator(),
   330  	}
   331  }
   332  
   333  func (k *hybridPrivateKey) decap(enc []byte) ([]byte, error) {
   334  	if len(enc) != k.kem.pqCiphertextSize+k.kem.curvePointSize {
   335  		return nil, errors.New("invalid encapsulated key size")
   336  	}
   337  	ctPQ, ctT := enc[:k.kem.pqCiphertextSize], enc[k.kem.pqCiphertextSize:]
   338  	ssPQ, err := k.pq.Decapsulate(ctPQ)
   339  	if err != nil {
   340  		return nil, err
   341  	}
   342  	var pub *ecdh.PublicKey
   343  	fips140.WithoutEnforcement(func() { // Hybrid of ML-KEM, which is Approved.
   344  		pub, err = k.t.Curve().NewPublicKey(ctT)
   345  	})
   346  	if err != nil {
   347  		return nil, err
   348  	}
   349  	var ssT []byte
   350  	fips140.WithoutEnforcement(func() {
   351  		ssT, err = k.t.ECDH(pub)
   352  	})
   353  	if err != nil {
   354  		return nil, err
   355  	}
   356  	ss := k.kem.sharedSecret(ssPQ, ssT, ctT, k.t.PublicKey().Bytes())
   357  	return ss, nil
   358  }
   359  
   360  var mlkem768 = &mlkemKEM{
   361  	id:             0x0041,
   362  	ciphertextSize: mlkem.CiphertextSize768,
   363  	newPublicKey: func(data []byte) (crypto.Encapsulator, error) {
   364  		return mlkem.NewEncapsulationKey768(data)
   365  	},
   366  	newPrivateKey: func(data []byte) (crypto.Decapsulator, error) {
   367  		return mlkem.NewDecapsulationKey768(data)
   368  	},
   369  	generateKey: func() (crypto.Decapsulator, error) {
   370  		return mlkem.GenerateKey768()
   371  	},
   372  }
   373  
   374  // MLKEM768 returns a KEM implementing ML-KEM-768 from draft-ietf-hpke-pq.
   375  func MLKEM768() KEM {
   376  	return mlkem768
   377  }
   378  
   379  var mlkem1024 = &mlkemKEM{
   380  	id:             0x0042,
   381  	ciphertextSize: mlkem.CiphertextSize1024,
   382  	newPublicKey: func(data []byte) (crypto.Encapsulator, error) {
   383  		return mlkem.NewEncapsulationKey1024(data)
   384  	},
   385  	newPrivateKey: func(data []byte) (crypto.Decapsulator, error) {
   386  		return mlkem.NewDecapsulationKey1024(data)
   387  	},
   388  	generateKey: func() (crypto.Decapsulator, error) {
   389  		return mlkem.GenerateKey1024()
   390  	},
   391  }
   392  
   393  // MLKEM1024 returns a KEM implementing ML-KEM-1024 from draft-ietf-hpke-pq.
   394  func MLKEM1024() KEM {
   395  	return mlkem1024
   396  }
   397  
   398  type mlkemKEM struct {
   399  	id             uint16
   400  	ciphertextSize int
   401  	newPublicKey   func(data []byte) (crypto.Encapsulator, error)
   402  	newPrivateKey  func(data []byte) (crypto.Decapsulator, error)
   403  	generateKey    func() (crypto.Decapsulator, error)
   404  }
   405  
   406  func (kem *mlkemKEM) ID() uint16 {
   407  	return kem.id
   408  }
   409  
   410  func (kem *mlkemKEM) encSize() int {
   411  	return kem.ciphertextSize
   412  }
   413  
   414  type mlkemPublicKey struct {
   415  	kem *mlkemKEM
   416  	pq  crypto.Encapsulator
   417  }
   418  
   419  // NewMLKEMPublicKey returns a KEMPublicKey implementing
   420  //
   421  //   - ML-KEM-768
   422  //   - ML-KEM-1024
   423  //
   424  // from draft-ietf-hpke-pq, depending on the type of pub
   425  // (*[mlkem.EncapsulationKey768] or *[mlkem.EncapsulationKey1024]).
   426  //
   427  // This function is meant for applications that already have an instantiated
   428  // crypto/mlkem public key. Otherwise, applications should use the
   429  // [KEM.NewPublicKey] method of e.g. [MLKEM768].
   430  func NewMLKEMPublicKey(pub crypto.Encapsulator) (PublicKey, error) {
   431  	switch pub.(type) {
   432  	case *mlkem.EncapsulationKey768:
   433  		return &mlkemPublicKey{mlkem768, pub}, nil
   434  	case *mlkem.EncapsulationKey1024:
   435  		return &mlkemPublicKey{mlkem1024, pub}, nil
   436  	default:
   437  		return nil, errors.New("unsupported public key type")
   438  	}
   439  }
   440  
   441  func (kem *mlkemKEM) NewPublicKey(data []byte) (PublicKey, error) {
   442  	pq, err := kem.newPublicKey(data)
   443  	if err != nil {
   444  		return nil, err
   445  	}
   446  	return NewMLKEMPublicKey(pq)
   447  }
   448  
   449  func (pk *mlkemPublicKey) KEM() KEM {
   450  	return pk.kem
   451  }
   452  
   453  func (pk *mlkemPublicKey) Bytes() []byte {
   454  	return pk.pq.Bytes()
   455  }
   456  
   457  func (pk *mlkemPublicKey) encap() (sharedSecret []byte, encapPub []byte, err error) {
   458  	ss, ct := pk.pq.Encapsulate()
   459  	if testingOnlyEncapsulate != nil {
   460  		ss, ct = testingOnlyEncapsulate()
   461  	}
   462  	return ss, ct, nil
   463  }
   464  
   465  type mlkemPrivateKey struct {
   466  	kem *mlkemKEM
   467  	pq  crypto.Decapsulator
   468  }
   469  
   470  // NewMLKEMPrivateKey returns a KEMPrivateKey implementing
   471  //
   472  //   - ML-KEM-768
   473  //   - ML-KEM-1024
   474  //
   475  // from draft-ietf-hpke-pq, depending on the type of priv.Encapsulator()
   476  // (either *[mlkem.EncapsulationKey768] or *[mlkem.EncapsulationKey1024]).
   477  //
   478  // This function is meant for applications that already have an instantiated
   479  // crypto/mlkem private key. Otherwise, applications should use the
   480  // [KEM.NewPrivateKey] method of e.g. [MLKEM768].
   481  func NewMLKEMPrivateKey(priv crypto.Decapsulator) (PrivateKey, error) {
   482  	switch priv.Encapsulator().(type) {
   483  	case *mlkem.EncapsulationKey768:
   484  		return &mlkemPrivateKey{mlkem768, priv}, nil
   485  	case *mlkem.EncapsulationKey1024:
   486  		return &mlkemPrivateKey{mlkem1024, priv}, nil
   487  	default:
   488  		return nil, errors.New("unsupported public key type")
   489  	}
   490  }
   491  
   492  func (kem *mlkemKEM) GenerateKey() (PrivateKey, error) {
   493  	pq, err := kem.generateKey()
   494  	if err != nil {
   495  		return nil, err
   496  	}
   497  	return NewMLKEMPrivateKey(pq)
   498  }
   499  
   500  func (kem *mlkemKEM) NewPrivateKey(priv []byte) (PrivateKey, error) {
   501  	pq, err := kem.newPrivateKey(priv)
   502  	if err != nil {
   503  		return nil, err
   504  	}
   505  	return NewMLKEMPrivateKey(pq)
   506  }
   507  
   508  func (kem *mlkemKEM) DeriveKeyPair(ikm []byte) (PrivateKey, error) {
   509  	suiteID := byteorder.BEAppendUint16([]byte("KEM"), kem.id)
   510  	dk, err := SHAKE256().labeledDerive(suiteID, ikm, "DeriveKeyPair", nil, 64)
   511  	if err != nil {
   512  		return nil, err
   513  	}
   514  	return kem.NewPrivateKey(dk)
   515  }
   516  
   517  func (k *mlkemPrivateKey) KEM() KEM {
   518  	return k.kem
   519  }
   520  
   521  func (k *mlkemPrivateKey) Bytes() ([]byte, error) {
   522  	pq, ok := k.pq.(interface {
   523  		Bytes() []byte
   524  	})
   525  	if !ok {
   526  		return nil, errors.New("private key seed not available")
   527  	}
   528  	return pq.Bytes(), nil
   529  }
   530  
   531  func (k *mlkemPrivateKey) PublicKey() PublicKey {
   532  	return &mlkemPublicKey{
   533  		kem: k.kem,
   534  		pq:  k.pq.Encapsulator(),
   535  	}
   536  }
   537  
   538  func (k *mlkemPrivateKey) decap(enc []byte) ([]byte, error) {
   539  	return k.pq.Decapsulate(enc)
   540  }
   541
View as plain text