From b296b36d2b2aaa2f81d26b3c133ace2714c58a7d Mon Sep 17 00:00:00 2001
From: Gustav Simonsson <gustav.simonsson@gmail.com>
Date: Wed, 25 Feb 2015 17:29:23 +0100
Subject: [PATCH] Add automatic locking / unlocking of accounts

* Change account signing API to two sign functions;
  Sign without passphrase - works if account is unlocked
  Sign with passphrase - always works and unlocks the account
* Account stays unlocked for X ms and is then automatically locked
---
 accounts/account_manager.go | 48 ++++++++++++++++++++++++++-------
 accounts/accounts_test.go   | 53 +++++++++++++++++++++++++++++++++++--
 2 files changed, 90 insertions(+), 11 deletions(-)

diff --git a/accounts/account_manager.go b/accounts/account_manager.go
index f7a7506ba2..4d63bd0f28 100644
--- a/accounts/account_manager.go
+++ b/accounts/account_manager.go
@@ -34,24 +34,33 @@ package accounts
 
 import (
 	crand "crypto/rand"
+	"errors"
 	"github.com/ethereum/go-ethereum/crypto"
+	"sync"
+	"time"
 )
 
+var ErrLocked = errors.New("account is locked; please request passphrase")
+
 // TODO: better name for this struct?
 type Account struct {
 	Address []byte
 }
 
 type AccountManager struct {
-	keyStore crypto.KeyStore2
+	keyStore             crypto.KeyStore2
+	unlockedKeys         map[string]crypto.Key
+	unlockedMilliSeconds int
+	mutex                sync.Mutex
 }
 
-// TODO: get key by addr - modify KeyStore2 GetKey to work with addr
-
-// TODO: pass through passphrase for APIs which require access to private key?
-func NewAccountManager(keyStore crypto.KeyStore2) AccountManager {
+func NewAccountManager(keyStore crypto.KeyStore2, unlockMilliSeconds int) AccountManager {
+	keysMap := make(map[string]crypto.Key)
 	am := &AccountManager{
-		keyStore: keyStore,
+		keyStore:             keyStore,
+		unlockedKeys:         keysMap,
+		unlockedMilliSeconds: unlockMilliSeconds,
+		mutex:                sync.Mutex{}, // for accessing unlockedKeys map
 	}
 	return *am
 }
@@ -60,11 +69,26 @@ func (am AccountManager) DeleteAccount(address []byte, auth string) error {
 	return am.keyStore.DeleteKey(address, auth)
 }
 
-func (am *AccountManager) Sign(fromAccount *Account, keyAuth string, toSign []byte) (signature []byte, err error) {
+func (am *AccountManager) Sign(fromAccount *Account, toSign []byte) (signature []byte, err error) {
+	am.mutex.Lock()
+	unlockedKey := am.unlockedKeys[string(fromAccount.Address)]
+	am.mutex.Unlock()
+	if unlockedKey.Address == nil {
+		return nil, ErrLocked
+	}
+	signature, err = crypto.Sign(toSign, unlockedKey.PrivateKey)
+	return signature, err
+}
+
+func (am *AccountManager) SignLocked(fromAccount *Account, keyAuth string, toSign []byte) (signature []byte, err error) {
 	key, err := am.keyStore.GetKey(fromAccount.Address, keyAuth)
 	if err != nil {
 		return nil, err
 	}
+	am.mutex.Lock()
+	am.unlockedKeys[string(fromAccount.Address)] = *key
+	am.mutex.Unlock()
+	go unlockLater(am, fromAccount.Address)
 	signature, err = crypto.Sign(toSign, key.PrivateKey)
 	return signature, err
 }
@@ -80,8 +104,6 @@ func (am AccountManager) NewAccount(auth string) (*Account, error) {
 	return ua, err
 }
 
-// set of accounts == set of keys in given key store
-// TODO: do we need persistence of accounts as well?
 func (am *AccountManager) Accounts() ([]Account, error) {
 	addresses, err := am.keyStore.GetKeyAddresses()
 	if err != nil {
@@ -97,3 +119,11 @@ func (am *AccountManager) Accounts() ([]Account, error) {
 	}
 	return accounts, err
 }
+
+func unlockLater(am *AccountManager, addr []byte) {
+	time.Sleep(time.Millisecond * time.Duration(am.unlockedMilliSeconds))
+	am.mutex.Lock()
+	// TODO: how do we know the key is actually gone from memory?
+	delete(am.unlockedKeys, string(addr))
+	am.mutex.Unlock()
+}
diff --git a/accounts/accounts_test.go b/accounts/accounts_test.go
index 4e97de5456..8f036fd1f1 100644
--- a/accounts/accounts_test.go
+++ b/accounts/accounts_test.go
@@ -6,19 +6,68 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/crypto/randentropy"
 	"github.com/ethereum/go-ethereum/ethutil"
+	"time"
 )
 
 func TestAccountManager(t *testing.T) {
 	ks := crypto.NewKeyStorePlain(ethutil.DefaultDataDir() + "/testaccounts")
-	am := NewAccountManager(ks)
+	am := NewAccountManager(ks, 100)
 	pass := "" // not used but required by API
 	a1, err := am.NewAccount(pass)
 	toSign := randentropy.GetEntropyCSPRNG(32)
-	_, err = am.Sign(a1, pass, toSign)
+	_, err = am.SignLocked(a1, pass, toSign)
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	// Cleanup
+	time.Sleep(time.Millisecond * time.Duration(150)) // wait for locking
+
+	accounts, err := am.Accounts()
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, account := range accounts {
+		err := am.DeleteAccount(account.Address, pass)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestAccountManagerLocking(t *testing.T) {
+	ks := crypto.NewKeyStorePassphrase(ethutil.DefaultDataDir() + "/testaccounts")
+	am := NewAccountManager(ks, 200)
+	pass := "foo"
+	a1, err := am.NewAccount(pass)
+	toSign := randentropy.GetEntropyCSPRNG(32)
+
+	// Signing without passphrase fails because account is locked
+	_, err = am.Sign(a1, toSign)
+	if err != ErrLocked {
+		t.Fatal(err)
+	}
+
+	// Signing with passphrase works
+	_, err = am.SignLocked(a1, pass, toSign)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Signing without passphrase works because account is temp unlocked
+	_, err = am.Sign(a1, toSign)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Signing without passphrase fails after automatic locking
+	time.Sleep(time.Millisecond * time.Duration(250))
+
+	_, err = am.Sign(a1, toSign)
+	if err != ErrLocked {
+		t.Fatal(err)
+	}
+
 	// Cleanup
 	accounts, err := am.Accounts()
 	if err != nil {