stash: fixup crypto.go and add a test.

pkg/crypto/crypto.go

@@ -1,7 +1,7 @@
 package crypto
 
 import (
-	"crypto/ed25519"
+	"crypto/ecdsa"
 	"errors"
 	"fmt"
 
@@ -13,7 +13,6 @@ import (
 var ErrInvalidPublicKey = errors.New("invalid public key")
 
 type PlaintextCredential struct {
-	Provider    string
 	Name        string
 	Description string
 	Schema      string
@@ -30,45 +29,67 @@ type PubKeyEncryptionManager struct {
 
 	keys       map[string]*jose.JSONWebKey
 	encrypters map[string]jose.Encrypter
+
+	fullSerialize bool
 }
 
 func (pkem *PubKeyEncryptionManager) Encrypt(cred *PlaintextCredential) ([]*v2.EncryptedData, error) {
-	encryptedDatas := make([]*v2.EncryptedData, len(pkem.configs))
+	encryptedDatas := make([]*v2.EncryptedData, 0, len(pkem.configs))
 
 	for keyId, encrypter := range pkem.encrypters {
 		jwe, err := encrypter.Encrypt(cred.Bytes)
 		if err != nil {
 			return nil, err
 		}
+
 		cypherText, err := jwe.CompactSerialize()
 		if err != nil {
 			return nil, err
 		}
+
 		encryptedData := &v2.EncryptedData{
-			Provider:       cred.Provider,
 			KeyId:          keyId,
+			Name:           cred.Name,
+			Description:    cred.Description,
+			Schema:         cred.Schema,
 			EncryptedBytes: []byte(cypherText),
 		}
 		encryptedDatas = append(encryptedDatas, encryptedData)
 	}
 	return encryptedDatas, nil
 }
 
-// parsePublicKey parses a public ed25519 JWK, all other key types return errors.
+// parsePublicKey parses a public ecdsa JWK, all other key types return errors.
 func parsePublicKey(input []byte) (*jose.JSONWebKey, error) {
 	npk := &jose.JSONWebKey{}
 	err := npk.UnmarshalJSON(input)
 	if err != nil {
 		return nil, fmt.Errorf("%w: failed unmarshalling public key: %w", ErrInvalidPublicKey, err)
 	}
 
-	if !npk.Valid() || !npk.IsPublic() {
+	if npk.KeyID == "" {
+		return nil, fmt.Errorf("%w: kid is required", ErrInvalidPublicKey)
+	}
+
+	if !npk.Valid() {
 		return nil, ErrInvalidPublicKey
 	}
 
-	_, ok := npk.Key.(ed25519.PrivateKey)
+	if !npk.IsPublic() {
+		return nil, fmt.Errorf("%w: key is not public", ErrInvalidPublicKey)
+	}
+
+	if npk.Use != "enc" {
+		return nil, fmt.Errorf("%w: invalid use (%s) - 'enc' is required", ErrInvalidPublicKey, npk.Use)
+	}
+
+	if npk.Algorithm != string(jose.ECDH_ES_A256KW) {
+		return nil, fmt.Errorf("%w: invalid algorithm (%s) - 'ECDH-ES+A256KW' is required", ErrInvalidPublicKey, npk.Algorithm)
+	}
+
+	_, ok := npk.Key.(*ecdsa.PublicKey)
 	if !ok {
-		return nil, ErrInvalidPublicKey
+		return nil, fmt.Errorf("%w: invalid key type - ecdsa is required", ErrInvalidPublicKey)
 	}
 
 	return npk, nil
@@ -99,16 +120,13 @@ func NewPubKeyEncryptionManager(co *v2.CredentialOptions, ec []*v2.EncryptionCon
 		if err != nil {
 			return nil, fmt.Errorf("public_key_encryption_manager: failed parsing public key %w", err)
 		}
-		if key.KeyID == "" {
-			return nil, fmt.Errorf("public_key_encryption_manager: kid is required")
-		}
 		_, ok := pkem.keys[key.KeyID]
 		if ok {
 			return nil, fmt.Errorf("public_key_encryption_manager: duplicate key id %s", key.KeyID)
 		}
 		pkem.keys[key.KeyID] = key
 
-		encryptor, err := jose.NewEncrypter(jose.A256GCM, jose.Recipient{Algorithm: jose.ED25519, Key: key, KeyID: key.KeyID}, nil)
+		encryptor, err := jose.NewEncrypter(jose.A256GCM, jose.Recipient{Algorithm: jose.KeyAlgorithm(key.Algorithm), Key: key, KeyID: key.KeyID}, nil)
 		if err != nil {
 			return nil, err
 		}

pkg/crypto/crypto_test.go

@@ -0,0 +1,85 @@
+package crypto
+
+import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"testing"
+
+	"github.com/go-jose/go-jose/v3"
+	"github.com/segmentio/ksuid"
+	"github.com/stretchr/testify/require"
+
+	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
+)
+
+// genKey generates a new ECDSA key and returns the private *ecdsa.PrivateKey, the *ecdsa.PublicKey as a JWK,
+// and the marshalled public key JWK.
+func genKey(t *testing.T) (*ecdsa.PrivateKey, *jose.JSONWebKey, []byte) {
+	key, err := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
+	require.NoError(t, err)
+
+	kid := ksuid.New().String()
+
+	jsonPubKey := &jose.JSONWebKey{
+		Key:       key.Public(),
+		KeyID:     kid,
+		Use:       "enc",
+		Algorithm: string(jose.ECDH_ES_A256KW),
+	}
+	marshalledPubKey, err := jsonPubKey.MarshalJSON()
+	require.NoError(t, err)
+
+	return key, jsonPubKey, marshalledPubKey
+}
+
+func TestNewPubKeyEncryptionManager(t *testing.T) {
+	// generate a keypair to encrypt to
+	privKey, pubKeyJWK, pubKeyJWKBytes := genKey(t)
+
+	// create an encryption manager
+	opts := &v2.CredentialOptions{}
+	config := []*v2.EncryptionConfig{
+		{
+			Config: &v2.EncryptionConfig_PublicKeyConfig_{
+				PublicKeyConfig: &v2.EncryptionConfig_PublicKeyConfig{
+					PubKey: pubKeyJWKBytes,
+				},
+			},
+		},
+	}
+	pkem, err := NewPubKeyEncryptionManager(opts, config)
+	require.NoError(t, err)
+
+	// encrypt a plaintext credential
+	cred := &PlaintextCredential{
+		Name:        "password",
+		Description: "this is the password",
+		Schema:      "string",
+		Bytes:       []byte("hunter2"),
+	}
+	encryptedValues, err := pkem.Encrypt(cred)
+	require.NoError(t, err)
+
+	// assert encrypt
+	require.Len(t, encryptedValues, 1)
+	encryptedValue := encryptedValues[0]
+	require.Equal(t, encryptedValue.Name, cred.Name)
+	require.Equal(t, encryptedValue.Description, cred.Description)
+	require.Equal(t, encryptedValue.Schema, cred.Schema)
+
+	require.Equal(t, encryptedValue.KeyId, pubKeyJWK.KeyID)
+	require.NotEmpty(t, encryptedValue.EncryptedBytes)
+
+	// assert we can decrypt with our private key
+	jwe, err := jose.ParseEncrypted(string(encryptedValue.EncryptedBytes))
+	require.NoError(t, err)
+	plaintext, err := jwe.Decrypt(privKey)
+	require.NoError(t, err)
+	require.Equal(t, []byte("hunter2"), plaintext)
+
+	// assert a different private key cannot decrypt
+	privKey2, _, _ := genKey(t)
+	_, err = jwe.Decrypt(privKey2)
+	require.ErrorContains(t, err, "go-jose/go-jose: error in cryptographic primitive")
+}