feat: replace private key with keystore (#16)

Co-authored-by: HAOYUatHZ <haoyu@protonmail.com>

.gitignore

@@ -1,2 +1,3 @@
+.idea
 assets/params*
 assets/seed

roller/.gitignore

@@ -12,3 +12,5 @@ core/prover/rust/target
 
 params/
 seed
+
+keystore

roller/config.toml

@@ -1,5 +1,6 @@
 roller_name = "my_roller"
-secret_key = "dcf2cbdd171a21c480aa7f53d77f31bb102282b3ff099c78e3118b37348c72f7"
+keystore_path = "keystore"
+keystore_password = "roller-pwd"
 scroll_url = "ws://localhost:9000"
 db_path = "bbolt_db"
 

roller/config/config.go

@@ -9,11 +9,12 @@ import (
 
 // Config loads roller configuration items.
 type Config struct {
-	RollerName string        `toml:"roller_name"`
-	SecretKey  string        `toml:"secret_key"`
-	ScrollURL  string        `toml:"scroll_url"`
-	Prover     *ProverConfig `toml:"prover"`
-	DBPath     string        `toml:"db_path"`
+	RollerName       string        `toml:"roller_name"`
+	KeystorePath     string        `toml:"keystore_path"`
+	KeystorePassword string        `toml:"keystore_password"`
+	ScrollURL        string        `toml:"scroll_url"`
+	Prover           *ProverConfig `toml:"prover"`
+	DBPath           string        `toml:"db_path"`
 }
 
 // ProverConfig loads zk roller configuration items.

roller/core/roller.go

@@ -1,14 +1,19 @@
 package core
 
 import (
+	"crypto/ecdsa"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
 	"strings"
 	"sync/atomic"
 	"time"
 
 	"github.com/gorilla/websocket"
+	"github.com/scroll-tech/go-ethereum/accounts/keystore"
 	"github.com/scroll-tech/go-ethereum/common"
 	"github.com/scroll-tech/go-ethereum/crypto"
 	"github.com/scroll-tech/go-ethereum/log"
@@ -87,9 +92,9 @@ func (r *Roller) Run() error {
 
 // Register registers Roller to the Scroll through Websocket.
 func (r *Roller) Register() error {
-	priv, err := crypto.HexToECDSA(r.cfg.SecretKey)
+	priv, err := r.loadOrCreateKey()
 	if err != nil {
-		return fmt.Errorf("generate private-key failed %v", err)
+		return err
 	}
 	authMsg := &AuthMessage{
 		Identity: Identity{
@@ -102,7 +107,7 @@ func (r *Roller) Register() error {
 
 	// Sign auth message
 	if err = authMsg.Sign(priv); err != nil {
-		return fmt.Errorf("Sign auth message failed %v", err)
+		return fmt.Errorf("sign auth message failed %v", err)
 	}
 
 	msgByt, err := MakeMsgByt(Register, authMsg)
@@ -263,6 +268,37 @@ func (r *Roller) persistTrace(byt []byte) error {
 	return r.stack.Push(traces)
 }
 
+func (r *Roller) loadOrCreateKey() (*ecdsa.PrivateKey, error) {
+	keystoreFilePath := r.cfg.KeystorePath
+	if _, err := os.Stat(r.cfg.KeystorePath); os.IsNotExist(err) {
+		// If there is no keystore, make a new one.
+		ks := keystore.NewKeyStore(r.cfg.KeystorePath, keystore.StandardScryptN, keystore.StandardScryptP)
+		account, err := ks.NewAccount(r.cfg.KeystorePassword)
+		if err != nil {
+			return nil, fmt.Errorf("generate crypto account failed %v", err)
+		}
+		log.Info("create a new account", "address", account.Address.Hex())
+
+		fis, err := ioutil.ReadDir(r.cfg.KeystorePath)
+		if err != nil {
+			return nil, err
+		}
+		keystoreFilePath = filepath.Join(r.cfg.KeystorePath, fis[0].Name())
+	} else {
+		return nil, err
+	}
+
+	keyjson, err := ioutil.ReadFile(keystoreFilePath)
+	if err != nil {
+		return nil, err
+	}
+	key, err := keystore.DecryptKey(keyjson, r.cfg.KeystorePassword)
+	if err != nil {
+		return nil, err
+	}
+	return key.PrivateKey, nil
+}
+
 // MakeMsgByt Marshals Msg to bytes.
 func MakeMsgByt(msgTyp MsgType, payloadVal interface{}) ([]byte, error) {
 	payload, err := json.Marshal(payloadVal)

roller/mock/mock_test.go

@@ -27,19 +27,22 @@ var (
 )
 
 func TestMain(m *testing.M) {
-	mockPath = "/tmp/roller_mock_test"
+	mockPath = "./roller_mock_test"
+
 	_ = os.RemoveAll(mockPath)
 	if err := os.Mkdir(mockPath, os.ModePerm); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	}
+
 	scrollPort = rand.Intn(9000)
 	cfg = &config.Config{
-		RollerName: "test-roller",
-		SecretKey:  "dcf2cbdd171a21c480aa7f53d77f31bb102282b3ff099c78e3118b37348c72f7",
-		ScrollURL:  fmt.Sprintf("ws://localhost:%d", scrollPort),
-		Prover:     &config.ProverConfig{MockMode: true},
-		DBPath:     filepath.Join(mockPath, "stack_db"),
+		RollerName:       "test-roller",
+		KeystorePath:     filepath.Join(mockPath, "roller-keystore"),
+		KeystorePassword: "mock_test",
+		ScrollURL:        fmt.Sprintf("ws://localhost:%d", scrollPort),
+		Prover:           &config.ProverConfig{MockMode: true},
+		DBPath:           filepath.Join(mockPath, "stack_db"),
 	}
 
 	os.Exit(m.Run())
@@ -50,9 +53,11 @@ func TestRoller(t *testing.T) {
 
 	r, err := core.NewRoller(cfg)
 	assert.NoError(t, err)
+
 	go r.Run()
 
 	<-time.NewTimer(2 * time.Second).C
+
 	r.Close()
 }
 
@@ -60,7 +65,7 @@ func mockScroll(t *testing.T) {
 	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
 		up := websocket.Upgrader{}
 		c, err := up.Upgrade(w, req, nil)
-		assert.NoError(t, err)
+		assert.NoError(t, err, "Upgrade WS")
 
 		var payload []byte
 		payload, err = func(c *websocket.Conn) ([]byte, error) {
@@ -76,22 +81,22 @@ func mockScroll(t *testing.T) {
 				}
 			}
 		}(c)
-		assert.NoError(t, err)
+		assert.NoError(t, err, "read payload")
 
 		msg := &Msg{}
 		err = json.Unmarshal(payload, msg)
-		assert.NoError(t, err)
+		assert.NoError(t, err, "json Unmarshal raw payload")
 
 		authMsg := &AuthMessage{}
 		err = json.Unmarshal(msg.Payload, authMsg)
-		assert.NoError(t, err)
+		assert.NoError(t, err, "json Unmarshal inner payload")
 
 		// Verify signature
 		hash, err := authMsg.Identity.Hash()
-		assert.NoError(t, err)
+		assert.NoError(t, err, "authMsg.Identity.Hash()")
 
 		if !secp256k1.VerifySignature(common.FromHex(authMsg.Identity.PublicKey), hash, common.FromHex(authMsg.Signature)[:64]) {
-			assert.NoError(t, err)
+			assert.NoError(t, err, "VerifySignature")
 		}
 		t.Log("signature verification successfully. Roller: ", authMsg.Identity.Name)
 		assert.Equal(t, cfg.RollerName, authMsg.Identity.Name)
@@ -101,10 +106,10 @@ func mockScroll(t *testing.T) {
 			Traces: nil,
 		}
 		msgByt, err := core.MakeMsgByt(BlockTrace, traces)
-		assert.NoError(t, err)
+		assert.NoError(t, err, "MakeMsgByt")
 
 		err = c.WriteMessage(websocket.BinaryMessage, msgByt)
-		assert.NoError(t, err)
+		assert.NoError(t, err, "WriteMessage")
 	})
 	http.ListenAndServe(fmt.Sprintf(":%d", scrollPort), nil)
 }