[CRE] Allow the Gateway to handle multiple dons

core/services/gateway/handlers/vault/aggregator.go

@@ -8,6 +8,7 @@ import (
 	"maps"
 	"slices"
 	"strconv"
+	"strings"
 
 	"github.com/ethereum/go-ethereum/common"
 
@@ -20,6 +21,14 @@ import (
 
 type baseAggregator struct {
 	capabilitiesRegistry capabilitiesRegistry
+	// vaultHandlerDonID scopes registry lookup when several vault DONs exist.
+	//
+	// Source: gateway job TOML [[gatewayConfig.ShardedDONs]] DonName (see deployment/cre/jobs/pkg/gateway_job.go),
+	// loaded as ShardedDONConfig.DonName and passed as DONConfig.DonId (handler_factory.shardedDONsToLegacy;
+	// DonId is a legacy field name for that string, not the on-chain uint32 id).
+	//
+	// Matching: capabilities.DON.Name when non-empty (v2), else decimal capabilities.DON.ID string (v1 sync).
+	vaultHandlerDonID string
 }
 
 func (a *baseAggregator) Aggregate(ctx context.Context, l logger.Logger, resps map[string]jsonrpc.Response[json.RawMessage], currResp *jsonrpc.Response[json.RawMessage]) (*jsonrpc.Response[json.RawMessage], error) {
@@ -47,15 +56,57 @@ func (a *baseAggregator) donForVaultCapability(ctx context.Context) (*capabiliti
 	if err != nil {
 		return nil, err
 	}
-	// TODO: Support multiple vault capabilities in the capability registry.
-	// For the initial Smartcon deployment there will be exactly one Vault capability
-	// split across both DON families.
-	if len(dons) != 1 {
-		return nil, fmt.Errorf("expected exactly one DON for vault capability, found %d", len(dons))
+	if len(dons) == 0 {
+		return nil, fmt.Errorf("no DON found for vault capability %s", vaultcommon.CapabilityID)
+	}
+	if len(dons) == 1 {
+		don := dons[0]
+		return &don, nil
+	}
+
+	handlerDonID := strings.TrimSpace(a.vaultHandlerDonID)
+	if handlerDonID == "" {
+		return nil, fmt.Errorf("multiple DONs (%d) host vault capability %s but vault handler DonId is empty; set ShardedDONConfig.DonName so DONConfig.DonId matches the vault DON name or id in the registry (%s)",
+			len(dons), vaultcommon.CapabilityID, summarizeVaultRegistryDONs(dons))
 	}
 
-	don := dons[0]
-	return &don, nil
+	var matches []capabilities.DONWithNodes
+	for i := range dons {
+		d := dons[i]
+		if vaultDONMatchesHandlerDonID(&d.DON, handlerDonID) {
+			matches = append(matches, d)
+		}
+	}
+	switch len(matches) {
+	case 0:
+		return nil, fmt.Errorf("multiple DONs (%d) host vault capability %s but none match vault handler DonId %q; registry has %s",
+			len(dons), vaultcommon.CapabilityID, a.vaultHandlerDonID, summarizeVaultRegistryDONs(dons))
+	case 1:
+		d := matches[0]
+		return &d, nil
+	default:
+		return nil, fmt.Errorf("%d DONs match vault handler DonId %q for vault capability %s", len(matches), a.vaultHandlerDonID, vaultcommon.CapabilityID)
+	}
+}
+
+// vaultDONMatchesHandlerDonID reports whether don is the vault DON this handler is configured for.
+// handlerDonID is vaultHandlerDonID (jobspec DonName / DONConfig.DonId); see struct comment.
+func vaultDONMatchesHandlerDonID(don *capabilities.DON, handlerDonID string) bool {
+	if don.Name != "" {
+		return don.Name == handlerDonID
+	}
+	return strconv.FormatUint(uint64(don.ID), 10) == handlerDonID
+}
+
+func summarizeVaultRegistryDONs(dons []capabilities.DONWithNodes) string {
+	var b strings.Builder
+	for i, d := range dons {
+		if i > 0 {
+			b.WriteString("; ")
+		}
+		_, _ = fmt.Fprintf(&b, "name=%q id=%d", d.DON.Name, d.DON.ID)
+	}
+	return b.String()
 }
 
 func (a *baseAggregator) validateUsingQuorum(don capabilities.DON, resps map[string]jsonrpc.Response[json.RawMessage], l logger.Logger) (*jsonrpc.Response[json.RawMessage], error) {

core/services/gateway/handlers/vault/aggregator_test.go

@@ -287,3 +287,119 @@ func TestAggregator_QuorumUnobtainable(t *testing.T) {
 	_, err := agg.Aggregate(t.Context(), logger.Test(t), responses, resp3)
 	require.ErrorContains(t, err, "failed to validate using quorum: quorum unobtainable")
 }
+
+func makeDONWithNodesForTest(t *testing.T, name string, id uint32, f uint8, memberOffset byte, nodeCount int) capabilities.DONWithNodes {
+	t.Helper()
+	nodes := make([]capabilities.Node, nodeCount)
+	members := make([]p2ptypes.PeerID, nodeCount)
+	for i := 0; i < nodeCount; i++ {
+		pid := p2ptypes.PeerID{}
+		pid[0] = memberOffset + byte(i)
+		pid[1] = byte(i)
+		nodes[i] = capabilities.Node{PeerID: &pid, Signer: [32]byte{}}
+		members[i] = pid
+	}
+	return capabilities.DONWithNodes{
+		DON: capabilities.DON{
+			Name:    name,
+			ID:      id,
+			F:       f,
+			Members: members,
+		},
+		Nodes: nodes,
+	}
+}
+
+func TestAggregator_MultipleRegistryDONs_SelectsByVaultHandlerDonName(t *testing.T) {
+	donOther := makeDONWithNodesForTest(t, "staging-vault", 1, 2, 0x10, 7)
+	donMine := makeDONWithNodesForTest(t, "cre-reliability-vault", 2, 1, 0x20, 4)
+	mcr := &mockCapabilitiesRegistry{DONs: []capabilities.DONWithNodes{donOther, donMine}}
+	agg := &baseAggregator{
+		capabilitiesRegistry: mcr,
+		vaultHandlerDonID:    "cre-reliability-vault",
+	}
+
+	rm := json.RawMessage([]byte(`{}`))
+	currResp := jsonrpc.Response[json.RawMessage]{
+		Version: jsonrpc.JsonRpcVersion,
+		ID:      "1",
+		Method:  vaulttypes.MethodSecretsCreate,
+		Result:  &rm,
+	}
+	responses := map[string]jsonrpc.Response[json.RawMessage]{
+		"a": currResp,
+		"b": currResp,
+		"c": currResp,
+	}
+	resp, err := agg.Aggregate(t.Context(), logger.Test(t), responses, &currResp)
+	require.NoError(t, err)
+	require.Equal(t, currResp.ID, resp.ID)
+}
+
+func TestAggregator_MultipleRegistryDONs_SelectsByIDWhenNameEmpty(t *testing.T) {
+	donOther := makeDONWithNodesForTest(t, "", 1, 2, 0x10, 7)
+	donMine := makeDONWithNodesForTest(t, "", 99, 1, 0x20, 4)
+	mcr := &mockCapabilitiesRegistry{DONs: []capabilities.DONWithNodes{donOther, donMine}}
+	agg := &baseAggregator{
+		capabilitiesRegistry: mcr,
+		vaultHandlerDonID:    "99",
+	}
+
+	rm := json.RawMessage([]byte(`{}`))
+	currResp := jsonrpc.Response[json.RawMessage]{
+		Version: jsonrpc.JsonRpcVersion,
+		ID:      "1",
+		Method:  vaulttypes.MethodSecretsCreate,
+		Result:  &rm,
+	}
+	responses := map[string]jsonrpc.Response[json.RawMessage]{
+		"a": currResp,
+		"b": currResp,
+		"c": currResp,
+	}
+	resp, err := agg.Aggregate(t.Context(), logger.Test(t), responses, &currResp)
+	require.NoError(t, err)
+	require.Equal(t, currResp.ID, resp.ID)
+}
+
+func TestAggregator_MultipleRegistryDONs_NoMatchingVaultHandlerDonId(t *testing.T) {
+	donA := makeDONWithNodesForTest(t, "don-a", 1, 1, 0x10, 4)
+	donB := makeDONWithNodesForTest(t, "don-b", 2, 1, 0x20, 4)
+	mcr := &mockCapabilitiesRegistry{DONs: []capabilities.DONWithNodes{donA, donB}}
+	agg := &baseAggregator{
+		capabilitiesRegistry: mcr,
+		vaultHandlerDonID:    "unknown-vault",
+	}
+
+	rm := json.RawMessage([]byte(`{}`))
+	currResp := jsonrpc.Response[json.RawMessage]{
+		Version: jsonrpc.JsonRpcVersion,
+		ID:      "1",
+		Method:  vaulttypes.MethodSecretsCreate,
+		Result:  &rm,
+	}
+	responses := map[string]jsonrpc.Response[json.RawMessage]{"a": currResp}
+	_, err := agg.Aggregate(t.Context(), logger.Test(t), responses, &currResp)
+	require.ErrorContains(t, err, "none match vault handler DonId")
+}
+
+func TestAggregator_MultipleRegistryDONs_AmbiguousMatchingVaultHandlerDonId(t *testing.T) {
+	donA := makeDONWithNodesForTest(t, "same-name", 1, 1, 0x10, 4)
+	donB := makeDONWithNodesForTest(t, "same-name", 2, 1, 0x20, 4)
+	mcr := &mockCapabilitiesRegistry{DONs: []capabilities.DONWithNodes{donA, donB}}
+	agg := &baseAggregator{
+		capabilitiesRegistry: mcr,
+		vaultHandlerDonID:    "same-name",
+	}
+
+	rm := json.RawMessage([]byte(`{}`))
+	currResp := jsonrpc.Response[json.RawMessage]{
+		Version: jsonrpc.JsonRpcVersion,
+		ID:      "1",
+		Method:  vaulttypes.MethodSecretsCreate,
+		Result:  &rm,
+	}
+	responses := map[string]jsonrpc.Response[json.RawMessage]{"a": currResp}
+	_, err := agg.Aggregate(t.Context(), logger.Test(t), responses, &currResp)
+	require.ErrorContains(t, err, "2 DONs match vault handler DonId")
+}

core/services/gateway/handlers/vault/handler.go

@@ -271,7 +271,10 @@ func newHandlerWithAuthorizer(methodConfig json.RawMessage, donConfig *config.DO
 		jwtAuth:                   jwtAuth,
 		stopCh:                    make(services.StopChan),
 		metrics:                   metrics,
-		aggregator:                &baseAggregator{capabilitiesRegistry: capabilitiesRegistry},
+		aggregator: &baseAggregator{
+			capabilitiesRegistry: capabilitiesRegistry,
+			vaultHandlerDonID:    donConfig.DonId,
+		},
 		clock:                     clock,
 		RequestValidator:          vaultcap.NewRequestValidator(limiter, ciphertextLimiter, idKeyLengthLimiter, idOwnerLengthLimiter, idNamespaceLengthLimiter),
 	}, nil

core/services/gateway/handlers/vault/handler_test.go

@@ -127,11 +127,16 @@ func (m *mockAggregator) Aggregate(_ context.Context, _ logger.Logger, _ map[str
 type mockCapabilitiesRegistry struct {
 	F     uint8
 	Nodes []capabilities.Node
+	// DONs, if set, is returned as-is from DONsForCapability (for multi-DON tests).
+	DONs []capabilities.DONWithNodes
 }
 
 var owner = "test_owner"
 
 func (m *mockCapabilitiesRegistry) DONsForCapability(_ context.Context, _ string) ([]capabilities.DONWithNodes, error) {
+	if len(m.DONs) > 0 {
+		return m.DONs, nil
+	}
 	members := make([]p2ptypes.PeerID, 0, len(m.Nodes))
 	for _, n := range m.Nodes {
 		members = append(members, *n.PeerID)
@@ -1009,6 +1014,7 @@ func TestVaultHandler_PublicKeyGet(t *testing.T) {
 	mcr := &mockCapabilitiesRegistry{F: 1, Nodes: nodes}
 	h.(*handler).aggregator = &baseAggregator{
 		capabilitiesRegistry: mcr,
+		vaultHandlerDonID:    h.(*handler).donConfig.DonId,
 	}
 
 	don.On("SendToNode", mock.Anything, mock.Anything, mock.Anything).Return(nil)