diff --git a/management/server/account.go b/management/server/account.go index f5d53161f..894a9832f 100644 --- a/management/server/account.go +++ b/management/server/account.go @@ -84,7 +84,7 @@ func (manager *AccountManager) GetAccount(accountId string) (*Account, error) { account, err := manager.Store.GetAccount(accountId) if err != nil { - return nil, status.Errorf(codes.Internal, "failed retrieving account") + return nil, status.Errorf(codes.NotFound, "account not found") } return account, nil diff --git a/management/server/account_test.go b/management/server/account_test.go index 3bc350133..81ed25fe4 100644 --- a/management/server/account_test.go +++ b/management/server/account_test.go @@ -1,14 +1,18 @@ package server import ( + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "net" "testing" ) func TestAccountManager_AddAccount(t *testing.T) { - store, err := createStore(t) + manager, err := createManager(t) if err != nil { t.Fatal(err) + return } expectedId := "test_account" @@ -19,7 +23,6 @@ func TestAccountManager_AddAccount(t *testing.T) { Mask: net.IPMask{255, 192, 0, 0}, } - manager := NewManager(store) account, err := manager.AddAccount(expectedId) if err != nil { t.Fatal(err) @@ -40,24 +43,158 @@ func TestAccountManager_AddAccount(t *testing.T) { if account.Network.Net.String() != expectedNetwork.String() { t.Errorf("expected account to have Network = %v, got %v", expectedNetwork.String(), account.Network.Net.String()) } +} + +func TestAccountManager_GetOrCreateAccount(t *testing.T) { + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + expectedId := "test_account" + + //make sure account doesn't exist + account, err := manager.GetAccount(expectedId) + if err != nil { + errStatus, ok := status.FromError(err) + if !(ok && errStatus.Code() == codes.NotFound) { + t.Fatal(err) + } + } + if account != nil { + t.Fatal("expecting empty account") + } + + account, err = manager.GetOrCreateAccount(expectedId) + if err != nil { + t.Fatal(err) + } + + if account.Id != expectedId { + t.Fatalf("expected to create an account, got wrong account") + } + + account, err = manager.GetOrCreateAccount(expectedId) + if err != nil { + t.Errorf("expected to get existing account after creation, failed") + } + +} + +func TestAccountManager_AccountExists(t *testing.T) { + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + expectedId := "test_account" + _, err = manager.AddAccount(expectedId) + if err != nil { + t.Fatal(err) + } + + exists, err := manager.AccountExists(expectedId) + if err != nil { + t.Fatal(err) + } + + if !*exists { + t.Errorf("expected account to exist after creation, got false") + } + +} + +func TestAccountManager_GetAccount(t *testing.T) { + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + expectedId := "test_account" + account, err := manager.AddAccount(expectedId) + if err != nil { + t.Fatal(err) + } + + //AddAccount has been already tested so we can assume it is correct and compare results + getAccount, err := manager.GetAccount(expectedId) + if err != nil { + t.Fatal(err) + return + } + + if account.Id != getAccount.Id { + t.Errorf("expected account.ID %s, got %s", account.Id, getAccount.Id) + } + + for _, peer := range account.Peers { + if _, ok := getAccount.Peers[peer.Key]; !ok { + t.Errorf("expected account to have peer %s, not found", peer.Key) + } + } + + for _, key := range account.SetupKeys { + if _, ok := getAccount.SetupKeys[key.Key]; !ok { + t.Errorf("expected account to have setup key %s, not found", key.Key) + } + } } func TestAccountManager_AddPeer(t *testing.T) { + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + account, err := manager.AddAccount("test_account") + if err != nil { + t.Fatal(err) + } + + var setupKey *SetupKey + for _, key := range account.SetupKeys { + setupKey = key + } + + if setupKey == nil { + t.Errorf("expecting account to have a default setup key") + return + } + + key, err := wgtypes.GenerateKey() + if err != nil { + t.Fatal(err) + return + } + expectedPeerKey := key.PublicKey().String() + expectedPeerIP := "100.64.0.1" + + peer, err := manager.AddPeer(setupKey.Key, expectedPeerKey) + if err != nil { + t.Errorf("expecting peer to be added, got failure %v", err) + return + } + + if peer.Key != expectedPeerKey { + t.Errorf("expecting just added peer to have key = %s, got %s", expectedPeerKey, peer.Key) + } + + if peer.Key != expectedPeerKey { + t.Errorf("expecting just added peer to have IP = %s, got %s", expectedPeerIP, peer.IP.String()) + } + +} +func createManager(t *testing.T) (*AccountManager, error) { store, err := createStore(t) if err != nil { - t.Fatal(err) + return nil, err } - - manager := NewManager(store) - - _, err = manager.AddAccount("test_account") - if err != nil { - t.Fatal(err) - } - - //manager.AddPeer(account.SetupKeys[0].Key, "peer-key") + return NewManager(store), nil } func createStore(t *testing.T) (Store, error) { diff --git a/management/server/file_store.go b/management/server/file_store.go index a561c2e30..0be51e558 100644 --- a/management/server/file_store.go +++ b/management/server/file_store.go @@ -116,7 +116,7 @@ func (s *FileStore) SaveAccount(account *Account) error { // 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 + s.SetupKeyId2AccountId[strings.ToUpper(keyId)] = account.Id } for _, peer := range account.Peers {