adding age plugin wrapping default X25519 Identity/Recipient implementation with hooks to more efficiently lookup private keys given their respective public key

Author: eriktate

lib/auth/recordingencryption/age.go

@@ -0,0 +1,115 @@
+// Teleport
+// Copyright (C) 2025 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package recordingencryption
+
+import (
+	"filippo.io/age"
+	"github.com/gravitational/trace"
+
+	"github.com/gravitational/teleport/api/types"
+)
+
+// X25519Stanza is the default stanza type used by age.
+const X25519Stanza = "X25519"
+
+// RecordingStanza is the type used for the identifying stanza added by RecordingRecipient.
+const RecordingStanza = "Recording-X25519"
+
+// DecryptionKeyFinder returns an EncryptionKeyPair related to at least one of the given public keys to be used
+// for file key unwrapping.
+type DecryptionKeyFinder interface {
+	FindDecryptionKey(publicKeys ...[]byte) (*types.EncryptionKeyPair, error)
+}
+
+// RecordingIdentity removes public keys from stanzas and passes the unwrap call to the default
+// age.X25519Identity.
+type RecordingIdentity struct {
+	keyFinder DecryptionKeyFinder
+}
+
+// NewRecordingIdentity returns a RecordingIdentity that will use the given DecryptionKeyFinder in order to facilitate
+// file key unwrapping.
+func NewRecordingIdentity(keyFinder DecryptionKeyFinder) *RecordingIdentity {
+	return &RecordingIdentity{
+		keyFinder: keyFinder,
+	}
+}
+
+// Unwrap uses the additional stanzas added by RecordingRecipient.Wrap in order to find a matching X25519 identity.
+func (i *RecordingIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
+	var publicKeys [][]byte
+	for _, stanza := range stanzas {
+		if stanza.Type != RecordingStanza {
+			continue
+		}
+
+		if len(stanza.Args) != 1 {
+			continue
+		}
+
+		publicKeys = append(publicKeys, []byte(stanza.Args[0]))
+	}
+
+	pair, err := i.keyFinder.FindDecryptionKey(publicKeys...)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	identity, err := age.ParseX25519Identity(string(pair.PrivateKey))
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return identity.Unwrap(stanzas)
+}
+
+// RecordingRecipient adds the public key to the stanzas generated by the default age.X25519Recipient
+type RecordingRecipient struct {
+	*age.X25519Recipient
+}
+
+// ParseRecordingRecipient parses an Bech32 encoded age X25519 public key into a RecordingRecipient.
+func ParseRecordingRecipient(s string) (*RecordingRecipient, error) {
+	recipient, err := age.ParseX25519Recipient(s)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return &RecordingRecipient{X25519Recipient: recipient}, nil
+}
+
+// Wrap a fileKey using the wrapped X2519Recipient. An additional stanza containing the bech32 encoded X25519
+// public key will be created to enable lookups during Unwrap.
+func (r *RecordingRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
+	stanzas, err := r.X25519Recipient.Wrap(fileKey)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	// a new stanza has to be added because modifying the original stanza and returning it to "normal" during
+	// Unwrap fails due to MAC errors
+	for _, stanza := range stanzas {
+		if stanza.Type == X25519Stanza {
+			stanzas = append(stanzas, &age.Stanza{
+				Type: RecordingStanza,
+				Args: []string{r.String()},
+			})
+		}
+	}
+
+	return stanzas, nil
+}

lib/auth/recordingencryption/age_test.go

@@ -0,0 +1,109 @@
+// Teleport
+// Copyright (C) 2025 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package recordingencryption_test
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"testing"
+
+	"filippo.io/age"
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/auth/recordingencryption"
+)
+
+func TestRecordingAgePlugin(t *testing.T) {
+	keyStore := newFakeKeyStore()
+	recordingIdentity := recordingencryption.NewRecordingIdentity(keyStore)
+
+	ident, err := keyStore.generateIdentity()
+	require.NoError(t, err)
+
+	recipient, err := recordingencryption.ParseRecordingRecipient(ident.Recipient().String())
+	require.NoError(t, err)
+
+	out := bytes.NewBuffer(nil)
+	writer, err := age.Encrypt(out, recipient)
+	require.NoError(t, err)
+
+	msg := []byte("testing age plugin for session recordings")
+	_, err = writer.Write(msg)
+	require.NoError(t, err)
+
+	// writer must be closed to ensure data is flushed
+	err = writer.Close()
+	require.NoError(t, err)
+
+	reader, err := age.Decrypt(out, recordingIdentity)
+	require.NoError(t, err)
+	plaintext, err := io.ReadAll(reader)
+	require.NoError(t, err)
+
+	require.Equal(t, msg, plaintext)
+
+	// running the same test with the raw recipient should fail because the
+	// the extra stanza added by RecordingRecipient won't be present and
+	// the private key won't be found
+	out.Reset()
+	writer, err = age.Encrypt(out, ident.Recipient())
+	require.NoError(t, err)
+	_, err = writer.Write(msg)
+	require.NoError(t, err)
+	err = writer.Close()
+	require.NoError(t, err)
+	_, err = age.Decrypt(out, recordingIdentity)
+	require.Error(t, err)
+}
+
+type fakeKeyStore struct {
+	keys map[string]string
+}
+
+func newFakeKeyStore() *fakeKeyStore {
+	return &fakeKeyStore{
+		keys: make(map[string]string),
+	}
+}
+
+func (f *fakeKeyStore) FindDecryptionKey(publicKeys ...[]byte) (*types.EncryptionKeyPair, error) {
+	for _, pubKey := range publicKeys {
+		key, ok := f.keys[string(pubKey)]
+		if !ok {
+			continue
+		}
+
+		return &types.EncryptionKeyPair{
+			PrivateKey: []byte(key),
+			PublicKey:  pubKey,
+		}, nil
+	}
+
+	return nil, errors.New("no accessible decryption key found")
+}
+
+func (f *fakeKeyStore) generateIdentity() (*age.X25519Identity, error) {
+	ident, err := age.GenerateX25519Identity()
+	if err != nil {
+		return nil, err
+	}
+
+	f.keys[ident.Recipient().String()] = ident.String()
+	return ident, nil
+}