feat(policy): Add platform key indexer.

c-r33d · c-r33d · commit 79971f558d18 · 2025-05-07T13:27:37.000-05:00
remove comments.

linting.

linting.
diff --git a/service/trust/platform_key_indexer.go b/service/trust/platform_key_indexer.go
@@ -0,0 +1,255 @@
+package trust
+
+import (
+	"context"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/json"
+	"encoding/pem"
+	"errors"
+	"fmt"
+
+	"github.com/lestrrat-go/jwx/v2/jwk"
+	"github.com/opentdf/platform/protocol/go/policy"
+	"github.com/opentdf/platform/protocol/go/policy/kasregistry"
+	"github.com/opentdf/platform/sdk"
+	"github.com/opentdf/platform/service/logger"
+)
+
+var ErrNoActiveKeyForAlgorithm = errors.New("no active key found for specified algorithm")
+
+// Used for reaching out to platform to get keys
+type PlatformKeyIndexer struct {
+	// KeyIndex is the key index used to manage keys
+	KeyIndex
+	// SDK is the SDK instance used to interact with the platform
+	sdk *sdk.SDK
+	// KasURI
+	kasURI string
+	// Logger is the logger instance used for logging
+	log *logger.Logger
+}
+
+// platformKeyAdapter is an adapter for KeyDetails, where keys come from the platform
+type KasKeyAdapter struct {
+	key *policy.KasKey
+	log *logger.Logger
+}
+
+func NewPlatformKeyIndexer(sdk *sdk.SDK, kasURI string, l *logger.Logger) *PlatformKeyIndexer {
+	return &PlatformKeyIndexer{
+		sdk:    sdk,
+		kasURI: kasURI,
+		log:    l,
+	}
+}
+
+func convertAlgToEnum(alg string) (policy.Algorithm, error) {
+	switch alg {
+	case "rsa:2048":
+		return policy.Algorithm_ALGORITHM_RSA_2048, nil
+	case "rsa:4096":
+		return policy.Algorithm_ALGORITHM_RSA_4096, nil
+	case "ec:secp256r1":
+		return policy.Algorithm_ALGORITHM_EC_P256, nil
+	case "ec:secp384r1":
+		return policy.Algorithm_ALGORITHM_EC_P384, nil
+	case "ec:secp521r1":
+		return policy.Algorithm_ALGORITHM_EC_P521, nil
+	default:
+		return policy.Algorithm_ALGORITHM_UNSPECIFIED, errors.New("unsupported algorithm")
+	}
+}
+
+func (p *PlatformKeyIndexer) FindKeyByAlgorithm(ctx context.Context, algorithm string, _ bool) (KeyDetails, error) {
+	alg, err := convertAlgToEnum(algorithm)
+	if err != nil {
+		return nil, err
+	}
+
+	req := &kasregistry.ListKeysRequest{
+		KeyAlgorithm: alg,
+		KasFilter: &kasregistry.ListKeysRequest_KasUri{
+			KasUri: p.kasURI,
+		},
+	}
+	resp, err := p.sdk.KeyAccessServerRegistry.ListKeys(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	// Find active key.
+	var activeKey *policy.KasKey
+	for _, key := range resp.GetKasKeys() {
+		if key.GetKey().GetKeyStatus() == policy.KeyStatus_KEY_STATUS_ACTIVE {
+			activeKey = key
+			break
+		}
+	}
+	if activeKey == nil {
+		return nil, ErrNoActiveKeyForAlgorithm
+	}
+
+	return &KasKeyAdapter{
+		key: activeKey,
+		log: p.log,
+	}, nil
+}
+
+func (p *PlatformKeyIndexer) FindKeyByID(ctx context.Context, id KeyIdentifier) (KeyDetails, error) {
+	req := &kasregistry.GetKeyRequest{
+		Identifier: &kasregistry.GetKeyRequest_Key{
+			Key: &kasregistry.KasKeyIdentifier{
+				Identifier: &kasregistry.KasKeyIdentifier_Uri{
+					Uri: p.kasURI,
+				},
+				Kid: string(id),
+			},
+		},
+	}
+
+	resp, err := p.sdk.KeyAccessServerRegistry.GetKey(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	return &KasKeyAdapter{
+		key: resp.GetKasKey(),
+		log: p.log,
+	}, nil
+}
+
+func (p *PlatformKeyIndexer) ListKeys(ctx context.Context) ([]KeyDetails, error) {
+	req := &kasregistry.ListKeysRequest{
+		KasFilter: &kasregistry.ListKeysRequest_KasUri{
+			KasUri: p.kasURI,
+		},
+	}
+	resp, err := p.sdk.KeyAccessServerRegistry.ListKeys(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	keys := make([]KeyDetails, len(resp.GetKasKeys()))
+	for i, key := range resp.GetKasKeys() {
+		keys[i] = &KasKeyAdapter{
+			key: key,
+			log: p.log,
+		}
+	}
+
+	return keys, nil
+}
+
+func (p *KasKeyAdapter) ID() KeyIdentifier {
+	return KeyIdentifier(p.key.GetKey().GetKeyId())
+}
+
+// Might need to convert this to a standard format
+func (p *KasKeyAdapter) Algorithm() string {
+	return p.key.GetKey().GetKeyAlgorithm().String()
+}
+
+func (p *KasKeyAdapter) IsLegacy() bool {
+	return false
+}
+
+// This will point to the correct "manager"
+func (p *KasKeyAdapter) System() string {
+	var mode string
+	if p.key.GetKey().GetProviderConfig() != nil {
+		mode = p.key.GetKey().GetProviderConfig().GetName()
+	}
+	return mode
+}
+
+func pemToPublicKey(publicPEM string) (*rsa.PublicKey, error) {
+	// Decode the PEM data
+	block, _ := pem.Decode([]byte(publicPEM))
+	if block == nil || block.Type != "PUBLIC KEY" {
+		return nil, fmt.Errorf("failed to decode PEM block or incorrect PEM type")
+	}
+
+	// Parse the public key
+	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse public key: %w", err)
+	}
+
+	// Assert type and return
+	rsaPub, ok := pub.(*rsa.PublicKey)
+	if !ok {
+		return nil, fmt.Errorf("not an RSA public key")
+	}
+
+	return rsaPub, nil
+}
+
+// Repurpose of the StandardCrypto function
+func rsaPublicKeyAsJSON(_ context.Context, publicPEM string) (string, error) {
+	pubKey, err := pemToPublicKey(publicPEM)
+	if err != nil {
+		return "", err
+	}
+
+	rsaPublicKeyJwk, err := jwk.FromRaw(pubKey)
+	if err != nil {
+		return "", fmt.Errorf("jwk.FromRaw: %w", err)
+	}
+
+	// Convert the public key to JSON format
+	pubKeyJSON, err := json.Marshal(rsaPublicKeyJwk)
+	if err != nil {
+		return "", err
+	}
+
+	return string(pubKeyJSON), nil
+}
+
+// Repurpose of the StandardCrypto function
+func convertPEMToJWK(_ string) (string, error) {
+	return "", errors.New("convertPEMToJWK function is not implemented")
+}
+
+func (p *KasKeyAdapter) ExportPublicKey(ctx context.Context, format KeyType) (string, error) {
+	publicKeyCtx := p.key.GetKey().GetPublicKeyCtx()
+	var pubKeyCtxMap map[string]any
+	if err := json.Unmarshal(publicKeyCtx, &pubKeyCtxMap); err != nil {
+		return "", err
+	}
+
+	pubKey, ok := pubKeyCtxMap["pubKey"].(string)
+	if !ok {
+		return "", errors.New("public key is not a string")
+	}
+	// Decode the base64-encoded public key
+	decodedPubKey, err := base64.StdEncoding.DecodeString(pubKey)
+	if err != nil {
+		return "", err
+	}
+
+	switch format {
+	case KeyTypeJWK:
+		// For JWK format (currently only supported for RSA)
+		if p.key.GetKey().GetKeyAlgorithm() == policy.Algorithm_ALGORITHM_RSA_2048 ||
+			p.key.GetKey().GetKeyAlgorithm() == policy.Algorithm_ALGORITHM_RSA_4096 {
+			return rsaPublicKeyAsJSON(ctx, string(decodedPubKey))
+		}
+		// For EC keys, we return the public key in PEM format
+		jwkKey, err := convertPEMToJWK(string(decodedPubKey))
+		if err != nil {
+			return "", err
+		}
+
+		return jwkKey, nil
+	case KeyTypePKCS8:
+		return string(decodedPubKey), nil
+	default:
+		return "", errors.New("unsupported key type")
+	}
+}
+
+func (p *KasKeyAdapter) ExportCertificate(_ context.Context) (string, error) {
+	return "", errors.New("not implemented")
+}
diff --git a/service/trust/platform_key_indexer_test.go b/service/trust/platform_key_indexer_test.go
@@ -0,0 +1,79 @@
+package trust
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/lestrrat-go/jwx/v2/jwk"
+	"github.com/opentdf/platform/lib/ocrypto"
+	"github.com/opentdf/platform/protocol/go/policy"
+	"github.com/stretchr/testify/suite"
+)
+
+type PlatformKeyIndexTestSuite struct {
+	suite.Suite
+	rsaKey KeyDetails
+}
+
+func (s *PlatformKeyIndexTestSuite) SetupTest() {
+	s.rsaKey = &KasKeyAdapter{
+		key: &policy.KasKey{
+			KasId: "test-kas-id",
+			Key: &policy.AsymmetricKey{
+				Id:           "test-key-id",
+				KeyId:        "test-key-id",
+				KeyAlgorithm: policy.Algorithm_ALGORITHM_RSA_2048,
+				KeyStatus:    policy.KeyStatus_KEY_STATUS_ACTIVE,
+				KeyMode:      policy.KeyMode_KEY_MODE_LOCAL,
+				PublicKeyCtx: []byte(`{"pubKey": "LS0tLS1CRUdJTiBQVUJMSUMgS0VZLS0tLS0KTUlJQklqQU5CZ2txaGtpRzl3MEJBUUVGQUFPQ0FROEFNSUlCQ2dLQ0FRRUF3SEw0TkVrOFpDa0JzNjZXQVpWagpIS3NseDRseWdmaXN3aW42RUx5OU9OczZLVDRYa1crRGxsdExtck14bHZkbzVRaDg1UmFZS01mWUdDTWtPM0dGCkFsK0JOeWFOM1kwa0N1QjNPU2ErTzdyMURhNVZteVVuaEJNbFBrYnVPY1Y0cjlLMUhOSGd3eDl2UFp3RjRpQW8KQStEY1VBcWFEeHlvYjV6enNGZ0hUNjJHLzdLdEtiZ2hYT1dCanRUYUl1ZHpsK2FaSjFPemY0U1RkOXhST2QrMQordVo2VG1ocmFEUm9zdDUrTTZUN0toL2lGWk40TTFUY2hwWXU1TDhKR2tVaG9YaEdZcHUrMGczSzlqYlh6RVh5CnpJU3VXN2d6SGRWYUxvcnBkQlNkRHpOWkNvTFVoL0U1T3d5TFZFQkNKaDZJVUtvdWJ5WHVucnIxQnJmK2tLbEsKeHdJREFRQUIKLS0tLS1FTkQgUFVCTElDIEtFWS0tLS0tCg=="}`),
+				ProviderConfig: &policy.KeyProviderConfig{
+					Id:   "test-provider-id",
+					Name: "openbao",
+				},
+			},
+		},
+	}
+}
+func (s *PlatformKeyIndexTestSuite) TearDownTest() {}
+
+func (s *PlatformKeyIndexTestSuite) TestKeyDetails() {
+	s.Equal("test-key-id", s.rsaKey.ID())
+	s.Equal("ALGORITHM_RSA_2048", s.rsaKey.Algorithm())
+	s.False(s.rsaKey.IsLegacy())
+	s.Equal("openbao", s.rsaKey.System())
+}
+
+func (s *PlatformKeyIndexTestSuite) TestKeyExportPublicKey_JWKFormat() {
+	// Export JWK format
+	jwkString, err := s.rsaKey.ExportPublicKey(context.Background(), KeyTypeJWK)
+	s.Require().NoError(err)
+	s.Require().NotEmpty(jwkString)
+
+	rsaKey, err := jwk.ParseKey([]byte(jwkString))
+	s.Require().NoError(err)
+	s.Require().NotNil(rsaKey)
+}
+
+func (s *PlatformKeyIndexTestSuite) TestKeyExportPublicKey_PKCSFormat() {
+	// Export JWK format
+	pem, err := s.rsaKey.ExportPublicKey(context.Background(), KeyTypePKCS8)
+	s.Require().NoError(err)
+	s.Require().NotEmpty(pem)
+
+	keyAdapter, ok := s.rsaKey.(*KasKeyAdapter)
+	s.Require().True(ok)
+	pubCtx := keyAdapter.key.GetKey().GetPublicKeyCtx()
+	s.Require().NotEmpty(pubCtx)
+	base64Pem := ocrypto.Base64Encode([]byte(pem))
+
+	var pubCtxMap map[string]interface{}
+	err = json.Unmarshal(pubCtx, &pubCtxMap)
+	s.Require().NoError(err)
+
+	s.Equal(pubCtxMap["pubKey"], string(base64Pem))
+}
+
+func TestNewPlatformKeyIndexTestSuite(t *testing.T) {
+	suite.Run(t, new(PlatformKeyIndexTestSuite))
+}
