Make vault client more robust by handling invariant violations (spiff…

…e#5058)

pkg/server/plugin/keymanager/hashicorpvault/hashicorp_vault.go

@@ -17,7 +17,6 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"os"
-	"strings"
 	"sync"
 )
 
@@ -214,29 +213,16 @@ func (p *Plugin) SignData(ctx context.Context, req *keymanagerv1.SignDataRequest
 		return nil, err
 	}
 
-	// TODO: Should this be done in SignData?
+	// TODO: Should the encoding be done in SignData?
 	encodedData := base64.StdEncoding.EncodeToString(req.Data)
-	signResp, err := p.vc.SignData(ctx, req.KeyId, encodedData, hashAlgo, signingAlgo)
 
+	signature, err := p.vc.SignData(ctx, req.KeyId, encodedData, hashAlgo, signingAlgo)
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to sign: %v", err)
-	}
-
-	// TODO: Should this be done in SignData?
-	sig := signResp.Data["signature"].(string)
-	cutSig, ok := strings.CutPrefix(sig, "vault:v1:")
-
-	if !ok {
-		return nil, status.Errorf(codes.Internal, "response should contain vault prefix: %v", err)
-	}
-
-	data, err := base64.StdEncoding.DecodeString(cutSig)
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "unable to base64 decode signature: %v", err)
+		return nil, err
 	}
 
 	return &keymanagerv1.SignDataResponse{
-		Signature:      data,
+		Signature:      signature,
 		KeyFingerprint: keyEntry.PublicKey.Fingerprint,
 	}, nil
 }

pkg/server/plugin/keymanager/hashicorpvault/vault_client.go

@@ -3,6 +3,7 @@ package hashicorpvault
 import (
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/base64"
 	"fmt"
 	"github.com/hashicorp/go-hclog"
 	vapi "github.com/hashicorp/vault/api"
@@ -12,6 +13,7 @@ import (
 	"google.golang.org/grpc/status"
 	"net/http"
 	"os"
+	"strings"
 
 	"github.com/spiffe/spire/pkg/common/pemutil"
 )
@@ -388,7 +390,7 @@ func (c *Client) GetKey(ctx context.Context, spireKeyID string) (*vapi.Secret, e
 	return c.vaultClient.Logical().ReadWithContext(ctx, fmt.Sprintf("/transit/keys/%s", spireKeyID))
 }
 
-func (c *Client) SignData(ctx context.Context, spireKeyID string, data string, hashAlgo TransitHashAlgorithm, signatureAlgo TransitSignatureAlgorithm) (*vapi.Secret, error) {
+func (c *Client) SignData(ctx context.Context, spireKeyID string, data string, hashAlgo TransitHashAlgorithm, signatureAlgo TransitSignatureAlgorithm) ([]byte, error) {
 	body := map[string]interface{}{
 		"key_version":           "0", // always use tha latest version
 		"input":                 data,
@@ -399,5 +401,30 @@ func (c *Client) SignData(ctx context.Context, spireKeyID string, data string, h
 
 	// TODO: Handle errors here
 	// TODO: Make the transit engine path configurable
-	return c.vaultClient.Logical().WriteWithContext(ctx, fmt.Sprintf("/transit/sign/%s/%s", spireKeyID, hashAlgo), body)
+	sigResp, err := c.vaultClient.Logical().WriteWithContext(ctx, fmt.Sprintf("/transit/sign/%s/%s", spireKeyID, hashAlgo), body)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "transit engine sign call failed: %v", err)
+	}
+
+	sig, ok := sigResp.Data["signature"]
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "transit engine sign call was successful but signature is missing: %v", err)
+	}
+
+	sigStr, ok := sig.(string)
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "expected signature data type %T but got %T", sigStr, sig)
+	}
+
+	cutSig, ok := strings.CutPrefix(sigStr, "vault:v1:")
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "signature is missing vault prefix: %v", err)
+	}
+
+	sigData, err := base64.StdEncoding.DecodeString(cutSig)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "unable to base64 decode signature: %v", err)
+	}
+
+	return sigData, nil
 }