Improve addition of new peers in Management service. (#56)

* Store refactoring
* Improve addition of new peers in Management service.

management/file_store.go

@@ -0,0 +1,125 @@
+package management
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/wiretrustee/wiretrustee/util"
+)
+
+// storeFileName Store file name. Stored in the datadir
+const storeFileName = "store.json"
+
+// Store represents an account storage
+type FileStore struct {
+	Accounts             map[string]*Account
+	SetupKeyId2AccountId map[string]string `json:"-"`
+
+	// mutex to synchronise Store read/write operations
+	mux       sync.Mutex `json:"-"`
+	storeFile string     `json:"-"`
+}
+
+// NewStore restores a store from the file located in the datadir
+func NewStore(dataDir string) (*FileStore, error) {
+	return restore(filepath.Join(dataDir, storeFileName))
+}
+
+// restore restores the state of the store from the file.
+// Creates a new empty store file if doesn't exist
+func restore(file string) (*FileStore, error) {
+
+	if _, err := os.Stat(file); os.IsNotExist(err) {
+		// create a new FileStore if previously didn't exist (e.g. first run)
+		s := &FileStore{
+			Accounts:  make(map[string]*Account),
+			mux:       sync.Mutex{},
+			storeFile: file,
+		}
+
+		err = s.persist(file)
+		if err != nil {
+			return nil, err
+		}
+
+		return s, nil
+	}
+
+	read, err := util.ReadJson(file, &FileStore{})
+	if err != nil {
+		return nil, err
+	}
+
+	store := read.(*FileStore)
+	store.storeFile = file
+	store.SetupKeyId2AccountId = make(map[string]string)
+	for accountId, account := range store.Accounts {
+		for setupKeyId := range account.SetupKeys {
+			store.SetupKeyId2AccountId[strings.ToLower(setupKeyId)] = accountId
+		}
+	}
+
+	return store, nil
+}
+
+// persist persists account data to a file
+// It is recommended to call it with locking FileStore.mux
+func (s *FileStore) persist(file string) error {
+	return util.WriteJson(file, s)
+}
+
+// AddPeer adds peer to the store and associates it with a Account and a SetupKey. Returns related Account
+// Each Account has a list of pre-authorised SetupKey and if no Account has a given key nil will be returned, meaning the key is invalid
+func (s *FileStore) AddPeer(setupKey string, peerKey string) error {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	accountId, accountIdFound := s.SetupKeyId2AccountId[strings.ToLower(setupKey)]
+	if !accountIdFound {
+		return status.Errorf(codes.NotFound, "Provided setup key doesn't exists")
+	}
+
+	account, accountFound := s.Accounts[accountId]
+	if !accountFound {
+		return status.Errorf(codes.Internal, "Invalid setup key")
+	}
+
+	key, setupKeyFound := account.SetupKeys[strings.ToLower(setupKey)]
+	if !setupKeyFound {
+		return status.Errorf(codes.Internal, "Invalid setup key")
+	}
+
+	account.Peers[peerKey] = &Peer{Key: peerKey, SetupKey: key}
+	err := s.persist(s.storeFile)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// AddAccount adds new account to the store.
+func (s *FileStore) AddAccount(account *Account) error {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	// todo will override, handle existing keys
+	s.Accounts[account.Id] = account
+
+	// todo check that account.Id and keyId are not exist already
+	// because if keyId exists for other accounts this can be bad
+	for keyId := range account.SetupKeys {
+		s.SetupKeyId2AccountId[strings.ToLower(keyId)] = account.Id
+	}
+
+	err := s.persist(s.storeFile)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

management/management_test.go

@@ -2,6 +2,13 @@ package management_test
 
 import (
 	"context"
+	"io/ioutil"
+	"net"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 	log "github.com/sirupsen/logrus"
@@ -11,12 +18,6 @@ import (
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
-	"io/ioutil"
-	"net"
-	"os"
-	"path/filepath"
-	"strings"
-	"time"
 )
 
 var _ = Describe("Client", func() {
@@ -114,11 +115,11 @@ var _ = Describe("Client", func() {
 
 				Expect(err).NotTo(HaveOccurred())
 
-				store, err := util.ReadJson(filepath.Join(dataDir, "store.json"), &mgmt.Store{})
+				store, err := util.ReadJson(filepath.Join(dataDir, "store.json"), &mgmt.FileStore{})
 				Expect(err).NotTo(HaveOccurred())
 
-				Expect(store.(*mgmt.Store)).NotTo(BeNil())
+				Expect(store.(*mgmt.FileStore)).NotTo(BeNil())
-				user := store.(*mgmt.Store).Accounts["bf1c8084-ba50-4ce7-9439-34653001fc3b"]
+				user := store.(*mgmt.FileStore).Accounts["bf1c8084-ba50-4ce7-9439-34653001fc3b"]
 				Expect(user.Peers[key.PublicKey().String()]).NotTo(BeNil())
 				Expect(user.SetupKeys[strings.ToLower(setupKey)]).NotTo(BeNil())
 

management/server.go

@@ -2,13 +2,14 @@ package management
 
 import (
 	"context"
+
 	"github.com/wiretrustee/wiretrustee/management/proto"
 	"google.golang.org/grpc/status"
 )
 
 // Server an instance of a Management server
 type Server struct {
-	Store *Store
+	Store *FileStore
 }
 
 // NewServer creates a new Management server

management/store.go

@@ -1,17 +1,5 @@
 package management
 
-import (
-	"fmt"
-	"github.com/wiretrustee/wiretrustee/util"
-	"os"
-	"path/filepath"
-	"strings"
-	"sync"
-)
-
-// storeFileName Store file name. Stored in the datadir
-const storeFileName = "store.json"
-
 // Account represents a unique account of the system
 type Account struct {
 	Id        string
@@ -34,87 +22,6 @@ type Peer struct {
 	SetupKey *SetupKey
 }
 
-// Store represents an account storage
+type Store interface {
-type Store struct {
+	AddPeer(setupKey string, peerKey string) error
-	Accounts map[string]*Account
-
-	// mutex to synchronise Store read/write operations
-	mux       sync.Mutex `json:"-"`
-	storeFile string     `json:"-"`
-}
-
-// NewStore restores a store from the file located in the datadir
-func NewStore(dataDir string) (*Store, error) {
-	return restore(filepath.Join(dataDir, storeFileName))
-}
-
-// restore restores the state of the store from the file.
-// Creates a new empty store file if doesn't exist
-func restore(file string) (*Store, error) {
-
-	if _, err := os.Stat(file); os.IsNotExist(err) {
-		// create a new Store if previously didn't exist (e.g. first run)
-		s := &Store{
-			Accounts:  make(map[string]*Account),
-			mux:       sync.Mutex{},
-			storeFile: file,
-		}
-
-		err = s.persist(file)
-		if err != nil {
-			return nil, err
-		}
-
-		return s, nil
-	}
-
-	read, err := util.ReadJson(file, &Store{})
-	if err != nil {
-		return nil, err
-	}
-	read.(*Store).storeFile = file
-
-	return read.(*Store), nil
-}
-
-// persist persists account data to a file
-// It is recommended to call it with locking Store,mux
-func (s *Store) persist(file string) error {
-	return util.WriteJson(file, s)
-}
-
-// AddPeer adds peer to the store and associates it with a Account and a SetupKey. Returns related Account
-// Each Account has a list of pre-authorised SetupKey and if no Account has a given key nil will be returned, meaning the key is invalid
-func (s *Store) AddPeer(setupKey string, peerKey string) error {
-	s.mux.Lock()
-	defer s.mux.Unlock()
-
-	for _, u := range s.Accounts {
-		for _, key := range u.SetupKeys {
-			if key.Key == strings.ToLower(setupKey) {
-				u.Peers[peerKey] = &Peer{Key: peerKey, SetupKey: key}
-				err := s.persist(s.storeFile)
-				if err != nil {
-					return err
-				}
-				return nil
-			}
-		}
-	}
-
-	return fmt.Errorf("invalid setup key")
-}
-
-// AddAccount adds new account to the store.
-func (s *Store) AddAccount(account *Account) error {
-	s.mux.Lock()
-	defer s.mux.Unlock()
-	// todo will override, handle existing keys
-	s.Accounts[account.Id] = account
-	err := s.persist(s.storeFile)
-	if err != nil {
-		return err
-	}
-
-	return nil
 }