[kms]: Add KeyMaterial to avoid hand-rolled framing

api/ext/v1/key_material.proto

@@ -0,0 +1,41 @@
+syntax = "proto3";
+
+package api.ext.v1;
+
+option go_package = "github.com/temporalio/temporal-proxy/pkg/api/ext/v1;ext";
+
+// KeyMaterial is an optional framing for the ciphertext an extension server
+// returns from Encrypt. Decrypt is handed nothing but that ciphertext, so
+// anything the server needs to find the wrapping key again has to travel inside
+// it. Rather than hand-roll a binary frame, a server can marshal this message as
+// its ciphertext and unmarshal it on the way back.
+//
+// The fields other than encrypted_dek are metadata in the clear: this is
+// framing, not encryption. A namespace is not a secret (it already travels in
+// gRPC metadata), but a server that would rather not expose one can leave
+// namespace empty and identify its key through opaque instead.
+//
+// Every byte of this message rides in the metadata of each payload it seals,
+// and stays there for as long as the Workflow retention period. Carry what
+// decrypt needs to find the key, and nothing else.
+message KeyMaterial {
+  // encrypted_dek is the wrapped DEK, as produced by whatever wrapping the
+  // extension server performs. It is the one field that must be set.
+  bytes encrypted_dek = 1;
+  // version identifies which version of the wrapping key sealed
+  // encrypted_dek, so a server that rotates key material can select the same
+  // version again on decrypt. Empty means the server does not version its
+  // keys.
+  string version = 2;
+  // namespace is the pre-translation (local) namespace the DEK belongs to,
+  // copied from EncryptRequest.namespace. A server that derives a key per
+  // namespace needs it to derive the same key on decrypt, where the request no
+  // longer carries it. Empty means the server does not key off the namespace.
+  string namespace = 3;
+  // opaque belongs to the extension server. Neither the proxy nor this message
+  // gives it any meaning: it is round-tripped untouched, so a server that needs
+  // more than version and namespace to find its key can put its own encoding
+  // here instead of replacing this framing wholesale. Whatever goes in owns its
+  // own compatibility, since the proxy cannot migrate what it cannot read.
+  bytes opaque = 4;
+}

examples/kms/server/keyring.go

@@ -7,31 +7,28 @@ import (
 	"crypto/hkdf"
 	"crypto/rand"
 	"crypto/sha256"
-	"encoding/binary"
 	"errors"
 	"fmt"
-	"math"
+
+	extv1 "github.com/temporalio/temporal-proxy/pkg/api/ext/v1"
 )
 
 const (
-	// formatVersion prefixes every ciphertext so a later change to the framing is
-	// rejected rather than mis-parsed.
-	formatVersion = 0x01
-
-	// headerSize is the fixed part of the frame: the version byte plus the uint16
-	// namespace length that follows it.
-	headerSize = 3
+	// currentVersion is the wrapping key version Wrap stamps into new key
+	// material. Unwrap derives from whatever version it is handed, so bumping
+	// this rotates new payloads without stranding the ones already sealed.
+	currentVersion = "1"
 
 	// keySize selects AES-256.
 	keySize = 32
 
 	// infoPrefix domain-separates these derived keys from any other use of the
 	// same master secret.
-	infoPrefix = "temporal-proxy-kek/v1/"
+	infoPrefix = "temporal-proxy-kek"
 )
 
-// keyring derives one wrapping key per namespace from a master secret, so a
-// compromise of one namespace's key does not hand over the others.
+// keyring derives one wrapping key per version and namespace from a master
+// secret, so a compromise of one namespace's key does not hand over the others.
 type keyring struct {
 	secret []byte
 }
@@ -52,78 +49,83 @@ func newKeyring(secret []byte) (*keyring, error) {
 	return &keyring{secret: secret}, nil
 }
 
-// Wrap seals dek under the key derived for namespace.
+// Wrap seals dek under the key derived for namespace and returns it as
+// [extv1.KeyMaterial], which carries everything Unwrap needs to derive the same
+// key again.
 //
