mirror of
https://github.com/netbirdio/netbird.git
synced 2025-05-31 07:07:42 +02:00
Improve addition of new peers in Management service. (#56)
* Store refactoring * Improve addition of new peers in Management service.
This commit is contained in:
andpar83
GitHub
parent
f9c3ed784f
commit
7b52049333
125
management/file_store.go
Normal file
125
management/file_store.go
Normal file
@ -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
|
||||
}
|
||||
@ -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())
|
||||
user := store.(*mgmt.Store).Accounts["bf1c8084-ba50-4ce7-9439-34653001fc3b"]
|
||||
Expect(store.(*mgmt.FileStore)).NotTo(BeNil())
|
||||
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())
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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 struct {
|
||||
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
|
||||
type Store interface {
|
||||
AddPeer(setupKey string, peerKey string) error
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user