Enable Google KMS and Azure Key Vault for publisher login tool (#696)

## Motivation and Context This allows the private key used for DNS or HTTP based authentication to be stored securely in a cloud key management system. These services often provide HSM storage which makes it very hard to leak the private key. The full context is described in my design document here: https://github.com/modelcontextprotocol/registry/issues/482#issuecomment-3437117171 ## How Has This Been Tested? I have tested it against PROD using an Ed25519 key stored in Google KMS. I have added unit tests. I have tested both ECDSA P-384 and Ed25519 with Az KV and Google KMS against a locally running server. ## Breaking Changes None intended. ## Types of changes  - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [x] Documentation update ## Checklist  - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [x] New and existing tests pass locally - [x] I have added appropriate error handling - [x] I have added or updated documentation as needed ## Additional context  --------- Co-authored-by: adam jones <domdomegg+git@gmail.com>
2025-10-29 09:23:20 -04:00 · 2025-10-29 09:23:20 -04:00 · 88c2e1be16
commit 88c2e1be16
parent 3d375e6544
5 changed files with 641 additions and 91 deletions
--- a/cmd/publisher/auth/azurekeyvault/common.go
+++ b/cmd/publisher/auth/azurekeyvault/common.go
@ -0,0 +1,85 @@
+package azurekeyvault
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/sha512"
+	"fmt"
+	"math/big"
+	"os"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys"
+	"github.com/modelcontextprotocol/registry/cmd/publisher/auth"
+)
+
+func GetSignatureProvider(vaultName, keyName string) (auth.Signer, error) {
+	if vaultName == "" {
+		return nil, fmt.Errorf("--vault option (vault name) is required")
+	}
+
+	if keyName == "" {
+		return nil, fmt.Errorf("--key option (key name) is required")
+	}
+
+	return Signer{
+		vaultName: vaultName,
+		keyName:   keyName,
+	}, nil
+}
+
+type Signer struct {
+	vaultName string
+	keyName   string
+}
+
+func (d Signer) GetSignedTimestamp(ctx context.Context) (*string, []byte, error) {
+	fmt.Fprintf(os.Stdout, "Signing using Azure Key Vault %s and key %s\n", d.vaultName, d.keyName)
+
+	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	if err != nil {
+		return nil, nil, fmt.Errorf("authentication to Azure failed: %w", err)
+	}
+
+	vaultURL := fmt.Sprintf("https://%s.vault.azure.net/", d.vaultName)
+	client, err := azkeys.NewClient(vaultURL, cred, nil)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create Key Vault client: %w", err)
+	}
+
+	keyResp, err := client.GetKey(ctx, d.keyName, "", nil)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to retrieve key for public parameters: %w", err)
+	}
+
+	if *keyResp.Key.Kty != azkeys.KeyTypeEC && *keyResp.Key.Kty != azkeys.KeyTypeECHSM {
+		return nil, nil, fmt.Errorf("unsupported key type: kty: %s (only EC or EC-HSM keys are supported)", *keyResp.Key.Kty)
+	}
+
+	if *keyResp.Key.Crv != azkeys.CurveNameP384 {
+		return nil, nil, fmt.Errorf("unsupported curve: %s (only P-384 is supported)", *keyResp.Key.Crv)
+	}
+
+	fmt.Fprintln(os.Stdout, "Successfully read the public key from Key Vault.")
+	auth.PrintEcdsaP384KeyInfo(ecdsa.PublicKey{
+		Curve: elliptic.P384(),
+		X:     new(big.Int).SetBytes(keyResp.Key.X),
+		Y:     new(big.Int).SetBytes(keyResp.Key.Y),
+	})
+
+	timestamp := auth.GetTimestamp()
+	digest := sha512.Sum384([]byte(timestamp))
+	alg := azkeys.SignatureAlgorithmES384
+	fmt.Fprintln(os.Stdout, "Executing the sign request...")
+	signResp, err := client.Sign(ctx, d.keyName, "", azkeys.SignParameters{
+		Algorithm: &alg,
+		Value:     digest[:],
+	}, nil)
+
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to sign message: %w", err)
+	}
+
+	return &timestamp, signResp.Result, nil
+}
--- a/cmd/publisher/auth/common.go
+++ b/cmd/publisher/auth/common.go
@ -8,12 +8,14 @@ import (
 	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/sha512"
+	"encoding/base64"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
 	"math/big"
 	"net/http"
+	"os"
 	"time"
 )