-// The namespace is written into the frame in the clear because Unwrap is handed
-// nothing but ciphertext and has to derive the same key again. It doubles as the
-// GCM additional data, so a ciphertext relabelled with another namespace fails
-// to open. A namespace is not a secret (it already travels in gRPC metadata),
-// but a provider that would rather not expose one should carry an opaque key
-// identifier here and resolve it internally.
+// The namespace travels in the clear, and doubles as the GCM additional data so
+// key material relabelled with another namespace fails to open. A namespace is
+// not a secret (it already travels in gRPC metadata), but a provider that would
+// rather not expose one can leave the field empty and identify its key through
+// opaque instead.
 func (k *keyring) Wrap(_ context.Context, namespace string, dek []byte) ([]byte, error) {
-	if len(namespace) > math.MaxUint16 {
-		return nil, fmt.Errorf("server: namespace is too long to frame: %d bytes", len(namespace))
-	}
-
-	gcm, err := k.cipher(namespace)
+	gcm, err := k.cipher(currentVersion, namespace)
 	if err != nil {
 		return nil, err
 	}
 
-	out := make([]byte, 0, headerSize+len(namespace)+gcm.NonceSize()+len(dek)+gcm.Overhead())
-	out = append(out, formatVersion)
-	out = binary.BigEndian.AppendUint16(out, uint16(len(namespace)))
-	out = append(out, namespace...)
-
 	nonce := make([]byte, gcm.NonceSize())
 	// crypto/rand.Read never returns an error; it crashes the program instead.
 	_, _ = rand.Read(nonce)
-	out = append(out, nonce...)
-
-	return gcm.Seal(out, nonce, dek, []byte(namespace)), nil
-}
 
