Skip to content

Commit 83b0a77

Browse files
joshua-kimjoshua-kim
committed
nits
1 parent 1b346c1 commit 83b0a77

File tree

10 files changed

+223
-155
lines changed

10 files changed

+223
-155
lines changed

‎snow/validators/gvalidators/validator_state_client.go

Lines changed: 9 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,16 @@ package gvalidators
55

66
import (
77
"context"
8-
"errors"
98

109
"google.golang.org/protobuf/types/known/emptypb"
1110

1211
"github.com/ava-labs/avalanchego/ids"
1312
"github.com/ava-labs/avalanchego/snow/validators"
14-
"github.com/ava-labs/avalanchego/utils/crypto/bls"
1513

1614
pb "github.com/ava-labs/avalanchego/proto/pb/validatorstate"
1715
)
1816

19-
var (
20-
_ validators.State = (*Client)(nil)
21-
errFailedPublicKeyDeserialize = errors.New("couldn't deserialize public key")
22-
)
17+
var _ validators.State = (*Client)(nil)
2318

2419
type Client struct {
2520
client pb.ValidatorStateClient
@@ -74,26 +69,16 @@ func (c *Client) GetValidatorSet(
7469
if err != nil {
7570
return nil, err
7671
}
77-
var publicKey *bls.PublicKey
78-
if len(validator.PublicKey) > 0 {
79-
// This is a performance optimization to avoid the cost of compression
80-
// and key re-verification with PublicKeyFromBytes. We can safely
81-
// assume that the BLS Public Keys are verified before being added
82-
// to the P-Chain and served by the gRPC server.
83-
publicKey = new(bls.PublicKey).Deserialize(validator.PublicKey)
84-
if publicKey == nil {
85-
return nil, errFailedPublicKeyDeserialize
86-
}
87-
}
88-
vdrs[nodeID] = &validators.GetValidatorOutput{
89-
NodeID: nodeID,
72+
var publicKey validators.PublicKey
73+
if len(validator.PublicKey) != 0 {
9074
// As an optimization, we pre-populate the cached deserialized key
9175
// since it was given to us off gRPC.
92-
PublicKey: validators.PublicKey{
93-
PublicKey: publicKey,
94-
Bytes: validator.PublicKey,
95-
},
96-
Weight: validator.Weight,
76+
publicKey = validators.NewPublicKeyFromBytes(validator.GetPublicKey())
77+
}
78+
vdrs[nodeID] = &validators.GetValidatorOutput{
79+
NodeID: nodeID,
80+
PublicKey: publicKey,
81+
Weight: validator.Weight,
9782
}
9883
}
9984
return vdrs, nil

‎snow/validators/gvalidators/validator_state_server.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,10 @@ func (s *Server) GetValidatorSet(ctx context.Context, req *pb.GetValidatorSetReq
6868
NodeId: vdr.NodeID[:],
6969
Weight: vdr.Weight,
7070
}
71-
if vdr.PublicKey.PublicKey != nil {
71+
if vdr.PublicKey != nil {
7272
// This is a performance optimization to avoid the cost of compression
7373
// from PublicKeyToBytes.
74-
vdrPB.PublicKey = vdr.PublicKey.Serialize()
74+
vdrPB.PublicKey = vdr.PublicKey.Bytes()
7575
}
7676
resp.Validators[i] = vdrPB
7777
i++

‎snow/validators/gvalidators/validator_state_test.go

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -148,29 +148,23 @@ func TestGetValidatorSet(t *testing.T) {
148148
require.NoError(err)
149149
pk0 := bls.PublicFromSecretKey(sk0)
150150
vdr0 := &validators.GetValidatorOutput{
151-
NodeID: ids.GenerateTestNodeID(),
152-
PublicKey: validators.PublicKey{
153-
PublicKey: pk0,
154-
Bytes: pk0.Serialize(),
155-
},
156-
Weight: 1,
151+
NodeID: ids.GenerateTestNodeID(),
152+
PublicKey: validators.NewPublicKeyFromBytes(pk0.Serialize()),
153+
Weight: 1,
157154
}
158155

159156
sk1, err := bls.NewSecretKey()
160157
require.NoError(err)
161158
pk1 := bls.PublicFromSecretKey(sk1)
162159
vdr1 := &validators.GetValidatorOutput{
163-
NodeID: ids.GenerateTestNodeID(),
164-
PublicKey: validators.PublicKey{
165-
PublicKey: pk1,
166-
Bytes: pk1.Serialize(),
167-
},
168-
Weight: 2,
160+
NodeID: ids.GenerateTestNodeID(),
161+
PublicKey: validators.NewPublicKeyFromBytes(pk1.Serialize()),
162+
Weight: 2,
169163
}
170164

171165
vdr2 := &validators.GetValidatorOutput{
172166
NodeID: ids.GenerateTestNodeID(),
173-
PublicKey: validators.PublicKey{},
167+
PublicKey: nil,
174168
Weight: 3,
175169
}
176170

@@ -249,12 +243,9 @@ func setupValidatorSet(b *testing.B, size int) map[ids.NodeID]*validators.GetVal
249243
for i := 0; i < size; i++ {
250244
id := ids.GenerateTestNodeID()
251245
set[id] = &validators.GetValidatorOutput{
252-
NodeID: id,
253-
PublicKey: validators.PublicKey{
254-
PublicKey: pk,
255-
Bytes: nil,
256-
},
257-
Weight: uint64(i),
246+
NodeID: id,
247+
PublicKey: validators.NewPublicKey(pk),
248+
Weight: uint64(i),
258249
}
259250
}
260251
return set

‎snow/validators/key_test.go

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package validators
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
8+
"github.com/ava-labs/avalanchego/utils/crypto/bls"
9+
)
10+
11+
func TestPublicKey_GetKey(t *testing.T) {
12+
sk, err := bls.NewSecretKey()
13+
require.NoError(t, err)
14+
pk := bls.PublicFromSecretKey(sk)
15+
pkBytes := pk.Serialize()
16+
17+
type expected struct {
18+
key *bls.PublicKey
19+
err error
20+
}
21+
22+
tests := []struct {
23+
name string
24+
Key PublicKey
25+
expected expected
26+
}{
27+
{
28+
name: "nil key, populated Bytes",
29+
Key: NewPublicKeyFromBytes(pkBytes),
30+
expected: expected{
31+
key: pk,
32+
},
33+
},
34+
{
35+
name: "populated key, nil Bytes",
36+
Key: NewPublicKey(pk),
37+
expected: expected{
38+
key: pk,
39+
},
40+
},
41+
{
42+
name: "invalid key",
43+
Key: NewPublicKeyFromBytes([]byte("foobar")),
44+
expected: expected{
45+
err: errFailedPublicKeyDeserialize,
46+
},
47+
},
48+
}
49+
50+
for _, test := range tests {
51+
t.Run(test.name, func(t *testing.T) {
52+
r := require.New(t)
53+
54+
actual, err := test.Key.Key()
55+
r.Equal(test.expected.key, actual)
56+
r.Equal(test.expected.err, err)
57+
})
58+
}
59+
}
60+
func TestPublicKey_GetBytes(t *testing.T) {
61+
sk, err := bls.NewSecretKey()
62+
require.NoError(t, err)
63+
pk := bls.PublicFromSecretKey(sk)
64+
pkBytes := pk.Serialize()
65+
66+
tests := []struct {
67+
name string
68+
key PublicKey
69+
expected []byte
70+
}{
71+
{
72+
name: "nil key, populated Bytes",
73+
key: NewPublicKeyFromBytes(pkBytes),
74+
expected: pkBytes,
75+
},
76+
{
77+
name: "populated key, nil Bytes",
78+
key: NewPublicKey(pk),
79+
expected: pkBytes,
80+
},
81+
}
82+
83+
for _, test := range tests {
84+
t.Run(test.name, func(t *testing.T) {
85+
r := require.New(t)
86+
87+
actual := test.key.Bytes()
88+
r.Equal(test.expected, actual)
89+
})
90+
}
91+
}

‎snow/validators/validator.go

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,14 @@
44
package validators
55

66
import (
7+
"errors"
8+
79
"github.com/ava-labs/avalanchego/ids"
810
"github.com/ava-labs/avalanchego/utils/crypto/bls"
911
)
1012

13+
var errFailedPublicKeyDeserialize = errors.New("couldn't deserialize public key")
14+
1115
// Validator is a struct that contains the base values representing a validator
1216
// of the Avalanche Network.
1317
type Validator struct {
@@ -30,15 +34,47 @@ type GetValidatorOutput struct {
3034
Weight uint64
3135
}
3236

33-
type PublicKey struct {
34-
*bls.PublicKey
35-
Bytes []byte
37+
type PublicKey interface {
38+
Key() (*bls.PublicKey, error)
39+
Bytes() []byte
40+
}
41+
42+
func NewPublicKey(key *bls.PublicKey) *CachedPublicKey {
43+
return &CachedPublicKey{
44+
key: key,
45+
bytes: nil,
46+
}
47+
}
48+
func NewPublicKeyFromBytes(bytes []byte) *CachedPublicKey {
49+
return &CachedPublicKey{
50+
key: nil,
51+
bytes: bytes,
52+
}
53+
}
54+
55+
type CachedPublicKey struct {
56+
// at least one of (Key, bytes) must be provided
57+
key *bls.PublicKey
58+
bytes []byte
59+
}
60+
61+
func (k *CachedPublicKey) Key() (*bls.PublicKey, error) {
62+
if k.key == nil {
63+
pk := new(bls.PublicKey).Deserialize(k.bytes)
64+
if pk == nil {
65+
return nil, errFailedPublicKeyDeserialize
66+
}
67+
68+
k.key = pk
69+
}
70+
71+
return k.key, nil
3672
}
3773

38-
func (pk PublicKey) Serialize() []byte {
39-
if pk.Bytes == nil {
40-
pk.Bytes = pk.PublicKey.Serialize()
74+
func (k *CachedPublicKey) Bytes() []byte {
75+
if len(k.bytes) == 0 {
76+
k.bytes = k.key.Serialize()
4177
}
4278

43-
return pk.Bytes
79+
return k.bytes
4480
}

‎vms/platformvm/vm.go

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -525,11 +525,9 @@ func (vm *VM) GetValidatorSet(ctx context.Context, height uint64, subnetID ids.I
525525
return nil, fmt.Errorf("%w: %s", errMissingValidator, vdr.NodeID)
526526
}
527527
vdrSet[vdr.NodeID] = &validators.GetValidatorOutput{
528-
NodeID: vdr.NodeID,
529-
PublicKey: validators.PublicKey{
530-
PublicKey: primaryVdr.PublicKey,
531-
},
532-
Weight: vdr.Weight,
528+
NodeID: vdr.NodeID,
529+
PublicKey: validators.NewPublicKey(primaryVdr.PublicKey),
530+
Weight: vdr.Weight,
533531
}
534532
}
535533

@@ -586,9 +584,7 @@ func (vm *VM) GetValidatorSet(ctx context.Context, height uint64, subnetID ids.I
586584
if vdr, ok := vdrSet[nodeID]; ok {
587585
// The validator's public key was removed at this block, so it
588586
// was in the validator set before.
589-
vdr.PublicKey = validators.PublicKey{
590-
PublicKey: pk,
591-
}
587+
vdr.PublicKey = validators.NewPublicKey(pk)
592588
}
593589
}
594590
}

0 commit comments

Comments
 (0)