cleaned up keystore API + improved batching

api/keystore/service.go

@@ -11,6 +11,7 @@ import (
 
 	"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/prefixdb"
@@ -202,7 +203,7 @@ type ExportUserArgs struct {
 
 // ExportUserReply is the reply from ExportUser
 type ExportUserReply struct {
-	User string `json:"user"`
+	User formatting.CB58 `json:"user"`
 }
 
 // ExportUser exports a serialized encoding of a user's information complete with encrypted database values
@@ -242,16 +243,15 @@ func (ks *Keystore) ExportUser(_ *http.Request, args *ExportUserArgs, reply *Exp
 	if err != nil {
 		return err
 	}
-	cb58 := formatting.CB58{Bytes: b}
-	reply.User = cb58.String()
+	reply.User.Bytes = b
 	return nil
 }
 
 // ImportUserArgs are arguments for ImportUser
 type ImportUserArgs struct {
-	Username string `json:"username"`
-	Password string `json:"password"`
-	User     string `json:"user"`
+	Username string          `json:"username"`
+	Password string          `json:"password"`
+	User     formatting.CB58 `json:"user"`
 }
 
 // ImportUserReply is the response for ImportUser
@@ -270,13 +270,8 @@ func (ks *Keystore) ImportUser(r *http.Request, args *ImportUserArgs, reply *Imp
 		return fmt.Errorf("user already exists: %s", args.Username)
 	}
 
-	cb58 := formatting.CB58{}
-	if err := cb58.FromString(args.User); err != nil {
-		return err
-	}
-
 	userData := UserDB{}
-	if err := ks.codec.Unmarshal(cb58.Bytes, &userData); err != nil {
+	if err := ks.codec.Unmarshal(args.User.Bytes, &userData); err != nil {
 		return err
 	}
 
@@ -285,22 +280,25 @@ func (ks *Keystore) ImportUser(r *http.Request, args *ImportUserArgs, reply *Imp
 		return err
 	}
 
-	// TODO: Should add batching to prevent creating a user without importing
-	// the account
-	if err := ks.userDB.Put([]byte(args.Username), usrBytes); err != nil {
+	userBatch := ks.userDB.NewBatch()
+	if err := userBatch.Put([]byte(args.Username), usrBytes); err != nil {
 		return err
 	}
-	ks.users[args.Username] = &userData.User
-
-	userDB := prefixdb.New([]byte(args.Username), ks.bcDB)
-	batch := userDB.NewBatch()
 
+	userDataDB := prefixdb.New([]byte(args.Username), ks.bcDB)
+	dataBatch := userDataDB.NewBatch()
 	for _, kvp := range userData.Data {
-		batch.Put(kvp.Key, kvp.Value)
+		dataBatch.Put(kvp.Key, kvp.Value)
+	}
+
+	if err := atomic.WriteAll(dataBatch, userBatch); err != nil {
+		return err
 	}
 
+	ks.users[args.Username] = &userData.User
+
 	reply.Success = true
-	return batch.Write()
+	return nil
 }
 
 // NewBlockchainKeyStore ...