diff --git a/storage/localstate_mgr.go b/storage/localstate_mgr.go index 2d045f1..9c077c5 100644 --- a/storage/localstate_mgr.go +++ b/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,7 +63,7 @@ func NewFileStateMgr(folder string) (*FileStateMgr, error) { } } encryptMode := true - key, err := getFragmentSeed() + key, err := getFragmentSeed(password) if err != nil { encryptMode = false } @@ -70,7 +71,8 @@ func NewFileStateMgr(folder string) (*FileStateMgr, error) { 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) diff --git a/storage/localstate_mgr_test.go b/storage/localstate_mgr_test.go index 567cef6..27b8199 100644 --- a/storage/localstate_mgr_test.go +++ b/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) diff --git a/tss/tss.go b/tss/tss.go index 9ab15b8..aabfeb2 100644 --- a/tss/tss.go +++ b/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") }