Added password to storage module

storage/localstate_mgr.go

@@ -48,11 +48,12 @@ type FileStateMgr struct {
 	folder      string
 	writeLock   *sync.RWMutex
 	encryptMode bool
-	key         []byte
+	passkey     []byte
+	keyGen      *KeygenLocalState
 }
 
 // NewFileStateMgr create a new instance of the FileStateMgr which implements LocalStateManager
-func NewFileStateMgr(folder string) (*FileStateMgr, error) {
+func NewFileStateMgr(folder string, password string) (*FileStateMgr, error) {
 	if len(folder) > 0 {
 		_, err := os.Stat(folder)
 		if err != nil && os.IsNotExist(err) {
@@ -62,15 +63,16 @@ func NewFileStateMgr(folder string) (*FileStateMgr, error) {
 		}
 	}
 	encryptMode := true
-	key, err := getFragmentSeed()
+	key, err := getFragmentSeed(password)
 	if err != nil {
 		encryptMode = false
 	}
 	return &FileStateMgr{
 		folder:      folder,
 		writeLock:   &sync.RWMutex{},
 		encryptMode: encryptMode,
-		key:         key,
+		passkey:     key,
+		keyGen:      nil,
 	}, nil
 }
 
@@ -109,6 +111,9 @@ func (fsm *FileStateMgr) SaveLocalState(state KeygenLocalState) error {
 
 // GetLocalState read the local state from file system
 func (fsm *FileStateMgr) GetLocalState(pubKey string) (KeygenLocalState, error) {
+	if fsm.keyGen != nil {
+		return *fsm.keyGen, nil
+	}
 	if len(pubKey) == 0 {
 		return KeygenLocalState{}, errors.New("pub key is empty")
 	}
@@ -133,6 +138,7 @@ func (fsm *FileStateMgr) GetLocalState(pubKey string) (KeygenLocalState, error)
 	if err := json.Unmarshal(pt, &localState); nil != err {
 		return KeygenLocalState{}, fmt.Errorf("fail to unmarshal KeygenLocalState:%x %w", pt, err)
 	}
+	fsm.keyGen = &localState
 	return localState, nil
 }
 
@@ -199,7 +205,7 @@ func (fsm *FileStateMgr) encryptFragment(plainText []byte) ([]byte, error) {
 	if !fsm.encryptMode {
 		return plainText, nil
 	}
-	block, err := aes.NewCipher(fsm.key)
+	block, err := aes.NewCipher(fsm.passkey)
 	if err != nil {
 		return nil, err
 	}
@@ -221,7 +227,7 @@ func (fsm *FileStateMgr) decryptFragment(buf []byte) ([]byte, error) {
 	if !fsm.encryptMode {
 		return buf, nil
 	}
-	block, err := aes.NewCipher(fsm.key)
+	block, err := aes.NewCipher(fsm.passkey)
 	if err != nil {
 		return nil, err
 	}
@@ -240,11 +246,15 @@ func (fsm *FileStateMgr) decryptFragment(buf []byte) ([]byte, error) {
 	return plainText, nil
 }
 
-func getFragmentSeed() ([]byte, error) {
+func getFragmentSeed(password string) ([]byte, error) {
 	seedStr := os.Getenv(keyFragmentSeed)
 	if seedStr == "" {
-		return nil, errors.New("empty fragment seed, please populate env variable: " + keyFragmentSeed)
+		if password == "" {
+			return nil, errors.New("empty fragment seed, please check password: " + password)
+		}
+		seedStr = password
 	}
+
 	h := sha256.New()
 	h.Write([]byte(seedStr))
 	seed := h.Sum(nil)

storage/localstate_mgr_test.go

@@ -34,7 +34,7 @@ func (s *FileStateMgrTestSuite) TestNewFileStateMgr(c *C) {
 		err := os.RemoveAll(f)
 		c.Assert(err, IsNil)
 	}()
-	fsm, err := NewFileStateMgr(f)
+	fsm, err := NewFileStateMgr(f, "password")
 	c.Assert(err, IsNil)
 	c.Assert(fsm, NotNil)
 	_, err = os.Stat(f)
@@ -61,7 +61,7 @@ func (s *FileStateMgrTestSuite) TestSaveLocalState(c *C) {
 		err := os.RemoveAll(f)
 		c.Assert(err, IsNil)
 	}()
-	fsm, err := NewFileStateMgr(f)
+	fsm, err := NewFileStateMgr(f, "password")
 	c.Assert(err, IsNil)
 	c.Assert(fsm, NotNil)
 	c.Assert(fsm.SaveLocalState(stateItem), NotNil)
@@ -93,7 +93,7 @@ func (s *FileStateMgrTestSuite) TestSaveAddressBook(c *C) {
 		err := os.RemoveAll(f)
 		c.Assert(err, IsNil)
 	}()
-	fsm, err := NewFileStateMgr(f)
+	fsm, err := NewFileStateMgr(f, "password")
 	c.Assert(err, IsNil)
 	c.Assert(fsm, NotNil)
 	c.Assert(fsm.SaveAddressBook(testAddresses), IsNil)
@@ -112,7 +112,7 @@ func (s *FileStateMgrTestSuite) TestEncryption(c *C) {
 		err := os.RemoveAll(f)
 		c.Assert(err, IsNil)
 	}()
-	fsm, err := NewFileStateMgr(f)
+	fsm, err := NewFileStateMgr(f, "password")
 	c.Assert(err, IsNil)
 	c.Assert(fsm, NotNil)
 

tss/tss.go

@@ -51,6 +51,7 @@ func NewTss(
 	conf common.TssConfig,
 	preParams *bkeygen.LocalPreParams,
 	externalIP string,
+	tssPassword string,
 ) (*TssServer, error) {
 	pk := coskey.PubKey{
 		Key: priKey.PubKey().Bytes()[:],
@@ -61,7 +62,7 @@ func NewTss(
 		return nil, fmt.Errorf("fail to genearte the key: %w", err)
 	}
 
-	stateManager, err := storage.NewFileStateMgr(baseFolder)
+	stateManager, err := storage.NewFileStateMgr(baseFolder, tssPassword)
 	if err != nil {
 		return nil, fmt.Errorf("fail to create file state manager")
 	}