feat(sdk): resolve key splits via GetKeyMappingsByFqns

sdk/granter.go

@@ -10,6 +10,8 @@ import (
 	"sort"
 	"strings"
 
+	"connectrpc.com/connect"
+
 	"github.com/opentdf/platform/lib/ocrypto"
 	"github.com/opentdf/platform/protocol/go/policy"
 	"github.com/opentdf/platform/protocol/go/policy/attributes"
@@ -419,55 +421,124 @@ func (r granter) byAttribute(fqn AttributeValueFQN) *keyAccessGrant {
 	return r.grantTable[fqn.key]
 }
 
-// Gets a list of directory of KAS grants for a list of attribute FQNs
+// Gets a list of directory of KAS grants for a list of attribute FQNs.
+//
+// Resolution uses GetKeyMappingsByFqns, which returns the effective mapped KAS
+// keys per value (value > definition > namespace precedence, resolved
+// server-side). Values that resolve to no mapped keys (e.g. configured only with
+// legacy KAS grants) fall back to GetAttributeValuesByFqns so their grants still
+// resolve.
 func newGranterFromService(ctx context.Context, logger *slog.Logger, keyCache *kasKeyCache, as sdkconnect.AttributesServiceClient, fqns ...AttributeValueFQN) (granter, error) {
 	fqnsStr := make([]string, len(fqns))
 	for i, v := range fqns {
 		fqnsStr[i] = v.String()
 	}
 
-	av, err := as.GetAttributeValuesByFqns(ctx, &attributes.GetAttributeValuesByFqnsRequest{
-		Fqns: fqnsStr,
-	})
-	if err != nil {
-		return granter{}, err
-	}
-
 	grants := granter{
 		logger:     logger,
 		tags:       fqns,
 		grantTable: make(map[string]*keyAccessGrant),
 		keyCache:   &rlKeyCache{c: make(map[ResourceLocator]*policy.SimpleKasKey)},
 	}
+
+	km, err := as.GetKeyMappingsByFqns(ctx, &attributes.GetKeyMappingsByFqnsRequest{Fqns: fqnsStr})
+	if err != nil {
+		// Older platforms do not expose GetKeyMappingsByFqns; fall back to the full
+		// attribute lookup for all values.
+		if connect.CodeOf(err) == connect.CodeUnimplemented {
+			if e := grants.addGrantsFromAttributeValues(ctx, logger, keyCache, as, fqnsStr); e != nil {
+				return granter{}, e
+			}
+			return grants, nil
+		}
+		return granter{}, err
+	}
+
+	// Iterate the requested FQNs (not the response map) so a value the server
+	// omits still falls back rather than being silently dropped. The response is
+	// keyed by normalized FQN, and the protobuf getters are nil-safe, so an absent
+	// or nil entry yields empty keys and resolves via the fallback below.
+	mappings := km.GetFqnKeyMappings()
+	var fallback []string
+	for _, fqn := range fqns {
+		m := mappings[fqn.key]
+		keys := m.GetKeys()
+		if len(keys) == 0 {
+			// No mapped keys (legacy-grant-only, unconfigured, or omitted); resolve
+			// via the full attribute lookup below.
+			fallback = append(fallback, fqn.key)
+			continue
+		}
+		// Only the rule and FQN are read downstream (constructAttributeBoolean /
+		// ruleToOperator), so a minimal attribute is sufficient.
+		attr := &policy.Attribute{Fqn: fqn.Prefix().String(), Rule: m.GetRule()}
+		for _, sk := range keys {
+			if err := grants.addMappedKey(fqn, sk); err != nil {
+				logger.Warn("failed to add mapped key",
+					slog.Any("fqn", fqn),
+					slog.String("kas", sk.GetKasUri()),
+					slog.Any("error", err),
+				)
+				continue
+			}
+			grants.typ = mappedFound
+			if _, present := grants.grantTable[fqn.key]; !present {
+				grants.grantTable[fqn.key] = &keyAccessGrant{attr, []string{sk.GetKasUri()}}
+			} else {
+				grants.grantTable[fqn.key].kases = append(grants.grantTable[fqn.key].kases, sk.GetKasUri())
+			}
+			// Populate the SDK key cache for wrapping (addMappedKey already filled
+			// the granter's rlKeyCache).
+			storeKeysToCache(logger, nil, []*policy.SimpleKasKey{sk}, keyCache, nil)
+		}
+	}
+
+	if len(fallback) > 0 {
+		if err := grants.addGrantsFromAttributeValues(ctx, logger, keyCache, as, fallback); err != nil {
+			return granter{}, err
+		}
+	}
+
+	return grants, nil
+}
+
+// addGrantsFromAttributeValues resolves grants for the given value FQNs via the
+// full GetAttributeValuesByFqns lookup (value > definition > namespace), used as
+// the fallback for values with no mapped keys.
+func (r *granter) addGrantsFromAttributeValues(ctx context.Context, logger *slog.Logger, keyCache *kasKeyCache, as sdkconnect.AttributesServiceClient, fqns []string) error {
+	av, err := as.GetAttributeValuesByFqns(ctx, &attributes.GetAttributeValuesByFqnsRequest{Fqns: fqns})
+	if err != nil {
+		return err
+	}
 	for fqnstr, pair := range av.GetFqnAttributeValues() {
 		fqn, err := NewAttributeValueFQN(fqnstr)
 		if err != nil {
-			return grants, err
+			return err
 		}
 		def := pair.GetAttribute()
 
 		if def != nil {
-			storeKeysToCache(logger, def.GetGrants(), def.GetKasKeys(), keyCache, grants.keyCache)
+			storeKeysToCache(logger, def.GetGrants(), def.GetKasKeys(), keyCache, r.keyCache)
 		}
 		v := pair.GetValue()
 		gType := noKeysFound
 		if v != nil {
-			gType = grants.addAllGrants(fqn, v, def)
-			storeKeysToCache(logger, v.GetGrants(), v.GetKasKeys(), keyCache, grants.keyCache)
+			gType = r.addAllGrants(fqn, v, def)
+			storeKeysToCache(logger, v.GetGrants(), v.GetKasKeys(), keyCache, r.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(logger, def.GetGrants(), def.GetKasKeys(), keyCache, grants.keyCache)
+			gType = r.addAllGrants(fqn, def, def)
+			storeKeysToCache(logger, def.GetGrants(), def.GetKasKeys(), keyCache, r.keyCache)
 		}
 		if gType == noKeysFound && def.GetNamespace() != nil {
-			grants.addAllGrants(fqn, def.GetNamespace(), def)
-			storeKeysToCache(logger, def.GetNamespace().GetGrants(), def.GetNamespace().GetKasKeys(), keyCache, grants.keyCache)
+			r.addAllGrants(fqn, def.GetNamespace(), def)
+			storeKeysToCache(logger, def.GetNamespace().GetGrants(), def.GetNamespace().GetKasKeys(), keyCache, r.keyCache)
 		}
 	}
 
-	return grants, nil
+	return nil
 }
 
 func algProto2String(e policy.KasPublicKeyAlgEnum) string {

sdk/granter_test.go

@@ -11,6 +11,8 @@ import (
 	"strings"
 	"testing"
 
+	"connectrpc.com/connect"
+
 	"github.com/opentdf/platform/lib/ocrypto"
 	"github.com/opentdf/platform/protocol/go/policy"
 	"github.com/opentdf/platform/protocol/go/policy/attributes"
@@ -812,6 +814,115 @@ func (*mockAttributesClient) GetAttributeValuesByFqns(_ context.Context, req *at
 	}, nil
 }
 
+// effectiveMockKasKeys mirrors the server-side resolveEffectiveKasKeys: the
+// most-specific mapped keys with value > definition > namespace precedence.
+func effectiveMockKasKeys(v *policy.Value, a *policy.Attribute) []*policy.SimpleKasKey {
+	if k := v.GetKasKeys(); len(k) > 0 {
+		return k
+	}
+	if k := a.GetKasKeys(); len(k) > 0 {
+		return k
+	}
+	if k := a.GetNamespace().GetKasKeys(); len(k) > 0 {
+		return k
+	}
+	return nil
+}
+
+func (*mockAttributesClient) GetKeyMappingsByFqns(_ context.Context, req *attributes.GetKeyMappingsByFqnsRequest) (*attributes.GetKeyMappingsByFqnsResponse, error) {
+	out := make(map[string]*attributes.GetKeyMappingsByFqnsResponse_AttributeKeyMapping)
+	for _, v := range req.GetFqns() {
+		vfqn, err := NewAttributeValueFQN(v)
+		if err != nil {
+			return nil, err
+		}
+		val := mockValueFor(vfqn)
+		attr := val.GetAttribute()
+		out[v] = &attributes.GetKeyMappingsByFqnsResponse_AttributeKeyMapping{
+			Rule: attr.GetRule(),
+			Keys: effectiveMockKasKeys(val, attr),
+		}
+	}
+	return &attributes.GetKeyMappingsByFqnsResponse{FqnKeyMappings: out}, nil
+}
+
+// unimplementedKeyMappingsClient models an older platform that does not expose
+// GetKeyMappingsByFqns; the granter must fall back to GetAttributeValuesByFqns.
+type unimplementedKeyMappingsClient struct {
+	mockAttributesClient
+}
+
+func (*unimplementedKeyMappingsClient) GetKeyMappingsByFqns(_ context.Context, _ *attributes.GetKeyMappingsByFqnsRequest) (*attributes.GetKeyMappingsByFqnsResponse, error) {
+	return nil, connect.NewError(connect.CodeUnimplemented, errors.New("not implemented"))
+}
+
+// emptyKeyMappingsClient models a server that succeeds but omits the requested
+// FQNs from the response map; the granter must fall back per value rather than
+// silently drop them.
+type emptyKeyMappingsClient struct {
+	mockAttributesClient
+}
+
+func (*emptyKeyMappingsClient) GetKeyMappingsByFqns(_ context.Context, _ *attributes.GetKeyMappingsByFqnsRequest) (*attributes.GetKeyMappingsByFqnsResponse, error) {
+	return &attributes.GetKeyMappingsByFqnsResponse{}, nil
+}
+
+// TestReasonerKeyMappingFallback verifies that grant-only values and older
+// platforms (no GetKeyMappingsByFqns) resolve to the same plan as the key-mapping
+// path, via the GetAttributeValuesByFqns fallback.
+func TestReasonerKeyMappingFallback(t *testing.T) {
+	policy := []AttributeValueFQN{spk2spk} // grant-based value (no mapped keys)
+	defaults := []string{kasUs}
+	want := []keySplitStep{{evenMoreSpecificKas, ""}}
+	noSplit := func() string { return "" }
+
+	for _, tc := range []struct {
+		name string
+		as   sdkconnect.AttributesServiceClient
+	}{
+		{"per-value fallback (empty key set)", &mockAttributesClient{}},
+		{"full fallback (RPC unimplemented)", &unimplementedKeyMappingsClient{}},
+		{"per-value fallback (FQN omitted from response)", &emptyKeyMappingsClient{}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			reasoner, err := newGranterFromService(t.Context(), slog.Default(), newKasKeyCache(), tc.as, policy...)
+			require.NoError(t, err)
+			plan, err := reasoner.plan(defaults, noSplit)
+			require.NoError(t, err)
+			assert.ElementsMatch(t, want, plan)
+		})
+	}
+}
+
+// A single request that mixes mapped-key values (MP: mapped kas_keys) with
+// grant-configured values (REL: legacy grants) must resolve every attribute and
+// combine them into one plan. The server returns keys for both (cached-key grants
+// are converted alongside mapped keys), so the plan should carry the mapped KAS
+// and both grant KASes.
+func TestReasonerMixedMappingsAndGrants(t *testing.T) {
+	policy := []AttributeValueFQN{mpa, mpb, rel2aus, rel2can}
+	reasoner, err := newGranterFromService(t.Context(), slog.Default(), newKasKeyCache(), &mockAttributesClient{}, policy...)
+	require.NoError(t, err)
+
+	i := 0
+	plan, err := reasoner.plan([]string{kasUs}, func() string {
+		i++
+		return strconv.Itoa(i)
+	})
+	require.NoError(t, err)
+
+	// The mapped attribute (MP) forms its own split; the grant attribute (REL,
+	// ANY_OF) shares a split across its two KASes. Assert the full plan (KAS plus
+	// split id) rather than a KAS set, so mapped-key splits stay distinct instead
+	// of collapsing by KAS.
+	want := []keySplitStep{
+		{evenMoreSpecificKas, "1"},
+		{kasAu, "2"},
+		{kasCa, "2"},
+	}
+	assert.ElementsMatch(t, want, plan)
+}
+
 // Tests titles are written in the form [{attr}.{value}] => [{resulting kas boolean exp}]
 // where the left hand side is the list of attributes passed in and the right
 // is the resulting split steps