feat(sdk): Allow key splits with same algo

sdk/basekey.go

@@ -77,12 +77,12 @@ func getBaseKey(ctx context.Context, s SDK) (*policy.SimpleKasKey, error) {
 
 	baseKeyStructure, ok := configMap[baseKeyWellKnown]
 	if !ok {
-		return nil, errBaseKeyNotFound
+		return nil, ErrBaseKeyNotFound
 	}
 
 	baseKeyMap, ok := baseKeyStructure.(map[string]interface{})
 	if !ok {
-		return nil, errBaseKeyInvalidFormat
+		return nil, ErrBaseKeyInvalidFormat
 	}
 
 	simpleKasKey, err := parseSimpleKasKey(baseKeyMap)
@@ -97,28 +97,28 @@ func parseSimpleKasKey(baseKeyMap map[string]interface{}) (*policy.SimpleKasKey,
 	simpleKasKey := &policy.SimpleKasKey{}
 
 	if len(baseKeyMap) == 0 {
-		return nil, errBaseKeyEmpty
+		return nil, ErrBaseKeyEmpty
 	}
 
 	publicKey, ok := baseKeyMap[baseKeyPublicKey].(map[string]interface{})
 	if !ok {
-		return nil, errBaseKeyInvalidFormat
+		return nil, ErrBaseKeyInvalidFormat
 	}
 
 	alg, ok := publicKey[baseKeyAlg].(string)
 	if !ok {
-		return nil, errBaseKeyInvalidFormat
+		return nil, ErrBaseKeyInvalidFormat
 	}
 	publicKey[baseKeyAlg] = getKasKeyAlg(alg)
 	baseKeyMap[baseKeyPublicKey] = publicKey
 	configJSON, err := json.Marshal(baseKeyMap)
 	if err != nil {
-		return nil, errors.Join(errMarshalBaseKeyFailed, err)
+		return nil, errors.Join(ErrMarshalBaseKeyFailed, err)
 	}
 
 	err = protojson.Unmarshal(configJSON, simpleKasKey)
 	if err != nil {
-		return nil, errors.Join(errUnmarshalBaseKeyFailed, err)
+		return nil, errors.Join(ErrUnmarshalBaseKeyFailed, err)
 	}
 	return simpleKasKey, nil
 }

sdk/basekey_test.go

@@ -249,7 +249,7 @@ func (s *BaseKeyTestSuite) TestGetBaseKeyMissingBaseKey() {
 	s.Require().True(mockService.called)
 	s.Require().Error(err)
 	s.Require().Nil(baseKey)
-	s.Require().Contains(err.Error(), errBaseKeyNotFound.Error())
+	s.Require().ErrorIs(err, ErrBaseKeyNotFound)
 }
 
 func (s *BaseKeyTestSuite) TestGetBaseKeyInvalidBaseKeyFormat() {
@@ -267,7 +267,7 @@ func (s *BaseKeyTestSuite) TestGetBaseKeyInvalidBaseKeyFormat() {
 	s.Require().True(mockService.called)
 	s.Require().Error(err)
 	s.Require().Nil(baseKey)
-	s.Require().ErrorContains(err, errBaseKeyInvalidFormat.Error())
+	s.Require().ErrorIs(err, ErrBaseKeyInvalidFormat)
 }
 
 func (s *BaseKeyTestSuite) TestGetBaseKeyEmptyBaseKey() {
@@ -286,7 +286,7 @@ func (s *BaseKeyTestSuite) TestGetBaseKeyEmptyBaseKey() {
 	s.Require().True(mockService.called)
 	s.Require().Error(err)
 	s.Require().Nil(baseKey)
-	s.Require().ErrorContains(err, errBaseKeyEmpty.Error())
+	s.Require().ErrorIs(err, ErrBaseKeyEmpty)
 }
 
 func (s *BaseKeyTestSuite) TestGetBaseKeyMissingPublicKey() {
@@ -308,7 +308,7 @@ func (s *BaseKeyTestSuite) TestGetBaseKeyMissingPublicKey() {
 	s.Require().True(mockService.called)
 	s.Require().Error(err)
 	s.Require().Nil(baseKey)
-	s.Require().ErrorContains(err, errBaseKeyInvalidFormat.Error())
+	s.Require().ErrorIs(err, ErrBaseKeyInvalidFormat)
 }
 
 func (s *BaseKeyTestSuite) TestGetBaseKeyInvalidPublicKey() {
@@ -330,5 +330,5 @@ func (s *BaseKeyTestSuite) TestGetBaseKeyInvalidPublicKey() {
 	s.Require().True(mockService.called)
 	s.Require().Error(err)
 	s.Require().Nil(baseKey)
-	s.Require().ErrorContains(err, errBaseKeyInvalidFormat.Error())
+	s.Require().ErrorIs(err, ErrBaseKeyInvalidFormat)
 }

sdk/basekeyerrors.go

@@ -3,9 +3,9 @@ package sdk
 import "errors"
 
 var (
-	errBaseKeyNotFound        = errors.New("base key not found in well-known configuration")
-	errBaseKeyInvalidFormat   = errors.New("base key has invalid format")
-	errBaseKeyEmpty           = errors.New("base key is empty or not provided")
-	errMarshalBaseKeyFailed   = errors.New("failed to marshal base key configuration")
-	errUnmarshalBaseKeyFailed = errors.New("failed to unmarshal base key configuration")
+	ErrBaseKeyNotFound        = errors.New("base key not found in well-known configuration")
+	ErrBaseKeyInvalidFormat   = errors.New("base key has invalid format")
+	ErrBaseKeyEmpty           = errors.New("base key is empty or not provided")
+	ErrMarshalBaseKeyFailed   = errors.New("failed to marshal base key configuration")
+	ErrUnmarshalBaseKeyFailed = errors.New("failed to unmarshal base key configuration")
 )

sdk/granter.go

@@ -388,23 +388,23 @@ func newGranterFromService(ctx context.Context, keyCache *kasKeyCache, as sdkcon
 		def := pair.GetAttribute()
 
 		if def != nil {
-			storeKeysToCache(def.GetGrants(), keyCache)
+			storeKeysToCache(def.GetGrants(), def.GetKasKeys(), keyCache)
 		}
 		v := pair.GetValue()
 		gType := noKeysFound
 		if v != nil {
 			gType = grants.addAllGrants(fqn, v, def)
-			storeKeysToCache(v.GetGrants(), keyCache)
+			storeKeysToCache(v.GetGrants(), v.GetKasKeys(), keyCache)
 		}
 
 		// If no more specific grant was found, then add the value grants
 		if gType == noKeysFound && def != nil {
 			gType = grants.addAllGrants(fqn, def, def)
-			storeKeysToCache(def.GetGrants(), keyCache)
+			storeKeysToCache(def.GetGrants(), def.GetKasKeys(), keyCache)
 		}
 		if gType == noKeysFound && def.GetNamespace() != nil {
 			grants.addAllGrants(fqn, def.GetNamespace(), def)
-			storeKeysToCache(def.GetNamespace().GetGrants(), keyCache)
+			storeKeysToCache(def.GetNamespace().GetGrants(), def.GetNamespace().GetKasKeys(), keyCache)
 		}
 	}
 
@@ -429,7 +429,7 @@ func algProto2String(e policy.KasPublicKeyAlgEnum) string {
 	return ""
 }
 
-func storeKeysToCache(kases []*policy.KeyAccessServer, c *kasKeyCache) {
+func storeKeysToCache(kases []*policy.KeyAccessServer, keys []*policy.SimpleKasKey, c *kasKeyCache) {
 	for _, kas := range kases {
 		keys := kas.GetPublicKey().GetCached().GetKeys()
 		if len(keys) == 0 {
@@ -445,6 +445,18 @@ func storeKeysToCache(kases []*policy.KeyAccessServer, c *kasKeyCache) {
 			})
 		}
 	}
+	for _, key := range keys {
+		alg, err := formatAlg(key.GetPublicKey().GetAlgorithm())
+		if err != nil {
+			continue
+		}
+		c.store(KASInfo{
+			URL:       key.GetKasUri(),
+			KID:       key.GetPublicKey().GetKid(),
+			Algorithm: alg,
+			PublicKey: key.GetPublicKey().GetPem(),
+		})
+	}
 }
 
 // Given a policy (list of data attributes or tags),
@@ -472,16 +484,16 @@ func newGranterFromAttributes(keyCache *kasKeyCache, attrs ...*policy.Value) (gr
 		}
 
 		if grants.addAllGrants(fqn, v, def) != noKeysFound {
-			storeKeysToCache(v.GetGrants(), keyCache)
+			storeKeysToCache(v.GetGrants(), v.GetKasKeys(), keyCache)
 			continue
 		}
 		// If no more specific grant was found, then add the attr grants
 		if grants.addAllGrants(fqn, def, def) != noKeysFound {
-			storeKeysToCache(def.GetGrants(), keyCache)
+			storeKeysToCache(def.GetGrants(), def.GetKasKeys(), keyCache)
 			continue
 		}
 		grants.addAllGrants(fqn, namespace, def)
-		storeKeysToCache(namespace.GetGrants(), keyCache)
+		storeKeysToCache(namespace.GetGrants(), namespace.GetKasKeys(), keyCache)
 	}
 
 	return grants, nil

sdk/granter_test.go

@@ -30,7 +30,7 @@ const (
 	specifiedKas        = "https://attr.kas.com/"
 	evenMoreSpecificKas = "https://value.kas.com/"
 	lessSpecificKas     = "https://namespace.kas.com/"
-	fakePem             = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQ...\n-----END PUBLIC KEY-----\n"
+	fakePem             = mockRSAPublicKey1
 )
 
 var (
@@ -74,6 +74,8 @@ var (
 	MP, _  = NewAttributeNameFQN("https://virtru.com/attr/mapped")
 	mpa, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/a")
 	mpb, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/b")
+	mpc, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/c")
+	mpd, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/d")
 	mpu, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/unspecified")
 )
 
@@ -248,7 +250,7 @@ func mockSimpleKasKey(kas, kid string) *policy.SimpleKasKey {
 	}
 	var alg policy.Algorithm
 	switch kid {
-	case "r1":
+	case "r0", "r1", "r3":
 		alg = policy.Algorithm_ALGORITHM_RSA_2048
 	case "r2":
 		alg = policy.Algorithm_ALGORITHM_RSA_4096
@@ -262,7 +264,7 @@ func mockSimpleKasKey(kas, kid string) *policy.SimpleKasKey {
 		PublicKey: &policy.SimpleKasPublicKey{
 			Algorithm: alg,
 			Kid:       kid,
-			Pem:       "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQ...\n-----END PUBLIC KEY-----\n",
+			Pem:       fakePem,
 		},
 	}
 }
@@ -339,17 +341,20 @@ func mockValueFor(fqn AttributeValueFQN) *policy.Value {
 		switch strings.ToLower(fqn.Value()) {
 		case "a":
 			p.KasKeys = make([]*policy.SimpleKasKey, 1)
-			p.Grants = make([]*policy.KeyAccessServer, 1)
 			p.KasKeys[0] = mockSimpleKasKey(evenMoreSpecificKas, "r2")
-			p.Grants[0] = mockGrant(evenMoreSpecificKas, "r2")
-			p.Grants[0].PublicKey = createPublicKey("r2", fakePem, policy.KasPublicKeyAlgEnum_KAS_PUBLIC_KEY_ALG_ENUM_RSA_2048)
 
 		case "b":
 			p.KasKeys = make([]*policy.SimpleKasKey, 1)
-			p.Grants = make([]*policy.KeyAccessServer, 1)
 			p.KasKeys[0] = mockSimpleKasKey(evenMoreSpecificKas, "e1")
-			p.Grants[0] = mockGrant(evenMoreSpecificKas, "e1")
-			p.Grants[0].PublicKey = createPublicKey("e1", fakePem, policy.KasPublicKeyAlgEnum_KAS_PUBLIC_KEY_ALG_ENUM_RSA_2048)
+
+		case "c":
+			p.KasKeys = make([]*policy.SimpleKasKey, 1)
+			p.KasKeys[0] = mockSimpleKasKey(evenMoreSpecificKas, "r0")
+
+		case "d":
+			p.KasKeys = make([]*policy.SimpleKasKey, 1)
+			p.KasKeys[0] = mockSimpleKasKey(evenMoreSpecificKas, "r3")
+
 		case "unspecified":
 			// defaults only
 		default:

sdk/kas_client.go

@@ -378,7 +378,7 @@ func (k *KASClient) getRewrapRequest(reqs []*kas.UnsignedRewrapRequest_WithPolic
 }
 
 type kasKeyRequest struct {
-	url, algorithm string
+	url, algorithm, kid string
 }
 
 type timeStampedKASInfo struct {
@@ -399,10 +399,19 @@ func (c *kasKeyCache) clear() {
 	c.c = make(map[kasKeyRequest]timeStampedKASInfo)
 }
 
-func (c *kasKeyCache) get(url, algorithm string) *KASInfo {
-	cacheKey := kasKeyRequest{url, algorithm}
+func (c *kasKeyCache) get(url, algorithm, kid string) *KASInfo {
+	cacheKey := kasKeyRequest{url, algorithm, kid}
 	now := time.Now()
 	cv, ok := c.c[cacheKey]
+	if !ok && kid == "" {
+		for k, v := range c.c {
+			if k.url == url && k.algorithm == algorithm {
+				cv = v
+				ok = true
+				break
+			}
+		}
+	}
 	if !ok {
 		return nil
 	}
@@ -415,13 +424,13 @@ func (c *kasKeyCache) get(url, algorithm string) *KASInfo {
 }
 
 func (c *kasKeyCache) store(ki KASInfo) {
-	cacheKey := kasKeyRequest{ki.URL, ki.Algorithm}
+	cacheKey := kasKeyRequest{ki.URL, ki.Algorithm, ki.KID}
 	c.c[cacheKey] = timeStampedKASInfo{ki, time.Now()}
 }
 
-func (s SDK) getPublicKey(ctx context.Context, kasurl, algorithm string) (*KASInfo, error) {
+func (s SDK) getPublicKey(ctx context.Context, kasurl, algorithm, kidToFind string) (*KASInfo, error) {
 	if s.kasKeyCache != nil {
-		if cachedValue := s.kasKeyCache.get(kasurl, algorithm); nil != cachedValue {
+		if cachedValue := s.kasKeyCache.get(kasurl, algorithm, kidToFind); nil != cachedValue {
 			return cachedValue, nil
 		}
 	}