Skip to content

Commit

Permalink
Address code review comments II
Browse files Browse the repository at this point in the history 
Signed-off-by: Yacov Manevich <[email protected]>
  • Loading branch information
@yacovm
yacovm committed Oct 30, 2024
1 parent feb2ce9 commit d08e41a
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 49 deletions.
18 changes: 6 additions & 12 deletions network/peer/tls_config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,15 @@ import (
"errors"
"io"

"golang.org/x/crypto/ed25519"

"github.com/ava-labs/avalanchego/staking"
)

var (
ErrNoCertsSent = errors.New("no certificates sent by peer")
ErrEmptyCert = errors.New("certificate sent by peer is empty")
ErrEmptyPublicKey = errors.New("no public key sent by peer")
ErrCurveMismatch = errors.New("only P256 is allowed for ECDSA")
ErrForbidden25519Key = errors.New("ed25519 is not allowed in this version")
ErrNoCertsSent = errors.New("no certificates sent by peer")
ErrEmptyCert = errors.New("certificate sent by peer is empty")
ErrEmptyPublicKey = errors.New("no public key sent by peer")
ErrCurveMismatch = errors.New("only P256 is allowed for ECDSA")
ErrUnsupportedKeyType = errors.New("key type is not supported")
)

// TLSConfig returns the TLS config that will allow secure connections to other
Expand Down Expand Up @@ -59,8 +57,6 @@ func ValidateCertificate(cs tls.ConnectionState) error {
pk := cs.PeerCertificates[0].PublicKey

switch key := pk.(type) {
case ed25519.PublicKey:
return ErrForbidden25519Key
case *ecdsa.PublicKey:
if key == nil {
return ErrEmptyPublicKey
Expand All @@ -71,9 +67,7 @@ func ValidateCertificate(cs tls.ConnectionState) error {
return nil
case *rsa.PublicKey:
return staking.ValidateRSAPublicKeyIsWellFormed(key)
case nil:
return ErrEmptyPublicKey
default:
return nil
return ErrUnsupportedKeyType
}
}
43 changes: 23 additions & 20 deletions network/peer/tls_config_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ func TestValidateCertificate(t *testing.T) {
input: func(t *testing.T) tls.ConnectionState {
key, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)
x509Cert := makeRSACertAndKey(t, key)
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{&x509Cert.cert}}
x509Cert := makeCert(t, key, &key.PublicKey)
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{x509Cert}}
},
},
{
Expand All @@ -54,9 +54,9 @@ func TestValidateCertificate(t *testing.T) {
key, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)

x509CertWithNilPK := makeRSACertAndKey(t, key)
x509CertWithNilPK.cert.PublicKey = (*rsa.PublicKey)(nil)
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{&x509CertWithNilPK.cert}}
x509CertWithNilPK := makeCert(t, key, &key.PublicKey)
x509CertWithNilPK.PublicKey = (*rsa.PublicKey)(nil)
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{x509CertWithNilPK}}
},
expectedErr: staking.ErrInvalidRSAPublicKey,
},
Expand All @@ -66,23 +66,20 @@ func TestValidateCertificate(t *testing.T) {
key, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)

x509CertWithNoPK := makeRSACertAndKey(t, key)
x509CertWithNoPK.cert.PublicKey = nil
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{&x509CertWithNoPK.cert}}
x509CertWithNoPK := makeCert(t, key, &key.PublicKey)
x509CertWithNoPK.PublicKey = nil
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{x509CertWithNoPK}}
},
expectedErr: peer.ErrEmptyPublicKey,
expectedErr: peer.ErrUnsupportedKeyType,
},
{
description: "EC cert",
input: func(t *testing.T) tls.ConnectionState {
ecKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)

basicCert := basicCert()
certBytes, err := x509.CreateCertificate(rand.Reader, basicCert, basicCert, &ecKey.PublicKey, ecKey)
require.NoError(t, err)
ecCert := makeCert(t, ecKey, &ecKey.PublicKey)

ecCert, err := x509.ParseCertificate(certBytes)
require.NoError(t, err)
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{ecCert}}
},
Expand All @@ -94,21 +91,27 @@ func TestValidateCertificate(t *testing.T) {
ecKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)

