Merge pull request #245 from aaronbuchwald/duplicates-avm-import-key

Prevent duplicated addresses in avm import key
This commit is contained in:
Stephen Buttolph 2020-06-15 15:22:40 -04:00 committed by GitHub
commit 7c60e8d966
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 175 additions and 52 deletions

View File

@ -8,12 +8,14 @@ import (
"fmt"
"net/http"
"sync"
"testing"
"github.com/gorilla/rpc/v2"
"github.com/ava-labs/gecko/chains/atomic"
"github.com/ava-labs/gecko/database"
"github.com/ava-labs/gecko/database/encdb"
"github.com/ava-labs/gecko/database/memdb"
"github.com/ava-labs/gecko/database/prefixdb"
"github.com/ava-labs/gecko/ids"
"github.com/ava-labs/gecko/snow/engine/common"
@ -137,35 +139,9 @@ func (ks *Keystore) CreateUser(_ *http.Request, args *CreateUserArgs, reply *Cre
ks.log.Verbo("CreateUser called with %.*s", maxUserPassLen, args.Username)
if len(args.Username) > maxUserPassLen || len(args.Password) > maxUserPassLen {
return errUserPassMaxLength
}
if args.Username == "" {
return errEmptyUsername
}
if usr, err := ks.getUser(args.Username); err == nil || usr != nil {
return fmt.Errorf("user already exists: %s", args.Username)
}
if zxcvbn.PasswordStrength(args.Password, nil).Score < requiredPassScore {
return errWeakPassword
}
usr := &User{}
if err := usr.Initialize(args.Password); err != nil {
if err := ks.AddUser(args.Username, args.Password); err != nil {
return err
}
usrBytes, err := ks.codec.Marshal(usr)
if err != nil {
return err
}
if err := ks.userDB.Put([]byte(args.Username), usrBytes); err != nil {
return err
}
ks.users[args.Username] = usr
reply.Success = true
return nil
}
@ -403,3 +379,43 @@ func (ks *Keystore) GetDatabase(bID ids.ID, username, password string) (database
return encDB, nil
}
func (ks *Keystore) AddUser(username, password string) error {
if len(username) > maxUserPassLen || len(password) > maxUserPassLen {
return errUserPassMaxLength
}
if username == "" {
return errEmptyUsername
}
if usr, err := ks.getUser(username); err == nil || usr != nil {
return fmt.Errorf("user already exists: %s", username)
}
if zxcvbn.PasswordStrength(password, nil).Score < requiredPassScore {
return errWeakPassword
}
usr := &User{}
if err := usr.Initialize(password); err != nil {
return err
}
usrBytes, err := ks.codec.Marshal(usr)
if err != nil {
return err
}
if err := ks.userDB.Put([]byte(username), usrBytes); err != nil {
return err
}
ks.users[username] = usr
return nil
}
func CreateTestKeystore(t *testing.T) *Keystore {
ks := &Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
return ks
}

View File

@ -10,9 +10,7 @@ import (
"reflect"
"testing"
"github.com/ava-labs/gecko/database/memdb"
"github.com/ava-labs/gecko/ids"
"github.com/ava-labs/gecko/utils/logging"
)
var (
@ -22,8 +20,7 @@ var (
)
func TestServiceListNoUsers(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
reply := ListUsersReply{}
if err := ks.ListUsers(nil, &ListUsersArgs{}, &reply); err != nil {
@ -35,8 +32,7 @@ func TestServiceListNoUsers(t *testing.T) {
}
func TestServiceCreateUser(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
{
reply := CreateUserReply{}
@ -75,8 +71,7 @@ func genStr(n int) string {
// TestServiceCreateUserArgsChecks generates excessively long usernames or
// passwords to assure the santity checks on string length are not exceeded
func TestServiceCreateUserArgsCheck(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
{
reply := CreateUserReply{}
@ -117,8 +112,7 @@ func TestServiceCreateUserArgsCheck(t *testing.T) {
// TestServiceCreateUserWeakPassword tests creating a new user with a weak
// password to ensure the password strength check is working
func TestServiceCreateUserWeakPassword(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
{
reply := CreateUserReply{}
@ -138,8 +132,7 @@ func TestServiceCreateUserWeakPassword(t *testing.T) {
}
func TestServiceCreateDuplicate(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
{
reply := CreateUserReply{}
@ -166,8 +159,7 @@ func TestServiceCreateDuplicate(t *testing.T) {
}
func TestServiceCreateUserNoName(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
reply := CreateUserReply{}
if err := ks.CreateUser(nil, &CreateUserArgs{
@ -178,8 +170,7 @@ func TestServiceCreateUserNoName(t *testing.T) {
}
func TestServiceUseBlockchainDB(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
{
reply := CreateUserReply{}
@ -218,8 +209,7 @@ func TestServiceUseBlockchainDB(t *testing.T) {
}
func TestServiceExportImport(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
{
reply := CreateUserReply{}
@ -252,8 +242,7 @@ func TestServiceExportImport(t *testing.T) {
t.Fatal(err)
}
newKS := Keystore{}
newKS.Initialize(logging.NoLog{}, memdb.New())
newKS := CreateTestKeystore(t)
{
reply := ImportUserReply{}
@ -358,11 +347,10 @@ func TestServiceDeleteUser(t *testing.T) {
for _, tt := range tests {
t.Run(tt.desc, func(t *testing.T) {
ks := Keystore{}
ks.Initialize(logging.NoLog{}, memdb.New())
ks := CreateTestKeystore(t)
if tt.setup != nil {
if err := tt.setup(&ks); err != nil {
if err := tt.setup(ks); err != nil {
t.Fatalf("failed to create user setup in keystore: %v", err)
}
}

View File

@ -666,13 +666,20 @@ func (service *Service) ImportKey(r *http.Request, args *ImportKeyArgs, reply *I
}
addresses, _ := user.Addresses(db)
addresses = append(addresses, sk.PublicKey().Address())
newAddress := sk.PublicKey().Address()
reply.Address = service.vm.Format(newAddress.Bytes())
for _, address := range addresses {
if newAddress.Equals(address) {
return nil
}
}
addresses = append(addresses, newAddress)
if err := user.SetAddresses(db, addresses); err != nil {
return fmt.Errorf("problem saving addresses: %w", err)
}
reply.Address = service.vm.Format(sk.PublicKey().Address().Bytes())
return nil
}

View File

@ -9,8 +9,10 @@ import (
"github.com/stretchr/testify/assert"
"github.com/ava-labs/gecko/api/keystore"
"github.com/ava-labs/gecko/ids"
"github.com/ava-labs/gecko/snow/choices"
"github.com/ava-labs/gecko/utils/crypto"
"github.com/ava-labs/gecko/utils/formatting"
)
@ -340,3 +342,113 @@ func TestCreateVariableCapAsset(t *testing.T) {
t.Fatalf("Wrong assetID returned from CreateFixedCapAsset %s", reply.AssetID)
}
}
func TestImportAvmKey(t *testing.T) {
_, vm, s := setup(t)
defer func() {
vm.Shutdown()
ctx.Lock.Unlock()
}()
userKeystore := keystore.CreateTestKeystore(t)
username := "bobby"
password := "StrnasfqewiurPasswdn56d"
if err := userKeystore.AddUser(username, password); err != nil {
t.Fatal(err)
}
vm.ctx.Keystore = userKeystore.NewBlockchainKeyStore(vm.ctx.ChainID)
_, err := vm.ctx.Keystore.GetDatabase(username, password)
if err != nil {
t.Fatal(err)
}
factory := crypto.FactorySECP256K1R{}
skIntf, err := factory.NewPrivateKey()
if err != nil {
t.Fatalf("problem generating private key: %w", err)
}
sk := skIntf.(*crypto.PrivateKeySECP256K1R)
args := ImportKeyArgs{
Username: username,
Password: password,
PrivateKey: formatting.CB58{Bytes: sk.Bytes()},
}
reply := ImportKeyReply{}
if err = s.ImportKey(nil, &args, &reply); err != nil {
t.Fatal(err)
}
}
func TestImportAvmKeyNoDuplicates(t *testing.T) {
_, vm, s := setup(t)
defer func() {
vm.Shutdown()
ctx.Lock.Unlock()
}()
userKeystore := keystore.CreateTestKeystore(t)
username := "bobby"
password := "StrnasfqewiurPasswdn56d"
if err := userKeystore.AddUser(username, password); err != nil {
t.Fatal(err)
}
vm.ctx.Keystore = userKeystore.NewBlockchainKeyStore(vm.ctx.ChainID)
_, err := vm.ctx.Keystore.GetDatabase(username, password)
if err != nil {
t.Fatal(err)
}
factory := crypto.FactorySECP256K1R{}
skIntf, err := factory.NewPrivateKey()
if err != nil {
t.Fatalf("problem generating private key: %w", err)
}
sk := skIntf.(*crypto.PrivateKeySECP256K1R)
args := ImportKeyArgs{
Username: username,
Password: password,
PrivateKey: formatting.CB58{Bytes: sk.Bytes()},
}
reply := ImportKeyReply{}
if err = s.ImportKey(nil, &args, &reply); err != nil {
t.Fatal(err)
}
expectedAddress := vm.Format(sk.PublicKey().Address().Bytes())
if reply.Address != expectedAddress {
t.Fatalf("Reply address: %s did not match expected address: %s", reply.Address, expectedAddress)
}
reply2 := ImportKeyReply{}
if err = s.ImportKey(nil, &args, &reply2); err != nil {
t.Fatal(err)
}
if reply2.Address != expectedAddress {
t.Fatalf("Reply address: %s did not match expected address: %s", reply2.Address, expectedAddress)
}
addrsArgs := ListAddressesArgs{
Username: username,
Password: password,
}
addrsReply := ListAddressesResponse{}
if err := s.ListAddresses(nil, &addrsArgs, &addrsReply); err != nil {
t.Fatal(err)
}
if len(addrsReply.Addresses) != 1 {
t.Fatal("Importing the same key twice created duplicate addresses")
}
if addrsReply.Addresses[0] != expectedAddress {
t.Fatal("List addresses returned an incorrect address")
}
}