Pure Go implementation of a P-256 age recipient compatible with age-plugin-yubikey

age_p256_recipient.go

// Copyright (c) 2025 Monzo Bank Limited
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package monzoage

import (
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"fmt"

	"filippo.io/age"
	"github.com/awnumar/agekd/bech32"
	"golang.org/x/crypto/chacha20poly1305"
)

// yubikeyRecipient implements age.Recipient for YubiKey P256 keys: https://pkg.go.dev/filippo.io/age#Recipient
// This allows encrypting to a Yubikey-backed age-plugin-yubikey recipient without requiring the rust binary or pcscd.
// This is a port of parts of age-plugin-yubikey: https://github.com/str4d/age-plugin-yubikey
//   - https://github.com/str4d/age-plugin-yubikey/issues/192#issuecomment-2456071304
//   - https://github.com/C2SP/C2SP/blob/main/age-plugin.md
type yubikeyRecipient struct {
	pubKey          *ecdh.PublicKey
	tag             []byte
	compressedBytes []byte
}

// parseYubikeyRecipient parses an age-plugin-yubikey recipient string, e.g. age1yubikey1qwla8v7cu3mx6mp79asgrh5ad2h52flwln7c66ydcyy50lg5uh0gxh4kmaz
// and from it initialises an age.Recipient that can be used with the age library to encrypt data.
func parseYubikeyRecipient(s string) (age.Recipient, error) {
	hrp, data, err := bech32.Decode(s)
	if err != nil {
		return nil, fmt.Errorf("decoding bech32: %w", err)
	}
	if hrp != "age1yubikey" {
		return nil, fmt.Errorf("this implementation only supports age1yubikey recipients, not '%s'", hrp)
	}

	// We need the uncompressed format for the curve point
	x, y := elliptic.UnmarshalCompressed(elliptic.P256(), data)
	if x == nil {
		return nil, fmt.Errorf("invalid compressed public key data: %x", data)
	}

	// elliptic.Marshal is depreciated so use the ecdsa package instead
	ecdsaKey := &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y}
	pubKey, err := ecdsaKey.ECDH()
	if err != nil {
		return nil, fmt.Errorf("converting to ECDH key: %w", err)
	}

	// The tag is calculated from the compressed point encoding (like Rust implementation)
	tag := sha256.Sum256(data)

	return &yubikeyRecipient{
		pubKey:          pubKey,
		tag:             tag[:4],
		compressedBytes: data,
	}, nil
}

// compressPublicKeyP256 turns an uncompressed P256 point representation into a compressed representation.
func compressPublicKeyP256(uncompressedPublicKey []byte) ([]byte, error) {
	ecdsaPubKey, err := ecdsa.ParseUncompressedPublicKey(elliptic.P256(), uncompressedPublicKey)
	if err != nil {
		return nil, fmt.Errorf("parsing uncompressed public key: %w", err)
	}
	return elliptic.MarshalCompressed(elliptic.P256(), ecdsaPubKey.X, ecdsaPubKey.Y), nil
}

// Wrap implements the age.Recipient interface, which allows us to use the age library for encryption.
//   - https://pkg.go.dev/filippo.io/age#Recipient
//   - https://pkg.go.dev/filippo.io/age#Encrypt
func (recipient *yubikeyRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
	const stanzaTag = "piv-p256"
	const stanzaKeyLabel = "piv-p256"

	ephemeralPrivateKey, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generating ephemeral private key: %w", err)
	}

	sharedSecret, err := ephemeralPrivateKey.ECDH(recipient.pubKey)
	if err != nil {
		return nil, fmt.Errorf("performing ECDH: %w", err)
	}

	ephemeralCompressed, err := compressPublicKeyP256(ephemeralPrivateKey.PublicKey().Bytes())
	if err != nil {
		return nil, fmt.Errorf("compressing ephemeral public key: %w", err)
	}

	// Embedding both public keys in the salt binds the derived key to both keys. This matches the Rust implementation.
	salt := make([]byte, 0, len(ephemeralCompressed)+len(recipient.compressedBytes))
	salt = append(salt, ephemeralCompressed...)
	salt = append(salt, recipient.compressedBytes...)

	encryptionKey, err := hkdf.Key(sha256.New, sharedSecret, salt, stanzaKeyLabel, 32)
	if err != nil {
		return nil, fmt.Errorf("hkdf: deriving encryption key: %w", err)
	}

	fileKeyCipher, err := chacha20poly1305.New(encryptionKey)
	if err != nil {
		return nil, fmt.Errorf("creating chacha20poly1305 cipher: %w", err)
	}
	nonce := make([]byte, fileKeyCipher.NonceSize()) // zero nonce as per Rust implementation

	encryptedFileKey := fileKeyCipher.Seal(nil, nonce, fileKey, nil)

	return []*age.Stanza{
		{
			Type: stanzaTag,
			Args: []string{
				base64.RawStdEncoding.EncodeToString(recipient.tag),
				base64.RawStdEncoding.EncodeToString(ephemeralCompressed),
			},
			Body: encryptedFileKey,
		},
	}, nil
}