-// Unwrap opens a ciphertext produced by Wrap, deriving the key from the
-// namespace the frame carries.
-func (k *keyring) Unwrap(_ context.Context, ciphertext []byte) ([]byte, error) {
-	if len(ciphertext) < headerSize {
-		return nil, errors.New("server: ciphertext is too short to hold a header")
+	material := &extv1.KeyMaterial{
+		EncryptedDek: gcm.Seal(nil, nonce, dek, []byte(namespace)),
+		Version:      currentVersion,
+		Namespace:    namespace,
+		// KeyMaterial has no field for a nonce, and needs none: opaque is where a
+		// server puts whatever its own wrapping requires. A nonce is not secret,
+		// only single-use.
+		Opaque: nonce,
 	}
 
-	if ciphertext[0] != formatVersion {
-		return nil, fmt.Errorf("server: unsupported ciphertext version: %#x", ciphertext[0])
+	packed, err := material.Marshal()
+	if err != nil {
+		return nil, fmt.Errorf("server: failed to pack key material: %w", err)
 	}
 
-	nsLen := int(binary.BigEndian.Uint16(ciphertext[1:headerSize]))
-	if len(ciphertext) < headerSize+nsLen {
-		return nil, errors.New("server: ciphertext is truncated inside its namespace")
+	return packed, nil
+}
+
+// Unwrap opens key material produced by Wrap, deriving the key from the version
+// and namespace it carries: the two together address one key, the way a lookup
+// against a real key service would.
+func (k *keyring) Unwrap(_ context.Context, ciphertext []byte) ([]byte, error) {
+	material, err := extv1.UnmarshalKeyMaterial(ciphertext)
+	if err != nil {
+		return nil, fmt.Errorf("server: %w", err)
 	}
 
-	namespace := string(ciphertext[headerSize : headerSize+nsLen])
-	sealed := ciphertext[headerSize+nsLen:]
+	// Unmarshal accepts bytes that set no fields at all, so anything arriving
+	// from outside gets checked before it is trusted.
+	if err := material.Validate(); err != nil {
+		return nil, fmt.Errorf("server: invalid key material: %w", err)
+	}
 
-	gcm, err := k.cipher(namespace)
+	gcm, err := k.cipher(material.GetVersion(), material.GetNamespace())
 	if err != nil {
 		return nil, err
 	}
 
-	if len(sealed) < gcm.NonceSize() {
-		return nil, errors.New("server: ciphertext is truncated inside its nonce")
+	nonce := material.GetOpaque()
+	if len(nonce) != gcm.NonceSize() {
+		return nil, fmt.Errorf("server: key material carries a %d-byte nonce, want %d", len(nonce), gcm.NonceSize())
 	}
 
-	dek, err := gcm.Open(nil, sealed[:gcm.NonceSize()], sealed[gcm.NonceSize():], []byte(namespace))
+	dek, err := gcm.Open(nil, nonce, material.GetEncryptedDek(), []byte(material.GetNamespace()))
 	if err != nil {
-		return nil, fmt.Errorf("server: failed to open ciphertext for namespace %q: %w", namespace, err)
+		return nil, fmt.Errorf("server: failed to open key material for namespace %q: %w", material.GetNamespace(), err)
 	}
 
 	return dek, nil
 }
 
-// cipher derives the wrapping key for namespace and returns a GCM cipher over
-// it. Derivation is deterministic, so a restarted provider still opens
-// ciphertexts sealed before the restart.
-func (k *keyring) cipher(namespace string) (cipher.AEAD, error) {
-	key, err := hkdf.Key(sha256.New, k.secret, nil, infoPrefix+namespace, keySize)
+// cipher derives the wrapping key for version and namespace and returns a GCM
+// cipher over it. Derivation is deterministic, so a restarted provider still
+// opens key material sealed before the restart.
+func (k *keyring) cipher(version, namespace string) (cipher.AEAD, error) {
+	info := fmt.Sprintf("%s/v%s/%s", infoPrefix, version, namespace)
+
+	key, err := hkdf.Key(sha256.New, k.secret, nil, info, keySize)
 	if err != nil {
 		return nil, fmt.Errorf("server: failed to derive a key for namespace %q: %w", namespace, err)
 	}

examples/kms/server/main.go

@@ -8,12 +8,13 @@
 // from. Both are required.
 //
 // Only the provider itself lives here, in keyring.go, which derives one
-// AES-256-GCM key per namespace from the master secret and frames the
-// ciphertext. The gRPC surface and the bearer token check come from
-// [github.com/temporalio/temporal-proxy/pkg/ext]: keyring's Wrap and Unwrap
-// satisfy [ext.KMS], and [ext.Serve] registers them, serves TLS, and shuts down
-// on a signal. That split is the point of the example. The interesting part of
-// writing one of these is the key handling, not the server around it.
+// AES-256-GCM key per version and namespace from the master secret and returns
+// the wrapped DEK as api.ext.v1.KeyMaterial. The gRPC surface and the bearer
+// token check come from [github.com/temporalio/temporal-proxy/pkg/ext]:
+// keyring's Wrap and Unwrap satisfy [ext.KMS], and [ext.Serve] registers them,
+// serves TLS, and shuts down on a signal. That split is the point of the
+// example. The interesting part of writing one of these is the key handling,
+// not the server around it.
 //
 // This is enough to show the shape of the contract and it is not a key manager:
 // the master secret sits in an environment variable, nothing is rotated, and

pkg/api/ext/v1/key_material.go

@@ -0,0 +1,54 @@
+package ext
+
+import (
+	"errors"
+	"fmt"
+
+	"google.golang.org/protobuf/proto"
+
+	"github.com/temporalio/temporal-proxy/pkg/validation"
+)
+
+// UnmarshalKeyMaterial decodes key material produced by [KeyMaterial.Marshal].
+// It does not check what it decoded: proto3 happily accepts bytes that set none
+// of the fields, so call [KeyMaterial.Validate] on anything whose framing you
+// did not produce yourself.
+func UnmarshalKeyMaterial(raw []byte) (*KeyMaterial, error) {
+	km := &KeyMaterial{}
+	if err := proto.Unmarshal(raw, km); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal key material: %w", err)
+	}
+
+	return km, nil
+}
+
+// Marshal validates km and returns its wire encoding, ready to hand back as an
+// EncryptResponse ciphertext.
+func (km *KeyMaterial) Marshal() ([]byte, error) {
+	if err := km.Validate(); err != nil {
+		return nil, err
+	}
+
+	packed, err := proto.Marshal(km)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal key material: %w", err)
+	}
+
+	return packed, nil
+}
+
+// Validate reports whether km carries a wrapped DEK, treating a nil km as one
+// that does not. Every other field is optional: an extension server may version
+// no keys, key off no namespace, and carry nothing of its own.
+func (km *KeyMaterial) Validate() error {
+	return validation.Validate(
+		"",
+		validation.Field("encrypted_dek", km.GetEncryptedDek(), func(v []byte) error {
+			if len(v) == 0 {
+				return errors.New("must not be empty")
+			}
+
+			return nil
+		}),
+	)
+}