From 1b27bc7ea24137196cb2f1168fdde32920de27fe Mon Sep 17 00:00:00 2001
From: Paul Flynn <paul@arkavo.com>
Date: Sat, 11 May 2024 19:22:51 -0400
Subject: [PATCH] Implement certificate caching in PublicKey Provider

Added certificate caching mechanism to the PublicKey Provider to reduce the number of calls to the CryptoProvider. Now, before performing an expensive certificate generation operation, the system will first check if a cached certificate for the requested key exists. If it does, it returns the cached certificate, otherwise, it generates a new one and stores it in the cache for future requests.
---
 service/kas/access/publicKey.go | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/service/kas/access/publicKey.go b/service/kas/access/publicKey.go
index 120e6188f..399212cdf 100644
--- a/service/kas/access/publicKey.go
+++ b/service/kas/access/publicKey.go
@@ -8,11 +8,12 @@ import (
 	"encoding/pem"
 	"errors"
 	"log/slog"
+	"sync"
 
 	kaspb "github.com/arkavo-org/opentdf-platform/protocol/go/kas"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
-	wrapperspb "google.golang.org/protobuf/types/known/wrapperspb"
+	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
 const (
@@ -21,9 +22,10 @@ const (
 	algorithmEc256       = "ec:secp256r1"
 )
 
+var ecCertCache sync.Map
+
 func (p *Provider) LegacyPublicKey(ctx context.Context, in *kaspb.LegacyPublicKeyRequest) (*wrapperspb.StringValue, error) {
 	algorithm := in.GetAlgorithm()
-	var cert string
 	var err error
 	if p.CryptoProvider == nil {
 		return nil, errors.Join(ErrConfig, status.Error(codes.Internal, "configuration error"))
@@ -34,17 +36,25 @@ func (p *Provider) LegacyPublicKey(ctx context.Context, in *kaspb.LegacyPublicKe
 		if !ok {
 			return nil, errors.New("services.kas.eccertid is not a string")
 		}
-		cert, err = p.CryptoProvider.ECCertificate(ecCertID)
-		if err != nil {
-			slog.ErrorContext(ctx, "CryptoProvider.ECPublicKey failed", "err", err)
-			return nil, errors.Join(ErrConfig, status.Error(codes.Internal, "configuration error"))
+		if cert, exists := ecCertCache.Load(ecCertID); exists {
+			return cert.(*wrapperspb.StringValue), nil
 		}
-	} else {
-		cert, err = p.CryptoProvider.RSAPublicKey("unknown")
+		cert, err := p.CryptoProvider.ECCertificate(ecCertID)
 		if err != nil {
-			slog.ErrorContext(ctx, "CryptoProvider.RSAPublicKey failed", "err", err)
+			slog.ErrorContext(ctx, "CryptoProvider.ECPublicKey failed", "err", err)
 			return nil, errors.Join(ErrConfig, status.Error(codes.Internal, "configuration error"))
 		}
+		// workaround for Error code 75497574.  [ec_key_pair.cpp:650] Failed to create X509 cert struct.error:04800066:PEM routines::bad end line
+		cert += "\n"
+		ecCertStringValue := &wrapperspb.StringValue{Value: cert}
+		// Store the certificate in the cache
+		ecCertCache.Store(ecCertID, ecCertStringValue)
+		return ecCertStringValue, nil
+	}
+	cert, err := p.CryptoProvider.RSAPublicKey("unknown")
+	if err != nil {
+		slog.ErrorContext(ctx, "CryptoProvider.RSAPublicKey failed", "err", err)
+		return nil, errors.Join(ErrConfig, status.Error(codes.Internal, "configuration error"))
 	}
 	// workaround for Error code 75497574.  [ec_key_pair.cpp:650] Failed to create X509 cert struct.error:04800066:PEM routines::bad end line
 	cert += "\n"