Add ML-KEM agreement

encryption.go

@@ -13,6 +13,15 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+// EncryptionKeyMaterial holds material for encryption post-derivation, so
+// regardless of X25519 or ML-KEM (or future types) this is the resulting
+// struct
+type EncryptionKeyMaterial struct {
+	KeyId     string
+	KeyType   uint
+	SharedKey []byte
+}
+
 // X25519KeyProducer is an interface that can be satisfied by an underlying type
 // that produces an encryption key via X25519, along with a key identifier used
 // for AAD and embedding in the wrapping data. If the ID is empty, it is simply
@@ -22,38 +31,33 @@ type X25519KeyProducer interface {
 	PreviousX25519EncryptionKey() (string, []byte, error)
 }
 
-// EncryptMessage takes any proto.Message and a valid key source that implements
-// X25519KeyProducer. Internally it uses an `aead` wrapper from go-kms-wrapping v2.
-// No options are currently supported but in the future non-AES-GCM encryption
-// types could be supported by the wrapper and chosen here.
-//
-// The resulting value from the wrapper is marshaled before being returned.
-//
-// Supported options: WithRandomReader
-func EncryptMessage(ctx context.Context, msg proto.Message, keySource X25519KeyProducer, opt ...Option) ([]byte, error) {
-	const op = "nodeenrollment.EncryptMessage"
-	switch {
-	case IsNil(msg):
-		return nil, fmt.Errorf("(%s) incoming message is nil", op)
-	case IsNil(keySource):
-		return nil, fmt.Errorf("(%s) incoming key source is nil", op)
-	}
+// KeyProducer exposes shared keys for current and previous encryption
+// material.
+type KeyProducer interface {
+	CurrentSharedEncryptionKey() (EncryptionKeyMaterial, error)
+	PreviousSharedEncryptionKey() (EncryptionKeyMaterial, error)
+}
 
-	opts, err := GetOpts(opt...)
+func encryptMessageWithKeyProducer(ctx context.Context, msg proto.Message, keySource KeyProducer, opt ...Option) ([]byte, error) {
+	const op = "nodeenrollment.encryptMessageWithKeyProducer"
+	selectedKey, err := keySource.CurrentSharedEncryptionKey()
 	if err != nil {
-		return nil, fmt.Errorf("(%s) error parsing options: %w", op, err)
+		return nil, fmt.Errorf("(%s) error deriving shared encryption key: %w", op, err)
+	}
+	if len(selectedKey.SharedKey) == 0 {
+		return nil, fmt.Errorf("(%s) no shared encryption keys available", op)
 	}
 
-	keyId, sharedKey, err := keySource.X25519EncryptionKey()
+	opts, err := GetOpts(opt...)
 	if err != nil {
-		return nil, fmt.Errorf("(%s) error deriving shared encryption key: %w", op, err)
+		return nil, fmt.Errorf("(%s) error parsing options: %w", op, err)
 	}
 
 	aeadWrapper := aead.NewWrapper()
 	if _, err := aeadWrapper.SetConfig(
 		ctx,
-		wrapping.WithKeyId(keyId),
-		aead.WithKey(sharedKey),
+		wrapping.WithKeyId(selectedKey.KeyId),
+		aead.WithKey(selectedKey.SharedKey),
 		aead.WithRandomReader(opts.WithRandomReader),
 	); err != nil {
 		return nil, fmt.Errorf("(%s) error instantiating aead wrapper: %w", op, err)
@@ -65,8 +69,8 @@ func EncryptMessage(ctx context.Context, msg proto.Message, keySource X25519KeyP
 	}
 
 	var aadOpt wrapping.Option
-	if keyId != "" {
-		aadOpt = wrapping.WithAad([]byte(keyId))
+	if selectedKey.KeyId != "" {
+		aadOpt = wrapping.WithAad([]byte(selectedKey.KeyId))
 	}
 	blobInfo, err := aeadWrapper.Encrypt(ctx, marshaledMsg, aadOpt)
 	if err != nil {
@@ -81,8 +85,53 @@ func EncryptMessage(ctx context.Context, msg proto.Message, keySource X25519KeyP
 	return marshaledBlob, nil
 }
 
-// DecryptMessage takes any a value encrypted with EncryptMessage and a valid
-// key source that implements X25519KeyProducer and decrypts the message into the
+func decryptMessageWithKeyProducer(ctx context.Context, ct []byte, keySource KeyProducer, result proto.Message) error {
+	const op = "nodeenrollment.decryptMessageWithKeyProducer"
+	currentKey, err := keySource.CurrentSharedEncryptionKey()
+	if err != nil {
+		return fmt.Errorf("(%s) error deriving shared encryption key: %w", op, err)
+	}
+
+	err = decryptWithKey(ctx, currentKey.KeyId, ct, currentKey.SharedKey, result)
+	if err == nil {
+		return nil
+	}
+
+	previousKey, prevErr := keySource.PreviousSharedEncryptionKey()
+	if prevErr != nil || len(previousKey.SharedKey) == 0 {
+		return err
+	}
+
+	prevErr = decryptWithKey(ctx, previousKey.KeyId, ct, previousKey.SharedKey, result)
+	if prevErr != nil {
+		return errors.Join(err, fmt.Errorf("(%s) error decrypting with previous key: %w", op, prevErr))
+	}
+
+	return nil
+}
+
+// EncryptMessage takes any proto.Message and a valid key source that implements
+// KeyProducer. Internally it uses an `aead` wrapper from go-kms-wrapping v2.
+// No options are currently supported but in the future non-AES-GCM encryption
+// types could be supported by the wrapper and chosen here.
+//
+// The resulting value from the wrapper is marshaled before being returned.
+//
+// Supported options: WithRandomReader
+func EncryptMessage(ctx context.Context, msg proto.Message, keySource KeyProducer, opt ...Option) ([]byte, error) {
+	const op = "nodeenrollment.EncryptMessage"
+	switch {
+	case IsNil(msg):
+		return nil, fmt.Errorf("(%s) incoming message is nil", op)
+	case IsNil(keySource):
+		return nil, fmt.Errorf("(%s) incoming key source is nil", op)
+	}
+
+	return encryptMessageWithKeyProducer(ctx, msg, keySource, opt...)
+}
+
+// DecryptMessage takes a value encrypted with EncryptMessage and a valid
+// key source that implements KeyProducer and decrypts the message into the
 // given proto.Message. Internally it uses an `aead` wrapper from
 // go-kms-wrapping v2. No options are currently supported but in the future
 // non-AES-GCM decryption types could be supported by the wrapper and chosen
@@ -95,7 +144,7 @@ func EncryptMessage(ctx context.Context, msg proto.Message, keySource X25519KeyP
 // If decryption fails with the current key, and a prior key is present,
 // use that to try and decrypt the message in the case an older key
 // was used to encrypt the incoming message
-func DecryptMessage(ctx context.Context, ct []byte, keySource X25519KeyProducer, result proto.Message, _ ...Option) error {
+func DecryptMessage(ctx context.Context, ct []byte, keySource KeyProducer, result proto.Message, _ ...Option) error {
 	const op = "nodeenrollment.DecryptMessage"
 	switch {
 	case len(ct) == 0:
@@ -106,27 +155,7 @@ func DecryptMessage(ctx context.Context, ct []byte, keySource X25519KeyProducer,
 		return fmt.Errorf("(%s) incoming result message is nil", op)
 	}
 
-	keyId, sharedKey, err := keySource.X25519EncryptionKey()
-	if err != nil {
-		return fmt.Errorf("(%s) error deriving shared encryption key: %w", op, err)
-	}
-
-	err = decryptWithKey(ctx, keyId, ct, sharedKey, result)
-
-	// If decryption fails with the current key, try with the previous key, if present
-	if err != nil {
-		prevId, previousKey, prevErr := keySource.PreviousX25519EncryptionKey()
-		if prevErr != nil || previousKey == nil {
-			return err
-		}
-		prevErr = decryptWithKey(ctx, prevId, ct, previousKey, result)
-		if prevErr != nil {
-			err = errors.Join(err, fmt.Errorf("(%s) error decrypting with previous key: %w", op, prevErr))
-			return err
-		}
-	}
-
-	return nil
+	return decryptMessageWithKeyProducer(ctx, ct, keySource, result)
 }
 
 func decryptWithKey(ctx context.Context, keyId string, ct []byte, sharedKey []byte, result proto.Message) error {

encryption_test.go

@@ -4,6 +4,7 @@
 package nodeenrollment
 
 import (
+	"bytes"
 	"context"
 	"crypto/ecdh"
 	"crypto/rand"
@@ -24,6 +25,11 @@ type testNode struct {
 	otherPub []byte
 }
 
+type orderedTestKeyProducer struct {
+	current  EncryptionKeyMaterial
+	previous EncryptionKeyMaterial
+}
+
 func (t testNode) X25519EncryptionKey() (string, []byte, error) {
 	privKey, err := ecdh.X25519().NewPrivateKey(t.priv)
 	if err != nil {
@@ -44,7 +50,45 @@ func (t testNode) PreviousX25519EncryptionKey() (string, []byte, error) {
 	return "", nil, nil
 }
 
-var _ X25519KeyProducer = (*testNode)(nil)
+func (t testNode) CurrentSharedEncryptionKey() (EncryptionKeyMaterial, error) {
+	keyId, sharedKey, err := t.X25519EncryptionKey()
+	if err != nil {
+		return EncryptionKeyMaterial{}, err
+	}
+
+	return EncryptionKeyMaterial{
+		KeyId:     keyId,
+		SharedKey: sharedKey,
+	}, nil
+}
+
+func (t testNode) PreviousSharedEncryptionKey() (EncryptionKeyMaterial, error) {
+	keyId, sharedKey, err := t.PreviousX25519EncryptionKey()
+	if err != nil {
+		return EncryptionKeyMaterial{}, err
+	}
+	if sharedKey == nil {
+		return EncryptionKeyMaterial{}, nil
+	}
+
+	return EncryptionKeyMaterial{
+		KeyId:     keyId,
+		SharedKey: sharedKey,
+	}, nil
+}
+
+func (t orderedTestKeyProducer) CurrentSharedEncryptionKey() (EncryptionKeyMaterial, error) {
+	return t.current, nil
+}
+
+func (t orderedTestKeyProducer) PreviousSharedEncryptionKey() (EncryptionKeyMaterial, error) {
+	return t.previous, nil
+}
+
+var (
+	_ X25519KeyProducer = (*testNode)(nil)
+	_ KeyProducer       = (*testNode)(nil)
+)
 
 func Test_EncryptionDecryption(t *testing.T) {
 	t.Parallel()
@@ -80,8 +124,8 @@ func Test_EncryptionDecryption(t *testing.T) {
 		decryptId          string
 		encryptMsg         proto.Message
 		decryptMsg         proto.Message
-		encryptKeySource   X25519KeyProducer
-		decryptKeySource   X25519KeyProducer
+		encryptKeySource   KeyProducer
+		decryptKeySource   KeyProducer
 		encDecWrapper      wrapping.Wrapper
 		wantErrContains    string
 		wantEncErrContains string
@@ -213,3 +257,48 @@ func Test_EncryptionDecryption(t *testing.T) {
 		})
 	}
 }
+
+func Test_encryptMessageWithKeyProducer_UsesCurrentKey(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	currentKey := bytes.Repeat([]byte{1}, 32)
+	msg := &wrapping.BlobInfo{
+		Ciphertext: []byte("foo"),
+		Iv:         []byte("bar"),
+		Hmac:       []byte("baz"),
+	}
+
+	ct, err := encryptMessageWithKeyProducer(ctx, msg, orderedTestKeyProducer{
+		current: EncryptionKeyMaterial{KeyId: "current", SharedKey: currentKey},
+	})
+	require.NoError(t, err)
+
+	decryptedMsg := new(wrapping.BlobInfo)
+	require.NoError(t, decryptWithKey(ctx, "current", ct, currentKey, decryptedMsg))
+	assert.Empty(t, cmp.Diff(msg, decryptedMsg, protocmp.Transform()))
+}
+
+func Test_decryptMessageWithKeyProducer_FallsBackToPreviousKeys(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	previousKey := bytes.Repeat([]byte{3}, 32)
+	currentKey := bytes.Repeat([]byte{4}, 32)
+	msg := &wrapping.BlobInfo{
+		Ciphertext: []byte("foo"),
+		Iv:         []byte("bar"),
+		Hmac:       []byte("baz"),
+	}
+
+	ct, err := encryptMessageWithKeyProducer(ctx, msg, orderedTestKeyProducer{
+		current: EncryptionKeyMaterial{KeyId: "previous", SharedKey: previousKey},
+	})
+	require.NoError(t, err)
+
+	decryptedMsg := new(wrapping.BlobInfo)
+	err = decryptMessageWithKeyProducer(ctx, ct, orderedTestKeyProducer{
+		current:  EncryptionKeyMaterial{KeyId: "current", SharedKey: currentKey},
+		previous: EncryptionKeyMaterial{KeyId: "previous", SharedKey: previousKey},
+	}, decryptedMsg)
+	require.NoError(t, err)
+	assert.Empty(t, cmp.Diff(msg, decryptedMsg, protocmp.Transform()))
+}

options.go

@@ -6,12 +6,14 @@ package nodeenrollment
 import (
 	"crypto/rand"
 	"crypto/x509"
+	"errors"
 	"fmt"
 	"io"
 	"time"
 
 	"github.com/hashicorp/go-hclog"
 	wrapping "github.com/hashicorp/go-kms-wrapping/v2"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -50,8 +52,9 @@ type Options struct {
 	WithExtraAlpnProtos                                   []string
 	WithReinitializeRoots                                 bool
 	WithActivationToken                                   string
-	WithPrivateKey                                        []byte
-	WithPrivateKeyType                                    uint
+	WithEncryptionPrivateKey                              []byte
+	WithEncryptionPrivateKeyType                          uint
+	WithMlkemParameters                                   proto.Message
 	WithoutRegistrationChallenge                          bool
 	WithMaximumServerLedActivationTokenLifetime           time.Duration
 	WithNativeConns                                       bool
@@ -245,11 +248,33 @@ func WithActivationToken(with string) Option {
 	}
 }
 
-// WithPrivateKey allows indicating a private key to be used for signing
-func WithPrivateKey(withKey []byte, withType uint) Option {
+// WithEncryptionPrivateKey allows indicating a private key to be used for encryption
+func WithEncryptionPrivateKey(withKey []byte, withType uint) Option {
 	return func(o *Options) error {
-		o.WithPrivateKey = withKey
-		o.WithPrivateKeyType = withType
+		o.WithEncryptionPrivateKey = withKey
+		o.WithEncryptionPrivateKeyType = withType
+		return nil
+	}
+}
+
+// WithEncryptionPrivateKeyType allows indicating a type of private key for generation
+func WithEncryptionPrivateKeyType(withType uint) Option {
+	return func(o *Options) error {
+		o.WithEncryptionPrivateKeyType = withType
+		return nil
+	}
+}
+
+// WithMlkemParameters allows passing existing MLKEM parameters
+func WithMlkemParameters(with proto.Message) Option {
+	return func(o *Options) error {
+		if IsNil(with) {
+			return errors.New("mlkem parameters cannot be nil")
+		}
+		if with.ProtoReflect().Descriptor().FullName() != "github.com.hashicorp.nodeenrollment.types.v1.MLKEMParameters" {
+			return fmt.Errorf("mlkem parameters must be github.com.hashicorp.nodeenrollment.types.v1.MLKEMParameters, got %s", with.ProtoReflect().Descriptor().FullName())
+		}
+		o.WithMlkemParameters = with
 		return nil
 	}
 }