mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-22 18:26:41 +00:00
Try PostgreSQL store
This commit is contained in:
Yury Gargay
353
management/server/postgresql_store_test.go
Normal file
353
management/server/postgresql_store_test.go
Normal file
@@ -0,0 +1,353 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"net"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
func TestPostgresql_NewStore(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStore(t)
|
||||
defer cleanup()
|
||||
|
||||
if len(store.GetAllAccounts()) != 0 {
|
||||
t.Errorf("expected to create a new empty Accounts map when creating a new FileStore")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresql_SaveAccount(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStore(t)
|
||||
defer cleanup()
|
||||
|
||||
account := newAccountWithId("account_id", "testuser", "")
|
||||
setupKey := GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
SetupKey: "peerkeysetupkey",
|
||||
IP: net.IP{127, 0, 0, 1},
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err := store.SaveAccount(account)
|
||||
require.NoError(t, err)
|
||||
|
||||
account2 := newAccountWithId("account_id2", "testuser2", "")
|
||||
setupKey = GenerateDefaultSetupKey()
|
||||
account2.SetupKeys[setupKey.Key] = setupKey
|
||||
account2.Peers["testpeer2"] = &nbpeer.Peer{
|
||||
Key: "peerkey2",
|
||||
SetupKey: "peerkeysetupkey2",
|
||||
IP: net.IP{127, 0, 0, 2},
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name 2",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(account2)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts()) != 2 {
|
||||
t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
a, err := store.GetAccount(account.Id)
|
||||
if a == nil {
|
||||
t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a != nil && len(a.Policies) != 1 {
|
||||
t.Errorf("expecting Account to have one policy stored after SaveAccount(), got %d", len(a.Policies))
|
||||
}
|
||||
|
||||
if a != nil && len(a.Policies[0].Rules) != 1 {
|
||||
t.Errorf("expecting Account to have one policy rule stored after SaveAccount(), got %d", len(a.Policies[0].Rules))
|
||||
return
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByPeerPubKey("peerkey"); a == nil {
|
||||
t.Errorf("expecting PeerKeyID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByUser("testuser"); a == nil {
|
||||
t.Errorf("expecting UserID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByPeerID("testpeer"); a == nil {
|
||||
t.Errorf("expecting PeerID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountBySetupKey(setupKey.Key); a == nil {
|
||||
t.Errorf("expecting SetupKeyID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresql_DeleteAccount(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStore(t)
|
||||
defer cleanup()
|
||||
|
||||
testUserID := "testuser"
|
||||
user := NewAdminUser(testUserID)
|
||||
user.PATs = map[string]*PersonalAccessToken{"testtoken": {
|
||||
ID: "testtoken",
|
||||
Name: "test token",
|
||||
}}
|
||||
|
||||
account := newAccountWithId("account_id", testUserID, "")
|
||||
setupKey := GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
SetupKey: "peerkeysetupkey",
|
||||
IP: net.IP{127, 0, 0, 1},
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
account.Users[testUserID] = user
|
||||
|
||||
err := store.SaveAccount(account)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts()) != 1 {
|
||||
t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
err = store.DeleteAccount(account)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts()) != 0 {
|
||||
t.Errorf("expecting 0 Accounts to be stored after DeleteAccount()")
|
||||
}
|
||||
|
||||
_, err = store.GetAccountByPeerPubKey("peerkey")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer public key")
|
||||
|
||||
_, err = store.GetAccountByUser("testuser")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by user")
|
||||
|
||||
_, err = store.GetAccountByPeerID("testpeer")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer id")
|
||||
|
||||
_, err = store.GetAccountBySetupKey(setupKey.Key)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by setup key")
|
||||
|
||||
_, err = store.GetAccount(account.Id)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by id")
|
||||
|
||||
for _, policy := range account.Policies {
|
||||
var rules []*PolicyRule
|
||||
err = store.db.Model(&PolicyRule{}).Find(&rules, "policy_id = ?", policy.ID).Error
|
||||
require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for policy rules")
|
||||
require.Len(t, rules, 0, "expecting no policy rules to be found after removing DeleteAccount")
|
||||
|
||||
}
|
||||
|
||||
for _, accountUser := range account.Users {
|
||||
var pats []*PersonalAccessToken
|
||||
err = store.db.Model(&PersonalAccessToken{}).Find(&pats, "user_id = ?", accountUser.Id).Error
|
||||
require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for personal access token")
|
||||
require.Len(t, pats, 0, "expecting no personal access token to be found after removing DeleteAccount")
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestPostgresql_SavePeerStatus(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStoreFromFile(t, "testdata/store.json")
|
||||
defer cleanup()
|
||||
|
||||
account, err := store.GetAccount("bf1c8084-ba50-4ce7-9439-34653001fc3b")
|
||||
require.NoError(t, err)
|
||||
|
||||
// save status of non-existing peer
|
||||
newStatus := nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()}
|
||||
err = store.SavePeerStatus(account.Id, "non-existing-peer", newStatus)
|
||||
assert.Error(t, err)
|
||||
|
||||
// save new status of existing peer
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
ID: "testpeer",
|
||||
SetupKey: "peerkeysetupkey",
|
||||
IP: net.IP{127, 0, 0, 1},
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: false, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(account)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = store.SavePeerStatus(account.Id, "testpeer", newStatus)
|
||||
require.NoError(t, err)
|
||||
|
||||
account, err = store.GetAccount(account.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
actual := account.Peers["testpeer"].Status
|
||||
assert.Equal(t, newStatus.Connected, actual.Connected)
|
||||
assert.True(t, newStatus.LastSeen.Equal(actual.LastSeen))
|
||||
}
|
||||
|
||||
func TestPostgresql_TestGetAccountByPrivateDomain(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStoreFromFile(t, "testdata/store.json")
|
||||
defer cleanup()
|
||||
|
||||
existingDomain := "test.com"
|
||||
|
||||
account, err := store.GetAccountByPrivateDomain(existingDomain)
|
||||
require.NoError(t, err, "should found account")
|
||||
require.Equal(t, existingDomain, account.Domain, "domains should match")
|
||||
|
||||
_, err = store.GetAccountByPrivateDomain("missing-domain.com")
|
||||
require.Error(t, err, "should return error on domain lookup")
|
||||
}
|
||||
|
||||
func TestPostgresql_GetTokenIDByHashedToken(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStoreFromFile(t, "testdata/store.json")
|
||||
defer cleanup()
|
||||
|
||||
hashed := "SoMeHaShEdToKeN"
|
||||
id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
token, err := store.GetTokenIDByHashedToken(hashed)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, id, token)
|
||||
}
|
||||
|
||||
func TestPostgresql_GetUserByTokenID(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The PostgreSQL store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
store, cleanup := newPostgresqlStoreFromFile(t, "testdata/store.json")
|
||||
defer cleanup()
|
||||
|
||||
id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
user, err := store.GetUserByTokenID(id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, id, user.PATs[id].ID)
|
||||
}
|
||||
|
||||
func newPostgresqlStore(t *testing.T) (*PostgresqlStore, func()) {
|
||||
t.Helper()
|
||||
|
||||
dbName := "store_" + randString(10)
|
||||
postgresDsn := "host=localhost user=postgres password=postgres port=5432 sslmode=disable"
|
||||
db, _ := gorm.Open(postgres.Open(postgresDsn), &gorm.Config{})
|
||||
result := db.Exec(fmt.Sprintf("CREATE DATABASE %s ENCODING = 'UTF8'", dbName))
|
||||
if result.Error != nil {
|
||||
t.Fatalf("could not initialize postgresql store: %s", result.Error)
|
||||
}
|
||||
postgresDsn = fmt.Sprintf("%s dbname=%s ", postgresDsn, dbName)
|
||||
cleanup := func() {
|
||||
db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", dbName))
|
||||
}
|
||||
|
||||
store, err := NewPostgresqlStore(postgresDsn, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("could not initialize postgresql store: %s", err)
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, store)
|
||||
|
||||
return store, cleanup
|
||||
}
|
||||
|
||||
func randString(n int) string {
|
||||
var letterRunes = []rune("abcdefghijklmnopqrstuvwxyz1234567890")
|
||||
b := make([]rune, n)
|
||||
for i := range b {
|
||||
b[i] = letterRunes[rand.Intn(len(letterRunes))]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func newPostgresqlStoreFromFile(t *testing.T, filename string) (*PostgresqlStore, func()) {
|
||||
t.Helper()
|
||||
|
||||
storeDir := t.TempDir()
|
||||
|
||||
err := util.CopyFileContents(filename, filepath.Join(storeDir, "store.json"))
|
||||
require.NoError(t, err)
|
||||
|
||||
fStore, err := NewFileStore(storeDir, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
dbName := "store_" + randString(10)
|
||||
postgresDsn := "host=localhost user=postgres password=postgres port=5432 sslmode=disable"
|
||||
db, _ := gorm.Open(postgres.Open(postgresDsn), &gorm.Config{})
|
||||
result := db.Exec(fmt.Sprintf("CREATE DATABASE %s ENCODING = 'UTF8'", dbName))
|
||||
if result.Error != nil {
|
||||
t.Fatalf("could not initialize postgresql store: %s", result.Error)
|
||||
}
|
||||
postgresDsn = fmt.Sprintf("%s dbname=%s ", postgresDsn, dbName)
|
||||
cleanup := func() {
|
||||
db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", dbName))
|
||||
}
|
||||
|
||||
store, err := NewPostgresqlStoreFromFileStore(fStore, postgresDsn, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, store)
|
||||
|
||||
return store, cleanup
|
||||
}
|
||||
|
||||
/*
|
||||
func newAccount(store Store, id int) error {
|
||||
str := fmt.Sprintf("%s-%d", uuid.New().String(), id)
|
||||
account := newAccountWithId(str, str+"-testuser", "example.com")
|
||||
setupKey := GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["p"+str] = &nbpeer.Peer{
|
||||
Key: "peerkey" + str,
|
||||
SetupKey: "peerkeysetupkey",
|
||||
IP: net.IP{127, 0, 0, 1},
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
return store.SaveAccount(account)
|
||||
}
|
||||
*/
|
||||
Reference in New Issue
Block a user