@ -30,11 +32,35 @@ const (

 // CryptoProvider provides common functionality for DNS and HTTP authentication
 type CryptoProvider struct {
-	registryURL     string
-	domain          string
-	privateKey      string
-	cryptoAlgorithm CryptoAlgorithm
-	authMethod      string
+	registryURL string
+	domain      string
+	signer      Signer
+	authMethod  string
+}
+
+type Signer interface {
+	GetSignedTimestamp(ctx context.Context) (*string, []byte, error)
+}
+
+func GetTimestamp() string {
+	return time.Now().UTC().Format(time.RFC3339)
+}
+
+func NewInProcessSigner(privateKey string, algorithm CryptoAlgorithm) (Signer, error) {
+	if privateKey == "" {
+		return nil, fmt.Errorf("%s private key (hex) is required", algorithm)
+	}
+
+	// Decode private key from hex
+	privateKeyBytes, err := hex.DecodeString(privateKey)
+	if err != nil {
+		return nil, fmt.Errorf("invalid hex private key format: %w", err)
+	}
+
+	return &InProcessSigner{
+		privateKey:      privateKeyBytes,
+		cryptoAlgorithm: algorithm,
+	}, nil
 }

 // GetToken retrieves the registry JWT token using cryptographic authentication
@ -43,26 +69,15 @@ func (c *CryptoProvider) GetToken(ctx context.Context) (string, error) {
 		return "", fmt.Errorf("%s domain is required", c.authMethod)
 	}

-	if c.privateKey == "" {
-		return "", fmt.Errorf("%s private key (hex) is required", c.authMethod)
-	}
-
-	// Decode private key from hex
-	privateKeyBytes, err := hex.DecodeString(c.privateKey)
-	if err != nil {
-		return "", fmt.Errorf("invalid hex private key format: %w", err)
-	}
-
 	// Generate current timestamp
-	timestamp := time.Now().UTC().Format(time.RFC3339)
-	signedTimestamp, err := c.signMessage(privateKeyBytes, []byte(timestamp))
+	timestamp, signedTimestamp, err := c.signer.GetSignedTimestamp(ctx)
 	if err != nil {
 		return "", fmt.Errorf("failed to sign timestamp: %w", err)
 	}
 	signedTimestampHex := hex.EncodeToString(signedTimestamp)

 	// Exchange signature for registry token
-	registryToken, err := c.exchangeTokenForRegistry(ctx, c.domain, timestamp, signedTimestampHex)
+	registryToken, err := c.exchangeTokenForRegistry(ctx, c.domain, *timestamp, signedTimestampHex)
 	if err != nil {
 		return "", fmt.Errorf("failed to exchange %s signature: %w", c.authMethod, err)
 	}
@ -70,38 +85,52 @@ func (c *CryptoProvider) GetToken(ctx context.Context) (string, error) {
 	return registryToken, nil
 }

-func (c *CryptoProvider) signMessage(privateKeyBytes []byte, message []byte) ([]byte, error) {
+type InProcessSigner struct {
+	privateKey      []byte
+	cryptoAlgorithm CryptoAlgorithm
+}
+
+func (c *InProcessSigner) GetSignedTimestamp(_ context.Context) (*string, []byte, error) {
+	fmt.Fprintf(os.Stdout, "Signing in process using key algorithm %s\n", c.cryptoAlgorithm)
+
+	timestamp := GetTimestamp()
+
 	switch c.cryptoAlgorithm {
 	case AlgorithmEd25519:
-		if len(privateKeyBytes) != ed25519.SeedSize {
-			return nil, fmt.Errorf("invalid seed length: expected %d bytes, got %d", ed25519.SeedSize, len(privateKeyBytes))
+		if len(c.privateKey) != ed25519.SeedSize {
+			return nil, nil, fmt.Errorf("invalid seed length: expected %d bytes, got %d", ed25519.SeedSize, len(c.privateKey))
 		}

-		privateKey := ed25519.NewKeyFromSeed(privateKeyBytes)
-		signature := ed25519.Sign(privateKey, message)
-		return signature, nil
+		privateKey := ed25519.NewKeyFromSeed(c.privateKey)
+
+		PrintEd25519KeyInfo(privateKey.Public().(ed25519.PublicKey))
+
+		signature := ed25519.Sign(privateKey, []byte(timestamp))
+		return &timestamp, signature, nil
 	case AlgorithmECDSAP384:
-		if len(privateKeyBytes) != 48 {
-			return nil, fmt.Errorf("invalid seed length for ECDSA P-384: expected 48 bytes, got %d", len(privateKeyBytes))
+		if len(c.privateKey) != 48 {
+			return nil, nil, fmt.Errorf("invalid seed length for ECDSA P-384: expected 48 bytes, got %d", len(c.privateKey))
 		}

-		digest := sha512.Sum384(message)
+		digest := sha512.Sum384([]byte(timestamp))
 		curve := elliptic.P384()

 		// Parse the raw private key (compatible with Go 1.24)
-		privateKey, err := parseRawPrivateKey(curve, privateKeyBytes)
+		privateKey, err := parseRawPrivateKey(curve, c.privateKey)
 		if err != nil {
-			return nil, fmt.Errorf("failed to parse ECDSA private key: %w", err)
+			return nil, nil, fmt.Errorf("failed to parse ECDSA private key: %w", err)
 		}

+		PrintEcdsaP384KeyInfo(privateKey.PublicKey)
+
 		r, s, err := ecdsa.Sign(rand.Reader, privateKey, digest[:])
 		if err != nil {
-			return nil, fmt.Errorf("failed to sign message: %w", err)
+			return nil, nil, fmt.Errorf("failed to sign message: %w", err)
 		}
 		signature := append(r.Bytes(), s.Bytes()...)
-		return signature, nil
+		return &timestamp, signature, nil
 	default:
-		return nil, fmt.Errorf("unsupported crypto algorithm: %s", c.cryptoAlgorithm)
+		return nil, nil, fmt.Errorf("unsupported crypto algorithm: %s", c.cryptoAlgorithm)
 	}
 }

@ -112,6 +141,11 @@ func parseRawPrivateKey(curve elliptic.Curve, privateKeyBytes []byte) (*ecdsa.Pr
 		return nil, fmt.Errorf("nil curve")
 	}

+	expectedBytes := (curve.Params().N.BitLen() + 7) / 8
+	if len(privateKeyBytes) != expectedBytes {
+		return nil, fmt.Errorf("invalid private key length: expected %d bytes, got %d", expectedBytes, len(privateKeyBytes))
+	}
+
 	// Only standard NIST curves supported
 	switch curve {
 	case elliptic.P224(), elliptic.P256(), elliptic.P384(), elliptic.P521():
@ -147,6 +181,23 @@ func (c *CryptoProvider) Login(_ context.Context) error {
 	return nil
 }

+func PrintEd25519KeyInfo(pubKey ed25519.PublicKey) {
+	pubKeyString := base64.StdEncoding.EncodeToString(pubKey)
+	fmt.Fprint(os.Stdout, "Expected proof record:\n")
+	fmt.Fprintf(os.Stdout, "v=MCPv1; k=ed25519; p=%s\n", pubKeyString)
+}
+
+func PrintEcdsaP384KeyInfo(pubKey ecdsa.PublicKey) {
+	printEcdsaKeyInfo("ecdsap384", pubKey)
+}
+
+func printEcdsaKeyInfo(k string, pubKey ecdsa.PublicKey) {
+	compressed := elliptic.MarshalCompressed(pubKey.Curve, pubKey.X, pubKey.Y)
+	pubKeyString := base64.StdEncoding.EncodeToString(compressed)
+	fmt.Fprint(os.Stdout, "Expected proof record:\n")
+	fmt.Fprintf(os.Stdout, "v=MCPv1; k=%s; p=%s\n", k, pubKeyString)
+}
+
 // exchangeTokenForRegistry exchanges signature for a registry JWT token
 func (c *CryptoProvider) exchangeTokenForRegistry(ctx context.Context, domain, timestamp, signedTimestamp string) (string, error) {
 	if c.registryURL == "" {
--- a/cmd/publisher/auth/dns.go
+++ b/cmd/publisher/auth/dns.go
@ -5,14 +5,13 @@ type DNSProvider struct {
 }

 // NewDNSProvider creates a new DNS-based auth provider
-func NewDNSProvider(registryURL, domain, privateKey string, cryptoAlgorithm CryptoAlgorithm) Provider {
+func NewDNSProvider(registryURL, domain string, signer *Signer) Provider {
 	return &DNSProvider{
 		CryptoProvider: &CryptoProvider{
-			registryURL:     registryURL,
-			domain:          domain,
-			privateKey:      privateKey,
-			cryptoAlgorithm: cryptoAlgorithm,
-			authMethod:      "dns",
+			registryURL: registryURL,
+			domain:      domain,
+			signer:      *signer,
+			authMethod:  "dns",
 		},
 	}
 }
--- a/cmd/publisher/auth/googlekms/common.go
+++ b/cmd/publisher/auth/googlekms/common.go
@ -0,0 +1,147 @@
+package googlekms
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/elliptic"
+	"crypto/sha512"
+	"crypto/x509"
+	"encoding/asn1"
+	"encoding/pem"
+	"errors"
+	"fmt"
+	"math/big"
+	"os"
+
+	kms "cloud.google.com/go/kms/apiv1"
+	"cloud.google.com/go/kms/apiv1/kmspb"
+	"github.com/modelcontextprotocol/registry/cmd/publisher/auth"
+)
+
+// GetSignatureProvider validates inputs and returns a GoogleKMSSigner implementing auth.Signer.
+func GetSignatureProvider(resourceName string) (auth.Signer, error) {
+	if resourceName == "" {
+		return nil, fmt.Errorf("--resource is required, e.g. projects/my-project/locations/global/keyRings/fellowship/cryptoKeys/bilbo/cryptoKeyVersions/1")
+	}
+
+	return &Signer{
+		resource: resourceName,
+	}, nil
+}
+
+type Signer struct {
+	resource string
+}
+
+type ecdsaSignature struct {
+	R, S *big.Int
+}
+
+func derToRS(der []byte, curve elliptic.Curve) ([]byte, error) {
+	var sig ecdsaSignature
+	if _, err := asn1.Unmarshal(der, &sig); err != nil {
+		return nil, fmt.Errorf("invalid DER ECDSA signature: %w", err)
+	}
+	if sig.R == nil || sig.S == nil || sig.R.Sign() <= 0 || sig.S.Sign() <= 0 {
+		return nil, fmt.Errorf("invalid ECDSA signature components")
+	}
+
+	size := (curve.Params().BitSize + 7) / 8
+	rBytes := sig.R.Bytes()
+	sBytes := sig.S.Bytes()
+	if len(rBytes) > size || len(sBytes) > size {
+		return nil, fmt.Errorf("ECDSA signature component too large")
+	}
+
+	out := make([]byte, 2*size)
+	copy(out[size-len(rBytes):size], rBytes)
+	copy(out[2*size-len(sBytes):], sBytes)
+	return out, nil
+}
+
+func (g *Signer) GetSignedTimestamp(ctx context.Context) (*string, []byte, error) {
+	client, err := kms.NewKeyManagementClient(ctx)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create KMS client: %w", err)
+	}
+	defer client.Close()
+
+	// Fetch public key (PEM) so we can output expected proof record.
+	algo, err := g.showPublicKeyAndGetAlgorithm(ctx, client)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	timestamp := auth.GetTimestamp()
+
+	fmt.Fprintln(os.Stdout, "Executing the sign request...")
+	switch algo {
+	case auth.AlgorithmEd25519:
+		signReq := &kmspb.AsymmetricSignRequest{
+			Name: g.resource,
+			Data: []byte(timestamp),
+		}
+		signResp, err := client.AsymmetricSign(ctx, signReq)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to sign with KMS: %w", err)
+		}
+
+		return &timestamp, signResp.Signature, nil
+	case auth.AlgorithmECDSAP384:
+		digest := sha512.Sum384([]byte(timestamp))
+		signReq := &kmspb.AsymmetricSignRequest{
+			Name:   g.resource,
+			Digest: &kmspb.Digest{Digest: &kmspb.Digest_Sha384{Sha384: digest[:]}},
+		}
+		signResp, err := client.AsymmetricSign(ctx, signReq)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to sign with KMS: %w", err)
+		}
+
+		sigBytes, err := derToRS(signResp.Signature, elliptic.P384())
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to convert DER signature: %w", err)
+		}
+
+		return &timestamp, sigBytes, nil
+	}
+
+	return nil, nil, fmt.Errorf("unsupported algorithm: %s", algo)
+}
+
+func (g *Signer) showPublicKeyAndGetAlgorithm(ctx context.Context, client *kms.KeyManagementClient) (auth.CryptoAlgorithm, error) {
+	pubResp, err := client.GetPublicKey(ctx, &kmspb.GetPublicKeyRequest{Name: g.resource})
+	if err != nil {
+		return "", fmt.Errorf("failed to get public key: %w", err)
+	}
+
+	block, _ := pem.Decode([]byte(pubResp.Pem))
+	if block == nil {
+		return "", errors.New("failed to decode PEM public key from KMS")
+	}
+
+	parsed, err := x509.ParsePKIXPublicKey(block.Bytes)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse public key: %w", err)
+	}
+
+	switch pubResp.Algorithm { //nolint:exhaustive
+	case kmspb.CryptoKeyVersion_EC_SIGN_ED25519:
+		pk, ok := parsed.(ed25519.PublicKey)
+		if !ok {
+			return "", errors.New("KMS reported ED25519 but parsed key is different type")
+		}
+		auth.PrintEd25519KeyInfo(pk)
+		return auth.AlgorithmEd25519, nil
+	case kmspb.CryptoKeyVersion_EC_SIGN_P384_SHA384:
+		pk, ok := parsed.(*ecdsa.PublicKey)
+		if !ok || pk.Curve != elliptic.P384() {
+			return "", errors.New("KMS reported P-384 but parsed key mismatch")
+		}
+		auth.PrintEcdsaP384KeyInfo(*pk)
+		return auth.AlgorithmECDSAP384, nil
+	default:
+		return "", fmt.Errorf("unsupported KMS key algorithm: %s", pubResp.Algorithm.String())
+	}
+}
--- a/cmd/publisher/auth/http.go
+++ b/cmd/publisher/auth/http.go
@ -5,14 +5,13 @@ type HTTPProvider struct {
 }

 // NewHTTPProvider creates a new HTTP-based auth provider
-func NewHTTPProvider(registryURL, domain, privateKey string, cryptoAlgorithm CryptoAlgorithm) Provider {
+func NewHTTPProvider(registryURL, domain string, signer *Signer) Provider {
 	return &HTTPProvider{
 		CryptoProvider: &CryptoProvider{
-			registryURL:     registryURL,
-			domain:          domain,
-			privateKey:      privateKey,
-			cryptoAlgorithm: cryptoAlgorithm,
-			authMethod:      "http",
+			registryURL: registryURL,
+			domain:      domain,
+			signer:      *signer,
+			authMethod:  "http",
 		},
 	}
 }