Merge pull request #15 from ava-labs/grpc-fds

cleaned up keystore apis

1 files changed, 34 insertions, 28 deletions

diff --git a/plugin/evm/service.go b/plugin/evm/service.go
index dbd6d46..fdbfd91 100644
--- a/plugin/evm/service.go
+++ b/plugin/evm/service.go
@@ -147,19 +147,25 @@ type ExportKeyReply struct {
 // ExportKey returns a private key from the provided user
 func (service *AvaxAPI) ExportKey(r *http.Request, args *ExportKeyArgs, reply *ExportKeyReply) error {
 	service.vm.ctx.Log.Info("Platform: ExportKey called")
+
+	address, err := service.vm.ParseEthAddress(args.Address)
+	if err != nil {
+		return fmt.Errorf("couldn't parse %s to address: %s", args.Address, err)
+	}
+
 	db, err := service.vm.ctx.Keystore.GetDatabase(args.Username, args.Password)
 	if err != nil {
 		return fmt.Errorf("problem retrieving user '%s': %w", args.Username, err)
 	}
+	defer db.Close()
+
 	user := user{db: db}
-	if address, err := service.vm.ParseEthAddress(args.Address); err != nil {
-		return fmt.Errorf("couldn't parse %s to address: %s", args.Address, err)
-	} else if sk, err := user.getKey(address); err != nil {
+	sk, err := user.getKey(address)
+	if err != nil {
 		return fmt.Errorf("problem retrieving private key: %w", err)
-	} else {
-		reply.PrivateKey = constants.SecretKeyPrefix + formatting.CB58{Bytes: sk.Bytes()}.String()
-		return nil
 	}
+	reply.PrivateKey = constants.SecretKeyPrefix + formatting.CB58{Bytes: sk.Bytes()}.String()
+	return nil
 }
 
 // ImportKeyArgs are arguments for ImportKey
@@ -171,43 +177,40 @@ type ImportKeyArgs struct {
 // ImportKey adds a private key to the provided user
 func (service *AvaxAPI) ImportKey(r *http.Request, args *ImportKeyArgs, reply *api.JsonAddress) error {
 	service.vm.ctx.Log.Info("Platform: ImportKey called for user '%s'", args.Username)
-	if service.vm.ctx.Keystore == nil {
-		return fmt.Errorf("oh no")
-	}
-	fmt.Sprintf("good")
-	db, err := service.vm.ctx.Keystore.GetDatabase(args.Username, args.Password)
-	if err != nil {
-		return fmt.Errorf("problem retrieving data: %w", err)
-	}
-
-	user := user{db: db}
-
-	factory := crypto.FactorySECP256K1R{}
 
 	if !strings.HasPrefix(args.PrivateKey, constants.SecretKeyPrefix) {
 		return fmt.Errorf("private key missing %s prefix", constants.SecretKeyPrefix)
 	}
+
 	trimmedPrivateKey := strings.TrimPrefix(args.PrivateKey, constants.SecretKeyPrefix)
 	formattedPrivateKey := formatting.CB58{}
 	if err := formattedPrivateKey.FromString(trimmedPrivateKey); err != nil {
 		return fmt.Errorf("problem parsing private key: %w", err)
 	}
 
+	factory := crypto.FactorySECP256K1R{}
 	skIntf, err := factory.ToPrivateKey(formattedPrivateKey.Bytes)
 	if err != nil {
 		return fmt.Errorf("problem parsing private key: %w", err)
 	}
 	sk := skIntf.(*crypto.PrivateKeySECP256K1R)
 
-	if err := user.putAddress(sk); err != nil {
-		return fmt.Errorf("problem saving key %w", err)
-	}
-
 	// TODO: return eth address here
 	reply.Address, err = service.vm.FormatEthAddress(GetEthAddress(sk))
 	if err != nil {
 		return fmt.Errorf("problem formatting address: %w", err)
 	}
+
+	db, err := service.vm.ctx.Keystore.GetDatabase(args.Username, args.Password)
+	if err != nil {
+		return fmt.Errorf("problem retrieving data: %w", err)
+	}
+	defer db.Close()
+
+	user := user{db: db}
+	if err := user.putAddress(sk); err != nil {
+		return fmt.Errorf("problem saving key %w", err)
+	}
 	return nil
 }
 
@@ -232,18 +235,19 @@ func (service *AvaxAPI) ImportAVAX(_ *http.Request, args *ImportAVAXArgs, respon
 		return fmt.Errorf("problem parsing chainID %q: %w", args.SourceChain, err)
 	}
 
+	to, err := service.vm.ParseEthAddress(args.To)
+	if err != nil { // Parse address
+		return fmt.Errorf("couldn't parse argument 'to' to an address: %w", err)
+	}
+
 	// Get the user's info
 	db, err := service.vm.ctx.Keystore.GetDatabase(args.Username, args.Password)
 	if err != nil {
 		return fmt.Errorf("couldn't get user '%s': %w", args.Username, err)
 	}
-	user := user{db: db}
-
-	to, err := service.vm.ParseEthAddress(args.To)
-	if err != nil { // Parse address
-		return fmt.Errorf("couldn't parse argument 'to' to an address: %w", err)
-	}
+	defer db.Close()
 
+	user := user{db: db}
 	privKeys, err := user.getKeys()
 	if err != nil { // Get keys
 		return fmt.Errorf("couldn't get keys controlled by the user: %w", err)
@@ -289,6 +293,8 @@ func (service *AvaxAPI) ExportAVAX(_ *http.Request, args *ExportAVAXArgs, respon
 	if err != nil {
 		return fmt.Errorf("problem retrieving user '%s': %w", args.Username, err)
 	}
+	defer db.Close()
+
 	user := user{db: db}
 	privKeys, err := user.getKeys()
 	if err != nil {