Skip to content

Commit

Permalink
go/registry: Add per-role runtime admission policy support
Browse files Browse the repository at this point in the history
  • Loading branch information
@kostko
kostko committed Sep 11, 2023
1 parent d917152 commit 9bd2a4f
Show file tree
Hide file tree
Showing 7 changed files with 461 additions and 131 deletions.
85 changes: 3 additions & 82 deletions go/consensus/cometbft/apps/registry/transactions.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,89 +235,10 @@ func (app *registryApplication) registerNode( // nolint: gocyclo
}
}

// Check runtime's whitelist.
// Verify admission policies for all node's runtimes are satisfied.
for _, rt := range paidRuntimes {
if rt.AdmissionPolicy.EntityWhitelist == nil {
continue
}
wcfg, entIsWhitelisted := rt.AdmissionPolicy.EntityWhitelist.Entities[newNode.EntityID]
if !entIsWhitelisted {
ctx.Logger().Debug("RegisterNode: node's entity not in a runtime's whitelist",
"entity_id", newNode.EntityID,
"runtime_id", rt.ID,
"node_id", newNode.ID,
)
return registry.ErrForbidden
}
if len(wcfg.MaxNodes) == 0 {
continue
}

// Map is present and non-empty, check per-role restrictions
// on the maximum number of nodes per entity.

// Iterate over all valid roles (each entry in the map can
// only have a single role).
for _, role := range node.Roles() {
if !newNode.HasRoles(role) {
// Skip unset roles.
continue
}

maxNodes, exists := wcfg.MaxNodes[role]
if !exists {
// No such role found in whitelist.
ctx.Logger().Debug("RegisterNode: runtime's whitelist does not allow nodes with given role",
"role", role.String(),
"runtime_id", rt.ID,
"node_id", newNode.ID,
)
return registry.ErrForbidden
}
if maxNodes == 0 {
// No nodes of this type are allowed.
ctx.Logger().Debug("RegisterNode: runtime's whitelist does not allow nodes with given role",
"role", role.String(),
"runtime_id", rt.ID,
)
return registry.ErrForbidden
}

// Count existing nodes owned by entity.
nodes, grr := state.GetEntityNodes(ctx, newNode.EntityID)
if grr != nil {
ctx.Logger().Error("RegisterNode: failed to query entity nodes",
"err", grr,
"entity_id", newNode.EntityID,
)
return grr
}
var curNodes uint16
for _, n := range nodes {
if n.ID.Equal(newNode.ID) || n.IsExpired(uint64(epoch)) || !n.HasRuntime(rt.ID) {
// Skip existing node when re-registering. Also skip
// expired nodes and nodes that haven't registered
// for the same runtime.
continue
}

if n.HasRoles(role) {
curNodes++
}

// The check is inside the for loop, so we can stop as
// soon as possible once we're over the limit.
if curNodes+1 > maxNodes {
// Too many nodes with given role already registered.
ctx.Logger().Error("RegisterNode: too many nodes with given role already registered for runtime",
"role", role.String(),
"runtime_id", rt.ID,
"node_id", newNode.ID,
"num_registered_nodes", curNodes,
)
return registry.ErrForbidden
}
}
if err = rt.AdmissionPolicy.Verify(ctx, state, newNode, rt, epoch); err != nil {
return err
}
}

Expand Down
161 changes: 161 additions & 0 deletions go/consensus/cometbft/apps/registry/transactions_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,167 @@ func TestRegisterNode(t *testing.T) {
false,
false,
},
// Compute node on whitelist.
{
"ComputeNodeOnWhitelist",
func(tcd *testCaseData) {
rt := registry.Runtime{
Versioned: cbor.NewVersioned(registry.LatestRuntimeDescriptorVersion),
ID: common.NewTestNamespaceFromSeed([]byte("consensus/cometbft/apps/registry: runtime: ComputeNodeOnWhitelist"), 0),
Kind: registry.KindCompute,
GovernanceModel: registry.GovernanceEntity,
AdmissionPolicy: registry.RuntimeAdmissionPolicy{
EntityWhitelist: &registry.EntityWhitelistRuntimeAdmissionPolicy{
Entities: map[signature.PublicKey]registry.EntityWhitelistConfig{
tcd.node.EntityID: {},
},
},
},
}
_ = state.SetRuntime(ctx, &rt, false)

tcd.node.AddRoles(node.RoleComputeWorker)
tcd.node.Runtimes = []*node.Runtime{
{ID: rt.ID},
}
},
nil,
true,
true,
},
// Compute node not on whitelist.
{
"ComputeNodeNotOnWhitelist",
func(tcd *testCaseData) {
// Generate a random entity ID.
sig := memorySigner.NewTestSigner("consensus/cometbft/apps/registry: random signer 1")

rt := registry.Runtime{
Versioned: cbor.NewVersioned(registry.LatestRuntimeDescriptorVersion),
ID: common.NewTestNamespaceFromSeed([]byte("consensus/cometbft/apps/registry: runtime: ComputeNodeNotOnWhitelist"), 0),
Kind: registry.KindCompute,
GovernanceModel: registry.GovernanceEntity,
AdmissionPolicy: registry.RuntimeAdmissionPolicy{
EntityWhitelist: &registry.EntityWhitelistRuntimeAdmissionPolicy{
Entities: map[signature.PublicKey]registry.EntityWhitelistConfig{
sig.Public(): {},
},
},
},
}
_ = state.SetRuntime(ctx, &rt, false)

tcd.node.AddRoles(node.RoleComputeWorker)
tcd.node.Runtimes = []*node.Runtime{
{ID: rt.ID},
}
},
nil,
false,
false,
},
// Observer node on per-role whitelist.
{
"ObserverNodeOnPerRoleWhitelist",
func(tcd *testCaseData) {
rt := registry.Runtime{
Versioned: cbor.NewVersioned(registry.LatestRuntimeDescriptorVersion),
ID: common.NewTestNamespaceFromSeed([]byte("consensus/cometbft/apps/registry: runtime: ObserverNodeOnPerRoleWhitelist"), 0),
Kind: registry.KindCompute,
GovernanceModel: registry.GovernanceEntity,
AdmissionPolicy: registry.RuntimeAdmissionPolicy{
PerRole: map[node.RolesMask]registry.PerRoleAdmissionPolicy{
node.RoleObserver: {
EntityWhitelist: &registry.EntityWhitelistRoleAdmissionPolicy{
Entities: map[signature.PublicKey]registry.EntityWhitelistRoleConfig{
tcd.node.EntityID: {},
},
},
},
},
},
}
_ = state.SetRuntime(ctx, &rt, false)

tcd.node.AddRoles(node.RoleObserver)
tcd.node.Runtimes = []*node.Runtime{
{ID: rt.ID},
}
},
nil,
true,
true,
},
// Observer node not on per-role whitelist.
{
"ObserverNodeNotOnPerRoleWhitelist",
func(tcd *testCaseData) {
// Generate a random entity ID.
sig := memorySigner.NewTestSigner("consensus/cometbft/apps/registry: random signer 1")

rt := registry.Runtime{
Versioned: cbor.NewVersioned(registry.LatestRuntimeDescriptorVersion),
ID: common.NewTestNamespaceFromSeed([]byte("consensus/cometbft/apps/registry: runtime: ObserverNodeNotOnPerRoleWhitelist"), 0),
Kind: registry.KindCompute,
GovernanceModel: registry.GovernanceEntity,
AdmissionPolicy: registry.RuntimeAdmissionPolicy{
PerRole: map[node.RolesMask]registry.PerRoleAdmissionPolicy{
node.RoleObserver: {
EntityWhitelist: &registry.EntityWhitelistRoleAdmissionPolicy{
Entities: map[signature.PublicKey]registry.EntityWhitelistRoleConfig{
sig.Public(): {},
},
},
},
},
},
}
_ = state.SetRuntime(ctx, &rt, false)

tcd.node.AddRoles(node.RoleObserver)
tcd.node.Runtimes = []*node.Runtime{
{ID: rt.ID},
}
},
nil,
false,
false,
},
// Compute node not on per-role whitelist, but per-role whitelist only set for observer nodes.
{
"ComputeNodeNotOnPerRoleWhitelist",
func(tcd *testCaseData) {
// Generate a random entity ID.
sig := memorySigner.NewTestSigner("consensus/cometbft/apps/registry: random signer 1")

rt := registry.Runtime{
Versioned: cbor.NewVersioned(registry.LatestRuntimeDescriptorVersion),
ID: common.NewTestNamespaceFromSeed([]byte("consensus/cometbft/apps/registry: runtime: ComputeNodeNotOnPerRoleWhitelist"), 0),
Kind: registry.KindCompute,
GovernanceModel: registry.GovernanceEntity,
AdmissionPolicy: registry.RuntimeAdmissionPolicy{
PerRole: map[node.RolesMask]registry.PerRoleAdmissionPolicy{
node.RoleObserver: { // Note: The node will register with compute role.
EntityWhitelist: &registry.EntityWhitelistRoleAdmissionPolicy{
Entities: map[signature.PublicKey]registry.EntityWhitelistRoleConfig{
sig.Public(): {},
},
},
},
},
},
}
_ = state.SetRuntime(ctx, &rt, false)

tcd.node.AddRoles(node.RoleComputeWorker)
tcd.node.Runtimes = []*node.Runtime{
{ID: rt.ID},
}
},
nil,
true,
true,
},
// Updating a node should be allowed.
{
"UpdateValidator",
Expand Down
Loading

0 comments on commit 9bd2a4f

Please sign in to comment.