basicCert := basicCert()
certBytes, err := x509.CreateCertificate(rand.Reader, basicCert, basicCert, &ecKey.PublicKey, ecKey)
require.NoError(t, err)
ecCert := makeCert(t, ecKey, &ecKey.PublicKey)
ecCert.PublicKey = (*ecdsa.PublicKey)(nil)

ecCert, err := x509.ParseCertificate(certBytes)
return tls.ConnectionState{PeerCertificates: []*x509.Certificate{ecCert}}
},
},
{
description: "EC cert with P384 curve",
expectedErr: peer.ErrCurveMismatch,
input: func(t *testing.T) tls.ConnectionState {
ecKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
require.NoError(t, err)

ecCert.PublicKey = nil
ecCert := makeCert(t, ecKey, &ecKey.PublicKey)

return tls.ConnectionState{PeerCertificates: []*x509.Certificate{ecCert}}
},
},
{
description: "EC cert with ed25519 key",
expectedErr: peer.ErrForbidden25519Key,
description: "EC cert with ed25519 key not supported",
expectedErr: peer.ErrUnsupportedKeyType,
input: func(t *testing.T) tls.ConnectionState {
pub, priv, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
Expand Down
22 changes: 6 additions & 16 deletions network/peer/upgrader_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,6 @@ func TestBlockClientsWithIncorrectRSAKeys(t *testing.T) {
// Initialize upgrader with a mock that fails when it's incremented.
failOnIncrementCounter := &mockPrometheusCounter{
Counter: c,
t: t,
onIncrement: func() {
require.FailNow(t, "should not have invoked")
},
Expand Down Expand Up @@ -168,10 +167,10 @@ func nonStandardRSAKey(t *testing.T) *rsa.PrivateKey {
}

func makeTLSCert(t *testing.T, privKey *rsa.PrivateKey) tls.Certificate {
x509Cert := makeRSACertAndKey(t, privKey)
x509Cert := makeCert(t, privKey, &privKey.PublicKey)

rawX509PEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: x509Cert.cert.Raw})
privateKeyInDER, err := x509.MarshalPKCS8PrivateKey(x509Cert.key)
rawX509PEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: x509Cert.Raw})
privateKeyInDER, err := x509.MarshalPKCS8PrivateKey(privKey)
require.NoError(t, err)

privateKeyInPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privateKeyInDER})
Expand All @@ -182,24 +181,16 @@ func makeTLSCert(t *testing.T, privKey *rsa.PrivateKey) tls.Certificate {
return tlsCertServer
}

type certAndKey struct {
cert x509.Certificate
key *rsa.PrivateKey
}

func makeRSACertAndKey(t *testing.T, privKey *rsa.PrivateKey) certAndKey {
func makeCert(t *testing.T, privateKey any, publicKey any) *x509.Certificate {
// Create a self-signed cert
basicCert := basicCert()
certBytes, err := x509.CreateCertificate(rand.Reader, basicCert, basicCert, &privKey.PublicKey, privKey)
certBytes, err := x509.CreateCertificate(rand.Reader, basicCert, basicCert, publicKey, privateKey)
require.NoError(t, err)

cert, err := x509.ParseCertificate(certBytes)
require.NoError(t, err)

return certAndKey{
cert: *cert,
key: privKey,
}
return cert
}

func basicCert() *x509.Certificate {
Expand All @@ -212,7 +203,6 @@ func basicCert() *x509.Certificate {
}

type mockPrometheusCounter struct {
t *testing.T
prometheus.Counter
onIncrement func()
}
Expand Down
2 changes: 1 addition & 1 deletion staking/parse_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ func TestValidateRSAPublicKeyIsWellFormed(t *testing.T) {
},
{
description: "even modulus",
expectErr: ErrRSAModulusIsEven,
getPK: func() rsa.PublicKey {
evenN := new(big.Int).Set(validKey.N)
evenN.Add(evenN, big.NewInt(1))
Expand All @@ -60,7 +61,6 @@ func TestValidateRSAPublicKeyIsWellFormed(t *testing.T) {
E: 65537,
}
},
expectErr: ErrRSAModulusIsEven,
},
{
description: "unsupported exponent",
Expand Down

0 comments on commit d08e41a

Please sign in to comment.