kem.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  	"crypto/ecdh"
     9  	"crypto/rand"
    10  	"errors"
    11  	"internal/byteorder"
    12  )
    13  
    14  // A KEM is a Key Encapsulation Mechanism, one of the three components of an
    15  // HPKE ciphersuite.
    16  type KEM interface {
    17  	// ID returns the HPKE KEM identifier.
    18  	ID() uint16
    19  
    20  	// GenerateKey generates a new key pair.
    21  	GenerateKey() (PrivateKey, error)
    22  
    23  	// NewPublicKey deserializes a public key from bytes.
    24  	//
    25  	// It implements DeserializePublicKey, as defined in RFC 9180.
    26  	NewPublicKey([]byte) (PublicKey, error)
    27  
    28  	// NewPrivateKey deserializes a private key from bytes.
    29  	//
    30  	// It implements DeserializePrivateKey, as defined in RFC 9180.
    31  	NewPrivateKey([]byte) (PrivateKey, error)
    32  
    33  	// DeriveKeyPair derives a key pair from the given input keying material.
    34  	//
    35  	// It implements DeriveKeyPair, as defined in RFC 9180.
    36  	DeriveKeyPair(ikm []byte) (PrivateKey, error)
    37  
    38  	encSize() int
    39  }
    40  
    41  // NewKEM returns the KEM implementation for the given KEM ID.
    42  //
    43  // Applications are encouraged to use specific implementations like [DHKEM] or
    44  // [MLKEM768X25519] instead, unless runtime agility is required.
    45  func NewKEM(id uint16) (KEM, error) {
    46  	switch id {
    47  	case 0x0010: // DHKEM(P-256, HKDF-SHA256)
    48  		return DHKEM(ecdh.P256()), nil
    49  	case 0x0011: // DHKEM(P-384, HKDF-SHA384)
    50  		return DHKEM(ecdh.P384()), nil
    51  	case 0x0012: // DHKEM(P-521, HKDF-SHA512)
    52  		return DHKEM(ecdh.P521()), nil
    53  	case 0x0020: // DHKEM(X25519, HKDF-SHA256)
    54  		return DHKEM(ecdh.X25519()), nil
    55  	case 0x0041: // ML-KEM-768
    56  		return MLKEM768(), nil
    57  	case 0x0042: // ML-KEM-1024
    58  		return MLKEM1024(), nil
    59  	case 0x647a: // MLKEM768-X25519
    60  		return MLKEM768X25519(), nil
    61  	case 0x0050: // MLKEM768-P256
    62  		return MLKEM768P256(), nil
    63  	case 0x0051: // MLKEM1024-P384
    64  		return MLKEM1024P384(), nil
    65  	default:
    66  		return nil, errors.New("unsupported KEM")
    67  	}
    68  }
    69  
    70  // A PublicKey is an instantiation of a KEM (one of the three components of an
    71  // HPKE ciphersuite) with an encapsulation key (i.e. the public key).
    72  //
    73  // A PublicKey is usually obtained from a method of the corresponding [KEM] or
    74  // [PrivateKey], such as [KEM.NewPublicKey] or [PrivateKey.PublicKey].
    75  type PublicKey interface {
    76  	// KEM returns the instantiated KEM.
    77  	KEM() KEM
    78  
    79  	// Bytes returns the public key as the output of SerializePublicKey.
    80  	Bytes() []byte
    81  
    82  	encap() (sharedSecret, enc []byte, err error)
    83  }
    84  
    85  // A PrivateKey is an instantiation of a KEM (one of the three components of
    86  // an HPKE ciphersuite) with a decapsulation key (i.e. the secret key).
    87  //
    88  // A PrivateKey is usually obtained from a method of the corresponding [KEM],
    89  // such as [KEM.GenerateKey] or [KEM.NewPrivateKey].
    90  type PrivateKey interface {
    91  	// KEM returns the instantiated KEM.
    92  	KEM() KEM
    93  
    94  	// Bytes returns the private key as the output of SerializePrivateKey, as
    95  	// defined in RFC 9180.
    96  	//
    97  	// Note that for X25519 this might not match the input to NewPrivateKey.
    98  	// This is a requirement of RFC 9180, Section 7.1.2.
    99  	Bytes() ([]byte, error)
   100  
   101  	// PublicKey returns the corresponding PublicKey.
   102  	PublicKey() PublicKey
   103  
   104  	decap(enc []byte) (sharedSecret []byte, err error)
   105  }
   106  
   107  type dhKEM struct {
   108  	kdf     KDF
   109  	id      uint16
   110  	curve   ecdh.Curve
   111  	Nsecret uint16
   112  	Nsk     uint16
   113  	Nenc    int
   114  }
   115  
   116  func (kem *dhKEM) extractAndExpand(dhKey, kemContext []byte) ([]byte, error) {
   117  	suiteID := byteorder.BEAppendUint16([]byte("KEM"), kem.id)
   118  	eaePRK, err := kem.kdf.labeledExtract(suiteID, nil, "eae_prk", dhKey)
   119  	if err != nil {
   120  		return nil, err
   121  	}
   122  	return kem.kdf.labeledExpand(suiteID, eaePRK, "shared_secret", kemContext, kem.Nsecret)
   123  }
   124  
   125  func (kem *dhKEM) ID() uint16 {
   126  	return kem.id
   127  }
   128  
   129  func (kem *dhKEM) encSize() int {
   130  	return kem.Nenc
   131  }
   132  
   133  var dhKEMP256 = &dhKEM{HKDFSHA256(), 0x0010, ecdh.P256(), 32, 32, 65}
   134  var dhKEMP384 = &dhKEM{HKDFSHA384(), 0x0011, ecdh.P384(), 48, 48, 97}
   135  var dhKEMP521 = &dhKEM{HKDFSHA512(), 0x0012, ecdh.P521(), 64, 66, 133}
   136  var dhKEMX25519 = &dhKEM{HKDFSHA256(), 0x0020, ecdh.X25519(), 32, 32, 32}
   137  
   138  // DHKEM returns a KEM implementing one of
   139  //
   140  //   - DHKEM(P-256, HKDF-SHA256)
   141  //   - DHKEM(P-384, HKDF-SHA384)
   142  //   - DHKEM(P-521, HKDF-SHA512)
   143  //   - DHKEM(X25519, HKDF-SHA256)
   144  //
   145  // depending on curve.
   146  func DHKEM(curve ecdh.Curve) KEM {
   147  	switch curve {
   148  	case ecdh.P256():
   149  		return dhKEMP256
   150  	case ecdh.P384():
   151  		return dhKEMP384
   152  	case ecdh.P521():
   153  		return dhKEMP521
   154  	case ecdh.X25519():
   155  		return dhKEMX25519
   156  	default:
   157  		// The set of ecdh.Curve implementations is closed, because the
   158  		// interface has unexported methods. Therefore, this default case is
   159  		// only hit if a new curve is added that DHKEM doesn't support.
   160  		return unsupportedCurveKEM{}
   161  	}
   162  }
   163  
   164  type unsupportedCurveKEM struct{}
   165  
   166  func (unsupportedCurveKEM) ID() uint16 {
   167  	return 0
   168  }
   169  func (unsupportedCurveKEM) GenerateKey() (PrivateKey, error) {
   170  	return nil, errors.New("unsupported curve")
   171  }
   172  func (unsupportedCurveKEM) NewPublicKey([]byte) (PublicKey, error) {
   173  	return nil, errors.New("unsupported curve")
   174  }
   175  func (unsupportedCurveKEM) NewPrivateKey([]byte) (PrivateKey, error) {
   176  	return nil, errors.New("unsupported curve")
   177  }
   178  func (unsupportedCurveKEM) DeriveKeyPair([]byte) (PrivateKey, error) {
   179  	return nil, errors.New("unsupported curve")
   180  }
   181  func (unsupportedCurveKEM) encSize() int {
   182  	return 0
   183  }
   184  
   185  type dhKEMPublicKey struct {
   186  	kem *dhKEM
   187  	pub *ecdh.PublicKey
   188  }
   189  
   190  // NewDHKEMPublicKey returns a PublicKey implementing
   191  //
   192  //   - DHKEM(P-256, HKDF-SHA256)
   193  //   - DHKEM(P-384, HKDF-SHA384)
   194  //   - DHKEM(P-521, HKDF-SHA512)
   195  //   - DHKEM(X25519, HKDF-SHA256)
   196  //
   197  // depending on the underlying curve of pub ([ecdh.X25519], [ecdh.P256],
   198  // [ecdh.P384], or [ecdh.P521]).
   199  //
   200  // This function is meant for applications that already have an instantiated
   201  // crypto/ecdh public key. Otherwise, applications should use the
   202  // [KEM.NewPublicKey] method of [DHKEM].
   203  func NewDHKEMPublicKey(pub *ecdh.PublicKey) (PublicKey, error) {
   204  	kem, ok := DHKEM(pub.Curve()).(*dhKEM)
   205  	if !ok {
   206  		return nil, errors.New("unsupported curve")
   207  	}
   208  	return &dhKEMPublicKey{
   209  		kem: kem,
   210  		pub: pub,
   211  	}, nil
   212  }
   213  
   214  func (kem *dhKEM) NewPublicKey(data []byte) (PublicKey, error) {
   215  	pub, err := kem.curve.NewPublicKey(data)
   216  	if err != nil {
   217  		return nil, err
   218  	}
   219  	return NewDHKEMPublicKey(pub)
   220  }
   221  
   222  func (pk *dhKEMPublicKey) KEM() KEM {
   223  	return pk.kem
   224  }
   225  
   226  func (pk *dhKEMPublicKey) Bytes() []byte {
   227  	return pk.pub.Bytes()
   228  }
   229  
   230  // testingOnlyGenerateKey is only used during testing, to provide
   231  // a fixed test key to use when checking the RFC 9180 vectors.
   232  var testingOnlyGenerateKey func() *ecdh.PrivateKey
   233  
   234  func (pk *dhKEMPublicKey) encap() (sharedSecret []byte, encapPub []byte, err error) {
   235  	privEph, err := pk.pub.Curve().GenerateKey(rand.Reader)
   236  	if err != nil {
   237  		return nil, nil, err
   238  	}
   239  	if testingOnlyGenerateKey != nil {
   240  		privEph = testingOnlyGenerateKey()
   241  	}
   242  	dhVal, err := privEph.ECDH(pk.pub)
   243  	if err != nil {
   244  		return nil, nil, err
   245  	}
   246  	encPubEph := privEph.PublicKey().Bytes()
   247  
   248  	encPubRecip := pk.pub.Bytes()
   249  	kemContext := append(encPubEph, encPubRecip...)
   250  	sharedSecret, err = pk.kem.extractAndExpand(dhVal, kemContext)
   251  	if err != nil {
   252  		return nil, nil, err
   253  	}
   254  	return sharedSecret, encPubEph, nil
   255  }
   256  
   257  type dhKEMPrivateKey struct {
   258  	kem  *dhKEM
   259  	priv ecdh.KeyExchanger
   260  }
   261  
   262  // NewDHKEMPrivateKey returns a PrivateKey implementing
   263  //
   264  //   - DHKEM(P-256, HKDF-SHA256)
   265  //   - DHKEM(P-384, HKDF-SHA384)
   266  //   - DHKEM(P-521, HKDF-SHA512)
   267  //   - DHKEM(X25519, HKDF-SHA256)
   268  //
   269  // depending on the underlying curve of priv ([ecdh.X25519], [ecdh.P256],
   270  // [ecdh.P384], or [ecdh.P521]).
   271  //
   272  // This function is meant for applications that already have an instantiated
   273  // crypto/ecdh private key, or another implementation of a [ecdh.KeyExchanger]
   274  // (e.g. a hardware key). Otherwise, applications should use the
   275  // [KEM.NewPrivateKey] method of [DHKEM].
   276  func NewDHKEMPrivateKey(priv ecdh.KeyExchanger) (PrivateKey, error) {
   277  	kem, ok := DHKEM(priv.Curve()).(*dhKEM)
   278  	if !ok {
   279  		return nil, errors.New("unsupported curve")
   280  	}
   281  	return &dhKEMPrivateKey{
   282  		kem:  kem,
   283  		priv: priv,
   284  	}, nil
   285  }
   286  
   287  func (kem *dhKEM) GenerateKey() (PrivateKey, error) {
   288  	priv, err := kem.curve.GenerateKey(rand.Reader)
   289  	if err != nil {
   290  		return nil, err
   291  	}
   292  	return NewDHKEMPrivateKey(priv)
   293  }
   294  
   295  func (kem *dhKEM) NewPrivateKey(ikm []byte) (PrivateKey, error) {
   296  	priv, err := kem.curve.NewPrivateKey(ikm)
   297  	if err != nil {
   298  		return nil, err
   299  	}
   300  	return NewDHKEMPrivateKey(priv)
   301  }
   302  
   303  func (kem *dhKEM) DeriveKeyPair(ikm []byte) (PrivateKey, error) {
   304  	// DeriveKeyPair from RFC 9180 Section 7.1.3.
   305  	suiteID := byteorder.BEAppendUint16([]byte("KEM"), kem.id)
   306  	prk, err := kem.kdf.labeledExtract(suiteID, nil, "dkp_prk", ikm)
   307  	if err != nil {
   308  		return nil, err
   309  	}
   310  	if kem == dhKEMX25519 {
   311  		s, err := kem.kdf.labeledExpand(suiteID, prk, "sk", nil, kem.Nsk)
   312  		if err != nil {
   313  			return nil, err
   314  		}
   315  		return kem.NewPrivateKey(s)
   316  	}
   317  	var counter uint8
   318  	for counter < 4 {
   319  		s, err := kem.kdf.labeledExpand(suiteID, prk, "candidate", []byte{counter}, kem.Nsk)
   320  		if err != nil {
   321  			return nil, err
   322  		}
   323  		if kem == dhKEMP521 {
   324  			s[0] &= 0x01
   325  		}
   326  		r, err := kem.NewPrivateKey(s)
   327  		if err != nil {
   328  			counter++
   329  			continue
   330  		}
   331  		return r, nil
   332  	}
   333  	panic("chance of four rejections is < 2^-128")
   334  }
   335  
   336  func (k *dhKEMPrivateKey) KEM() KEM {
   337  	return k.kem
   338  }
   339  
   340  func (k *dhKEMPrivateKey) Bytes() ([]byte, error) {
   341  	// Bizarrely, RFC 9180, Section 7.1.2 says SerializePrivateKey MUST clamp
   342  	// the output, which I thought we all agreed to instead do as part of the DH
   343  	// function, letting private keys be random bytes.
   344  	//
   345  	// At the same time, it says DeserializePrivateKey MUST also clamp, implying
   346  	// that the input doesn't have to be clamped, so Bytes by spec doesn't
   347  	// necessarily match the NewPrivateKey input.
   348  	//
   349  	// I'm sure this will not lead to any unexpected behavior or interop issue.
   350  	priv, ok := k.priv.(*ecdh.PrivateKey)
   351  	if !ok {
   352  		return nil, errors.New("ecdh: private key does not support Bytes")
   353  	}
   354  	if k.kem == dhKEMX25519 {
   355  		b := priv.Bytes()
   356  		b[0] &= 248
   357  		b[31] &= 127
   358  		b[31] |= 64
   359  		return b, nil
   360  	}
   361  	return priv.Bytes(), nil
   362  }
   363  
   364  func (k *dhKEMPrivateKey) PublicKey() PublicKey {
   365  	return &dhKEMPublicKey{
   366  		kem: k.kem,
   367  		pub: k.priv.PublicKey(),
   368  	}
   369  }
   370  
   371  func (k *dhKEMPrivateKey) decap(encPubEph []byte) ([]byte, error) {
   372  	pubEph, err := k.priv.Curve().NewPublicKey(encPubEph)
   373  	if err != nil {
   374  		return nil, err
   375  	}
   376  	dhVal, err := k.priv.ECDH(pubEph)
   377  	if err != nil {
   378  		return nil, err
   379  	}
   380  	kemContext := append(encPubEph, k.priv.PublicKey().Bytes()...)
   381  	return k.kem.extractAndExpand(dhVal, kemContext)
   382  }
   383
View as plain text