mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-25 08:09:07 +02:00
[management] split store by table (#7646)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,149 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// CreateAccessLog creates a new access log entry in the database
|
||||
func (s *SqlStore) CreateAccessLog(ctx context.Context, logEntry *accesslogs.AccessLogEntry) error {
|
||||
result := s.db.Create(logEntry)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).WithFields(log.Fields{
|
||||
"service_id": logEntry.ServiceID,
|
||||
"method": logEntry.Method,
|
||||
"host": logEntry.Host,
|
||||
"path": logEntry.Path,
|
||||
}).Errorf("failed to create access log entry in store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to create access log entry in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAccountAccessLogs retrieves access logs for a given account with pagination and filtering
|
||||
func (s *SqlStore) GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) {
|
||||
var logs []*accesslogs.AccessLogEntry
|
||||
var totalCount int64
|
||||
|
||||
baseQuery := s.db.
|
||||
Model(&accesslogs.AccessLogEntry{}).
|
||||
Where(accountIDCondition, accountID)
|
||||
|
||||
baseQuery = s.applyAccessLogFilters(baseQuery, filter)
|
||||
|
||||
if err := baseQuery.Count(&totalCount).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count access logs: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to count access logs")
|
||||
}
|
||||
|
||||
query := s.db.
|
||||
Where(accountIDCondition, accountID)
|
||||
|
||||
query = s.applyAccessLogFilters(query, filter)
|
||||
|
||||
sortColumns := filter.GetSortColumn()
|
||||
sortOrder := strings.ToUpper(filter.GetSortOrder())
|
||||
|
||||
var orderClauses []string
|
||||
for _, col := range strings.Split(sortColumns, ",") {
|
||||
col = strings.TrimSpace(col)
|
||||
if col != "" {
|
||||
orderClauses = append(orderClauses, col+" "+sortOrder)
|
||||
}
|
||||
}
|
||||
orderClause := strings.Join(orderClauses, ", ")
|
||||
|
||||
query = query.
|
||||
Order(orderClause).
|
||||
Limit(filter.GetLimit()).
|
||||
Offset(filter.GetOffset())
|
||||
|
||||
if lockStrength != LockingStrengthNone {
|
||||
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
result := query.Find(&logs)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get access logs from store: %v", result.Error)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to get access logs from store")
|
||||
}
|
||||
|
||||
return logs, totalCount, nil
|
||||
}
|
||||
|
||||
// DeleteOldAccessLogs deletes all access logs older than the specified time
|
||||
func (s *SqlStore) DeleteOldAccessLogs(ctx context.Context, olderThan time.Time) (int64, error) {
|
||||
result := s.db.
|
||||
Where("timestamp < ?", olderThan).
|
||||
Delete(&accesslogs.AccessLogEntry{})
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
|
||||
return 0, status.Errorf(status.Internal, "failed to delete old access logs")
|
||||
}
|
||||
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
// applyAccessLogFilters applies filter conditions to the query
|
||||
func (s *SqlStore) applyAccessLogFilters(query *gorm.DB, filter accesslogs.AccessLogFilter) *gorm.DB {
|
||||
if filter.Search != nil {
|
||||
searchPattern := "%" + *filter.Search + "%"
|
||||
query = query.Where(
|
||||
"id LIKE ? OR location_connection_ip LIKE ? OR host LIKE ? OR path LIKE ? OR CONCAT(host, path) LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)",
|
||||
searchPattern, searchPattern, searchPattern, searchPattern, searchPattern, searchPattern, searchPattern,
|
||||
)
|
||||
}
|
||||
|
||||
if filter.SourceIP != nil {
|
||||
query = query.Where("location_connection_ip = ?", *filter.SourceIP)
|
||||
}
|
||||
|
||||
if filter.Host != nil {
|
||||
query = query.Where("host = ?", *filter.Host)
|
||||
}
|
||||
|
||||
if filter.Path != nil {
|
||||
// Support LIKE pattern for path filtering
|
||||
query = query.Where("path LIKE ?", "%"+*filter.Path+"%")
|
||||
}
|
||||
|
||||
if filter.UserID != nil {
|
||||
query = query.Where("user_id = ?", *filter.UserID)
|
||||
}
|
||||
|
||||
if filter.Method != nil {
|
||||
query = query.Where("method = ?", *filter.Method)
|
||||
}
|
||||
|
||||
if filter.Status != nil {
|
||||
switch *filter.Status {
|
||||
case "success":
|
||||
query = query.Where("status_code >= ? AND status_code < ?", 200, 400)
|
||||
case "failed":
|
||||
query = query.Where("status_code < ? OR status_code >= ?", 200, 400)
|
||||
}
|
||||
}
|
||||
|
||||
if filter.StatusCode != nil {
|
||||
query = query.Where("status_code = ?", *filter.StatusCode)
|
||||
}
|
||||
|
||||
if filter.StartDate != nil {
|
||||
query = query.Where("timestamp >= ?", *filter.StartDate)
|
||||
}
|
||||
|
||||
if filter.EndDate != nil {
|
||||
query = query.Where("timestamp <= ?", *filter.EndDate)
|
||||
}
|
||||
|
||||
return query
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,70 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// GetAccountOnboarding retrieves the onboarding information for a specific account.
|
||||
func (s *SqlStore) GetAccountOnboarding(ctx context.Context, accountID string) (*types.AccountOnboarding, error) {
|
||||
var accountOnboarding types.AccountOnboarding
|
||||
result := s.db.Model(&accountOnboarding).Take(&accountOnboarding, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewAccountOnboardingNotFoundError(accountID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("error when getting account onboarding %s from the store: %s", accountID, result.Error)
|
||||
return nil, status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
return &accountOnboarding, nil
|
||||
}
|
||||
|
||||
// SaveAccountOnboarding updates the onboarding information for a specific account.
|
||||
func (s *SqlStore) SaveAccountOnboarding(ctx context.Context, onboarding *types.AccountOnboarding) error {
|
||||
result := s.db.Clauses(clause.OnConflict{UpdateAll: true}).Create(onboarding)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("error when saving account onboarding %s in the store: %s", onboarding.AccountID, result.Error)
|
||||
return status.Errorf(status.Internal, "error when saving account onboarding %s in the store: %s", onboarding.AccountID, result.Error)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) getAccountOnboarding(ctx context.Context, accountID string, account *types.Account) error {
|
||||
const query = `SELECT account_id, onboarding_flow_pending, signup_form_pending, created_at, updated_at FROM account_onboardings WHERE account_id = $1`
|
||||
var onboardingFlowPending, signupFormPending sql.NullBool
|
||||
var createdAt, updatedAt sql.NullTime
|
||||
err := s.pool.QueryRow(ctx, query, accountID).Scan(
|
||||
&account.Onboarding.AccountID,
|
||||
&onboardingFlowPending,
|
||||
&signupFormPending,
|
||||
&createdAt,
|
||||
&updatedAt,
|
||||
)
|
||||
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
if createdAt.Valid {
|
||||
account.Onboarding.CreatedAt = createdAt.Time
|
||||
}
|
||||
if updatedAt.Valid {
|
||||
account.Onboarding.UpdatedAt = updatedAt.Time
|
||||
}
|
||||
if onboardingFlowPending.Valid {
|
||||
account.Onboarding.OnboardingFlowPending = onboardingFlowPending.Bool
|
||||
}
|
||||
if signupFormPending.Valid {
|
||||
account.Onboarding.SignupFormPending = signupFormPending.Bool
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetAccountOnboarding(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "9439-34653001fc3b-bf1c8084-ba50-4ce7"
|
||||
a, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
t.Logf("Onboarding: %+v", a.Onboarding)
|
||||
err = store.SaveAccount(context.Background(), a)
|
||||
require.NoError(t, err)
|
||||
onboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, onboarding)
|
||||
require.Equal(t, accountID, onboarding.AccountID)
|
||||
require.Equal(t, time.Date(2024, time.October, 2, 14, 1, 38, 210000000, time.UTC), onboarding.CreatedAt.UTC())
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveAccountOnboarding(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
t.Run("New onboarding should be saved correctly", func(t *testing.T) {
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
onboarding := &types.AccountOnboarding{
|
||||
AccountID: accountID,
|
||||
SignupFormPending: true,
|
||||
OnboardingFlowPending: true,
|
||||
}
|
||||
|
||||
err = store.SaveAccountOnboarding(context.Background(), onboarding)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedOnboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, onboarding.SignupFormPending, savedOnboarding.SignupFormPending)
|
||||
require.Equal(t, onboarding.OnboardingFlowPending, savedOnboarding.OnboardingFlowPending)
|
||||
})
|
||||
|
||||
t.Run("Existing onboarding should be updated correctly", func(t *testing.T) {
|
||||
accountID := "9439-34653001fc3b-bf1c8084-ba50-4ce7"
|
||||
onboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
onboarding.OnboardingFlowPending = !onboarding.OnboardingFlowPending
|
||||
onboarding.SignupFormPending = !onboarding.SignupFormPending
|
||||
|
||||
err = store.SaveAccountOnboarding(context.Background(), onboarding)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedOnboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, onboarding.SignupFormPending, savedOnboarding.SignupFormPending)
|
||||
require.Equal(t, onboarding.OnboardingFlowPending, savedOnboarding.OnboardingFlowPending)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,965 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
proxydomain "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
|
||||
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
||||
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
|
||||
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
|
||||
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
"github.com/netbirdio/netbird/shared/testing_helpers"
|
||||
)
|
||||
|
||||
func Test_SaveAccount_Large(t *testing.T) {
|
||||
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
|
||||
t.Skip("skip CI tests on darwin and windows")
|
||||
}
|
||||
|
||||
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
|
||||
runLargeTest(t, store)
|
||||
})
|
||||
}
|
||||
|
||||
func runLargeTest(t *testing.T, store Store) {
|
||||
t.Helper()
|
||||
|
||||
account := newAccountWithId(context.Background(), "account_id", "testuser", "")
|
||||
groupALL, err := account.GetGroupAll()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
const numPerAccount = 6000
|
||||
for n := 0; n < numPerAccount; n++ {
|
||||
netIP := sequentialIPv4(n)
|
||||
peerID := fmt.Sprintf("%s-peer-%d", account.Id, n)
|
||||
addr, _ := netip.AddrFromSlice(netIP)
|
||||
|
||||
peer := &nbpeer.Peer{
|
||||
ID: peerID,
|
||||
Key: peerID,
|
||||
IP: addr.Unmap(),
|
||||
Name: peerID,
|
||||
DNSLabel: peerID,
|
||||
UserID: "testuser",
|
||||
Status: &nbpeer.PeerStatus{Connected: false, LastSeen: time.Now()},
|
||||
SSHEnabled: false,
|
||||
}
|
||||
account.Peers[peerID] = peer
|
||||
group, _ := account.GetGroupAll()
|
||||
group.Peers = append(group.Peers, peerID)
|
||||
user := &types.User{
|
||||
Id: fmt.Sprintf("%s-user-%d", account.Id, n),
|
||||
AccountID: account.Id,
|
||||
}
|
||||
account.Users[user.Id] = user
|
||||
route := &nbroute.Route{
|
||||
ID: nbroute.ID(fmt.Sprintf("network-id-%d", n)),
|
||||
Description: "base route",
|
||||
NetID: nbroute.NetID(fmt.Sprintf("network-id-%d", n)),
|
||||
Network: netip.MustParsePrefix(netIP.String() + "/24"),
|
||||
NetworkType: nbroute.IPv4Network,
|
||||
Metric: 9999,
|
||||
Masquerade: false,
|
||||
Enabled: true,
|
||||
Groups: []string{groupALL.ID},
|
||||
}
|
||||
account.Routes[route.ID] = route
|
||||
|
||||
group = &types.Group{
|
||||
ID: fmt.Sprintf("group-id-%d", n),
|
||||
AccountID: account.Id,
|
||||
Name: fmt.Sprintf("group-id-%d", n),
|
||||
Issued: "api",
|
||||
Peers: nil,
|
||||
}
|
||||
account.Groups[group.ID] = group
|
||||
|
||||
nameserver := &nbdns.NameServerGroup{
|
||||
ID: fmt.Sprintf("nameserver-id-%d", n),
|
||||
AccountID: account.Id,
|
||||
Name: fmt.Sprintf("nameserver-id-%d", n),
|
||||
Description: "",
|
||||
NameServers: []nbdns.NameServer{{IP: netip.MustParseAddr(netIP.String()), NSType: nbdns.UDPNameServerType}},
|
||||
Groups: []string{group.ID},
|
||||
Primary: false,
|
||||
Domains: nil,
|
||||
Enabled: false,
|
||||
SearchDomainsEnabled: false,
|
||||
}
|
||||
account.NameServerGroups[nameserver.ID] = nameserver
|
||||
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
_, exists := account.SetupKeys[setupKey.Key]
|
||||
if exists {
|
||||
t.Errorf("setup key already exists")
|
||||
}
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
}
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 1 {
|
||||
t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
a, err := store.GetAccount(context.Background(), 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 != nil && len(a.Peers) != numPerAccount {
|
||||
t.Errorf("expecting Account to have %d peers stored after SaveAccount(), got %d",
|
||||
numPerAccount, len(a.Peers))
|
||||
return
|
||||
}
|
||||
|
||||
if a != nil && len(a.Users) != numPerAccount+1 {
|
||||
t.Errorf("expecting Account to have %d users stored after SaveAccount(), got %d",
|
||||
numPerAccount+1, len(a.Users))
|
||||
return
|
||||
}
|
||||
|
||||
if a != nil && len(a.Routes) != numPerAccount {
|
||||
t.Errorf("expecting Account to have %d routes stored after SaveAccount(), got %d",
|
||||
numPerAccount, len(a.Routes))
|
||||
return
|
||||
}
|
||||
|
||||
if a != nil && len(a.NameServerGroups) != numPerAccount {
|
||||
t.Errorf("expecting Account to have %d NameServerGroups stored after SaveAccount(), got %d",
|
||||
numPerAccount, len(a.NameServerGroups))
|
||||
return
|
||||
}
|
||||
|
||||
if a != nil && len(a.NameServerGroups) != numPerAccount {
|
||||
t.Errorf("expecting Account to have %d NameServerGroups stored after SaveAccount(), got %d",
|
||||
numPerAccount, len(a.NameServerGroups))
|
||||
return
|
||||
}
|
||||
|
||||
if a != nil && len(a.SetupKeys) != numPerAccount+1 {
|
||||
t.Errorf("expecting Account to have %d SetupKeys stored after SaveAccount(), got %d",
|
||||
numPerAccount+1, len(a.SetupKeys))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// sequentialIPv4 returns a unique IPv4 address for the given index, avoiding
|
||||
// the random collisions that would otherwise violate the unique (account_id, ip)
|
||||
// index when generating a large number of peers.
|
||||
func sequentialIPv4(n int) net.IP {
|
||||
b := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(b, 0x0A000000+uint32(n))
|
||||
return net.IP(b)
|
||||
}
|
||||
|
||||
func Test_SaveAccount(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
|
||||
account := newAccountWithId(context.Background(), "account_id", "testuser", "")
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err := store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
account2 := newAccountWithId(context.Background(), "account_id2", "testuser2", "")
|
||||
setupKey, _ = types.GenerateDefaultSetupKey()
|
||||
account2.SetupKeys[setupKey.Key] = setupKey
|
||||
account2.Peers["testpeer2"] = &nbpeer.Peer{
|
||||
Key: "peerkey2",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 2}),
|
||||
IPv6: netip.MustParseAddr("fd00::2"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name 2",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(context.Background(), account2)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 2 {
|
||||
t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
a, err := store.GetAccount(context.Background(), 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(context.Background(), "peerkey"); a == nil {
|
||||
t.Errorf("expecting PeerKeyID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByUser(context.Background(), "testuser"); a == nil {
|
||||
t.Errorf("expecting UserID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByPeerID(context.Background(), "testpeer"); a == nil {
|
||||
t.Errorf("expecting PeerID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountBySetupKey(context.Background(), setupKey.Key); a == nil {
|
||||
t.Errorf("expecting SetupKeyID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func Test_AccountSettings_SaveAndRetrieve(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
populateFields := testing_helpers.NewPopulateFields().WithCustomFieldSetter(
|
||||
reflect.PointerTo(reflect.TypeOf(types.ExtraSettings{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
|
||||
es := types.ExtraSettings{}
|
||||
reflectedEs := reflect.ValueOf(&es).Elem()
|
||||
n, err := this.PopulateAll(reflectedEs)
|
||||
if err != nil {
|
||||
return n, err
|
||||
}
|
||||
field.Set(reflectedEs.Addr())
|
||||
return n, nil
|
||||
}).WithCustomFieldSetter(
|
||||
reflect.PointerTo(reflect.TypeOf(types.DashboardFeatures{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
|
||||
t := true
|
||||
df := types.DashboardFeatures{AgentNetwork: &t}
|
||||
reflectedDf := reflect.ValueOf(&df).Elem()
|
||||
field.Set(reflectedDf.Addr())
|
||||
return 1, nil
|
||||
}).WithSkippedTag("gorm", "-")
|
||||
|
||||
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
|
||||
account := newAccountWithId(context.Background(), "account_id", "testuser", "")
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
|
||||
settings := types.Settings{}
|
||||
numOfExportedFields, err := populateFields.PopulateAll(reflect.ValueOf(&settings).Elem())
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 27, numOfExportedFields)
|
||||
account.Settings = &settings
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
assert.NoError(t, err)
|
||||
|
||||
accountFromDb, err := store.GetAccount(context.Background(), account.Id)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, accountFromDb)
|
||||
assert.NotNil(t, accountFromDb.Settings)
|
||||
|
||||
assert.True(t, reflect.DeepEqual(&settings, accountFromDb.Settings), "created settings and settings retrieved from the db should match")
|
||||
})
|
||||
}
|
||||
|
||||
func TestSqlite_DeleteAccount(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
testUserID := "testuser"
|
||||
user := types.NewAdminUser(testUserID)
|
||||
user.PATs = map[string]*types.PersonalAccessToken{"testtoken": {
|
||||
ID: "testtoken",
|
||||
Name: "test token",
|
||||
}}
|
||||
|
||||
account := newAccountWithId(context.Background(), "account_id", testUserID, "")
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
account.Users[testUserID] = user
|
||||
account.Networks = []*networkTypes.Network{
|
||||
{
|
||||
ID: "network_id",
|
||||
AccountID: account.Id,
|
||||
Name: "network name",
|
||||
Description: "network description",
|
||||
},
|
||||
}
|
||||
account.NetworkRouters = []*routerTypes.NetworkRouter{
|
||||
{
|
||||
ID: "router_id",
|
||||
NetworkID: account.Networks[0].ID,
|
||||
AccountID: account.Id,
|
||||
PeerGroups: []string{"group_id"},
|
||||
Masquerade: true,
|
||||
Metric: 1,
|
||||
},
|
||||
}
|
||||
account.NetworkResources = []*resourceTypes.NetworkResource{
|
||||
{
|
||||
ID: "resource_id",
|
||||
NetworkID: account.Networks[0].ID,
|
||||
AccountID: account.Id,
|
||||
Name: "Name",
|
||||
Description: "Description",
|
||||
Type: "Domain",
|
||||
Address: "example.com",
|
||||
},
|
||||
}
|
||||
|
||||
account.Services = []*rpservice.Service{
|
||||
{
|
||||
ID: "service_id",
|
||||
AccountID: account.Id,
|
||||
Name: "test service",
|
||||
Domain: "svc.example.com",
|
||||
Enabled: true,
|
||||
Targets: []*rpservice.Target{
|
||||
{
|
||||
AccountID: account.Id,
|
||||
ServiceID: "service_id",
|
||||
Host: "localhost",
|
||||
Port: 8080,
|
||||
Protocol: "http",
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
account.Domains = []*proxydomain.Domain{
|
||||
{
|
||||
ID: "domain_id",
|
||||
Domain: "custom.example.com",
|
||||
AccountID: account.Id,
|
||||
Validated: true,
|
||||
},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 1 {
|
||||
t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
o, err := store.GetAccountOnboarding(context.Background(), account.Id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, o.AccountID, account.Id)
|
||||
|
||||
err = store.DeleteAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = store.GetAccountOnboarding(context.Background(), account.Id)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting onboarding")
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 0 {
|
||||
t.Errorf("expecting 0 Accounts to be stored after DeleteAccount()")
|
||||
}
|
||||
|
||||
_, err = store.GetAccountByPeerPubKey(context.Background(), "peerkey")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer public key")
|
||||
|
||||
_, err = store.GetAccountByUser(context.Background(), "testuser")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by user")
|
||||
|
||||
_, err = store.GetAccountByPeerID(context.Background(), "testpeer")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer id")
|
||||
|
||||
_, err = store.GetAccountBySetupKey(context.Background(), setupKey.Key)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by setup key")
|
||||
|
||||
_, err = store.GetAccount(context.Background(), account.Id)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by id")
|
||||
|
||||
for _, policy := range account.Policies {
|
||||
var rules []*types.PolicyRule
|
||||
err = store.(*SqlStore).db.Model(&types.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 []*types.PersonalAccessToken
|
||||
err = store.(*SqlStore).db.Model(&types.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")
|
||||
|
||||
}
|
||||
|
||||
for _, network := range account.Networks {
|
||||
routers, err := store.GetNetworkRoutersByNetID(context.Background(), LockingStrengthNone, account.Id, network.ID)
|
||||
require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for network routers")
|
||||
require.Len(t, routers, 0, "expecting no network routers to be found after DeleteAccount")
|
||||
|
||||
resources, err := store.GetNetworkResourcesByNetID(context.Background(), LockingStrengthNone, account.Id, network.ID)
|
||||
require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for network resources")
|
||||
require.Len(t, resources, 0, "expecting no network resources to be found after DeleteAccount")
|
||||
}
|
||||
|
||||
domains, err := store.ListCustomDomains(context.Background(), account.Id)
|
||||
require.NoError(t, err, "expecting no error after DeleteAccount when searching for custom domains")
|
||||
require.Len(t, domains, 0, "expecting no custom domains to be found after DeleteAccount")
|
||||
|
||||
var services []*rpservice.Service
|
||||
err = store.(*SqlStore).db.Model(&rpservice.Service{}).Find(&services, "account_id = ?", account.Id).Error
|
||||
require.NoError(t, err, "expecting no error after DeleteAccount when searching for services")
|
||||
require.Len(t, services, 0, "expecting no services to be found after DeleteAccount")
|
||||
|
||||
var targets []*rpservice.Target
|
||||
err = store.(*SqlStore).db.Model(&rpservice.Target{}).Find(&targets, "account_id = ?", account.Id).Error
|
||||
require.NoError(t, err, "expecting no error after DeleteAccount when searching for service targets")
|
||||
require.Len(t, targets, 0, "expecting no service targets to be found after DeleteAccount")
|
||||
}
|
||||
|
||||
func Test_GetAccount(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
|
||||
id := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
account, err := store.GetAccount(context.Background(), id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, id, account.Id, "account id should match")
|
||||
require.Equal(t, false, account.Onboarding.OnboardingFlowPending)
|
||||
|
||||
id = "9439-34653001fc3b-bf1c8084-ba50-4ce7"
|
||||
|
||||
account, err = store.GetAccount(context.Background(), id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, id, account.Id, "account id should match")
|
||||
require.Equal(t, true, account.Onboarding.OnboardingFlowPending)
|
||||
|
||||
_, err = store.GetAccount(context.Background(), "non-existing-account")
|
||||
assert.Error(t, err)
|
||||
parsedErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
|
||||
|
||||
})
|
||||
}
|
||||
|
||||
func Test_TestGetAccountByPrivateDomain(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
|
||||
existingDomain := "test.com"
|
||||
|
||||
account, err := store.GetAccountByPrivateDomain(context.Background(), existingDomain)
|
||||
require.NoError(t, err, "should found account")
|
||||
require.Equal(t, existingDomain, account.Domain, "domains should match")
|
||||
|
||||
_, err = store.GetAccountByPrivateDomain(context.Background(), "missing-domain.com")
|
||||
require.Error(t, err, "should return error on domain lookup")
|
||||
parsedErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
|
||||
})
|
||||
}
|
||||
|
||||
func TestPostgresql_SaveAccount(t *testing.T) {
|
||||
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
|
||||
t.Skip("skip CI tests on darwin and windows")
|
||||
}
|
||||
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
account := newAccountWithId(context.Background(), "account_id", "testuser", "")
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
account2 := newAccountWithId(context.Background(), "account_id2", "testuser2", "")
|
||||
setupKey, _ = types.GenerateDefaultSetupKey()
|
||||
account2.SetupKeys[setupKey.Key] = setupKey
|
||||
account2.Peers["testpeer2"] = &nbpeer.Peer{
|
||||
Key: "peerkey2",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 2}),
|
||||
IPv6: netip.MustParseAddr("fd00::2"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name 2",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(context.Background(), account2)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 2 {
|
||||
t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
a, err := store.GetAccount(context.Background(), 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(context.Background(), "peerkey"); a == nil {
|
||||
t.Errorf("expecting PeerKeyID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByUser(context.Background(), "testuser"); a == nil {
|
||||
t.Errorf("expecting UserID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountByPeerID(context.Background(), "testpeer"); a == nil {
|
||||
t.Errorf("expecting PeerID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
|
||||
if a, err := store.GetAccountBySetupKey(context.Background(), setupKey.Key); a == nil {
|
||||
t.Errorf("expecting SetupKeyID2AccountID index updated after SaveAccount(): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresql_DeleteAccount(t *testing.T) {
|
||||
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
|
||||
t.Skip("skip CI tests on darwin and windows")
|
||||
}
|
||||
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
testUserID := "testuser"
|
||||
user := types.NewAdminUser(testUserID)
|
||||
user.PATs = map[string]*types.PersonalAccessToken{"testtoken": {
|
||||
ID: "testtoken",
|
||||
Name: "test token",
|
||||
}}
|
||||
|
||||
account := newAccountWithId(context.Background(), "account_id", testUserID, "")
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
account.Users[testUserID] = user
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 1 {
|
||||
t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
|
||||
}
|
||||
|
||||
err = store.DeleteAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(store.GetAllAccounts(context.Background())) != 0 {
|
||||
t.Errorf("expecting 0 Accounts to be stored after DeleteAccount()")
|
||||
}
|
||||
|
||||
_, err = store.GetAccountByPeerPubKey(context.Background(), "peerkey")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer public key")
|
||||
|
||||
_, err = store.GetAccountByUser(context.Background(), "testuser")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by user")
|
||||
|
||||
_, err = store.GetAccountByPeerID(context.Background(), "testpeer")
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer id")
|
||||
|
||||
_, err = store.GetAccountBySetupKey(context.Background(), setupKey.Key)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by setup key")
|
||||
|
||||
_, err = store.GetAccount(context.Background(), account.Id)
|
||||
require.Error(t, err, "expecting error after removing DeleteAccount when getting account by id")
|
||||
|
||||
for _, policy := range account.Policies {
|
||||
var rules []*types.PolicyRule
|
||||
err = store.(*SqlStore).db.Model(&types.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 []*types.PersonalAccessToken
|
||||
err = store.(*SqlStore).db.Model(&types.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_TestGetAccountByPrivateDomain(t *testing.T) {
|
||||
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
|
||||
t.Skip("skip CI tests on darwin and windows")
|
||||
}
|
||||
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
existingDomain := "test.com"
|
||||
|
||||
account, err := store.GetAccountByPrivateDomain(context.Background(), existingDomain)
|
||||
require.NoError(t, err, "should found account")
|
||||
require.Equal(t, existingDomain, account.Domain, "domains should match")
|
||||
|
||||
_, err = store.GetAccountByPrivateDomain(context.Background(), "missing-domain.com")
|
||||
require.Error(t, err, "should return error on domain lookup")
|
||||
}
|
||||
|
||||
func TestSqlite_GetAccountNetwork(t *testing.T) {
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
_, err = store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
network, err := store.GetAccountNetwork(context.Background(), LockingStrengthNone, existingAccountID)
|
||||
require.NoError(t, err)
|
||||
ip := net.IP{100, 64, 0, 0}.To16()
|
||||
assert.Equal(t, ip, network.Net.IP)
|
||||
assert.Equal(t, net.IPMask{255, 255, 0, 0}, network.Net.Mask)
|
||||
assert.Equal(t, "", network.Dns)
|
||||
assert.Equal(t, "af1c8024-ha40-4ce2-9418-34653101fc3c", network.Identifier)
|
||||
assert.Equal(t, uint64(0), network.Serial)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveAccountPersistsAgentNetworkOnly(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
account, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.False(t, account.Settings.AgentNetworkOnly, "setting should default to false")
|
||||
|
||||
account.Settings.AgentNetworkOnly = true
|
||||
require.NoError(t, store.SaveAccount(context.Background(), account))
|
||||
|
||||
reloaded, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, reloaded.Settings.AgentNetworkOnly, "setting should survive a save/load round-trip")
|
||||
|
||||
reloaded.Settings.AgentNetworkOnly = false
|
||||
require.NoError(t, store.SaveAccount(context.Background(), reloaded))
|
||||
|
||||
disabled, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.False(t, disabled.Settings.AgentNetworkOnly, "disabling should persist")
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveAccountPersistsDashboardFeatures(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
account, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, account.Settings.DashboardFeatures, "dashboard features should default to unset")
|
||||
|
||||
agentNetwork := true
|
||||
account.Settings.DashboardFeatures = &types.DashboardFeatures{AgentNetwork: &agentNetwork}
|
||||
require.NoError(t, store.SaveAccount(context.Background(), account))
|
||||
|
||||
reloaded, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, reloaded.Settings.DashboardFeatures, "dashboard features should survive a save/load round-trip")
|
||||
require.NotNil(t, reloaded.Settings.DashboardFeatures.AgentNetwork, "agent network flag should be set")
|
||||
require.True(t, *reloaded.Settings.DashboardFeatures.AgentNetwork, "agent network flag should persist as true")
|
||||
|
||||
disabled := false
|
||||
reloaded.Settings.DashboardFeatures = &types.DashboardFeatures{AgentNetwork: &disabled}
|
||||
require.NoError(t, store.SaveAccount(context.Background(), reloaded))
|
||||
|
||||
reloadedDisabled, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, reloadedDisabled.Settings.DashboardFeatures.AgentNetwork, "agent network flag should remain set")
|
||||
require.False(t, *reloadedDisabled.Settings.DashboardFeatures.AgentNetwork, "explicit false should persist")
|
||||
}
|
||||
|
||||
func TestSqlStore_UpdateAccountDomainAttributes(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
t.Run("Should update attributes with public domain", func(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
domain := "example.com"
|
||||
category := "public"
|
||||
IsDomainPrimaryAccount := false
|
||||
err = store.UpdateAccountDomainAttributes(context.Background(), accountID, domain, category, IsDomainPrimaryAccount)
|
||||
require.NoError(t, err)
|
||||
account, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, domain, account.Domain)
|
||||
require.Equal(t, category, account.DomainCategory)
|
||||
require.Equal(t, IsDomainPrimaryAccount, account.IsDomainPrimaryAccount)
|
||||
})
|
||||
|
||||
t.Run("Should update attributes with private domain", func(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
domain := "test.com"
|
||||
category := "private"
|
||||
IsDomainPrimaryAccount := true
|
||||
err = store.UpdateAccountDomainAttributes(context.Background(), accountID, domain, category, IsDomainPrimaryAccount)
|
||||
require.NoError(t, err)
|
||||
account, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, domain, account.Domain)
|
||||
require.Equal(t, category, account.DomainCategory)
|
||||
require.Equal(t, IsDomainPrimaryAccount, account.IsDomainPrimaryAccount)
|
||||
})
|
||||
|
||||
t.Run("Should fail when account does not exist", func(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
domain := "test.com"
|
||||
category := "private"
|
||||
IsDomainPrimaryAccount := true
|
||||
err = store.UpdateAccountDomainAttributes(context.Background(), "non-existing-account-id", domain, category, IsDomainPrimaryAccount)
|
||||
require.Error(t, err)
|
||||
})
|
||||
|
||||
}
|
||||
|
||||
func TestSqlStore_GetDNSSettings(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing account dns settings",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing account dns settings",
|
||||
accountID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve dns settings with empty account ID",
|
||||
accountID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
dnsSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, dnsSettings)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, dnsSettings)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveDNSSettings(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
dnsSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
dnsSettings.DisabledManagementGroups = []string{"groupA", "groupB"}
|
||||
err = store.SaveDNSSettings(context.Background(), accountID, dnsSettings)
|
||||
require.NoError(t, err)
|
||||
|
||||
saveDNSSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, saveDNSSettings, dnsSettings)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountCreatedBy(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectError bool
|
||||
createdBy string
|
||||
}{
|
||||
{
|
||||
name: "existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectError: false,
|
||||
createdBy: "edafee4e-63fb-11ec-90d6-0242ac120003",
|
||||
},
|
||||
{
|
||||
name: "non-existing account ID",
|
||||
accountID: "nonexistent",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "empty account ID",
|
||||
accountID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
createdBy, err := store.GetAccountCreatedBy(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Empty(t, createdBy)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, createdBy)
|
||||
require.Equal(t, tt.createdBy, createdBy)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountMeta(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
accountMeta, err := store.GetAccountMeta(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, accountMeta)
|
||||
require.Equal(t, accountID, accountMeta.AccountID)
|
||||
require.Equal(t, "edafee4e-63fb-11ec-90d6-0242ac120003", accountMeta.CreatedBy)
|
||||
require.Equal(t, "test.com", accountMeta.Domain)
|
||||
require.Equal(t, "private", accountMeta.DomainCategory)
|
||||
require.Equal(t, time.Date(2024, time.October, 2, 14, 1, 38, 210000000, time.UTC), accountMeta.CreatedAt.UTC())
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAnyAccountID(t *testing.T) {
|
||||
t.Run("should return account ID when accounts exist", func(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID, err := store.GetAnyAccountID(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "bf1c8084-ba50-4ce7-9439-34653001fc3b", accountID)
|
||||
})
|
||||
|
||||
t.Run("should return error when no accounts exist", func(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID, err := store.GetAnyAccountID(context.Background())
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, sErr.Type(), status.NotFound)
|
||||
assert.Empty(t, accountID)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,268 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// CreateAgentNetworkAccessLog persists a flattened agent-network access-log
|
||||
// entry together with its authorising-group child rows in a single
|
||||
// transaction.
|
||||
func (s *SqlStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error {
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Idempotent on the log id / (log_id, group_id) so a proxy resend of the
|
||||
// same entry can't fail the request.
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(entry).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(groups) > 0 {
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).WithFields(log.Fields{
|
||||
"account_id": entry.AccountID,
|
||||
"service_id": entry.ServiceID,
|
||||
"model": entry.Model,
|
||||
}).Errorf("failed to create agent-network access log entry in store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to create agent-network access log entry in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteOldAgentNetworkAccessLogs deletes an account's access-log rows (and
|
||||
// their authorising-group child rows) older than the cutoff. Usage records are
|
||||
// untouched — they are the long-term aggregate. Returns the number of log rows
|
||||
// deleted.
|
||||
func (s *SqlStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) {
|
||||
var deleted int64
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Remove group child rows for the soon-to-be-deleted logs first.
|
||||
if err := tx.Exec(
|
||||
"DELETE FROM agent_network_access_log_group WHERE account_id = ? AND log_id IN (SELECT id FROM agent_network_access_log WHERE account_id = ? AND timestamp < ?)",
|
||||
accountID, accountID, olderThan,
|
||||
).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
res := tx.Where("account_id = ? AND timestamp < ?", accountID, olderThan).
|
||||
Delete(&agentNetworkTypes.AgentNetworkAccessLog{})
|
||||
if res.Error != nil {
|
||||
return res.Error
|
||||
}
|
||||
deleted = res.RowsAffected
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete old agent-network access logs for account %s: %v", accountID, err)
|
||||
return 0, status.Errorf(status.Internal, "failed to delete old agent-network access logs")
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkAccessLogs retrieves flattened agent-network access logs for
|
||||
// an account with server-side pagination, filtering and sorting. Authorising
|
||||
// group ids are hydrated from the group child table for the returned page.
|
||||
func (s *SqlStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) {
|
||||
var logs []*agentNetworkTypes.AgentNetworkAccessLog
|
||||
var totalCount int64
|
||||
|
||||
countQuery := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
)
|
||||
if err := countQuery.Count(&totalCount).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count agent-network access logs: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access logs")
|
||||
}
|
||||
|
||||
query := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).
|
||||
Order(filter.GetSortColumn() + " " + filter.GetSortOrder()).
|
||||
Limit(filter.GetLimit()).
|
||||
Offset(filter.GetOffset())
|
||||
|
||||
if lockStrength != LockingStrengthNone {
|
||||
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
if err := query.Find(&logs).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent-network access logs from store: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access logs from store")
|
||||
}
|
||||
|
||||
if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, logs); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return logs, totalCount, nil
|
||||
}
|
||||
|
||||
// applyAgentNetworkAccessLogFilters applies the filter conditions to a query.
|
||||
func (s *SqlStore) applyAgentNetworkAccessLogFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
|
||||
if filter.Search != nil {
|
||||
p := "%" + *filter.Search + "%"
|
||||
query = query.Where(
|
||||
"id LIKE ? OR host LIKE ? OR path LIKE ? OR model LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)",
|
||||
p, p, p, p, p, p,
|
||||
)
|
||||
}
|
||||
if filter.UserID != nil {
|
||||
query = query.Where("user_id = ?", *filter.UserID)
|
||||
}
|
||||
if filter.SessionID != nil {
|
||||
query = query.Where("session_id = ?", *filter.SessionID)
|
||||
}
|
||||
if filter.Decision != nil {
|
||||
query = query.Where("decision = ?", *filter.Decision)
|
||||
}
|
||||
if filter.PathPrefix != nil {
|
||||
query = query.Where("path LIKE ?", *filter.PathPrefix+"%")
|
||||
}
|
||||
if len(filter.ProviderIDs) > 0 {
|
||||
query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
|
||||
}
|
||||
if len(filter.Models) > 0 {
|
||||
query = query.Where("model IN ?", filter.Models)
|
||||
}
|
||||
if len(filter.GroupIDs) > 0 {
|
||||
query = query.Where(
|
||||
"id IN (SELECT log_id FROM agent_network_access_log_group WHERE group_id IN ?)",
|
||||
filter.GroupIDs,
|
||||
)
|
||||
}
|
||||
if filter.StartDate != nil {
|
||||
query = query.Where("timestamp >= ?", *filter.StartDate)
|
||||
}
|
||||
if filter.EndDate != nil {
|
||||
query = query.Where("timestamp <= ?", *filter.EndDate)
|
||||
}
|
||||
return query
|
||||
}
|
||||
|
||||
// hydrateAgentNetworkAccessLogGroups loads the authorising group ids for the
|
||||
// given page of entries and assigns them onto each entry's GroupIDs field.
|
||||
func (s *SqlStore) hydrateAgentNetworkAccessLogGroups(ctx context.Context, accountID string, logs []*agentNetworkTypes.AgentNetworkAccessLog) error {
|
||||
if len(logs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ids := make([]string, 0, len(logs))
|
||||
for _, l := range logs {
|
||||
ids = append(ids, l.ID)
|
||||
}
|
||||
|
||||
var rows []agentNetworkTypes.AgentNetworkAccessLogGroup
|
||||
if err := s.db.
|
||||
Where(accountIDCondition, accountID).
|
||||
Where("log_id IN ?", ids).
|
||||
Find(&rows).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to hydrate agent-network access log groups: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to hydrate agent-network access log groups")
|
||||
}
|
||||
|
||||
byLog := make(map[string][]string, len(logs))
|
||||
for _, r := range rows {
|
||||
byLog[r.LogID] = append(byLog[r.LogID], r.GroupID)
|
||||
}
|
||||
for _, l := range logs {
|
||||
l.GroupIDs = byLog[l.ID]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// agentNetworkSessionKeyExpr is the SQL group key for session-grouped access
|
||||
// logs: the row's session id, or — when the client sent none — the row id, so
|
||||
// session-less requests each form their own singleton group. COALESCE/NULLIF
|
||||
// are standard SQL, so this stays portable across SQLite and Postgres.
|
||||
const agentNetworkSessionKeyExpr = "COALESCE(NULLIF(session_id, ''), id)"
|
||||
|
||||
// GetAgentNetworkAccessLogSessions retrieves agent-network access logs grouped
|
||||
// by session, with server-side pagination, filtering and sorting at the session
|
||||
// level. It paginates over the distinct session keys (ordered by the requested
|
||||
// session-level aggregate), fetches every entry for the page's sessions, and
|
||||
// folds them into per-session summaries. The returned count is the number of
|
||||
// matching sessions. Filters apply to the entries, so a session's summary
|
||||
// reflects only its filter-matching requests.
|
||||
func (s *SqlStore) GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error) {
|
||||
// Count distinct sessions via a grouped subquery — portable and avoids
|
||||
// relying on COUNT(DISTINCT <expr>) quoting quirks.
|
||||
sessionsSubquery := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).
|
||||
Select(agentNetworkSessionKeyExpr + " AS session_key").
|
||||
Group(agentNetworkSessionKeyExpr)
|
||||
|
||||
var totalCount int64
|
||||
if err := s.db.Table("(?) AS sessions", sessionsSubquery).Count(&totalCount).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count agent-network access-log sessions: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access-log sessions")
|
||||
}
|
||||
|
||||
// The page of session keys, ordered by the session-level aggregate. The
|
||||
// session-key tiebreaker keeps pagination deterministic when the primary
|
||||
// aggregate ties.
|
||||
type sessionKeyRow struct {
|
||||
SessionKey string
|
||||
}
|
||||
var keyRows []sessionKeyRow
|
||||
keyQuery := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).
|
||||
Select(agentNetworkSessionKeyExpr + " AS session_key").
|
||||
Group(agentNetworkSessionKeyExpr).
|
||||
Order(filter.GetSessionSortExpr() + " " + filter.GetSortOrder()).
|
||||
Order("session_key ASC").
|
||||
Limit(filter.GetLimit()).
|
||||
Offset(filter.GetOffset())
|
||||
if err := keyQuery.Scan(&keyRows).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to list agent-network access-log session keys: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to list agent-network access-log session keys")
|
||||
}
|
||||
if len(keyRows) == 0 {
|
||||
return nil, totalCount, nil
|
||||
}
|
||||
|
||||
keys := make([]string, 0, len(keyRows))
|
||||
for _, r := range keyRows {
|
||||
keys = append(keys, r.SessionKey)
|
||||
}
|
||||
|
||||
// All entries for the page's sessions, contiguous per session and oldest
|
||||
// first within each — the fold relies on that ordering.
|
||||
var entries []*agentNetworkTypes.AgentNetworkAccessLog
|
||||
entriesQuery := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).
|
||||
Where(agentNetworkSessionKeyExpr+" IN ?", keys).
|
||||
Order(agentNetworkSessionKeyExpr + ", timestamp ASC")
|
||||
|
||||
if lockStrength != LockingStrengthNone {
|
||||
entriesQuery = entriesQuery.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
if err := entriesQuery.Find(&entries).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent-network access-log session entries: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access-log session entries")
|
||||
}
|
||||
|
||||
if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, entries); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return agentNetworkTypes.FoldAccessLogSessions(keys, entries), totalCount, nil
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// CreateAgentNetworkUsage persists a stripped agent-network usage record
|
||||
// together with its authorising-group child rows in a single transaction.
|
||||
func (s *SqlStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error {
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Idempotent on the usage id / (usage_id, group_id) so a proxy resend of
|
||||
// the same entry can't fail the request.
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(usage).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(groups) > 0 {
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).WithFields(log.Fields{
|
||||
"account_id": usage.AccountID,
|
||||
"model": usage.Model,
|
||||
}).Errorf("failed to create agent-network usage record in store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to create agent-network usage record in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkUsageRows returns the stripped usage rows for an account that
|
||||
// match the filter (date / user / group / provider / model). Aggregation into
|
||||
// time buckets happens in the manager so granularities stay engine-portable.
|
||||
func (s *SqlStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) {
|
||||
var rows []*agentNetworkTypes.AgentNetworkUsage
|
||||
|
||||
query := s.applyAgentNetworkUsageFilters(
|
||||
s.db.Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).Order("timestamp ASC")
|
||||
|
||||
if lockStrength != LockingStrengthNone {
|
||||
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
if err := query.Find(&rows).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent-network usage rows from store: %v", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent-network usage rows from store")
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
// applyAgentNetworkUsageFilters applies the shared access-log filter's
|
||||
// date/user/group/provider/model conditions to a usage-table query. Pagination,
|
||||
// sort and free-text search are ignored — the overview is an aggregate.
|
||||
func (s *SqlStore) applyAgentNetworkUsageFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
|
||||
if filter.UserID != nil {
|
||||
query = query.Where("user_id = ?", *filter.UserID)
|
||||
}
|
||||
if filter.SessionID != nil {
|
||||
query = query.Where("session_id = ?", *filter.SessionID)
|
||||
}
|
||||
if len(filter.ProviderIDs) > 0 {
|
||||
query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
|
||||
}
|
||||
if len(filter.Models) > 0 {
|
||||
query = query.Where("model IN ?", filter.Models)
|
||||
}
|
||||
if len(filter.GroupIDs) > 0 {
|
||||
query = query.Where(
|
||||
"id IN (SELECT usage_id FROM agent_network_request_usage_group WHERE group_id IN ?)",
|
||||
filter.GroupIDs,
|
||||
)
|
||||
}
|
||||
if filter.StartDate != nil {
|
||||
query = query.Where("timestamp >= ?", *filter.StartDate)
|
||||
}
|
||||
if filter.EndDate != nil {
|
||||
query = query.Where("timestamp <= ?", *filter.EndDate)
|
||||
}
|
||||
return query
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/rs/xid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// GetCustomDomainsCounts returns the total and validated custom domain counts.
|
||||
func (s *SqlStore) GetCustomDomainsCounts(ctx context.Context) (int64, int64, error) {
|
||||
var total, validated int64
|
||||
if err := s.db.Model(&domain.Domain{}).Count(&total).Error; err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if err := s.db.Model(&domain.Domain{}).Where("validated = ?", true).Count(&validated).Error; err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
return total, validated, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetCustomDomain(ctx context.Context, accountID string, domainID string) (*domain.Domain, error) {
|
||||
tx := s.db
|
||||
|
||||
customDomain := &domain.Domain{}
|
||||
result := tx.Take(&customDomain, accountAndIDQueryCondition, accountID, domainID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "custom domain %s not found", domainID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get custom domain from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get custom domain from store")
|
||||
}
|
||||
|
||||
return customDomain, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) ListFreeDomains(ctx context.Context, accountID string) ([]string, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) ListCustomDomains(ctx context.Context, accountID string) ([]*domain.Domain, error) {
|
||||
tx := s.db
|
||||
|
||||
var domains []*domain.Domain
|
||||
result := tx.Find(&domains, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get reverse proxy custom domains from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get reverse proxy custom domains from store")
|
||||
}
|
||||
|
||||
return domains, nil
|
||||
}
|
||||
|
||||
// GetCustomDomainByName returns the custom domain row holding the given name,
|
||||
// regardless of which account owns it.
|
||||
func (s *SqlStore) GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error) {
|
||||
customDomain := &domain.Domain{}
|
||||
result := s.db.Take(customDomain, "domain = ?", domainName)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "custom domain %s not found", domainName)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get custom domain by name from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get custom domain from store")
|
||||
}
|
||||
|
||||
return customDomain, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error) {
|
||||
newDomain := &domain.Domain{
|
||||
ID: xid.New().String(), // Generate our own ID because gorm doesn't always configure the database to handle this for us.
|
||||
Domain: domainName,
|
||||
AccountID: accountID,
|
||||
TargetCluster: targetCluster,
|
||||
Type: domain.TypeCustom,
|
||||
Validated: validated,
|
||||
}
|
||||
if !validated {
|
||||
expiresAt := time.Now().UTC().Add(domain.ValidationTTL)
|
||||
newDomain.ValidationExpiresAt = &expiresAt
|
||||
}
|
||||
result := s.db.Create(newDomain)
|
||||
if result.Error != nil {
|
||||
// The unique index is the last guard when two requests clear the
|
||||
// manager's availability check at the same time. The one that loses the
|
||||
// insert is a conflict, not an internal failure.
|
||||
var count int64
|
||||
if err := s.db.Model(&domain.Domain{}).Where("domain = ?", domainName).Count(&count).Error; err == nil && count > 0 {
|
||||
// The insert error is logged even on this path: the name being taken
|
||||
// is what the caller has to act on, but if the insert also failed for
|
||||
// an unrelated reason the operator still needs to see it.
|
||||
log.WithContext(ctx).Warnf("create reverse proxy custom domain %s rejected, name already registered: %v", domainName, result.Error)
|
||||
return nil, status.Errorf(status.AlreadyExists, "domain %s is already registered", domainName)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to create reverse proxy custom domain to store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to create reverse proxy custom domain to store")
|
||||
}
|
||||
|
||||
return newDomain, nil
|
||||
}
|
||||
|
||||
// UpdateCustomDomain completes validation only while the original registration is pending.
|
||||
func (s *SqlStore) UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error) {
|
||||
if !d.Validated {
|
||||
return nil, status.Errorf(status.InvalidArgument, "custom domain update must complete validation")
|
||||
}
|
||||
result := s.db.WithContext(ctx).Model(&domain.Domain{}).
|
||||
Where(accountAndIDQueryCondition, accountID, d.ID).
|
||||
Where("domain = ? AND target_cluster = ?", d.Domain, d.TargetCluster).
|
||||
Where("validated = ? AND validation_expires_at > ?", false, time.Now().UTC()).
|
||||
Update("validated", true)
|
||||
if result.Error != nil {
|
||||
return nil, fmt.Errorf("validate custom domain in store: %w", result.Error)
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return nil, status.Errorf(status.PreconditionFailed, "custom domain registration is no longer pending validation")
|
||||
}
|
||||
|
||||
return d, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteCustomDomain(ctx context.Context, accountID string, domainID string) error {
|
||||
result := s.db.Delete(domain.Domain{}, accountAndIDQueryCondition, accountID, domainID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete reverse proxy custom domain from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete reverse proxy custom domain from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "reverse proxy custom domain %s not found", domainID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) CreateDNSRecord(ctx context.Context, record *records.Record) error {
|
||||
result := s.db.Create(record)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to create dns record to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to create dns record to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) UpdateDNSRecord(ctx context.Context, record *records.Record) error {
|
||||
result := s.db.Select("*").Save(record)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to update dns record to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to update dns record to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteDNSRecord(ctx context.Context, accountID, zoneID, recordID string) error {
|
||||
result := s.db.Delete(&records.Record{}, "account_id = ? AND zone_id = ? AND id = ?", accountID, zoneID, recordID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete dns record from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete dns record from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewDNSRecordNotFoundError(recordID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetDNSRecordByID(ctx context.Context, lockStrength LockingStrength, accountID, zoneID, recordID string) (*records.Record, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var record *records.Record
|
||||
result := tx.Where("account_id = ? AND zone_id = ? AND id = ?", accountID, zoneID, recordID).Take(&record)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewDNSRecordNotFoundError(recordID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get dns record from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get dns record from store")
|
||||
}
|
||||
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetZoneDNSRecords(ctx context.Context, lockStrength LockingStrength, accountID, zoneID string) ([]*records.Record, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var recordsList []*records.Record
|
||||
result := tx.Where("account_id = ? AND zone_id = ?", accountID, zoneID).Find(&recordsList)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get zone dns records from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get zone dns records from store")
|
||||
}
|
||||
|
||||
return recordsList, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetZoneDNSRecordsByName(ctx context.Context, lockStrength LockingStrength, accountID, zoneID, name string) ([]*records.Record, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var recordsList []*records.Record
|
||||
result := tx.Where("account_id = ? AND zone_id = ? AND name = ?", accountID, zoneID, name).Find(&recordsList)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get zone dns records by name from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get zone dns records by name from store")
|
||||
}
|
||||
|
||||
return recordsList, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteZoneDNSRecords(ctx context.Context, accountID, zoneID string) error {
|
||||
result := s.db.Delete(&records.Record{}, "account_id = ? AND zone_id = ?", accountID, zoneID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete zone dns records from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete zone dns records from store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_CreateDNSRecord(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
|
||||
err = store.CreateDNSRecord(context.Background(), record)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, accountID, zone.ID, record.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, savedRecord)
|
||||
assert.Equal(t, record.ID, savedRecord.ID)
|
||||
assert.Equal(t, record.Name, savedRecord.Name)
|
||||
assert.Equal(t, record.Type, savedRecord.Type)
|
||||
assert.Equal(t, record.Content, savedRecord.Content)
|
||||
assert.Equal(t, record.TTL, savedRecord.TTL)
|
||||
assert.Equal(t, zone.ID, savedRecord.ZoneID)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetDNSRecordByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), record)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
zoneID string
|
||||
recordID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing record",
|
||||
accountID: accountID,
|
||||
zoneID: zone.ID,
|
||||
recordID: record.ID,
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing record",
|
||||
accountID: accountID,
|
||||
zoneID: zone.ID,
|
||||
recordID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty record ID",
|
||||
accountID: accountID,
|
||||
zoneID: zone.ID,
|
||||
recordID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
savedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, tt.accountID, tt.zoneID, tt.recordID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, savedRecord)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, savedRecord)
|
||||
assert.Equal(t, tt.recordID, savedRecord.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetZoneDNSRecords(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
recordA := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), recordA)
|
||||
require.NoError(t, err)
|
||||
|
||||
recordAAAA := records.NewRecord(accountID, zone.ID, "ipv6.example.com", records.RecordTypeAAAA, "2001:db8::1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), recordAAAA)
|
||||
require.NoError(t, err)
|
||||
|
||||
recordCNAME := records.NewRecord(accountID, zone.ID, "alias.example.com", records.RecordTypeCNAME, "www.example.com", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), recordCNAME)
|
||||
require.NoError(t, err)
|
||||
|
||||
allRecords, err := store.GetZoneDNSRecords(context.Background(), LockingStrengthNone, accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, allRecords)
|
||||
assert.Equal(t, 3, len(allRecords))
|
||||
|
||||
recordIDs := make(map[string]bool)
|
||||
for _, r := range allRecords {
|
||||
recordIDs[r.ID] = true
|
||||
}
|
||||
assert.True(t, recordIDs[recordA.ID])
|
||||
assert.True(t, recordIDs[recordAAAA.ID])
|
||||
assert.True(t, recordIDs[recordCNAME.ID])
|
||||
}
|
||||
|
||||
func TestSqlStore_GetZoneDNSRecordsByName(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
record1 := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), record1)
|
||||
require.NoError(t, err)
|
||||
|
||||
record2 := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeAAAA, "2001:db8::1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), record2)
|
||||
require.NoError(t, err)
|
||||
|
||||
record3 := records.NewRecord(accountID, zone.ID, "mail.example.com", records.RecordTypeA, "192.168.1.2", 600)
|
||||
err = store.CreateDNSRecord(context.Background(), record3)
|
||||
require.NoError(t, err)
|
||||
|
||||
recordsByName, err := store.GetZoneDNSRecordsByName(context.Background(), LockingStrengthNone, accountID, zone.ID, "www.example.com")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, recordsByName)
|
||||
assert.Equal(t, 2, len(recordsByName))
|
||||
|
||||
for _, r := range recordsByName {
|
||||
assert.Equal(t, "www.example.com", r.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_UpdateDNSRecord(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), record)
|
||||
require.NoError(t, err)
|
||||
|
||||
record.Name = "api.example.com"
|
||||
record.Content = "192.168.1.100"
|
||||
record.TTL = 600
|
||||
|
||||
err = store.UpdateDNSRecord(context.Background(), record)
|
||||
require.NoError(t, err)
|
||||
|
||||
updatedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, accountID, zone.ID, record.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, updatedRecord)
|
||||
assert.Equal(t, "api.example.com", updatedRecord.Name)
|
||||
assert.Equal(t, "192.168.1.100", updatedRecord.Content)
|
||||
assert.Equal(t, 600, updatedRecord.TTL)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteDNSRecord(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), record)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = store.DeleteDNSRecord(context.Background(), accountID, zone.ID, record.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
deletedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, accountID, zone.ID, record.ID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, deletedRecord)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteZoneDNSRecords(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
record1 := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
|
||||
err = store.CreateDNSRecord(context.Background(), record1)
|
||||
require.NoError(t, err)
|
||||
|
||||
record2 := records.NewRecord(accountID, zone.ID, "mail.example.com", records.RecordTypeA, "192.168.1.2", 600)
|
||||
err = store.CreateDNSRecord(context.Background(), record2)
|
||||
require.NoError(t, err)
|
||||
|
||||
allRecords, err := store.GetZoneDNSRecords(context.Background(), LockingStrengthNone, accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, len(allRecords))
|
||||
|
||||
err = store.DeleteZoneDNSRecords(context.Background(), accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
remainingRecords, err := store.GetZoneDNSRecords(context.Background(), LockingStrengthNone, accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, len(remainingRecords))
|
||||
}
|
||||
@@ -0,0 +1,356 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// CreateGroups creates the given list of groups to the database.
|
||||
// groupUpsertColumns is the explicit allowlist of columns that get updated when
|
||||
// CreateGroups / UpdateGroups hit a PK conflict. public_id is intentionally
|
||||
// omitted so a caller passing an entity with the zero value (e.g. an HTTP
|
||||
// handler-built struct) cannot reset the persisted public_id during an upsert.
|
||||
// Keep this in sync with the Group schema in management/server/types/group.go.
|
||||
func groupUpsertColumns() clause.Set {
|
||||
return clause.AssignmentColumns([]string{
|
||||
"account_id",
|
||||
"name",
|
||||
"issued",
|
||||
"integration_ref_id",
|
||||
"integration_ref_integration_type",
|
||||
"resources",
|
||||
})
|
||||
}
|
||||
|
||||
func (s *SqlStore) CreateGroups(ctx context.Context, accountID string, groups []*types.Group) error {
|
||||
if len(groups) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return s.db.Transaction(func(tx *gorm.DB) error {
|
||||
result := tx.
|
||||
Clauses(
|
||||
clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "id"}},
|
||||
Where: clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}},
|
||||
DoUpdates: groupUpsertColumns(),
|
||||
},
|
||||
).
|
||||
Omit(clause.Associations).
|
||||
Create(&groups)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save groups to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save groups to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateGroups updates the given list of groups to the database.
|
||||
func (s *SqlStore) UpdateGroups(ctx context.Context, accountID string, groups []*types.Group) error {
|
||||
if len(groups) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return s.db.Transaction(func(tx *gorm.DB) error {
|
||||
result := tx.
|
||||
Clauses(
|
||||
clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "id"}},
|
||||
Where: clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}},
|
||||
DoUpdates: groupUpsertColumns(),
|
||||
},
|
||||
).
|
||||
Omit(clause.Associations).
|
||||
Create(&groups)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save groups to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save groups to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.Group, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var groups []*types.Group
|
||||
result := tx.Preload(clause.Associations).Find(&groups, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get account groups from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get account groups from the store")
|
||||
}
|
||||
|
||||
for _, g := range groups {
|
||||
g.LoadGroupPeers()
|
||||
}
|
||||
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetResourceGroups(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) ([]*types.Group, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var groups []*types.Group
|
||||
|
||||
likePattern := `%"ID":"` + resourceID + `"%`
|
||||
|
||||
result := tx.
|
||||
Preload(clause.Associations).
|
||||
Where("resources LIKE ?", likePattern).
|
||||
Find(&groups)
|
||||
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, result.Error
|
||||
}
|
||||
|
||||
for _, g := range groups {
|
||||
g.LoadGroupPeers()
|
||||
}
|
||||
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) getGroups(ctx context.Context, accountID string) ([]*types.Group, error) {
|
||||
const query = `SELECT id, account_id, public_id, name, issued, resources, integration_ref_id, integration_ref_integration_type FROM groups WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groups, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*types.Group, error) {
|
||||
var g types.Group
|
||||
var resources []byte
|
||||
var refID sql.NullInt64
|
||||
var refType sql.NullString
|
||||
err := row.Scan(&g.ID, &g.AccountID, &g.PublicID, &g.Name, &g.Issued, &resources, &refID, &refType)
|
||||
if err == nil {
|
||||
if refID.Valid {
|
||||
g.IntegrationReference.ID = int(refID.Int64)
|
||||
}
|
||||
if refType.Valid {
|
||||
g.IntegrationReference.IntegrationType = refType.String
|
||||
}
|
||||
if resources != nil {
|
||||
_ = json.Unmarshal(resources, &g.Resources)
|
||||
} else {
|
||||
g.Resources = []types.Resource{}
|
||||
}
|
||||
g.GroupPeers = []types.GroupPeer{}
|
||||
g.Peers = []string{}
|
||||
}
|
||||
return &g, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
// AddResourceToGroup adds a resource to a group. Method always needs to run n a transaction
|
||||
func (s *SqlStore) AddResourceToGroup(ctx context.Context, accountId string, groupID string, resource *types.Resource) error {
|
||||
var group types.Group
|
||||
result := s.db.Where(accountAndIDQueryCondition, accountId, groupID).Take(&group)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return status.NewGroupNotFoundError(groupID)
|
||||
}
|
||||
|
||||
return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
|
||||
}
|
||||
|
||||
for _, res := range group.Resources {
|
||||
if res.ID == resource.ID {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
group.Resources = append(group.Resources, *resource)
|
||||
|
||||
if err := s.db.Save(&group).Error; err != nil {
|
||||
return status.Errorf(status.Internal, "issue updating group: %s", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemoveResourceFromGroup removes a resource from a group. Method always needs to run in a transaction
|
||||
func (s *SqlStore) RemoveResourceFromGroup(ctx context.Context, accountId string, groupID string, resourceID string) error {
|
||||
var group types.Group
|
||||
result := s.db.Where(accountAndIDQueryCondition, accountId, groupID).Take(&group)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return status.NewGroupNotFoundError(groupID)
|
||||
}
|
||||
|
||||
return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
|
||||
}
|
||||
|
||||
for i, res := range group.Resources {
|
||||
if res.ID == resourceID {
|
||||
group.Resources = append(group.Resources[:i], group.Resources[i+1:]...)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.db.Save(&group).Error; err != nil {
|
||||
return status.Errorf(status.Internal, "issue updating group: %s", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetGroupByID retrieves a group by ID and account ID.
|
||||
func (s *SqlStore) GetGroupByID(ctx context.Context, lockStrength LockingStrength, accountID, groupID string) (*types.Group, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var group *types.Group
|
||||
result := tx.Preload(clause.Associations).Take(&group, accountAndIDQueryCondition, accountID, groupID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewGroupNotFoundError(groupID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get group from store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get group from store")
|
||||
}
|
||||
|
||||
group.LoadGroupPeers()
|
||||
|
||||
return group, nil
|
||||
}
|
||||
|
||||
// GetGroupByName retrieves a group by name and account ID.
|
||||
func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStrength, accountID, groupName string) (*types.Group, error) {
|
||||
tx := s.db
|
||||
|
||||
var group types.Group
|
||||
|
||||
// TODO: This fix is accepted for now, but if we need to handle this more frequently
|
||||
// we may need to reconsider changing the types.
|
||||
query := tx.Preload(clause.Associations)
|
||||
|
||||
result := query.
|
||||
Model(&types.Group{}).
|
||||
Joins("LEFT JOIN group_peers ON group_peers.group_id = groups.id").
|
||||
Where("groups.account_id = ? AND groups.name = ?", accountID, groupName).
|
||||
Group("groups.id").
|
||||
Order("COUNT(group_peers.peer_id) DESC").
|
||||
Limit(1).
|
||||
First(&group)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewGroupNotFoundError(groupName)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get group by name from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get group by name from store")
|
||||
}
|
||||
|
||||
group.LoadGroupPeers()
|
||||
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
// GetGroupsByIDs retrieves groups by their IDs and account ID.
|
||||
func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*types.Group, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var groups []*types.Group
|
||||
result := tx.Preload(clause.Associations).Find(&groups, accountAndIDsQueryCondition, accountID, groupIDs)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get groups by ID's from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get groups by ID's from store")
|
||||
}
|
||||
|
||||
groupsMap := make(map[string]*types.Group)
|
||||
for _, group := range groups {
|
||||
group.LoadGroupPeers()
|
||||
groupsMap[group.ID] = group
|
||||
}
|
||||
|
||||
return groupsMap, nil
|
||||
}
|
||||
|
||||
// CreateGroup creates a group in the store.
|
||||
func (s *SqlStore) CreateGroup(ctx context.Context, group *types.Group) error {
|
||||
if group == nil {
|
||||
return status.Errorf(status.InvalidArgument, "group is nil")
|
||||
}
|
||||
|
||||
if err := s.db.Omit(clause.Associations).Create(group).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save group to store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to save group to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateGroup updates a group in the store.
|
||||
func (s *SqlStore) UpdateGroup(ctx context.Context, group *types.Group) error {
|
||||
if group == nil {
|
||||
return status.Errorf(status.InvalidArgument, "group is nil")
|
||||
}
|
||||
|
||||
if err := s.db.Omit(clause.Associations, "public_id").Save(group).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save group to store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to save group to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteGroup deletes a group from the database.
|
||||
func (s *SqlStore) DeleteGroup(ctx context.Context, accountID, groupID string) error {
|
||||
result := s.db.Select(clause.Associations).
|
||||
Delete(&types.Group{}, accountAndIDQueryCondition, accountID, groupID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete group from store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete group from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewGroupNotFoundError(groupID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteGroups deletes groups from the database.
|
||||
func (s *SqlStore) DeleteGroups(ctx context.Context, accountID string, groupIDs []string) error {
|
||||
result := s.db.Select(clause.Associations).
|
||||
Delete(&types.Group{}, accountAndIDsQueryCondition, accountID, groupIDs)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete groups from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete groups from store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getGroupPeers(ctx context.Context, groupIDs []string) ([]types.GroupPeer, error) {
|
||||
if len(groupIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
const query = `SELECT account_id, group_id, peer_id FROM group_peers WHERE group_id = ANY($1)`
|
||||
rows, err := s.pool.Query(ctx, query, groupIDs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groupPeers, err := pgx.CollectRows(rows, pgx.RowToStructByName[types.GroupPeer])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return groupPeers, nil
|
||||
}
|
||||
|
||||
// AddPeerToAllGroup adds a peer to the 'All' group. Method always needs to run in a transaction
|
||||
func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error {
|
||||
var groupID string
|
||||
_ = s.db.Model(types.Group{}).
|
||||
Select("id").
|
||||
Where("account_id = ? AND name = ?", accountID, "All").
|
||||
Limit(1).
|
||||
Scan(&groupID)
|
||||
|
||||
if groupID == "" {
|
||||
return status.Errorf(status.NotFound, "group 'All' not found for account %s", accountID)
|
||||
}
|
||||
|
||||
err := s.db.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "group_id"}, {Name: "peer_id"}},
|
||||
DoNothing: true,
|
||||
}).Create(&types.GroupPeer{
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
PeerID: peerID,
|
||||
}).Error
|
||||
if err != nil {
|
||||
return status.Errorf(status.Internal, "error adding peer to group 'All': %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddPeerToGroup adds a peer to a group
|
||||
func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountID, peerID, groupID string) error {
|
||||
peer := &types.GroupPeer{
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
PeerID: peerID,
|
||||
}
|
||||
|
||||
err := s.db.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "group_id"}, {Name: "peer_id"}},
|
||||
DoNothing: true,
|
||||
}).Create(peer).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to add peer %s to group %s for account %s: %v", peerID, groupID, accountID, err)
|
||||
return status.Errorf(status.Internal, "failed to add peer to group")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemovePeerFromGroup removes a peer from a group
|
||||
func (s *SqlStore) RemovePeerFromGroup(ctx context.Context, peerID string, groupID string) error {
|
||||
err := s.db.
|
||||
Delete(&types.GroupPeer{}, "group_id = ? AND peer_id = ?", groupID, peerID).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to remove peer %s from group %s: %v", peerID, groupID, err)
|
||||
return status.Errorf(status.Internal, "failed to remove peer from group")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemovePeerFromAllGroups removes a peer from all groups
|
||||
func (s *SqlStore) RemovePeerFromAllGroups(ctx context.Context, peerID string) error {
|
||||
err := s.db.
|
||||
Delete(&types.GroupPeer{}, "peer_id = ?", peerID).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to remove peer %s from all groups: %v", peerID, err)
|
||||
return status.Errorf(status.Internal, "failed to remove peer from all groups")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetPeerGroups retrieves all groups assigned to a specific peer in a given account.
|
||||
func (s *SqlStore) GetPeerGroups(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string) ([]*types.Group, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var groups []*types.Group
|
||||
query := tx.
|
||||
Joins("JOIN group_peers ON group_peers.group_id = groups.id").
|
||||
Where("groups.account_id = ? AND group_peers.peer_id = ?", accountId, peerId).
|
||||
Preload(clause.Associations).
|
||||
Find(&groups)
|
||||
|
||||
if query.Error != nil {
|
||||
return nil, query.Error
|
||||
}
|
||||
|
||||
for _, group := range groups {
|
||||
group.LoadGroupPeers()
|
||||
}
|
||||
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
// GetPeerGroupIDs retrieves all group IDs assigned to a specific peer in a given account.
|
||||
func (s *SqlStore) GetPeerGroupIDs(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string) ([]string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var groupIDs []string
|
||||
query := tx.
|
||||
Model(&types.GroupPeer{}).
|
||||
Where("account_id = ? AND peer_id = ?", accountId, peerId).
|
||||
Pluck("group_id", &groupIDs)
|
||||
|
||||
if query.Error != nil {
|
||||
if errors.Is(query.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "no groups found for peer %s in account %s", peerId, accountId)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get group IDs for peer %s in account %s: %v", peerId, accountId, query.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get group IDs for peer from store")
|
||||
}
|
||||
|
||||
return groupIDs, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountGroupPeers(ctx context.Context, lockStrength LockingStrength, accountID string) (map[string]map[string]struct{}, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peers []types.GroupPeer
|
||||
result := tx.Find(&peers, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get account group peers from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get account group peers from store")
|
||||
}
|
||||
|
||||
groupPeers := make(map[string]map[string]struct{})
|
||||
for _, peer := range peers {
|
||||
if _, exists := groupPeers[peer.GroupID]; !exists {
|
||||
groupPeers[peer.GroupID] = make(map[string]struct{})
|
||||
}
|
||||
groupPeers[peer.GroupID][peer.PeerID] = struct{}{}
|
||||
}
|
||||
|
||||
return groupPeers, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeersByGroupIDs(ctx context.Context, accountID string, groupIDs []string) ([]*nbpeer.Peer, error) {
|
||||
if len(groupIDs) == 0 {
|
||||
return []*nbpeer.Peer{}, nil
|
||||
}
|
||||
|
||||
var peers []*nbpeer.Peer
|
||||
peerIDsSubquery := s.db.Model(&types.GroupPeer{}).
|
||||
Select("DISTINCT peer_id").
|
||||
Where("account_id = ? AND group_id IN ?", accountID, groupIDs)
|
||||
|
||||
result := s.db.Where("account_id = ? AND id IN (?)", accountID, peerIDsSubquery).Find(&peers)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peers by group IDs: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get peers by group IDs")
|
||||
}
|
||||
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeerIDsByGroups(ctx context.Context, accountID string, groupIDs []string) ([]string, error) {
|
||||
if len(groupIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var peerIDs []string
|
||||
result := s.db.Model(&types.GroupPeer{}).
|
||||
Select("DISTINCT peer_id").
|
||||
Where("account_id = ? AND group_id IN ?", accountID, groupIDs).
|
||||
Pluck("peer_id", &peerIDs)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get peer IDs by groups: %s", result.Error)
|
||||
}
|
||||
|
||||
return peerIDs, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetGroupIDsByPeerIDs(ctx context.Context, accountID string, peerIDs []string) ([]string, error) {
|
||||
if len(peerIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var groupIDs []string
|
||||
result := s.db.Model(&types.GroupPeer{}).
|
||||
Select("DISTINCT group_id").
|
||||
Where("account_id = ? AND peer_id IN ?", accountID, peerIDs).
|
||||
Pluck("group_id", &groupIDs)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get group IDs by peers: %s", result.Error)
|
||||
}
|
||||
|
||||
return groupIDs, nil
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
func TestSqlStore_AddPeerToGroup(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
peerID := "cfefqs706sqkneg59g4g"
|
||||
groupID := "cfefqs706sqkneg59g4h"
|
||||
|
||||
group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.NoError(t, err, "failed to get group")
|
||||
require.Len(t, group.Peers, 0, "group should have 0 peers")
|
||||
|
||||
err = store.AddPeerToGroup(context.Background(), accountID, peerID, groupID)
|
||||
require.NoError(t, err, "failed to add peer to group")
|
||||
|
||||
group, err = store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.NoError(t, err, "failed to get group")
|
||||
require.Len(t, group.Peers, 1, "group should have 1 peers")
|
||||
require.Contains(t, group.Peers, peerID)
|
||||
}
|
||||
|
||||
func TestSqlStore_AddPeerToAllGroup(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
groupID := "cfefqs706sqkneg59g3g"
|
||||
|
||||
peer := &nbpeer.Peer{
|
||||
ID: "peer1",
|
||||
AccountID: accountID,
|
||||
DNSLabel: "peer1.domain.test",
|
||||
}
|
||||
|
||||
group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.NoError(t, err, "failed to get group")
|
||||
require.Len(t, group.Peers, 2, "group should have 2 peers")
|
||||
require.NotContains(t, group.Peers, peer.ID)
|
||||
|
||||
err = store.AddPeerToAccount(context.Background(), peer)
|
||||
require.NoError(t, err, "failed to add peer to account")
|
||||
|
||||
err = store.AddPeerToAllGroup(context.Background(), accountID, peer.ID)
|
||||
require.NoError(t, err, "failed to add peer to all group")
|
||||
|
||||
group, err = store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.NoError(t, err, "failed to get group")
|
||||
require.Len(t, group.Peers, 3, "group should have peers")
|
||||
require.Contains(t, group.Peers, peer.ID)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPeerGroups(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
peerID := "cfefqs706sqkneg59g4g"
|
||||
|
||||
groups, err := store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peerID)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, groups, 1)
|
||||
assert.Equal(t, groups[0].Name, "All")
|
||||
|
||||
err = store.AddPeerToGroup(context.Background(), accountID, peerID, "cfefqs706sqkneg59g4h")
|
||||
require.NoError(t, err)
|
||||
|
||||
groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peerID)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, groups, 2)
|
||||
|
||||
foreignPeerID := "foreign-peer"
|
||||
err = store.AddPeerToGroup(context.Background(), accountID, foreignPeerID, "cfefqs706sqkneg59g4h")
|
||||
require.NoError(t, err)
|
||||
|
||||
groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, "other-account", foreignPeerID)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, groups, "groups of another account must not be returned")
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPeersByGroupIDs(t *testing.T) {
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
group1ID := "test-group-1"
|
||||
group2ID := "test-group-2"
|
||||
emptyGroupID := "empty-group"
|
||||
|
||||
peer1 := "cfefqs706sqkneg59g4g"
|
||||
peer2 := "cfeg6sf06sqkneg59g50"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
groupIDs []string
|
||||
expectedPeers []string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve peers from single group with multiple peers",
|
||||
groupIDs: []string{group1ID},
|
||||
expectedPeers: []string{peer1, peer2},
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "retrieve peers from single group with one peer",
|
||||
groupIDs: []string{group2ID},
|
||||
expectedPeers: []string{peer1},
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "retrieve peers from multiple groups (with overlap)",
|
||||
groupIDs: []string{group1ID, group2ID},
|
||||
expectedPeers: []string{peer1, peer2}, // should deduplicate
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "retrieve peers from existing 'All' group",
|
||||
groupIDs: []string{"cfefqs706sqkneg59g3g"}, // All group from test data
|
||||
expectedPeers: []string{peer1, peer2},
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "retrieve peers from empty group",
|
||||
groupIDs: []string{emptyGroupID},
|
||||
expectedPeers: []string{},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "retrieve peers from non-existing group",
|
||||
groupIDs: []string{"non-existing-group"},
|
||||
expectedPeers: []string{},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "empty group IDs list",
|
||||
groupIDs: []string{},
|
||||
expectedPeers: []string{},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "mix of existing and non-existing groups",
|
||||
groupIDs: []string{group1ID, "non-existing-group"},
|
||||
expectedPeers: []string{peer1, peer2},
|
||||
expectedCount: 2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
groups := []*types.Group{
|
||||
{
|
||||
ID: group1ID,
|
||||
AccountID: accountID,
|
||||
},
|
||||
{
|
||||
ID: group2ID,
|
||||
AccountID: accountID,
|
||||
},
|
||||
}
|
||||
require.NoError(t, store.CreateGroups(ctx, accountID, groups))
|
||||
|
||||
otherAccount := newAccountWithId(ctx, "other-account", "other-user", "")
|
||||
require.NoError(t, store.SaveAccount(ctx, otherAccount))
|
||||
foreignPeer := &nbpeer.Peer{ID: "foreign-peer", AccountID: otherAccount.Id}
|
||||
require.NoError(t, store.AddPeerToAccount(ctx, foreignPeer))
|
||||
|
||||
require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group1ID))
|
||||
require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer2, group1ID))
|
||||
require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group2ID))
|
||||
require.NoError(t, store.AddPeerToGroup(ctx, accountID, foreignPeer.ID, group1ID))
|
||||
|
||||
peers, err := store.GetPeersByGroupIDs(ctx, accountID, tt.groupIDs)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
|
||||
if tt.expectedCount > 0 {
|
||||
actualPeerIDs := make([]string, len(peers))
|
||||
for i, peer := range peers {
|
||||
actualPeerIDs[i] = peer.ID
|
||||
}
|
||||
assert.ElementsMatch(t, tt.expectedPeers, actualPeerIDs)
|
||||
|
||||
// Verify all returned peers belong to the correct account
|
||||
for _, peer := range peers {
|
||||
assert.Equal(t, accountID, peer.AccountID)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlite_GetGroupByName(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
group, err := store.GetGroupByName(context.Background(), LockingStrengthNone, accountID, "All")
|
||||
require.NoError(t, err)
|
||||
require.True(t, group.IsGroupAll())
|
||||
}
|
||||
|
||||
func TestSqlStore_GetGroupsByIDs(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
groupIDs []string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing groups by existing IDs",
|
||||
groupIDs: []string{"cfefqs706sqkneg59g4g", "cfefqs706sqkneg59g3g"},
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "empty group IDs list",
|
||||
groupIDs: []string{},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "non-existing group IDs",
|
||||
groupIDs: []string{"nonexistent1", "nonexistent2"},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "mixed existing and non-existing group IDs",
|
||||
groupIDs: []string{"cfefqs706sqkneg59g4g", "nonexistent"},
|
||||
expectedCount: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
groups, err := store.GetGroupsByIDs(context.Background(), LockingStrengthNone, accountID, tt.groupIDs)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, groups, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_CreateGroup(t *testing.T) {
|
||||
if os.Getenv("CI") == "true" {
|
||||
t.Log("Skipping MySQL test on CI")
|
||||
}
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.MysqlStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
group := &types.Group{
|
||||
ID: "group-id",
|
||||
AccountID: accountID,
|
||||
Issued: "api",
|
||||
Peers: []string{},
|
||||
Resources: []types.Resource{},
|
||||
GroupPeers: []types.GroupPeer{},
|
||||
}
|
||||
err = store.CreateGroup(context.Background(), group)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedGroup, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, "group-id")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, savedGroup, group)
|
||||
}
|
||||
|
||||
func TestSqlStore_CreateUpdateGroups(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
groups := []*types.Group{
|
||||
{
|
||||
ID: "group-1",
|
||||
AccountID: accountID,
|
||||
Issued: "api",
|
||||
Peers: []string{},
|
||||
Resources: []types.Resource{},
|
||||
GroupPeers: []types.GroupPeer{},
|
||||
},
|
||||
{
|
||||
ID: "group-2",
|
||||
AccountID: accountID,
|
||||
Issued: "integration",
|
||||
Peers: []string{},
|
||||
Resources: []types.Resource{},
|
||||
GroupPeers: []types.GroupPeer{},
|
||||
},
|
||||
}
|
||||
err = store.CreateGroups(context.Background(), accountID, groups)
|
||||
require.NoError(t, err)
|
||||
|
||||
groups[1].Peers = []string{}
|
||||
err = store.UpdateGroups(context.Background(), accountID, groups)
|
||||
require.NoError(t, err)
|
||||
|
||||
group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groups[1].ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, groups[1], group)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteGroup(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
groupID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "delete existing group",
|
||||
groupID: "cfefqs706sqkneg59g4g",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "delete non-existing group",
|
||||
groupID: "non-existing-group-id",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "delete with empty group ID",
|
||||
groupID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := store.DeleteGroup(context.Background(), accountID, tt.groupID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
|
||||
group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, tt.groupID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, group)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteGroups(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
groupIDs []string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "delete multiple existing groups",
|
||||
groupIDs: []string{"cfefqs706sqkneg59g4g", "cfefqs706sqkneg59g3g"},
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "delete non-existing groups",
|
||||
groupIDs: []string{"non-existing-id-1", "non-existing-id-2"},
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "delete with empty group IDs list",
|
||||
groupIDs: []string{},
|
||||
expectError: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := store.DeleteGroups(context.Background(), accountID, tt.groupIDs)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, groupID := range tt.groupIDs {
|
||||
group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, group)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_AddAndRemoveResourceFromGroup(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(cleanup)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
resourceId := "ctc4nci7qv9061u6ilfg"
|
||||
groupID := "cs1tnh0hhcjnqoiuebeg"
|
||||
|
||||
res := &types.Resource{
|
||||
ID: resourceId,
|
||||
Type: "host",
|
||||
}
|
||||
err = store.AddResourceToGroup(context.Background(), accountID, groupID, res)
|
||||
require.NoError(t, err)
|
||||
|
||||
group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, group.Resources, *res)
|
||||
|
||||
groups, err := store.GetResourceGroups(context.Background(), LockingStrengthNone, accountID, resourceId)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, groups, 1)
|
||||
|
||||
err = store.RemoveResourceFromGroup(context.Background(), accountID, groupID, res.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
group, err = store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, group.Resources, *res)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveGroups_LargeBatch(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
accountGroups, err := store.GetAccountGroups(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accountGroups, 3)
|
||||
|
||||
groupsToSave := make([]*types.Group, 0)
|
||||
|
||||
for i := 1; i <= 8000; i++ {
|
||||
groupsToSave = append(groupsToSave, &types.Group{
|
||||
ID: fmt.Sprintf("%d", i),
|
||||
AccountID: accountID,
|
||||
Name: fmt.Sprintf("group-%d", i),
|
||||
})
|
||||
}
|
||||
|
||||
err = store.CreateGroups(context.Background(), accountID, groupsToSave)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountGroups, err = store.GetAccountGroups(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 8003, len(accountGroups))
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type installation struct {
|
||||
ID uint `gorm:"primaryKey"`
|
||||
InstallationIDValue string
|
||||
}
|
||||
|
||||
func (s *SqlStore) SaveInstallationID(_ context.Context, ID string) error {
|
||||
installation := installation{InstallationIDValue: ID}
|
||||
installation.ID = uint(s.installationPK)
|
||||
|
||||
return s.db.Clauses(clause.OnConflict{UpdateAll: true}).Create(&installation).Error
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetInstallationID() string {
|
||||
var installation installation
|
||||
|
||||
if result := s.db.Take(&installation, idQueryCondition, s.installationPK); result.Error != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return installation.InstallationIDValue
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// SaveJob persists a job in DB
|
||||
func (s *SqlStore) CreatePeerJob(ctx context.Context, job *types.Job) error {
|
||||
result := s.db.Create(job)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to create job in store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to create job in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) CompletePeerJob(ctx context.Context, job *types.Job) error {
|
||||
result := s.db.
|
||||
Model(&types.Job{}).
|
||||
Where(idQueryCondition, job.ID).
|
||||
Updates(job)
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to update job in store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to update job in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// job was pending for too long and has been cancelled
|
||||
func (s *SqlStore) MarkPendingJobsAsFailed(ctx context.Context, accountID, peerID, jobID, reason string) error {
|
||||
now := time.Now().UTC()
|
||||
result := s.db.
|
||||
Model(&types.Job{}).
|
||||
Where(accountAndPeerIDQueryCondition+" AND id = ?"+" AND status = ?", accountID, peerID, jobID, types.JobStatusPending).
|
||||
Updates(types.Job{
|
||||
Status: types.JobStatusFailed,
|
||||
FailedReason: reason,
|
||||
CompletedAt: &now,
|
||||
})
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to mark pending jobs as Failed job in store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to mark pending job as Failed in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// job was pending for too long and has been cancelled
|
||||
func (s *SqlStore) MarkAllPendingJobsAsFailed(ctx context.Context, accountID, peerID, reason string) error {
|
||||
now := time.Now().UTC()
|
||||
result := s.db.
|
||||
Model(&types.Job{}).
|
||||
Where(accountAndPeerIDQueryCondition+" AND status = ?", accountID, peerID, types.JobStatusPending).
|
||||
Updates(types.Job{
|
||||
Status: types.JobStatusFailed,
|
||||
FailedReason: reason,
|
||||
CompletedAt: &now,
|
||||
})
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to mark pending jobs as Failed job in store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to mark pending job as Failed in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetJobByID fetches job by ID
|
||||
func (s *SqlStore) GetPeerJobByID(ctx context.Context, accountID, jobID string) (*types.Job, error) {
|
||||
var job types.Job
|
||||
err := s.db.
|
||||
Where(accountAndIDQueryCondition, accountID, jobID).
|
||||
First(&job).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "job %s not found", jobID)
|
||||
}
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to fetch job from store: %s", err)
|
||||
return nil, err
|
||||
}
|
||||
return &job, nil
|
||||
}
|
||||
|
||||
// get all jobs
|
||||
func (s *SqlStore) GetPeerJobs(ctx context.Context, accountID, peerID string) ([]*types.Job, error) {
|
||||
var jobs []*types.Job
|
||||
err := s.db.
|
||||
Where(accountAndPeerIDQueryCondition, accountID, peerID).
|
||||
Order("created_at DESC").
|
||||
Find(&jobs).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to fetch jobs from store: %s", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return jobs, nil
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getNameServerGroups(ctx context.Context, accountID string) ([]nbdns.NameServerGroup, error) {
|
||||
const query = `SELECT id, account_id, public_id, name, description, name_servers, groups, "primary", domains, enabled, search_domains_enabled FROM name_server_groups WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
nsgs, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nbdns.NameServerGroup, error) {
|
||||
var n nbdns.NameServerGroup
|
||||
var ns, groups, domains []byte
|
||||
var primary, enabled, searchDomainsEnabled sql.NullBool
|
||||
err := row.Scan(&n.ID, &n.AccountID, &n.PublicID, &n.Name, &n.Description, &ns, &groups, &primary, &domains, &enabled, &searchDomainsEnabled)
|
||||
if err == nil {
|
||||
if primary.Valid {
|
||||
n.Primary = primary.Bool
|
||||
}
|
||||
if enabled.Valid {
|
||||
n.Enabled = enabled.Bool
|
||||
}
|
||||
if searchDomainsEnabled.Valid {
|
||||
n.SearchDomainsEnabled = searchDomainsEnabled.Bool
|
||||
}
|
||||
if ns != nil {
|
||||
_ = json.Unmarshal(ns, &n.NameServers)
|
||||
} else {
|
||||
n.NameServers = []nbdns.NameServer{}
|
||||
}
|
||||
if groups != nil {
|
||||
_ = json.Unmarshal(groups, &n.Groups)
|
||||
} else {
|
||||
n.Groups = []string{}
|
||||
}
|
||||
if domains != nil {
|
||||
_ = json.Unmarshal(domains, &n.Domains)
|
||||
} else {
|
||||
n.Domains = []string{}
|
||||
}
|
||||
}
|
||||
return n, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nsgs, nil
|
||||
}
|
||||
|
||||
// GetAccountNameServerGroups retrieves name server groups for an account.
|
||||
func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbdns.NameServerGroup, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var nsGroups []*nbdns.NameServerGroup
|
||||
result := tx.Find(&nsGroups, accountIDCondition, accountID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get name server groups from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get name server groups from store")
|
||||
}
|
||||
|
||||
return nsGroups, nil
|
||||
}
|
||||
|
||||
// GetNameServerGroupByID retrieves a name server group by its ID and account ID.
|
||||
func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength LockingStrength, accountID, nsGroupID string) (*nbdns.NameServerGroup, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var nsGroup *nbdns.NameServerGroup
|
||||
result := tx.
|
||||
Take(&nsGroup, accountAndIDQueryCondition, accountID, nsGroupID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewNameServerGroupNotFoundError(nsGroupID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get name server group from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get name server group from store")
|
||||
}
|
||||
|
||||
return nsGroup, nil
|
||||
}
|
||||
|
||||
// SaveNameServerGroup saves a name server group to the database.
|
||||
func (s *SqlStore) SaveNameServerGroup(ctx context.Context, nameServerGroup *nbdns.NameServerGroup) error {
|
||||
result := s.db.Save(nameServerGroup)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save name server group to the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to save name server group to store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteNameServerGroup deletes a name server group from the database.
|
||||
func (s *SqlStore) DeleteNameServerGroup(ctx context.Context, accountID, nsGroupID string) error {
|
||||
result := s.db.Delete(&nbdns.NameServerGroup{}, accountAndIDQueryCondition, accountID, nsGroupID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete name server group from the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to delete name server group from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewNameServerGroupNotFoundError(nsGroupID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetAccountNameServerGroups(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve name server groups by existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "non-existing account ID",
|
||||
accountID: "nonexistent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "empty account ID",
|
||||
accountID: "",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peers, err := store.GetAccountNameServerGroups(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestSqlStore_GetNameServerByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
nsGroupID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing nameserver group",
|
||||
nsGroupID: "csqdelq7qv97ncu7d9t0",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing nameserver group",
|
||||
nsGroupID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty nameserver group ID",
|
||||
nsGroupID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
nsGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthNone, accountID, tt.nsGroupID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, nsGroup)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, nsGroup)
|
||||
require.Equal(t, tt.nsGroupID, nsGroup.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveNameServerGroup(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
nsGroup := &nbdns.NameServerGroup{
|
||||
ID: "ns-group-id",
|
||||
AccountID: accountID,
|
||||
Name: "NS Group",
|
||||
NameServers: []nbdns.NameServer{
|
||||
{
|
||||
IP: netip.MustParseAddr("8.8.8.8"),
|
||||
NSType: 1,
|
||||
Port: 53,
|
||||
},
|
||||
},
|
||||
Groups: []string{"groupA"},
|
||||
Primary: true,
|
||||
Enabled: true,
|
||||
SearchDomainsEnabled: false,
|
||||
}
|
||||
|
||||
err = store.SaveNameServerGroup(context.Background(), nsGroup)
|
||||
require.NoError(t, err)
|
||||
|
||||
saveNSGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthNone, accountID, nsGroup.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, saveNSGroup, nsGroup)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteNameServerGroup(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
nsGroupID := "csqdelq7qv97ncu7d9t0"
|
||||
|
||||
err = store.DeleteNameServerGroup(context.Background(), accountID, nsGroupID)
|
||||
require.NoError(t, err)
|
||||
|
||||
nsGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthNone, accountID, nsGroupID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, nsGroup)
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getNetworks(ctx context.Context, accountID string) ([]*networkTypes.Network, error) {
|
||||
const query = `SELECT id, account_id, public_id, name, description FROM networks WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
networks, err := pgx.CollectRows(rows, pgx.RowToStructByName[networkTypes.Network])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]*networkTypes.Network, len(networks))
|
||||
for i := range networks {
|
||||
result[i] = &networks[i]
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountNetworks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*networkTypes.Network, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var networks []*networkTypes.Network
|
||||
result := tx.Find(&networks, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get networks from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get networks from store")
|
||||
}
|
||||
|
||||
return networks, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkByID(ctx context.Context, lockStrength LockingStrength, accountID, networkID string) (*networkTypes.Network, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var network *networkTypes.Network
|
||||
result := tx.Take(&network, accountAndIDQueryCondition, accountID, networkID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewNetworkNotFoundError(networkID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get network from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network from store")
|
||||
}
|
||||
|
||||
return network, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) SaveNetwork(ctx context.Context, network *networkTypes.Network) error {
|
||||
result := s.db.Save(network)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save network to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save network to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteNetwork(ctx context.Context, accountID, networkID string) error {
|
||||
result := s.db.Delete(&networkTypes.Network{}, accountAndIDQueryCondition, accountID, networkID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete network from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete network from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewNetworkNotFoundError(networkID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getNetworkResources(ctx context.Context, accountID string) ([]*resourceTypes.NetworkResource, error) {
|
||||
const query = `SELECT id, network_id, account_id, public_id, name, description, type, domain, prefix, enabled FROM network_resources WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resources, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (resourceTypes.NetworkResource, error) {
|
||||
var r resourceTypes.NetworkResource
|
||||
var prefix []byte
|
||||
var enabled sql.NullBool
|
||||
err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.PublicID, &r.Name, &r.Description, &r.Type, &r.Domain, &prefix, &enabled)
|
||||
if err == nil {
|
||||
if enabled.Valid {
|
||||
r.Enabled = enabled.Bool
|
||||
}
|
||||
if prefix != nil {
|
||||
_ = json.Unmarshal(prefix, &r.Prefix)
|
||||
}
|
||||
}
|
||||
return r, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]*resourceTypes.NetworkResource, len(resources))
|
||||
for i := range resources {
|
||||
result[i] = &resources[i]
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkResourcesByNetID(ctx context.Context, lockStrength LockingStrength, accountID, networkID string) ([]*resourceTypes.NetworkResource, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netResources []*resourceTypes.NetworkResource
|
||||
result := tx.
|
||||
Find(&netResources, "account_id = ? AND network_id = ?", accountID, networkID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get network resources from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network resources from store")
|
||||
}
|
||||
|
||||
return netResources, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkResourcesByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*resourceTypes.NetworkResource, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netResources []*resourceTypes.NetworkResource
|
||||
result := tx.
|
||||
Find(&netResources, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get network resources from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network resources from store")
|
||||
}
|
||||
|
||||
return netResources, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkResourceByID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netResources *resourceTypes.NetworkResource
|
||||
result := tx.
|
||||
Take(&netResources, accountAndIDQueryCondition, accountID, resourceID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewNetworkResourceNotFoundError(resourceID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network resource from store")
|
||||
}
|
||||
|
||||
return netResources, nil
|
||||
}
|
||||
|
||||
// GetNetworkResourceByIDOrPublicID retrieves a network resource by either its ID or its
|
||||
// PublicID. See GetPolicyByIDOrPublicID for why peer-reported references need both.
|
||||
func (s *SqlStore) GetNetworkResourceByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netResources *resourceTypes.NetworkResource
|
||||
result := tx.
|
||||
Take(&netResources, accountAndAnyIDQueryCondition, accountID, resourceID, resourceID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewNetworkResourceNotFoundError(resourceID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network resource from store")
|
||||
}
|
||||
|
||||
return netResources, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkResourceByName(ctx context.Context, lockStrength LockingStrength, accountID, resourceName string) (*resourceTypes.NetworkResource, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netResources *resourceTypes.NetworkResource
|
||||
result := tx.
|
||||
Take(&netResources, "account_id = ? AND name = ?", accountID, resourceName)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewNetworkResourceNotFoundError(resourceName)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network resource from store")
|
||||
}
|
||||
|
||||
return netResources, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) SaveNetworkResource(ctx context.Context, resource *resourceTypes.NetworkResource) error {
|
||||
result := s.db.Save(resource)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save network resource to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save network resource to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteNetworkResource(ctx context.Context, accountID, resourceID string) error {
|
||||
result := s.db.Delete(&resourceTypes.NetworkResource{}, accountAndIDQueryCondition, accountID, resourceID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete network resource from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete network resource from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewNetworkResourceNotFoundError(resourceID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetNetworkResourcesByNetID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
networkID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve resources by existing network ID",
|
||||
networkID: "ct286bi7qv930dsrrug0",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "retrieve resources by non-existing network ID",
|
||||
networkID: "non-existent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
netResources, err := store.GetNetworkResourcesByNetID(context.Background(), LockingStrengthNone, accountID, tt.networkID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, netResources, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetNetworkResourceByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
netResourceID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing network resource ID",
|
||||
netResourceID: "ctc4nci7qv9061u6ilfg",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing network resource ID",
|
||||
netResourceID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve network with empty resource ID",
|
||||
netResourceID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
netResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, tt.netResourceID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, netResource)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, netResource)
|
||||
require.Equal(t, tt.netResourceID, netResource.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetNetworkResourceByIDOrPublicID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
netResourceID := "ctc4nci7qv9061u6ilfg"
|
||||
|
||||
netResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, netResourceID)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, netResource.PublicID)
|
||||
|
||||
for _, id := range []string{netResourceID, netResource.PublicID} {
|
||||
netResource, err := store.GetNetworkResourceByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, netResourceID, netResource.ID)
|
||||
}
|
||||
|
||||
netResource, err = store.GetNetworkResourceByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, netResource)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveNetworkResource(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
networkID := "ct286bi7qv930dsrrug0"
|
||||
|
||||
netResource, err := resourceTypes.NewNetworkResource(accountID, networkID, "resource-name", "", "example.com", []string{}, true)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = store.SaveNetworkResource(context.Background(), netResource)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedNetResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, netResource.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, netResource.ID, savedNetResource.ID)
|
||||
require.Equal(t, netResource.Name, savedNetResource.Name)
|
||||
require.Equal(t, netResource.NetworkID, savedNetResource.NetworkID)
|
||||
require.Equal(t, netResource.Type, resourceTypes.NetworkResourceType("domain"))
|
||||
require.Equal(t, netResource.Domain, "example.com")
|
||||
require.Equal(t, netResource.AccountID, savedNetResource.AccountID)
|
||||
require.Equal(t, netResource.Prefix, netip.Prefix{})
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteNetworkResource(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
netResourceID := "ctc4nci7qv9061u6ilfg"
|
||||
|
||||
err = store.DeleteNetworkResource(context.Background(), accountID, netResourceID)
|
||||
require.NoError(t, err)
|
||||
|
||||
netResource, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, netResourceID)
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, sErr.Type())
|
||||
require.Nil(t, netResource)
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
|
||||
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getNetworkRouters(ctx context.Context, accountID string) ([]*routerTypes.NetworkRouter, error) {
|
||||
const query = `SELECT id, network_id, account_id, public_id, peer, peer_groups, masquerade, metric, enabled FROM network_routers WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
routers, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (routerTypes.NetworkRouter, error) {
|
||||
var r routerTypes.NetworkRouter
|
||||
var peerGroups []byte
|
||||
var masquerade, enabled sql.NullBool
|
||||
var metric sql.NullInt64
|
||||
err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.PublicID, &r.Peer, &peerGroups, &masquerade, &metric, &enabled)
|
||||
if err == nil {
|
||||
if masquerade.Valid {
|
||||
r.Masquerade = masquerade.Bool
|
||||
}
|
||||
if enabled.Valid {
|
||||
r.Enabled = enabled.Bool
|
||||
}
|
||||
if metric.Valid {
|
||||
r.Metric = int(metric.Int64)
|
||||
}
|
||||
if peerGroups != nil {
|
||||
_ = json.Unmarshal(peerGroups, &r.PeerGroups)
|
||||
}
|
||||
}
|
||||
return r, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]*routerTypes.NetworkRouter, len(routers))
|
||||
for i := range routers {
|
||||
result[i] = &routers[i]
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkRoutersByNetID(ctx context.Context, lockStrength LockingStrength, accountID, netID string) ([]*routerTypes.NetworkRouter, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netRouters []*routerTypes.NetworkRouter
|
||||
result := tx.
|
||||
Find(&netRouters, "account_id = ? AND network_id = ?", accountID, netID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get network routers from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network routers from store")
|
||||
}
|
||||
|
||||
return netRouters, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkRoutersByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*routerTypes.NetworkRouter, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netRouters []*routerTypes.NetworkRouter
|
||||
result := tx.
|
||||
Find(&netRouters, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get network routers from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network routers from store")
|
||||
}
|
||||
|
||||
return netRouters, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetNetworkRouterByID(ctx context.Context, lockStrength LockingStrength, accountID, routerID string) (*routerTypes.NetworkRouter, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var netRouter *routerTypes.NetworkRouter
|
||||
result := tx.
|
||||
Take(&netRouter, accountAndIDQueryCondition, accountID, routerID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewNetworkRouterNotFoundError(routerID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get network router from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get network router from store")
|
||||
}
|
||||
|
||||
return netRouter, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) CreateNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error {
|
||||
if err := s.db.Create(router).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to create network router in store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to create network router in store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) UpdateNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error {
|
||||
result := s.db.
|
||||
Select("*").
|
||||
Where(accountAndIDQueryCondition, router.AccountID, router.ID).
|
||||
Updates(router)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to update network router in store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to update network router in store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewNetworkRouterNotFoundError(router.ID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteNetworkRouter(ctx context.Context, accountID, routerID string) error {
|
||||
result := s.db.Delete(&routerTypes.NetworkRouter{}, accountAndIDQueryCondition, accountID, routerID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete network router from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete network router from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewNetworkRouterNotFoundError(routerID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetRoutingPeerNetworks returns the distinct network names where the peer is assigned as a routing peer
|
||||
// in an enabled network router, either directly or via peer groups.
|
||||
func (s *SqlStore) GetRoutingPeerNetworks(_ context.Context, accountID, peerID string) ([]string, error) {
|
||||
var routers []*routerTypes.NetworkRouter
|
||||
if err := s.db.Select("peer, peer_groups, network_id").Where("account_id = ? AND enabled = true", accountID).Find(&routers).Error; err != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get enabled routers: %v", err)
|
||||
}
|
||||
|
||||
if len(routers) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var groupPeers []types.GroupPeer
|
||||
if err := s.db.Select("group_id").Where("account_id = ? AND peer_id = ?", accountID, peerID).Find(&groupPeers).Error; err != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get peer group memberships: %v", err)
|
||||
}
|
||||
|
||||
groupSet := make(map[string]struct{}, len(groupPeers))
|
||||
for _, gp := range groupPeers {
|
||||
groupSet[gp.GroupID] = struct{}{}
|
||||
}
|
||||
|
||||
networkIDs := make(map[string]struct{})
|
||||
for _, r := range routers {
|
||||
if r.Peer == peerID {
|
||||
networkIDs[r.NetworkID] = struct{}{}
|
||||
} else if r.Peer == "" {
|
||||
for _, pg := range r.PeerGroups {
|
||||
if _, ok := groupSet[pg]; ok {
|
||||
networkIDs[r.NetworkID] = struct{}{}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(networkIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
ids := make([]string, 0, len(networkIDs))
|
||||
for id := range networkIDs {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
var networks []*networkTypes.Network
|
||||
if err := s.db.Select("name").Where("account_id = ? AND id IN ?", accountID, ids).Find(&networks).Error; err != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get networks: %v", err)
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(networks))
|
||||
for _, n := range networks {
|
||||
names = append(names, n.Name)
|
||||
}
|
||||
|
||||
return names, nil
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetNetworkRoutersByNetID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
networkID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve routers by existing network ID",
|
||||
networkID: "ct286bi7qv930dsrrug0",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "retrieve routers by non-existing network ID",
|
||||
networkID: "non-existent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
routers, err := store.GetNetworkRoutersByNetID(context.Background(), LockingStrengthNone, accountID, tt.networkID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, routers, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetNetworkRouterByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
networkRouterID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing network router ID",
|
||||
networkRouterID: "ctc20ji7qv9ck2sebc80",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing network router ID",
|
||||
networkRouterID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve network with empty router ID",
|
||||
networkRouterID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
networkRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, tt.networkRouterID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, networkRouter)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, networkRouter)
|
||||
require.Equal(t, tt.networkRouterID, networkRouter.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_CreateNetworkRouter(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
networkID := "ct286bi7qv930dsrrug0"
|
||||
|
||||
netRouter, err := routerTypes.NewNetworkRouter(accountID, networkID, "", []string{"net-router-grp"}, true, 0, true)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = store.CreateNetworkRouter(context.Background(), netRouter)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedNetRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, netRouter.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, netRouter, savedNetRouter)
|
||||
}
|
||||
|
||||
func TestSqlStore_UpdateNetworkRouter(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
networkID := "ct286bi7qv930dsrrug0"
|
||||
routerID := "ctc20ji7qv9ck2sebc80"
|
||||
|
||||
netRouter := &routerTypes.NetworkRouter{
|
||||
ID: routerID,
|
||||
AccountID: accountID,
|
||||
NetworkID: networkID,
|
||||
Peer: "",
|
||||
PeerGroups: []string{"net-router-grp"},
|
||||
Masquerade: true,
|
||||
Metric: 42,
|
||||
Enabled: true,
|
||||
}
|
||||
|
||||
err = store.UpdateNetworkRouter(context.Background(), netRouter)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedNetRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, routerID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, netRouter, savedNetRouter)
|
||||
|
||||
// Updating a router under a different account must not match any row.
|
||||
netRouter.AccountID = "non-existent-account"
|
||||
err = store.UpdateNetworkRouter(context.Background(), netRouter)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteNetworkRouter(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
netRouterID := "ctc20ji7qv9ck2sebc80"
|
||||
|
||||
err = store.DeleteNetworkRouter(context.Background(), accountID, netRouterID)
|
||||
require.NoError(t, err)
|
||||
|
||||
netRouter, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, netRouterID)
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, sErr.Type())
|
||||
require.Nil(t, netRouter)
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetAccountNetworks(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve networks by existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectedCount: 1,
|
||||
},
|
||||
|
||||
{
|
||||
name: "retrieve networks by non-existing account ID",
|
||||
accountID: "non-existent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
networks, err := store.GetAccountNetworks(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, networks, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetNetworkByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
networkID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing network ID",
|
||||
networkID: "ct286bi7qv930dsrrug0",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing network ID",
|
||||
networkID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve network with empty ID",
|
||||
networkID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
network, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, tt.networkID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, network)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, network)
|
||||
require.Equal(t, tt.networkID, network.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveNetwork(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
network := &networkTypes.Network{
|
||||
ID: "net-id",
|
||||
AccountID: accountID,
|
||||
Name: "net",
|
||||
}
|
||||
|
||||
err = store.SaveNetwork(context.Background(), network)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedNet, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, network.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, network, savedNet)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteNetwork(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
networkID := "ct286bi7qv930dsrrug0"
|
||||
|
||||
err = store.DeleteNetwork(context.Background(), accountID, networkID)
|
||||
require.NoError(t, err)
|
||||
|
||||
network, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, networkID)
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, sErr.Type())
|
||||
require.Nil(t, network)
|
||||
}
|
||||
@@ -0,0 +1,781 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
|
||||
// To maintain data integrity, we create a copy of the peer's to prevent unintended updates to other fields.
|
||||
peerCopy := peer.Copy()
|
||||
peerCopy.AccountID = accountID
|
||||
|
||||
err := s.transaction(func(tx *gorm.DB) error {
|
||||
// check if peer exists before saving
|
||||
var peerID string
|
||||
result := tx.Model(&nbpeer.Peer{}).Select("id").Take(&peerID, accountAndIDQueryCondition, accountID, peer.ID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return status.Errorf(status.NotFound, peerNotFoundFMT, peer.ID)
|
||||
}
|
||||
return result.Error
|
||||
}
|
||||
|
||||
if peerID == "" {
|
||||
return status.Errorf(status.NotFound, peerNotFoundFMT, peer.ID)
|
||||
}
|
||||
|
||||
result = tx.Model(&nbpeer.Peer{}).Where(accountAndIDQueryCondition, accountID, peer.ID).Save(peerCopy)
|
||||
if result.Error != nil {
|
||||
return status.Errorf(status.Internal, "failed to save peer to store: %v", result.Error)
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) SavePeerStatus(ctx context.Context, accountID, peerID string, peerStatus nbpeer.PeerStatus) error {
|
||||
var peerCopy nbpeer.Peer
|
||||
peerCopy.Status = &peerStatus
|
||||
|
||||
fieldsToUpdate := []string{
|
||||
"peer_status_last_seen", "peer_status_session_started_at",
|
||||
"peer_status_connected", "peer_status_login_expired",
|
||||
"peer_status_requires_approval",
|
||||
}
|
||||
result := s.db.Model(&nbpeer.Peer{}).
|
||||
Select(fieldsToUpdate).
|
||||
Where(accountAndIDQueryCondition, accountID, peerID).
|
||||
Updates(&peerCopy)
|
||||
if result.Error != nil {
|
||||
return status.Errorf(status.Internal, "failed to save peer status to store: %v", result.Error)
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, peerNotFoundFMT, peerID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarkPeerConnectedIfNewerSession is an atomic optimistic-locked update.
|
||||
// The peer is marked connected with the given session token only when
|
||||
// the stored SessionStartedAt is strictly smaller than the incoming
|
||||
// one — equivalently, when no newer stream has already taken ownership.
|
||||
// The sentinel zero (set on peer creation or after a disconnect) counts
|
||||
// as the smallest possible token. This is the write half of the
|
||||
// fencing protocol described on PeerStatus.SessionStartedAt.
|
||||
//
|
||||
// The post-write side effects in the caller — geo lookup,
|
||||
// schedulePeerLoginExpiration, checkAndSchedulePeerInactivityExpiration,
|
||||
// OnPeersUpdated — all run AFTER this method returns and are deliberately
|
||||
// outside the database write so they cannot extend the row-lock window.
|
||||
//
|
||||
// LastSeen is set to the database's clock (CURRENT_TIMESTAMP) at the
|
||||
// moment the row is written. The caller never supplies LastSeen because
|
||||
// the value would otherwise drift under lock contention — a Go-side
|
||||
// time.Now() taken before the write can land minutes later than the
|
||||
// actual UPDATE under load, which previously caused real ordering bugs.
|
||||
func (s *SqlStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountID, peerID string, newSessionStartedAt int64) (bool, error) {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&nbpeer.Peer{}).
|
||||
Where(accountAndIDQueryCondition, accountID, peerID).
|
||||
Where("peer_status_session_started_at < ?", newSessionStartedAt).
|
||||
Updates(map[string]any{
|
||||
"peer_status_connected": true,
|
||||
"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
|
||||
"peer_status_session_started_at": newSessionStartedAt,
|
||||
"peer_status_login_expired": false,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return false, status.Errorf(status.Internal, "mark peer connected: %v", result.Error)
|
||||
}
|
||||
return result.RowsAffected > 0, nil
|
||||
}
|
||||
|
||||
// MarkPeerDisconnectedIfSameSession is an atomic optimistic-locked update.
|
||||
// The peer is marked disconnected only when the stored SessionStartedAt
|
||||
// matches the incoming token — meaning the stream that owns the current
|
||||
// session is the one ending. If a newer stream has already replaced the
|
||||
// session, the update is skipped. LastSeen is set to CURRENT_TIMESTAMP at
|
||||
// write time; see MarkPeerConnectedIfNewerSession for the rationale.
|
||||
//
|
||||
// A zero sessionStartedAt is rejected at the call site; the underlying
|
||||
// WHERE on equality would otherwise match every never-connected peer.
|
||||
func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error) {
|
||||
if sessionStartedAt == 0 {
|
||||
return false, nil
|
||||
}
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&nbpeer.Peer{}).
|
||||
Where(accountAndIDQueryCondition, accountID, peerID).
|
||||
Where("peer_status_session_started_at = ?", sessionStartedAt).
|
||||
Updates(map[string]any{
|
||||
"peer_status_connected": false,
|
||||
"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
|
||||
"peer_status_session_started_at": int64(0),
|
||||
})
|
||||
if result.Error != nil {
|
||||
return false, status.Errorf(status.Internal, "mark peer disconnected: %v", result.Error)
|
||||
}
|
||||
return result.RowsAffected > 0, nil
|
||||
}
|
||||
|
||||
// ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
|
||||
func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
|
||||
result := s.db.Model(&nbpeer.Peer{}).
|
||||
Where("account_id = ? AND peer_status_requires_approval = ?", accountID, true).
|
||||
Update("peer_status_requires_approval", false)
|
||||
if result.Error != nil {
|
||||
return 0, status.Errorf(status.Internal, "failed to approve pending account peers: %v", result.Error)
|
||||
}
|
||||
|
||||
return int(result.RowsAffected), nil
|
||||
}
|
||||
|
||||
// RefreshPeerLastSeen updates only peer_status_last_seen. Every other status
|
||||
// column is left untouched: peer_status_connected and
|
||||
// peer_status_session_started_at belong to the sync stream that owns the
|
||||
// session, and a blind write here would corrupt the fencing
|
||||
// MarkPeerConnectedIfNewerSession relies on.
|
||||
//
|
||||
// LastSeen comes from the database clock for the same reason it does there: a
|
||||
// Go-side timestamp is taken before the write and can land after a connect that
|
||||
// used CURRENT_TIMESTAMP, dragging the column backwards.
|
||||
//
|
||||
// staleBefore carries the caller's throttle into the same statement, so
|
||||
// concurrent requests for one peer collapse into a single write instead of
|
||||
// each racing on its own stale read. The column is nullable — Status is an
|
||||
// embedded pointer, so a peer stored without one leaves it NULL — and NULL
|
||||
// loses every comparison, hence the explicit branch for a peer never seen.
|
||||
func (s *SqlStore) RefreshPeerLastSeen(ctx context.Context, accountID, peerID string, staleBefore time.Time) (bool, error) {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&nbpeer.Peer{}).
|
||||
Where(accountAndIDQueryCondition, accountID, peerID).
|
||||
Where("(peer_status_last_seen IS NULL OR peer_status_last_seen < ?)", staleBefore).
|
||||
Update("peer_status_last_seen", gorm.Expr("CURRENT_TIMESTAMP"))
|
||||
if result.Error != nil {
|
||||
return false, status.Errorf(status.Internal, "refresh peer last seen: %v", result.Error)
|
||||
}
|
||||
|
||||
return result.RowsAffected > 0, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Peer, error) {
|
||||
const query = `SELECT id, account_id, key, ip, name, dns_label, user_id, ssh_key, ssh_enabled, login_expiration_enabled,
|
||||
inactivity_expiration_enabled, last_login, created_at, ephemeral, extra_dns_labels, allow_extra_dns_labels, meta_hostname,
|
||||
meta_go_os, meta_kernel, meta_core, meta_platform, meta_os, meta_os_version, meta_wt_version, meta_ui_version,
|
||||
meta_kernel_version, meta_network_addresses, meta_system_serial_number, meta_system_product_name, meta_system_manufacturer,
|
||||
meta_environment, meta_flags, meta_files, meta_capabilities, peer_status_last_seen, peer_status_session_started_at,
|
||||
peer_status_connected, peer_status_login_expired, peer_status_requires_approval, location_connection_ip,
|
||||
location_country_code, location_city_name, location_geo_name_id, proxy_meta_embedded, proxy_meta_cluster, ipv6, meta_sync_message_version
|
||||
FROM peers WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
peers, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nbpeer.Peer, error) {
|
||||
var p nbpeer.Peer
|
||||
p.Status = &nbpeer.PeerStatus{}
|
||||
var (
|
||||
lastLogin, createdAt sql.NullTime
|
||||
sshEnabled, loginExpirationEnabled, inactivityExpirationEnabled, ephemeral, allowExtraDNSLabels sql.NullBool
|
||||
peerStatusLastSeen sql.NullTime
|
||||
peerStatusSessionStartedAt sql.NullInt64
|
||||
peerStatusConnected, peerStatusLoginExpired, peerStatusRequiresApproval, proxyEmbedded sql.NullBool
|
||||
ip, extraDNS, netAddr, env, flags, files, capabilities, connIP, ipv6 []byte
|
||||
metaHostname, metaGoOS, metaKernel, metaCore, metaPlatform sql.NullString
|
||||
metaOS, metaOSVersion, metaWtVersion, metaUIVersion, metaKernelVersion sql.NullString
|
||||
metaSystemSerialNumber, metaSystemProductName, metaSystemManufacturer sql.NullString
|
||||
locationCountryCode, locationCityName, proxyCluster sql.NullString
|
||||
locationGeoNameID sql.NullInt64
|
||||
metaSyncMessageVersion sql.NullInt32
|
||||
)
|
||||
|
||||
err := row.Scan(&p.ID, &p.AccountID, &p.Key, &ip, &p.Name, &p.DNSLabel, &p.UserID, &p.SSHKey, &sshEnabled,
|
||||
&loginExpirationEnabled, &inactivityExpirationEnabled, &lastLogin, &createdAt, &ephemeral, &extraDNS,
|
||||
&allowExtraDNSLabels, &metaHostname, &metaGoOS, &metaKernel, &metaCore, &metaPlatform,
|
||||
&metaOS, &metaOSVersion, &metaWtVersion, &metaUIVersion, &metaKernelVersion, &netAddr,
|
||||
&metaSystemSerialNumber, &metaSystemProductName, &metaSystemManufacturer, &env, &flags, &files, &capabilities,
|
||||
&peerStatusLastSeen, &peerStatusSessionStartedAt, &peerStatusConnected, &peerStatusLoginExpired,
|
||||
&peerStatusRequiresApproval, &connIP, &locationCountryCode, &locationCityName, &locationGeoNameID,
|
||||
&proxyEmbedded, &proxyCluster, &ipv6, &metaSyncMessageVersion)
|
||||
|
||||
if err == nil {
|
||||
if lastLogin.Valid {
|
||||
p.LastLogin = &lastLogin.Time
|
||||
}
|
||||
if createdAt.Valid {
|
||||
p.CreatedAt = createdAt.Time
|
||||
}
|
||||
if sshEnabled.Valid {
|
||||
p.SSHEnabled = sshEnabled.Bool
|
||||
}
|
||||
if loginExpirationEnabled.Valid {
|
||||
p.LoginExpirationEnabled = loginExpirationEnabled.Bool
|
||||
}
|
||||
if inactivityExpirationEnabled.Valid {
|
||||
p.InactivityExpirationEnabled = inactivityExpirationEnabled.Bool
|
||||
}
|
||||
if ephemeral.Valid {
|
||||
p.Ephemeral = ephemeral.Bool
|
||||
}
|
||||
if allowExtraDNSLabels.Valid {
|
||||
p.AllowExtraDNSLabels = allowExtraDNSLabels.Bool
|
||||
}
|
||||
if peerStatusLastSeen.Valid {
|
||||
p.Status.LastSeen = peerStatusLastSeen.Time
|
||||
}
|
||||
if peerStatusSessionStartedAt.Valid {
|
||||
p.Status.SessionStartedAt = peerStatusSessionStartedAt.Int64
|
||||
}
|
||||
if peerStatusConnected.Valid {
|
||||
p.Status.Connected = peerStatusConnected.Bool
|
||||
}
|
||||
if peerStatusLoginExpired.Valid {
|
||||
p.Status.LoginExpired = peerStatusLoginExpired.Bool
|
||||
}
|
||||
if peerStatusRequiresApproval.Valid {
|
||||
p.Status.RequiresApproval = peerStatusRequiresApproval.Bool
|
||||
}
|
||||
if metaHostname.Valid {
|
||||
p.Meta.Hostname = metaHostname.String
|
||||
}
|
||||
if metaGoOS.Valid {
|
||||
p.Meta.GoOS = metaGoOS.String
|
||||
}
|
||||
if metaKernel.Valid {
|
||||
p.Meta.Kernel = metaKernel.String
|
||||
}
|
||||
if metaCore.Valid {
|
||||
p.Meta.Core = metaCore.String
|
||||
}
|
||||
if metaPlatform.Valid {
|
||||
p.Meta.Platform = metaPlatform.String
|
||||
}
|
||||
if metaOS.Valid {
|
||||
p.Meta.OS = metaOS.String
|
||||
}
|
||||
if metaOSVersion.Valid {
|
||||
p.Meta.OSVersion = metaOSVersion.String
|
||||
}
|
||||
if metaWtVersion.Valid {
|
||||
p.Meta.WtVersion = metaWtVersion.String
|
||||
}
|
||||
if metaUIVersion.Valid {
|
||||
p.Meta.UIVersion = metaUIVersion.String
|
||||
}
|
||||
if metaKernelVersion.Valid {
|
||||
p.Meta.KernelVersion = metaKernelVersion.String
|
||||
}
|
||||
if metaSystemSerialNumber.Valid {
|
||||
p.Meta.SystemSerialNumber = metaSystemSerialNumber.String
|
||||
}
|
||||
if metaSystemProductName.Valid {
|
||||
p.Meta.SystemProductName = metaSystemProductName.String
|
||||
}
|
||||
if metaSystemManufacturer.Valid {
|
||||
p.Meta.SystemManufacturer = metaSystemManufacturer.String
|
||||
}
|
||||
if locationCountryCode.Valid {
|
||||
p.Location.CountryCode = locationCountryCode.String
|
||||
}
|
||||
if locationCityName.Valid {
|
||||
p.Location.CityName = locationCityName.String
|
||||
}
|
||||
if locationGeoNameID.Valid {
|
||||
p.Location.GeoNameID = uint(locationGeoNameID.Int64)
|
||||
}
|
||||
if proxyEmbedded.Valid {
|
||||
p.ProxyMeta.Embedded = proxyEmbedded.Bool
|
||||
}
|
||||
if proxyCluster.Valid {
|
||||
p.ProxyMeta.Cluster = proxyCluster.String
|
||||
}
|
||||
if ip != nil {
|
||||
_ = json.Unmarshal(ip, &p.IP)
|
||||
}
|
||||
if ipv6 != nil {
|
||||
_ = json.Unmarshal(ipv6, &p.IPv6)
|
||||
}
|
||||
if extraDNS != nil {
|
||||
_ = json.Unmarshal(extraDNS, &p.ExtraDNSLabels)
|
||||
}
|
||||
if netAddr != nil {
|
||||
_ = json.Unmarshal(netAddr, &p.Meta.NetworkAddresses)
|
||||
}
|
||||
if env != nil {
|
||||
_ = json.Unmarshal(env, &p.Meta.Environment)
|
||||
}
|
||||
if flags != nil {
|
||||
_ = json.Unmarshal(flags, &p.Meta.Flags)
|
||||
}
|
||||
if files != nil {
|
||||
_ = json.Unmarshal(files, &p.Meta.Files)
|
||||
}
|
||||
if capabilities != nil {
|
||||
_ = json.Unmarshal(capabilities, &p.Meta.Capabilities)
|
||||
}
|
||||
if connIP != nil {
|
||||
_ = json.Unmarshal(connIP, &p.Location.ConnectionIP)
|
||||
}
|
||||
if metaSyncMessageVersion.Valid {
|
||||
p.Meta.SyncMessageVersion = int(metaSyncMessageVersion.Int32)
|
||||
}
|
||||
}
|
||||
return p, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountByPeerID(ctx context.Context, peerID string) (*types.Account, error) {
|
||||
var peer nbpeer.Peer
|
||||
result := s.db.Select("account_id").Take(&peer, idQueryCondition, peerID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
return nil, status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
if peer.AccountID == "" {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
|
||||
return s.GetAccount(ctx, peer.AccountID)
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (*types.Account, error) {
|
||||
var peer nbpeer.Peer
|
||||
result := s.db.Select("account_id").Take(&peer, GetKeyQueryCondition(s), peerKey)
|
||||
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
return nil, status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
if peer.AccountID == "" {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
|
||||
return s.GetAccount(ctx, peer.AccountID)
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountIDByPeerPubKey(ctx context.Context, peerKey string) (string, error) {
|
||||
var peer nbpeer.Peer
|
||||
var accountID string
|
||||
result := s.db.Model(&peer).Select("account_id").Where(GetKeyQueryCondition(s), peerKey).Take(&accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
return "", status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
return accountID, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountIDByPeerID(ctx context.Context, lockStrength LockingStrength, peerID string) (string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var accountID string
|
||||
result := tx.Model(&nbpeer.Peer{}).
|
||||
Select("account_id").Where(idQueryCondition, peerID).Take(&accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return "", status.Errorf(status.NotFound, "peer %s account not found", peerID)
|
||||
}
|
||||
return "", status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
return accountID, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength, accountID string) ([]netip.Addr, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var ipJSONStrings []string
|
||||
|
||||
result := tx.Model(&nbpeer.Peer{}).
|
||||
Where("account_id = ?", accountID).
|
||||
Pluck("ip", &ipJSONStrings)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "no peers found for the account")
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "issue getting IPs from store: %s", result.Error)
|
||||
}
|
||||
|
||||
ips := make([]netip.Addr, len(ipJSONStrings))
|
||||
for i, ipJSON := range ipJSONStrings {
|
||||
var ip netip.Addr
|
||||
if err := json.Unmarshal([]byte(ipJSON), &ip); err != nil {
|
||||
return nil, status.Errorf(status.Internal, "issue parsing IP JSON from store")
|
||||
}
|
||||
ips[i] = ip.Unmap()
|
||||
}
|
||||
|
||||
return ips, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountID string, dnsLabel string) ([]string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var labels []string
|
||||
result := tx.Model(&nbpeer.Peer{}).
|
||||
Where("account_id = ? AND dns_label LIKE ?", accountID, dnsLabel+"%").
|
||||
Pluck("dns_label", &labels)
|
||||
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "no peers found for the account")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("error when getting dns labels from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "issue getting dns labels from store: %s", result.Error)
|
||||
}
|
||||
|
||||
return labels, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peer nbpeer.Peer
|
||||
result := tx.Take(&peer, GetKeyQueryCondition(s), peerKey)
|
||||
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPeerNotFoundError(peerKey)
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "issue getting peer from store: %s", result.Error)
|
||||
}
|
||||
|
||||
return &peer, nil
|
||||
}
|
||||
|
||||
// GetAccountPeers retrieves peers for an account.
|
||||
func (s *SqlStore) GetAccountPeers(ctx context.Context, lockStrength LockingStrength, accountID, nameFilter, ipFilter string) ([]*nbpeer.Peer, error) {
|
||||
var peers []*nbpeer.Peer
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
query := tx.Where(accountIDCondition, accountID)
|
||||
|
||||
if nameFilter != "" {
|
||||
query = query.Where("name LIKE ?", "%"+nameFilter+"%")
|
||||
}
|
||||
if ipFilter != "" {
|
||||
query = query.Where("ip LIKE ? OR ipv6 LIKE ?", "%"+ipFilter+"%", "%"+ipFilter+"%")
|
||||
}
|
||||
|
||||
if err := query.Find(&peers).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peers from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get peers from store")
|
||||
}
|
||||
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
// GetUserPeers retrieves peers for a user.
|
||||
func (s *SqlStore) GetUserPeers(ctx context.Context, lockStrength LockingStrength, accountID, userID string) ([]*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peers []*nbpeer.Peer
|
||||
|
||||
// Exclude peers added via setup keys, as they are not user-specific and have an empty user_id.
|
||||
if userID == "" {
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
result := tx.
|
||||
Find(&peers, "account_id = ? AND user_id = ?", accountID, userID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peers from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get peers from store")
|
||||
}
|
||||
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) error {
|
||||
if err := s.db.Create(peer).Error; err != nil {
|
||||
return status.Errorf(status.Internal, "issue adding peer to account: %s", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetPeerByID retrieves a peer by its ID and account ID.
|
||||
func (s *SqlStore) GetPeerByID(ctx context.Context, lockStrength LockingStrength, accountID, peerID string) (*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peer *nbpeer.Peer
|
||||
result := tx.
|
||||
Take(&peer, accountAndIDQueryCondition, accountID, peerID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPeerNotFoundError(peerID)
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "failed to get peer from store")
|
||||
}
|
||||
|
||||
return peer, nil
|
||||
}
|
||||
|
||||
// GetPeersByIDs retrieves peers by their IDs and account ID.
|
||||
func (s *SqlStore) GetPeersByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, peerIDs []string) (map[string]*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peers []*nbpeer.Peer
|
||||
result := tx.Find(&peers, accountAndIDsQueryCondition, accountID, peerIDs)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peers by ID's from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get peers by ID's from the store")
|
||||
}
|
||||
|
||||
peersMap := make(map[string]*nbpeer.Peer)
|
||||
for _, peer := range peers {
|
||||
peersMap[peer.ID] = peer
|
||||
}
|
||||
|
||||
return peersMap, nil
|
||||
}
|
||||
|
||||
// GetAccountPeersWithExpiration retrieves a list of peers that have login expiration enabled and added by a user.
|
||||
func (s *SqlStore) GetAccountPeersWithExpiration(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peers []*nbpeer.Peer
|
||||
result := tx.
|
||||
Where("login_expiration_enabled = ? AND peer_status_login_expired != ? AND user_id IS NOT NULL AND user_id != ''", true, true).
|
||||
Find(&peers, accountIDCondition, accountID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peers with expiration from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get peers with expiration from store")
|
||||
}
|
||||
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
// GetAccountPeersWithInactivity retrieves a list of peers that have login expiration enabled and added by a user.
|
||||
func (s *SqlStore) GetAccountPeersWithInactivity(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peers []*nbpeer.Peer
|
||||
result := tx.
|
||||
Where("inactivity_expiration_enabled = ? AND user_id IS NOT NULL AND user_id != ''", true).
|
||||
Find(&peers, accountIDCondition, accountID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peers with inactivity from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get peers with inactivity from store")
|
||||
}
|
||||
|
||||
return peers, nil
|
||||
}
|
||||
|
||||
// GetAllEphemeralPeers retrieves all peers with Ephemeral set to true across all accounts, optimized for batch processing.
|
||||
func (s *SqlStore) GetAllEphemeralPeers(ctx context.Context, lockStrength LockingStrength) ([]*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var allEphemeralPeers, batchPeers []*nbpeer.Peer
|
||||
result := tx.
|
||||
Where("ephemeral = ?", true).
|
||||
FindInBatches(&batchPeers, 1000, func(tx *gorm.DB, batch int) error {
|
||||
allEphemeralPeers = append(allEphemeralPeers, batchPeers...)
|
||||
return nil
|
||||
})
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to retrieve ephemeral peers: %s", result.Error)
|
||||
return nil, fmt.Errorf("failed to retrieve ephemeral peers")
|
||||
}
|
||||
|
||||
return allEphemeralPeers, nil
|
||||
}
|
||||
|
||||
// DeletePeer removes a peer from the store.
|
||||
func (s *SqlStore) DeletePeer(ctx context.Context, accountID string, peerID string) error {
|
||||
result := s.db.Delete(&nbpeer.Peer{}, accountAndIDQueryCondition, accountID, peerID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete peer from the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to delete peer from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewPeerNotFoundError(peerID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeerByIP(ctx context.Context, lockStrength LockingStrength, accountID string, ip net.IP) (*nbpeer.Peer, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
column := "ip"
|
||||
if ip.To4() == nil {
|
||||
column = "ipv6"
|
||||
}
|
||||
jsonValue := fmt.Sprintf(`"%s"`, ip.String())
|
||||
|
||||
var peer nbpeer.Peer
|
||||
result := tx.
|
||||
Take(&peer, fmt.Sprintf("account_id = ? AND %s = ?", column), accountID, jsonValue)
|
||||
if result.Error != nil {
|
||||
// A tunnel-IP miss is an expected outcome (e.g. the proxy's
|
||||
// ValidateTunnelPeer probing an address that isn't in the
|
||||
// account roster); surface it as NotFound so callers can tell
|
||||
// it apart from a real store failure.
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "peer with ip %s not found", ip.String())
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "failed to get peer from store")
|
||||
}
|
||||
|
||||
return &peer, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeerIdByLabel(ctx context.Context, lockStrength LockingStrength, accountID string, hostname string) (string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peerID string
|
||||
result := tx.Model(&nbpeer.Peer{}).
|
||||
Select("id").
|
||||
// Where(" = ?", hostname).
|
||||
Where("account_id = ? AND dns_label = ?", accountID, hostname).
|
||||
Limit(1).
|
||||
Scan(&peerID)
|
||||
|
||||
if peerID == "" {
|
||||
return "", gorm.ErrRecordNotFound
|
||||
}
|
||||
|
||||
return peerID, result.Error
|
||||
}
|
||||
|
||||
// GetEmbeddedProxyPeerIDsByCluster returns peer IDs of all embedded proxy peers
|
||||
// in the account, grouped by their ProxyCluster. The map is nil when no embedded
|
||||
// proxy peers exist.
|
||||
func (s *SqlStore) GetEmbeddedProxyPeerIDsByCluster(ctx context.Context, accountID string) (map[string][]string, error) {
|
||||
type row struct {
|
||||
ID string
|
||||
Cluster string
|
||||
}
|
||||
var rows []row
|
||||
result := s.db.Model(&nbpeer.Peer{}).
|
||||
Select("id, proxy_meta_cluster AS cluster").
|
||||
Where("account_id = ? AND proxy_meta_embedded = ?", accountID, true).
|
||||
Scan(&rows)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get embedded proxy peers: %s", result.Error)
|
||||
}
|
||||
|
||||
out := make(map[string][]string, len(rows))
|
||||
for _, r := range rows {
|
||||
out[r.Cluster] = append(out[r.Cluster], r.ID)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetUserIDByPeerKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var userID string
|
||||
result := tx.Model(&nbpeer.Peer{}).
|
||||
Select("user_id").
|
||||
Take(&userID, GetKeyQueryCondition(s), peerKey)
|
||||
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return "", status.Errorf(status.NotFound, "peer not found: index lookup failed")
|
||||
}
|
||||
return "", status.Errorf(status.Internal, "failed to get user ID by peer key")
|
||||
}
|
||||
|
||||
return userID, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPeerIDByKey(ctx context.Context, lockStrength LockingStrength, key string) (string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var peerID string
|
||||
result := tx.Model(&nbpeer.Peer{}).
|
||||
Select("id").
|
||||
Where(GetKeyQueryCondition(s), key).
|
||||
Limit(1).
|
||||
Scan(&peerID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get peer ID by key: %s", result.Error)
|
||||
return "", status.Errorf(status.Internal, "failed to get peer ID by key")
|
||||
}
|
||||
|
||||
return peerID, nil
|
||||
}
|
||||
@@ -0,0 +1,901 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"reflect"
|
||||
"sort"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/management/server/util"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
"github.com/netbirdio/netbird/shared/testing_helpers"
|
||||
)
|
||||
|
||||
// TestSqlStore_GetPeerByIP_NotFound pins the not-found semantics the
|
||||
// proxy's ValidateTunnelPeer relies on: a tunnel-IP that isn't in the
|
||||
// account roster must surface as a NotFound error (not a generic
|
||||
// Internal) so callers can distinguish an expected miss from a real
|
||||
// store failure. A known IP still resolves.
|
||||
func TestSqlStore_GetPeerByIP_NotFound(t *testing.T) {
|
||||
runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
|
||||
const accountID = "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
peer, err := store.GetPeerByIP(context.Background(), LockingStrengthNone, accountID, net.ParseIP("192.168.0.0"))
|
||||
require.NoError(t, err, "known tunnel IP must resolve")
|
||||
require.NotNil(t, peer)
|
||||
|
||||
_, err = store.GetPeerByIP(context.Background(), LockingStrengthNone, accountID, net.ParseIP("100.65.0.99"))
|
||||
require.Error(t, err, "unknown tunnel IP must error")
|
||||
parsedErr, ok := status.FromError(err)
|
||||
require.True(t, ok, "error must be a status error")
|
||||
require.Equal(t, status.NotFound, parsedErr.Type(), "tunnel-IP miss must be NotFound, not Internal")
|
||||
})
|
||||
}
|
||||
|
||||
func TestSqlStore_SavePeer(t *testing.T) {
|
||||
populateFields := testing_helpers.NewPopulateFields()
|
||||
|
||||
runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
|
||||
account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
|
||||
require.NoError(t, err)
|
||||
|
||||
metadata := nbpeer.PeerSystemMeta{}
|
||||
reflectedMetadata := reflect.ValueOf(&metadata).Elem()
|
||||
|
||||
numOfFields, err := populateFields.PopulateAll(reflectedMetadata)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 33, numOfFields)
|
||||
|
||||
// save status of non-existing peer
|
||||
peer := &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
ID: "testpeer",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Meta: metadata, //nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
ctx := context.Background()
|
||||
err = store.SavePeer(ctx, account.Id, peer)
|
||||
assert.Error(t, err)
|
||||
parsedErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
|
||||
|
||||
// save new status of existing peer
|
||||
account.Peers[peer.ID] = peer
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
updatedPeer := peer.Copy()
|
||||
updatedPeer.Status.Connected = false
|
||||
updatedPeer.Meta.Hostname = "updatedpeer"
|
||||
|
||||
err = store.SavePeer(ctx, account.Id, updatedPeer)
|
||||
require.NoError(t, err)
|
||||
|
||||
account, err = store.GetAccount(context.Background(), account.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
actual := account.Peers[peer.ID]
|
||||
assert.Equal(t, updatedPeer.Meta, actual.Meta)
|
||||
assert.Equal(t, updatedPeer.Status.Connected, actual.Status.Connected)
|
||||
assert.Equal(t, updatedPeer.Status.LoginExpired, actual.Status.LoginExpired)
|
||||
assert.Equal(t, updatedPeer.Status.RequiresApproval, actual.Status.RequiresApproval)
|
||||
assert.WithinDurationf(t, updatedPeer.Status.LastSeen, actual.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
|
||||
})
|
||||
}
|
||||
|
||||
func TestSqlStore_SavePeerStatus(t *testing.T) {
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
|
||||
require.NoError(t, err)
|
||||
|
||||
// save status of non-existing peer
|
||||
newStatus := nbpeer.PeerStatus{Connected: false, LastSeen: time.Now().UTC()}
|
||||
err = store.SavePeerStatus(context.Background(), account.Id, "non-existing-peer", newStatus)
|
||||
assert.Error(t, err)
|
||||
parsedErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
|
||||
|
||||
// save new status of existing peer
|
||||
account.Peers["testpeer"] = &nbpeer.Peer{
|
||||
Key: "peerkey",
|
||||
ID: "testpeer",
|
||||
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Meta: nbpeer.PeerSystemMeta{},
|
||||
Name: "peer name",
|
||||
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
|
||||
}
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = store.SavePeerStatus(context.Background(), account.Id, "testpeer", newStatus)
|
||||
require.NoError(t, err)
|
||||
|
||||
account, err = store.GetAccount(context.Background(), account.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
actual := account.Peers["testpeer"].Status
|
||||
assert.Equal(t, newStatus.Connected, actual.Connected)
|
||||
assert.Equal(t, newStatus.LoginExpired, actual.LoginExpired)
|
||||
assert.Equal(t, newStatus.RequiresApproval, actual.RequiresApproval)
|
||||
assert.WithinDurationf(t, newStatus.LastSeen, actual.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
|
||||
|
||||
newStatus.Connected = true
|
||||
|
||||
err = store.SavePeerStatus(context.Background(), account.Id, "testpeer", newStatus)
|
||||
require.NoError(t, err)
|
||||
|
||||
account, err = store.GetAccount(context.Background(), account.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
actual = account.Peers["testpeer"].Status
|
||||
assert.Equal(t, newStatus.Connected, actual.Connected)
|
||||
assert.Equal(t, newStatus.LoginExpired, actual.LoginExpired)
|
||||
assert.Equal(t, newStatus.RequiresApproval, actual.RequiresApproval)
|
||||
assert.WithinDurationf(t, newStatus.LastSeen, actual.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
|
||||
}
|
||||
|
||||
func TestSqlite_GetTakenIPs(t *testing.T) {
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
defer cleanup()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
_, err = store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
takenIPs, err := store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []netip.Addr{}, takenIPs)
|
||||
|
||||
peer1 := &nbpeer.Peer{
|
||||
ID: "peer1",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key1",
|
||||
DNSLabel: "peer1",
|
||||
IP: netip.AddrFrom4([4]byte{1, 1, 1, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1:1:1:1"),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer1)
|
||||
require.NoError(t, err)
|
||||
|
||||
takenIPs, err = store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
|
||||
require.NoError(t, err)
|
||||
ip1 := netip.AddrFrom4([4]byte{1, 1, 1, 1})
|
||||
assert.Equal(t, []netip.Addr{ip1}, takenIPs)
|
||||
|
||||
peer2 := &nbpeer.Peer{
|
||||
ID: "peer1second",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key2",
|
||||
DNSLabel: "peer1-1",
|
||||
IP: netip.AddrFrom4([4]byte{2, 2, 2, 2}),
|
||||
IPv6: netip.MustParseAddr("fd00::2:2:2:2"),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer2)
|
||||
require.NoError(t, err)
|
||||
|
||||
takenIPs, err = store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
|
||||
require.NoError(t, err)
|
||||
ip2 := netip.AddrFrom4([4]byte{2, 2, 2, 2})
|
||||
assert.Equal(t, []netip.Addr{ip1, ip2}, takenIPs)
|
||||
}
|
||||
|
||||
func TestSqlite_GetPeerLabelsInAccount(t *testing.T) {
|
||||
runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
peerHostname := "peer1"
|
||||
|
||||
_, err := store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
labels, err := store.GetPeerLabelsInAccount(context.Background(), LockingStrengthNone, existingAccountID, peerHostname)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{}, labels)
|
||||
|
||||
peer1 := &nbpeer.Peer{
|
||||
ID: "peer1",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key1",
|
||||
DNSLabel: "peer1",
|
||||
IP: netip.AddrFrom4([4]byte{1, 1, 1, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1:1:1:1"),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer1)
|
||||
require.NoError(t, err)
|
||||
|
||||
labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthNone, existingAccountID, peerHostname)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"peer1"}, labels)
|
||||
|
||||
peer2 := &nbpeer.Peer{
|
||||
ID: "peer1second",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key2",
|
||||
DNSLabel: "peer1-1",
|
||||
IP: netip.AddrFrom4([4]byte{2, 2, 2, 2}),
|
||||
IPv6: netip.MustParseAddr("fd00::2:2:2:2"),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer2)
|
||||
require.NoError(t, err)
|
||||
|
||||
labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthNone, existingAccountID, peerHostname)
|
||||
require.NoError(t, err)
|
||||
|
||||
expected := []string{"peer1", "peer1-1"}
|
||||
sort.Strings(expected)
|
||||
sort.Strings(labels)
|
||||
assert.Equal(t, expected, labels)
|
||||
})
|
||||
}
|
||||
|
||||
func Test_AddPeerWithSameDnsLabel(t *testing.T) {
|
||||
runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
_, err := store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
peer1 := &nbpeer.Peer{
|
||||
ID: "peer1",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key1",
|
||||
DNSLabel: "peer1.domain.test",
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer1)
|
||||
require.NoError(t, err)
|
||||
|
||||
peer2 := &nbpeer.Peer{
|
||||
ID: "peer1second",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key2",
|
||||
DNSLabel: "peer1.domain.test",
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer2)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func Test_AddPeerWithSameIP(t *testing.T) {
|
||||
runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
_, err := store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
peer1 := &nbpeer.Peer{
|
||||
ID: "peer1",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key1",
|
||||
IP: netip.AddrFrom4([4]byte{1, 1, 1, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1:1:1:1"),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer1)
|
||||
require.NoError(t, err)
|
||||
|
||||
peer2 := &nbpeer.Peer{
|
||||
ID: "peer1second",
|
||||
AccountID: existingAccountID,
|
||||
Key: "key2",
|
||||
IP: netip.AddrFrom4([4]byte{1, 1, 1, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::2:2:2:2"),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer2)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPeerByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
peerID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing peer",
|
||||
peerID: "cfefqs706sqkneg59g4g",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing peer",
|
||||
peerID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty peer ID",
|
||||
peerID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, accountID, tt.peerID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, peer)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, peer)
|
||||
require.Equal(t, tt.peerID, peer.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPeersByIDs(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
peerIDs []string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing peers by existing IDs",
|
||||
peerIDs: []string{"cfefqs706sqkneg59g4g", "cfeg6sf06sqkneg59g50"},
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "empty peer IDs list",
|
||||
peerIDs: []string{},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "non-existing peer IDs",
|
||||
peerIDs: []string{"nonexistent1", "nonexistent2"},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "mixed existing and non-existing peer IDs",
|
||||
peerIDs: []string{"cfeg6sf06sqkneg59g50", "nonexistent"},
|
||||
expectedCount: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peers, err := store.GetPeersByIDs(context.Background(), LockingStrengthNone, accountID, tt.peerIDs)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_AddPeerToAccount(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
peer := &nbpeer.Peer{
|
||||
ID: "peer1",
|
||||
AccountID: accountID,
|
||||
Key: "key",
|
||||
IP: netip.AddrFrom4([4]byte{1, 1, 1, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::1:1:1:1"),
|
||||
Meta: nbpeer.PeerSystemMeta{
|
||||
Hostname: "hostname",
|
||||
GoOS: "linux",
|
||||
Kernel: "Linux",
|
||||
Core: "21.04",
|
||||
Platform: "x86_64",
|
||||
OS: "Ubuntu",
|
||||
WtVersion: "development",
|
||||
UIVersion: "development",
|
||||
},
|
||||
Name: "peer.test",
|
||||
DNSLabel: "peer",
|
||||
Status: &nbpeer.PeerStatus{
|
||||
LastSeen: time.Now().UTC(),
|
||||
Connected: true,
|
||||
LoginExpired: false,
|
||||
RequiresApproval: false,
|
||||
},
|
||||
SSHKey: "ssh-key",
|
||||
SSHEnabled: false,
|
||||
LoginExpirationEnabled: true,
|
||||
InactivityExpirationEnabled: false,
|
||||
LastLogin: util.ToPtr(time.Now().UTC()),
|
||||
CreatedAt: time.Now().UTC(),
|
||||
Ephemeral: true,
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer)
|
||||
require.NoError(t, err, "failed to add peer to account")
|
||||
|
||||
storedPeer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, accountID, peer.ID)
|
||||
require.NoError(t, err, "failed to get peer")
|
||||
|
||||
assert.Equal(t, peer.ID, storedPeer.ID)
|
||||
assert.Equal(t, peer.AccountID, storedPeer.AccountID)
|
||||
assert.Equal(t, peer.Key, storedPeer.Key)
|
||||
assert.Equal(t, peer.IP.String(), storedPeer.IP.String())
|
||||
assert.Equal(t, peer.Meta, storedPeer.Meta)
|
||||
assert.Equal(t, peer.Name, storedPeer.Name)
|
||||
assert.Equal(t, peer.DNSLabel, storedPeer.DNSLabel)
|
||||
assert.Equal(t, peer.SSHKey, storedPeer.SSHKey)
|
||||
assert.Equal(t, peer.SSHEnabled, storedPeer.SSHEnabled)
|
||||
assert.Equal(t, peer.LoginExpirationEnabled, storedPeer.LoginExpirationEnabled)
|
||||
assert.Equal(t, peer.InactivityExpirationEnabled, storedPeer.InactivityExpirationEnabled)
|
||||
assert.WithinDurationf(t, peer.GetLastLogin(), storedPeer.GetLastLogin().UTC(), time.Millisecond, "LastLogin should be equal")
|
||||
assert.WithinDurationf(t, peer.CreatedAt, storedPeer.CreatedAt.UTC(), time.Millisecond, "CreatedAt should be equal")
|
||||
assert.Equal(t, peer.Ephemeral, storedPeer.Ephemeral)
|
||||
assert.Equal(t, peer.Status.Connected, storedPeer.Status.Connected)
|
||||
assert.Equal(t, peer.Status.LoginExpired, storedPeer.Status.LoginExpired)
|
||||
assert.Equal(t, peer.Status.RequiresApproval, storedPeer.Status.RequiresApproval)
|
||||
assert.WithinDurationf(t, peer.Status.LastSeen, storedPeer.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountPeers(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
nameFilter string
|
||||
ipFilter string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "should retrieve peers for an existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectedCount: 5,
|
||||
},
|
||||
{
|
||||
name: "should return no peers for a non-existing account ID",
|
||||
accountID: "nonexistent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "should return no peers for an empty account ID",
|
||||
accountID: "",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "should filter peers by name",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
nameFilter: "expiredhost",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "should filter peers by partial name",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
nameFilter: "host",
|
||||
expectedCount: 4,
|
||||
},
|
||||
{
|
||||
name: "should filter peers by ip",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
ipFilter: "100.64.39.54",
|
||||
expectedCount: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peers, err := store.GetAccountPeers(context.Background(), LockingStrengthNone, tt.accountID, tt.nameFilter, tt.ipFilter)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountPeersWithExpiration(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectedCount int
|
||||
expectedPeerIDs []string
|
||||
}{
|
||||
{
|
||||
name: "should retrieve only non-expired peers with expiration enabled",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectedCount: 1,
|
||||
expectedPeerIDs: []string{"notexpired01"},
|
||||
},
|
||||
{
|
||||
name: "should return no peers with expiration for a non-existing account ID",
|
||||
accountID: "nonexistent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "should return no peers with expiration for a empty account ID",
|
||||
accountID: "",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peers, err := store.GetAccountPeersWithExpiration(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
for i, peer := range peers {
|
||||
assert.Equal(t, tt.expectedPeerIDs[i], peer.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountPeersWithExpiration_ExcludesAlreadyExpired(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
peers, err := store.GetAccountPeersWithExpiration(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify the already-expired peer (cg05lnblo1hkg2j514p0) is not returned
|
||||
for _, peer := range peers {
|
||||
assert.NotEqual(t, "cg05lnblo1hkg2j514p0", peer.ID, "already expired peer should not be returned")
|
||||
assert.False(t, peer.Status.LoginExpired, "returned peers should not have LoginExpired set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountPeersWithInactivity(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "should retrieve peers with inactivity for an existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "should return no peers with inactivity for a non-existing account ID",
|
||||
accountID: "nonexistent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "should return no peers with inactivity for an empty account ID",
|
||||
accountID: "",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peers, err := store.GetAccountPeersWithInactivity(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAllEphemeralPeers(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/storev1.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
peers, err := store.GetAllEphemeralPeers(context.Background(), LockingStrengthNone)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, 1)
|
||||
require.True(t, peers[0].Ephemeral)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserPeers(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
userID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "should retrieve peers for existing account ID and user ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
userID: "f4f6d672-63fb-11ec-90d6-0242ac120003",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "should return no peers for non-existing account ID with existing user ID",
|
||||
accountID: "nonexistent",
|
||||
userID: "f4f6d672-63fb-11ec-90d6-0242ac120003",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "should return no peers for non-existing user ID with existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
userID: "nonexistent_user",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "should retrieve peers for another valid account ID and user ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
userID: "edafee4e-63fb-11ec-90d6-0242ac120003",
|
||||
expectedCount: 3,
|
||||
},
|
||||
{
|
||||
name: "should return no peers for existing account ID with empty user ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
userID: "",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
peers, err := store.GetUserPeers(context.Background(), LockingStrengthNone, tt.accountID, tt.userID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, peers, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_DeletePeer(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
peerID := "csrnkiq7qv9d8aitqd50"
|
||||
|
||||
err = store.DeletePeer(context.Background(), accountID, peerID)
|
||||
require.NoError(t, err)
|
||||
|
||||
peer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, accountID, peerID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, peer)
|
||||
}
|
||||
|
||||
func BenchmarkGetAccountPeers(b *testing.B) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", b.TempDir())
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Cleanup(cleanup)
|
||||
|
||||
numberOfPeers := 1000
|
||||
numberOfGroups := 200
|
||||
numberOfPeersPerGroup := 500
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
peers := make([]*nbpeer.Peer, 0, numberOfPeers)
|
||||
for i := 0; i < numberOfPeers; i++ {
|
||||
peer := &nbpeer.Peer{
|
||||
ID: fmt.Sprintf("peer-%d", i),
|
||||
AccountID: accountID,
|
||||
Key: fmt.Sprintf("key-%d", i),
|
||||
DNSLabel: fmt.Sprintf("peer%d.example.com", i),
|
||||
IP: intToIPv4(uint32(i)),
|
||||
}
|
||||
err = store.AddPeerToAccount(context.Background(), peer)
|
||||
if err != nil {
|
||||
b.Fatalf("Failed to add peer: %v", err)
|
||||
}
|
||||
peers = append(peers, peer)
|
||||
}
|
||||
|
||||
for i := 0; i < numberOfGroups; i++ {
|
||||
groupID := fmt.Sprintf("group-%d", i)
|
||||
group := &types.Group{
|
||||
ID: groupID,
|
||||
AccountID: accountID,
|
||||
}
|
||||
err = store.CreateGroup(context.Background(), group)
|
||||
if err != nil {
|
||||
b.Fatalf("Failed to create group: %v", err)
|
||||
}
|
||||
for j := 0; j < numberOfPeersPerGroup; j++ {
|
||||
peerIndex := (i*numberOfPeersPerGroup + j) % numberOfPeers
|
||||
err = store.AddPeerToGroup(context.Background(), accountID, peers[peerIndex].ID, groupID)
|
||||
if err != nil {
|
||||
b.Fatalf("Failed to add peer to group: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peers[i%numberOfPeers].ID)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func intToIPv4(n uint32) netip.Addr {
|
||||
var b [4]byte
|
||||
binary.BigEndian.PutUint32(b[:], n)
|
||||
return netip.AddrFrom4(b)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserIDByPeerKey(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
userID := "test-user-123"
|
||||
peerKey := "peer-key-abc"
|
||||
|
||||
peer := &nbpeer.Peer{
|
||||
ID: "test-peer-1",
|
||||
Key: peerKey,
|
||||
AccountID: existingAccountID,
|
||||
UserID: userID,
|
||||
IP: netip.AddrFrom4([4]byte{10, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::a00:1"),
|
||||
DNSLabel: "test-peer-1",
|
||||
}
|
||||
|
||||
err = store.AddPeerToAccount(context.Background(), peer)
|
||||
require.NoError(t, err)
|
||||
|
||||
retrievedUserID, err := store.GetUserIDByPeerKey(context.Background(), LockingStrengthNone, peerKey)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, userID, retrievedUserID)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserIDByPeerKey_NotFound(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
nonExistentPeerKey := "non-existent-peer-key"
|
||||
|
||||
userID, err := store.GetUserIDByPeerKey(context.Background(), LockingStrengthNone, nonExistentPeerKey)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, "", userID)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserIDByPeerKey_NoUserID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
peerKey := "peer-key-abc"
|
||||
|
||||
peer := &nbpeer.Peer{
|
||||
ID: "test-peer-1",
|
||||
Key: peerKey,
|
||||
AccountID: existingAccountID,
|
||||
UserID: "",
|
||||
IP: netip.AddrFrom4([4]byte{10, 0, 0, 1}),
|
||||
IPv6: netip.MustParseAddr("fd00::a00:1"),
|
||||
DNSLabel: "test-peer-1",
|
||||
}
|
||||
|
||||
err = store.AddPeerToAccount(context.Background(), peer)
|
||||
require.NoError(t, err)
|
||||
|
||||
retrievedUserID, err := store.GetUserIDByPeerKey(context.Background(), LockingStrengthNone, peerKey)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "", retrievedUserID)
|
||||
}
|
||||
|
||||
func TestSqlStore_ApproveAccountPeers(t *testing.T) {
|
||||
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
|
||||
accountID := "test-account"
|
||||
ctx := context.Background()
|
||||
|
||||
account := newAccountWithId(ctx, accountID, "testuser", "example.com")
|
||||
err := store.SaveAccount(ctx, account)
|
||||
require.NoError(t, err)
|
||||
|
||||
peers := []*nbpeer.Peer{
|
||||
{
|
||||
ID: "peer1",
|
||||
AccountID: accountID,
|
||||
DNSLabel: "peer1.netbird.cloud",
|
||||
Key: "peer1-key",
|
||||
IP: netip.MustParseAddr("100.64.0.1"),
|
||||
IPv6: netip.MustParseAddr("fd00::1"),
|
||||
Status: &nbpeer.PeerStatus{
|
||||
RequiresApproval: true,
|
||||
LastSeen: time.Now().UTC(),
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "peer2",
|
||||
AccountID: accountID,
|
||||
DNSLabel: "peer2.netbird.cloud",
|
||||
Key: "peer2-key",
|
||||
IP: netip.MustParseAddr("100.64.0.2"),
|
||||
IPv6: netip.MustParseAddr("fd00::2"),
|
||||
Status: &nbpeer.PeerStatus{
|
||||
RequiresApproval: true,
|
||||
LastSeen: time.Now().UTC(),
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "peer3",
|
||||
AccountID: accountID,
|
||||
DNSLabel: "peer3.netbird.cloud",
|
||||
Key: "peer3-key",
|
||||
IP: netip.MustParseAddr("100.64.0.3"),
|
||||
IPv6: netip.MustParseAddr("fd00::3"),
|
||||
Status: &nbpeer.PeerStatus{
|
||||
RequiresApproval: false,
|
||||
LastSeen: time.Now().UTC(),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, peer := range peers {
|
||||
err = store.AddPeerToAccount(ctx, peer)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
t.Run("approve all pending peers", func(t *testing.T) {
|
||||
count, err := store.ApproveAccountPeers(ctx, accountID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, count)
|
||||
|
||||
allPeers, err := store.GetAccountPeers(ctx, LockingStrengthNone, accountID, "", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, peer := range allPeers {
|
||||
assert.False(t, peer.Status.RequiresApproval, "peer %s should not require approval", peer.ID)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no peers to approve", func(t *testing.T) {
|
||||
count, err := store.ApproveAccountPeers(ctx, accountID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count)
|
||||
})
|
||||
|
||||
t.Run("non-existent account", func(t *testing.T) {
|
||||
count, err := store.ApproveAccountPeers(ctx, "non-existent")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count)
|
||||
})
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/management/server/util"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// DeleteHashedPAT2TokenIDIndex is noop in SqlStore
|
||||
func (s *SqlStore) DeleteHashedPAT2TokenIDIndex(hashedToken string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteTokenID2UserIDIndex is noop in SqlStore
|
||||
func (s *SqlStore) DeleteTokenID2UserIDIndex(tokenID string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetTokenIDByHashedToken(ctx context.Context, hashedToken string) (string, error) {
|
||||
var token types.PersonalAccessToken
|
||||
result := s.db.Take(&token, "hashed_token = ?", hashedToken)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("error when getting token from the store: %s", result.Error)
|
||||
return "", status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
return token.ID, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) getPersonalAccessTokens(ctx context.Context, userIDs []string) ([]types.PersonalAccessToken, error) {
|
||||
if len(userIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
const query = `SELECT id, user_id, name, hashed_token, expiration_date, created_by, created_at, last_used FROM personal_access_tokens WHERE user_id = ANY($1)`
|
||||
rows, err := s.pool.Query(ctx, query, userIDs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pats, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (types.PersonalAccessToken, error) {
|
||||
var pat types.PersonalAccessToken
|
||||
var expirationDate, lastUsed, createdAt sql.NullTime
|
||||
err := row.Scan(&pat.ID, &pat.UserID, &pat.Name, &pat.HashedToken, &expirationDate, &pat.CreatedBy, &createdAt, &lastUsed)
|
||||
if err == nil {
|
||||
if expirationDate.Valid {
|
||||
pat.ExpirationDate = &expirationDate.Time
|
||||
}
|
||||
if createdAt.Valid {
|
||||
pat.CreatedAt = createdAt.Time
|
||||
}
|
||||
if lastUsed.Valid {
|
||||
pat.LastUsed = &lastUsed.Time
|
||||
}
|
||||
}
|
||||
return pat, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return pats, nil
|
||||
}
|
||||
|
||||
// GetPATByHashedToken returns a PersonalAccessToken by its hashed token.
|
||||
func (s *SqlStore) GetPATByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken string) (*types.PersonalAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var pat types.PersonalAccessToken
|
||||
result := tx.Take(&pat, "hashed_token = ?", hashedToken)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPATNotFoundError(hashedToken)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get pat by hash from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get pat by hash from store")
|
||||
}
|
||||
|
||||
return &pat, nil
|
||||
}
|
||||
|
||||
// GetPATByID retrieves a personal access token by its ID and user ID.
|
||||
func (s *SqlStore) GetPATByID(ctx context.Context, lockStrength LockingStrength, userID string, patID string) (*types.PersonalAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var pat types.PersonalAccessToken
|
||||
result := tx.
|
||||
Take(&pat, "id = ? AND user_id = ?", patID, userID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPATNotFoundError(patID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get pat from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get pat from store")
|
||||
}
|
||||
|
||||
return &pat, nil
|
||||
}
|
||||
|
||||
// GetUserPATs retrieves personal access tokens for a user.
|
||||
func (s *SqlStore) GetUserPATs(ctx context.Context, lockStrength LockingStrength, userID string) ([]*types.PersonalAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var pats []*types.PersonalAccessToken
|
||||
result := tx.Find(&pats, "user_id = ?", userID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get user pat's from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get user pat's from store")
|
||||
}
|
||||
|
||||
return pats, nil
|
||||
}
|
||||
|
||||
// MarkPATUsed marks a personal access token as used.
|
||||
func (s *SqlStore) MarkPATUsed(ctx context.Context, patID string) error {
|
||||
patCopy := types.PersonalAccessToken{
|
||||
LastUsed: util.ToPtr(time.Now().UTC()),
|
||||
}
|
||||
|
||||
fieldsToUpdate := []string{"last_used"}
|
||||
result := s.db.Select(fieldsToUpdate).
|
||||
Where(idQueryCondition, patID).Updates(&patCopy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to mark pat as used: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to mark pat as used")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewPATNotFoundError(patID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// SavePAT saves a personal access token to the database.
|
||||
func (s *SqlStore) SavePAT(ctx context.Context, pat *types.PersonalAccessToken) error {
|
||||
result := s.db.Save(pat)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save pat to the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to save pat to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeletePAT deletes a personal access token from the database.
|
||||
func (s *SqlStore) DeletePAT(ctx context.Context, userID, patID string) error {
|
||||
result := s.db.Delete(&types.PersonalAccessToken{}, "user_id = ? AND id = ?", userID, patID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete pat from the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to delete pat from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewPATNotFoundError(patID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"runtime"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/management/server/util"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func Test_GetTokenIDByHashedToken(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
|
||||
hashed := "SoMeHaShEdToKeN"
|
||||
id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
token, err := store.GetTokenIDByHashedToken(context.Background(), hashed)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, id, token)
|
||||
|
||||
_, err = store.GetTokenIDByHashedToken(context.Background(), "non-existing-hash")
|
||||
require.Error(t, err)
|
||||
parsedErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
|
||||
})
|
||||
}
|
||||
|
||||
func TestPostgresql_GetTokenIDByHashedToken(t *testing.T) {
|
||||
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
|
||||
t.Skip("skip CI tests on darwin and windows")
|
||||
}
|
||||
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
hashed := "SoMeHaShEdToKeN"
|
||||
id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
token, err := store.GetTokenIDByHashedToken(context.Background(), hashed)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, id, token)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPATByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
userID := "f4f6d672-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
patID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing PAT",
|
||||
patID: "9dj38s35-63fb-11ec-90d6-0242ac120003",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing PAT",
|
||||
patID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty PAT ID",
|
||||
patID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
pat, err := store.GetPATByID(context.Background(), LockingStrengthNone, userID, tt.patID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, pat)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, pat)
|
||||
require.Equal(t, tt.patID, pat.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserPATs(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
userPATs, err := store.GetUserPATs(context.Background(), LockingStrengthNone, "f4f6d672-63fb-11ec-90d6-0242ac120003")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, userPATs, 1)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPATByHashedToken(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
pat, err := store.GetPATByHashedToken(context.Background(), LockingStrengthNone, "SoMeHaShEdToKeN")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "9dj38s35-63fb-11ec-90d6-0242ac120003", pat.ID)
|
||||
}
|
||||
|
||||
func TestSqlStore_MarkPATUsed(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
userID := "f4f6d672-63fb-11ec-90d6-0242ac120003"
|
||||
patID := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
err = store.MarkPATUsed(context.Background(), patID)
|
||||
require.NoError(t, err)
|
||||
|
||||
pat, err := store.GetPATByID(context.Background(), LockingStrengthNone, userID, patID)
|
||||
require.NoError(t, err)
|
||||
now := time.Now().UTC()
|
||||
require.WithinRange(t, pat.LastUsed.UTC(), now.Add(-15*time.Second), now, "LastUsed should be within 1 second of now")
|
||||
}
|
||||
|
||||
func TestSqlStore_SavePAT(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
userID := "edafee4e-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
pat := &types.PersonalAccessToken{
|
||||
ID: "pat-id",
|
||||
UserID: userID,
|
||||
Name: "token",
|
||||
HashedToken: "SoMeHaShEdToKeN",
|
||||
ExpirationDate: util.ToPtr(time.Now().UTC().Add(12 * time.Hour)),
|
||||
CreatedBy: userID,
|
||||
CreatedAt: time.Now().UTC().Add(time.Hour),
|
||||
LastUsed: util.ToPtr(time.Now().UTC().Add(-15 * time.Minute)),
|
||||
}
|
||||
err = store.SavePAT(context.Background(), pat)
|
||||
require.NoError(t, err)
|
||||
|
||||
savePAT, err := store.GetPATByID(context.Background(), LockingStrengthNone, userID, pat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, pat.ID, savePAT.ID)
|
||||
require.Equal(t, pat.UserID, savePAT.UserID)
|
||||
require.Equal(t, pat.HashedToken, savePAT.HashedToken)
|
||||
require.Equal(t, pat.CreatedBy, savePAT.CreatedBy)
|
||||
require.WithinDurationf(t, pat.GetExpirationDate(), savePAT.ExpirationDate.UTC(), time.Millisecond, "ExpirationDate should be equal")
|
||||
require.WithinDurationf(t, pat.CreatedAt, savePAT.CreatedAt.UTC(), time.Millisecond, "CreatedAt should be equal")
|
||||
require.WithinDurationf(t, pat.GetLastUsed(), savePAT.LastUsed.UTC(), time.Millisecond, "LastUsed should be equal")
|
||||
}
|
||||
|
||||
func TestSqlStore_DeletePAT(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
userID := "f4f6d672-63fb-11ec-90d6-0242ac120003"
|
||||
patID := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
err = store.DeletePAT(context.Background(), userID, patID)
|
||||
require.NoError(t, err)
|
||||
|
||||
pat, err := store.GetPATByID(context.Background(), LockingStrengthNone, userID, patID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, pat)
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getPolicies(ctx context.Context, accountID string) ([]*types.Policy, error) {
|
||||
const query = `SELECT id, account_id, public_id, name, description, enabled, source_posture_checks FROM policies WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
policies, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*types.Policy, error) {
|
||||
var p types.Policy
|
||||
var checks []byte
|
||||
var enabled sql.NullBool
|
||||
err := row.Scan(&p.ID, &p.AccountID, &p.PublicID, &p.Name, &p.Description, &enabled, &checks)
|
||||
if err == nil {
|
||||
if enabled.Valid {
|
||||
p.Enabled = enabled.Bool
|
||||
}
|
||||
if checks != nil {
|
||||
_ = json.Unmarshal(checks, &p.SourcePostureChecks)
|
||||
}
|
||||
}
|
||||
return &p, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return policies, nil
|
||||
}
|
||||
|
||||
// GetAccountPolicies retrieves policies for an account.
|
||||
func (s *SqlStore) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.Policy, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var policies []*types.Policy
|
||||
result := tx.
|
||||
Preload(clause.Associations).Find(&policies, accountIDCondition, accountID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get policies from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get policies from store")
|
||||
}
|
||||
|
||||
return policies, nil
|
||||
}
|
||||
|
||||
// GetPolicyByID retrieves a policy by its ID and account ID.
|
||||
func (s *SqlStore) GetPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*types.Policy, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var policy *types.Policy
|
||||
|
||||
result := tx.Preload(clause.Associations).
|
||||
Take(&policy, accountAndIDQueryCondition, accountID, policyID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPolicyNotFoundError(policyID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get policy from store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get policy from store")
|
||||
}
|
||||
|
||||
return policy, nil
|
||||
}
|
||||
|
||||
// GetPolicyByIDOrPublicID retrieves a policy by either its ID or its PublicID. Peers report
|
||||
// whichever of the two the network map they were served carries, so callers resolving a
|
||||
// peer-reported reference cannot know upfront which namespace it belongs to.
|
||||
func (s *SqlStore) GetPolicyByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*types.Policy, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var policy *types.Policy
|
||||
|
||||
result := tx.Preload(clause.Associations).
|
||||
Take(&policy, accountAndAnyIDQueryCondition, accountID, policyID, policyID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPolicyNotFoundError(policyID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get policy from store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get policy from store")
|
||||
}
|
||||
|
||||
return policy, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) CreatePolicy(ctx context.Context, policy *types.Policy) error {
|
||||
result := s.db.Create(policy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to create policy in store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to create policy in store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// SavePolicy saves a policy to the database.
|
||||
func (s *SqlStore) SavePolicy(ctx context.Context, policy *types.Policy) error {
|
||||
result := s.db.Session(&gorm.Session{FullSaveAssociations: true}).Omit("public_id").Save(policy)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save policy to the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to save policy to store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeletePolicy(ctx context.Context, accountID, policyID string) error {
|
||||
return s.transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("policy_id = ?", policyID).Delete(&types.PolicyRule{}).Error; err != nil {
|
||||
return fmt.Errorf("delete policy rules: %w", err)
|
||||
}
|
||||
|
||||
result := tx.
|
||||
Where(accountAndIDQueryCondition, accountID, policyID).
|
||||
Delete(&types.Policy{})
|
||||
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to delete policy from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewPolicyNotFoundError(policyID)
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getPolicyRules(ctx context.Context, policyIDs []string) ([]*types.PolicyRule, error) {
|
||||
if len(policyIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
const query = `SELECT id, policy_id, name, description, enabled, action, destinations, destination_resource, sources, source_resource, bidirectional, protocol, ports, port_ranges, authorized_groups, authorized_user FROM policy_rules WHERE policy_id = ANY($1)`
|
||||
rows, err := s.pool.Query(ctx, query, policyIDs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rules, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*types.PolicyRule, error) {
|
||||
var r types.PolicyRule
|
||||
var dest, destRes, sources, sourceRes, ports, portRanges, authorizedGroups []byte
|
||||
var enabled, bidirectional sql.NullBool
|
||||
var authorizedUser sql.NullString
|
||||
err := row.Scan(&r.ID, &r.PolicyID, &r.Name, &r.Description, &enabled, &r.Action, &dest, &destRes, &sources, &sourceRes, &bidirectional, &r.Protocol, &ports, &portRanges, &authorizedGroups, &authorizedUser)
|
||||
if err == nil {
|
||||
if enabled.Valid {
|
||||
r.Enabled = enabled.Bool
|
||||
}
|
||||
if bidirectional.Valid {
|
||||
r.Bidirectional = bidirectional.Bool
|
||||
}
|
||||
if dest != nil {
|
||||
_ = json.Unmarshal(dest, &r.Destinations)
|
||||
}
|
||||
if destRes != nil {
|
||||
_ = json.Unmarshal(destRes, &r.DestinationResource)
|
||||
}
|
||||
if sources != nil {
|
||||
_ = json.Unmarshal(sources, &r.Sources)
|
||||
}
|
||||
if sourceRes != nil {
|
||||
_ = json.Unmarshal(sourceRes, &r.SourceResource)
|
||||
}
|
||||
if ports != nil {
|
||||
_ = json.Unmarshal(ports, &r.Ports)
|
||||
}
|
||||
if portRanges != nil {
|
||||
_ = json.Unmarshal(portRanges, &r.PortRanges)
|
||||
}
|
||||
if authorizedGroups != nil {
|
||||
_ = json.Unmarshal(authorizedGroups, &r.AuthorizedGroups)
|
||||
}
|
||||
if authorizedUser.Valid {
|
||||
r.AuthorizedUser = authorizedUser.String
|
||||
}
|
||||
}
|
||||
return &r, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rules, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPolicyRulesByResourceID(ctx context.Context, lockStrength LockingStrength, accountID string, resourceID string) ([]*types.PolicyRule, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var policyRules []*types.PolicyRule
|
||||
resourceIDPattern := `%"ID":"` + resourceID + `"%`
|
||||
result := tx.Where("source_resource LIKE ? OR destination_resource LIKE ?", resourceIDPattern, resourceIDPattern).
|
||||
Find(&policyRules)
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get policy rules for resource id from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get policy rules for resource id from store")
|
||||
}
|
||||
|
||||
return policyRules, nil
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetPolicyByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
policyID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing policy",
|
||||
policyID: "cs1tnh0hhcjnqoiuebf0",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing policy checks",
|
||||
policyID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty policy ID",
|
||||
policyID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
policy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, tt.policyID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, policy)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, policy)
|
||||
require.Equal(t, tt.policyID, policy.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPolicyByIDOrPublicID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
policyID := "cs1tnh0hhcjnqoiuebf0"
|
||||
|
||||
policy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, policyID)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, policy.PublicID)
|
||||
|
||||
for _, id := range []string{policyID, policy.PublicID} {
|
||||
policy, err := store.GetPolicyByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, policyID, policy.ID)
|
||||
}
|
||||
|
||||
policy, err = store.GetPolicyByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, policy)
|
||||
}
|
||||
|
||||
func TestSqlStore_CreatePolicy(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
policy := &types.Policy{
|
||||
ID: "policy-id",
|
||||
AccountID: accountID,
|
||||
Enabled: true,
|
||||
Rules: []*types.PolicyRule{
|
||||
{
|
||||
Enabled: true,
|
||||
Sources: []string{"groupA"},
|
||||
Destinations: []string{"groupC"},
|
||||
Bidirectional: true,
|
||||
Action: types.PolicyTrafficActionAccept,
|
||||
},
|
||||
},
|
||||
}
|
||||
err = store.CreatePolicy(context.Background(), policy)
|
||||
require.NoError(t, err)
|
||||
|
||||
savePolicy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, policy.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, savePolicy, policy)
|
||||
|
||||
}
|
||||
|
||||
func TestSqlStore_SavePolicy(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
policyID := "cs1tnh0hhcjnqoiuebf0"
|
||||
|
||||
policy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, policyID)
|
||||
require.NoError(t, err)
|
||||
|
||||
policy.Enabled = false
|
||||
policy.Description = "policy"
|
||||
policy.Rules[0].Sources = []string{"group"}
|
||||
policy.Rules[0].Ports = []string{"80", "443"}
|
||||
err = store.SavePolicy(context.Background(), policy)
|
||||
require.NoError(t, err)
|
||||
|
||||
savePolicy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, policy.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, savePolicy, policy)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeletePolicy(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
policyID := "cs1tnh0hhcjnqoiuebf0"
|
||||
|
||||
err = store.DeletePolicy(context.Background(), accountID, policyID)
|
||||
require.NoError(t, err)
|
||||
|
||||
policy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, policyID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, policy)
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/posture"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getPostureChecks(ctx context.Context, accountID string) ([]*posture.Checks, error) {
|
||||
const query = `SELECT id, account_id, public_id, name, description, checks FROM posture_checks WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
checks, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*posture.Checks, error) {
|
||||
var c posture.Checks
|
||||
var checksDef []byte
|
||||
err := row.Scan(&c.ID, &c.AccountID, &c.PublicID, &c.Name, &c.Description, &checksDef)
|
||||
if err == nil && checksDef != nil {
|
||||
_ = json.Unmarshal(checksDef, &c.Checks)
|
||||
}
|
||||
return &c, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return checks, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
|
||||
definitionJSON, err := json.Marshal(checks)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var postureCheck posture.Checks
|
||||
err = s.db.Where("account_id = ? AND checks = ?", accountID, string(definitionJSON)).Take(&postureCheck).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &postureCheck, nil
|
||||
}
|
||||
|
||||
// GetAccountPostureChecks retrieves posture checks for an account.
|
||||
func (s *SqlStore) GetAccountPostureChecks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*posture.Checks, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var postureChecks []*posture.Checks
|
||||
result := tx.Find(&postureChecks, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get posture checks from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get posture checks from store")
|
||||
}
|
||||
|
||||
return postureChecks, nil
|
||||
}
|
||||
|
||||
// GetPostureChecksByID retrieves posture checks by their ID and account ID.
|
||||
func (s *SqlStore) GetPostureChecksByID(ctx context.Context, lockStrength LockingStrength, accountID, postureChecksID string) (*posture.Checks, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var postureCheck *posture.Checks
|
||||
result := tx.
|
||||
Take(&postureCheck, accountAndIDQueryCondition, accountID, postureChecksID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPostureChecksNotFoundError(postureChecksID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get posture check from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get posture check from store")
|
||||
}
|
||||
|
||||
return postureCheck, nil
|
||||
}
|
||||
|
||||
// GetPostureChecksByIDs retrieves posture checks by their IDs and account ID.
|
||||
func (s *SqlStore) GetPostureChecksByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, postureChecksIDs []string) (map[string]*posture.Checks, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var postureChecks []*posture.Checks
|
||||
result := tx.Find(&postureChecks, accountAndIDsQueryCondition, accountID, postureChecksIDs)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get posture checks by ID's from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get posture checks by ID's from store")
|
||||
}
|
||||
|
||||
postureChecksMap := make(map[string]*posture.Checks)
|
||||
for _, postureCheck := range postureChecks {
|
||||
postureChecksMap[postureCheck.ID] = postureCheck
|
||||
}
|
||||
|
||||
return postureChecksMap, nil
|
||||
}
|
||||
|
||||
// SavePostureChecks saves a posture checks to the database.
|
||||
func (s *SqlStore) SavePostureChecks(ctx context.Context, postureCheck *posture.Checks) error {
|
||||
result := s.db.Save(postureCheck)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save posture checks to store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save posture checks to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeletePostureChecks deletes a posture checks from the database.
|
||||
func (s *SqlStore) DeletePostureChecks(ctx context.Context, accountID, postureChecksID string) error {
|
||||
result := s.db.Delete(&posture.Checks{}, accountAndIDQueryCondition, accountID, postureChecksID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete posture checks from store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete posture checks from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewPostureChecksNotFoundError(postureChecksID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/posture"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetPostureChecksByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
postureChecksID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing posture checks",
|
||||
postureChecksID: "csplshq7qv948l48f7t0",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing posture checks",
|
||||
postureChecksID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty posture checks ID",
|
||||
postureChecksID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
postureChecks, err := store.GetPostureChecksByID(context.Background(), LockingStrengthNone, accountID, tt.postureChecksID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, postureChecks)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, postureChecks)
|
||||
require.Equal(t, tt.postureChecksID, postureChecks.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetPostureChecksByIDs(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
postureCheckIDs []string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing posture checks by existing IDs",
|
||||
postureCheckIDs: []string{"csplshq7qv948l48f7t0", "cspnllq7qv95uq1r4k90"},
|
||||
expectedCount: 2,
|
||||
},
|
||||
{
|
||||
name: "empty posture check IDs list",
|
||||
postureCheckIDs: []string{},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "non-existing posture check IDs",
|
||||
postureCheckIDs: []string{"nonexistent1", "nonexistent2"},
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "mixed existing and non-existing posture check IDs",
|
||||
postureCheckIDs: []string{"cspnllq7qv95uq1r4k90", "nonexistent"},
|
||||
expectedCount: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
groups, err := store.GetPostureChecksByIDs(context.Background(), LockingStrengthNone, accountID, tt.postureCheckIDs)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, groups, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_SavePostureChecks(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
postureChecks := &posture.Checks{
|
||||
ID: "posture-checks-id",
|
||||
AccountID: accountID,
|
||||
Checks: posture.ChecksDefinition{
|
||||
NBVersionCheck: &posture.NBVersionCheck{
|
||||
MinVersion: "0.31.0",
|
||||
},
|
||||
OSVersionCheck: &posture.OSVersionCheck{
|
||||
Ios: &posture.MinVersionCheck{
|
||||
MinVersion: "13.0.1",
|
||||
},
|
||||
Linux: &posture.MinKernelVersionCheck{
|
||||
MinKernelVersion: "5.3.3-dev",
|
||||
},
|
||||
},
|
||||
GeoLocationCheck: &posture.GeoLocationCheck{
|
||||
Locations: []posture.Location{
|
||||
{
|
||||
CountryCode: "DE",
|
||||
CityName: "Berlin",
|
||||
},
|
||||
},
|
||||
Action: posture.CheckActionAllow,
|
||||
},
|
||||
},
|
||||
}
|
||||
err = store.SavePostureChecks(context.Background(), postureChecks)
|
||||
require.NoError(t, err)
|
||||
|
||||
savePostureChecks, err := store.GetPostureChecksByID(context.Background(), LockingStrengthNone, accountID, "posture-checks-id")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, savePostureChecks, postureChecks)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeletePostureChecks(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
postureChecksID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "delete existing posture checks",
|
||||
postureChecksID: "csplshq7qv948l48f7t0",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "delete non-existing posture checks",
|
||||
postureChecksID: "non-existing-posture-checks-id",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "delete with empty posture checks ID",
|
||||
postureChecksID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err = store.DeletePostureChecks(context.Background(), accountID, tt.postureChecksID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
group, err := store.GetPostureChecksByID(context.Background(), LockingStrengthNone, accountID, tt.postureChecksID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, group)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// GetProxyMetrics aggregates per-cluster + per-proxy counts for the
|
||||
// self-hosted telemetry payload. Single round-trip via conditional
|
||||
// aggregations so a large proxies table doesn't fan out into multiple
|
||||
// queries.
|
||||
func (s *SqlStore) GetProxyMetrics(ctx context.Context) (ProxyMetrics, error) {
|
||||
var m ProxyMetrics
|
||||
activeCutoff := time.Now().Add(-proxyActiveThreshold)
|
||||
|
||||
// COUNT(DISTINCT ... CASE WHEN ...) is portable across sqlite/postgres
|
||||
// (MySQL too) and keeps the round-trip to one. proxy.StatusConnected
|
||||
// is the same string the cluster-capability queries use; the active
|
||||
// window matches the cluster-capability semantics (only proxies
|
||||
// heartbeating within ~2 * heartbeat interval count as connected).
|
||||
row := s.db.WithContext(ctx).
|
||||
Model(&proxy.Proxy{}).
|
||||
Select(
|
||||
"COUNT(DISTINCT cluster_address) AS clusters, "+
|
||||
"COUNT(DISTINCT CASE WHEN account_id IS NOT NULL THEN cluster_address END) AS clusters_byop, "+
|
||||
"COUNT(DISTINCT CASE WHEN private = ? THEN cluster_address END) AS clusters_private, "+
|
||||
"COUNT(*) AS proxies, "+
|
||||
"COUNT(CASE WHEN status = ? AND last_seen > ? THEN 1 END) AS proxies_connected",
|
||||
true,
|
||||
proxy.StatusConnected,
|
||||
activeCutoff,
|
||||
).
|
||||
Row()
|
||||
if err := row.Scan(&m.Clusters, &m.ClustersBYOP, &m.ClustersPrivate, &m.Proxies, &m.ProxiesConnected); err != nil {
|
||||
return ProxyMetrics{}, fmt.Errorf("scan proxy metrics: %w", err)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// SaveProxy saves or updates a proxy in the database
|
||||
func (s *SqlStore) SaveProxy(ctx context.Context, p *proxy.Proxy) error {
|
||||
result := s.db.Save(p)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save proxy: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save proxy")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DisconnectProxy marks a proxy as disconnected only if the session ID matches.
|
||||
// This prevents a slow-to-close old session from overwriting a newer reconnection.
|
||||
func (s *SqlStore) DisconnectProxy(ctx context.Context, proxyID, sessionID string) error {
|
||||
now := time.Now()
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("id = ? AND session_id = ?", proxyID, sessionID).
|
||||
Updates(map[string]any{
|
||||
"status": proxy.StatusDisconnected,
|
||||
"disconnected_at": now,
|
||||
"last_seen": now,
|
||||
})
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to disconnect proxy %s session %s: %v", proxyID, sessionID, result.Error)
|
||||
return status.Errorf(status.Internal, "failed to disconnect proxy")
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
log.WithContext(ctx).Debugf("proxy %s session %s: no row updated (superseded by newer session)", proxyID, sessionID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAllProxies returns all reverse proxy instance rows.
|
||||
func (s *SqlStore) GetAllProxies(ctx context.Context) ([]*proxy.Proxy, error) {
|
||||
var proxies []*proxy.Proxy
|
||||
result := s.db.Order("cluster_address, id").Find(&proxies)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get proxies: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get proxies")
|
||||
}
|
||||
return proxies, nil
|
||||
}
|
||||
|
||||
// DisconnectAllProxies force-marks every proxy that is not already disconnected
|
||||
// as disconnected, regardless of session ID. Unlike DisconnectProxy it is not
|
||||
// session-guarded: it is an administrative repair helper, not part of the
|
||||
// connection lifecycle. last_seen is left untouched so the stale-proxy reaper
|
||||
// keeps working off the real last heartbeat. Returns the number of proxies updated.
|
||||
func (s *SqlStore) DisconnectAllProxies(ctx context.Context) (int64, error) {
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("status != ?", proxy.StatusDisconnected).
|
||||
Updates(map[string]any{
|
||||
"status": proxy.StatusDisconnected,
|
||||
"disconnected_at": time.Now(),
|
||||
})
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to disconnect all proxies: %v", result.Error)
|
||||
return 0, status.Errorf(status.Internal, "failed to disconnect all proxies")
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
// UpdateProxyHeartbeat updates the last_seen timestamp for the proxy's current session.
|
||||
func (s *SqlStore) UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error {
|
||||
now := time.Now()
|
||||
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("id = ? AND session_id = ?", p.ID, p.SessionID).
|
||||
Updates(map[string]any{
|
||||
"last_seen": now,
|
||||
"status": proxy.StatusConnected,
|
||||
"disconnected_at": nil,
|
||||
})
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to update proxy heartbeat: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to update proxy heartbeat")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
p.LastSeen = now
|
||||
p.ConnectedAt = &now
|
||||
p.Status = proxy.StatusConnected
|
||||
if err := s.db.Create(p).Error; err != nil {
|
||||
log.WithContext(ctx).Debugf("proxy %s session %s: heartbeat fallback insert skipped: %v", p.ID, p.SessionID, err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetActiveProxyClusterAddresses returns the unique cluster addresses of active
|
||||
// shared proxies (those without an account scope). BYOP cluster addresses are
|
||||
// excluded; use GetActiveProxyClusterAddressesForAccount to retrieve them.
|
||||
func (s *SqlStore) GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error) {
|
||||
var addresses []string
|
||||
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("account_id IS NULL AND status = ? AND last_seen > ?", proxy.StatusConnected, time.Now().Add(-proxyActiveThreshold)).
|
||||
Distinct("cluster_address").
|
||||
Pluck("cluster_address", &addresses)
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get active proxy cluster addresses: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get active proxy cluster addresses")
|
||||
}
|
||||
|
||||
return addresses, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetActiveProxyClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error) {
|
||||
var addresses []string
|
||||
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("account_id = ? AND status = ? AND last_seen > ?", accountID, proxy.StatusConnected, time.Now().Add(-proxyActiveThreshold)).
|
||||
Distinct("cluster_address").
|
||||
Pluck("cluster_address", &addresses)
|
||||
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "failed to get active proxy cluster addresses for account")
|
||||
}
|
||||
|
||||
return addresses, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetProxyByAccountID(ctx context.Context, accountID string) (*proxy.Proxy, error) {
|
||||
var p proxy.Proxy
|
||||
result := s.db.Where("account_id = ?", accountID).Take(&p)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "proxy not found for account")
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "get proxy by account ID: %v", result.Error)
|
||||
}
|
||||
return &p, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error) {
|
||||
var count int64
|
||||
result := s.db.Model(&proxy.Proxy{}).Where("account_id = ?", accountID).Count(&count)
|
||||
if result.Error != nil {
|
||||
return 0, status.Errorf(status.Internal, "count proxies by account ID: %v", result.Error)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// HasActiveProxyAtClusterAddress reports whether any proxy — shared or
|
||||
// account-scoped — is currently active at the given cluster address, using
|
||||
// the same connected-within-threshold window as the other active-proxy
|
||||
// queries. Backs the agent-network settings delete guard: settings cannot be
|
||||
// deleted while a proxy declares the endpoint hostname as its address.
|
||||
//
|
||||
// The comparison folds case on both sides: the caller passes a normalized
|
||||
// (lowercase) hostname, but proxies declare their cluster address verbatim
|
||||
// and Connect stores it unchanged, so on case-sensitive collations a proxy
|
||||
// declaring "GW.Example.com" would otherwise slip past the guard. Hostnames
|
||||
// are case-insensitive per RFC 4343; the guard must be too.
|
||||
func (s *SqlStore) HasActiveProxyAtClusterAddress(ctx context.Context, clusterAddress string) (bool, error) {
|
||||
var count int64
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("LOWER(cluster_address) = LOWER(?) AND status = ? AND last_seen > ?", clusterAddress, proxy.StatusConnected, time.Now().Add(-proxyActiveThreshold)).
|
||||
Count(&count)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count active proxies at cluster address: %v", result.Error)
|
||||
return false, status.Errorf(status.Internal, "failed to count active proxies at cluster address")
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error) {
|
||||
var count int64
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("cluster_address = ? AND (account_id IS NULL OR account_id != ?)", clusterAddress, accountID).
|
||||
Count(&count)
|
||||
if result.Error != nil {
|
||||
return false, status.Errorf(status.Internal, "check cluster address conflict: %v", result.Error)
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
// HasForeignAccountProxyAtHost reports whether a proxy owned by a different
|
||||
// account declares this host. Shared proxies (account_id IS NULL) are not
|
||||
// foreign: a shared cluster is what most accounts pin their agent network
|
||||
// gateway to. The match folds case because proxies declare their address as
|
||||
// the operator spelled it while the caller's host is normalised; that costs a
|
||||
// scan of the proxies table, taken once per account when its gateway is
|
||||
// bootstrapped, not on the per-connect path IsClusterAddressConflicting serves.
|
||||
func (s *SqlStore) HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error) {
|
||||
var count int64
|
||||
result := s.db.
|
||||
Model(&proxy.Proxy{}).
|
||||
Where("LOWER(cluster_address) = LOWER(?) AND account_id IS NOT NULL AND account_id != ?", host, accountID).
|
||||
Count(&count)
|
||||
if result.Error != nil {
|
||||
return false, status.Errorf(status.Internal, "check proxy host ownership: %v", result.Error)
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error {
|
||||
result := s.db.
|
||||
Where("cluster_address = ? AND account_id = ?", clusterAddress, accountID).
|
||||
Delete(&proxy.Proxy{})
|
||||
if result.Error != nil {
|
||||
return status.Errorf(status.Internal, "delete account cluster: %v", result.Error)
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "cluster not found")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetProxyClusters returns every cluster the account can see (shared
|
||||
// plus its own BYOP), regardless of whether any proxy in the cluster
|
||||
// is currently heartbeating. Online and ConnectedProxies are derived
|
||||
// from the 2-min active window so the dashboard can render offline
|
||||
// clusters distinctly; the 1-hour heartbeat reaper still removes rows
|
||||
// that go quiet for too long.
|
||||
//
|
||||
// AccountOwned is determined by whether any proxy row in the group
|
||||
// carries a non-NULL account_id; the caller maps that to Cluster.Type.
|
||||
// Capability flags are NOT filled here — the handler enriches them via
|
||||
// the per-cluster capability lookups.
|
||||
func (s *SqlStore) GetProxyClusters(ctx context.Context, accountID string) ([]proxy.Cluster, error) {
|
||||
activeCutoff := time.Now().Add(-proxyActiveThreshold)
|
||||
|
||||
type clusterRow struct {
|
||||
ID string
|
||||
Address string
|
||||
ConnectedProxies int
|
||||
Online bool
|
||||
AccountOwned bool
|
||||
}
|
||||
|
||||
var rows []clusterRow
|
||||
result := s.db.Model(&proxy.Proxy{}).
|
||||
Select(
|
||||
"MIN(id) AS id, "+
|
||||
"cluster_address AS address, "+
|
||||
// COUNT(CASE WHEN ... THEN 1 END) counts only non-NULL — i.e. only
|
||||
// rows that satisfy the predicate — so it works portably across
|
||||
// sqlite/postgres/mysql without dialect-specific FILTER syntax.
|
||||
"COUNT(CASE WHEN status = ? AND last_seen > ? THEN 1 END) AS connected_proxies, "+
|
||||
// MAX(CASE …) > 0 expresses BOOL_OR in a way Postgres tolerates
|
||||
// (Postgres can't MAX a boolean column).
|
||||
"MAX(CASE WHEN status = ? AND last_seen > ? THEN 1 ELSE 0 END) > 0 AS online, "+
|
||||
"MAX(CASE WHEN account_id IS NOT NULL THEN 1 ELSE 0 END) > 0 AS account_owned",
|
||||
proxy.StatusConnected, activeCutoff,
|
||||
proxy.StatusConnected, activeCutoff,
|
||||
).
|
||||
Where("account_id IS NULL OR account_id = ?", accountID).
|
||||
Group("cluster_address").
|
||||
Scan(&rows)
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get proxy clusters: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "get proxy clusters")
|
||||
}
|
||||
|
||||
clusters := make([]proxy.Cluster, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
c := proxy.Cluster{
|
||||
ID: r.ID,
|
||||
Address: r.Address,
|
||||
Online: r.Online,
|
||||
ConnectedProxies: r.ConnectedProxies,
|
||||
}
|
||||
if r.AccountOwned {
|
||||
c.Type = proxy.ClusterTypeAccount
|
||||
} else {
|
||||
c.Type = proxy.ClusterTypeShared
|
||||
}
|
||||
clusters = append(clusters, c)
|
||||
}
|
||||
|
||||
return clusters, nil
|
||||
}
|
||||
|
||||
// proxyActiveThreshold is the maximum age of a heartbeat for a proxy to be
|
||||
// considered active. Must be at least 2x the heartbeat interval (1 min).
|
||||
const proxyActiveThreshold = 2 * time.Minute
|
||||
|
||||
var validCapabilityColumns = map[string]struct{}{
|
||||
"supports_custom_ports": {},
|
||||
"require_subdomain": {},
|
||||
"supports_crowdsec": {},
|
||||
"private": {},
|
||||
}
|
||||
|
||||
// GetClusterSupportsCustomPorts returns whether any active proxy in the cluster
|
||||
// supports custom ports. Returns nil when no proxy reported the capability.
|
||||
func (s *SqlStore) GetClusterSupportsCustomPorts(ctx context.Context, clusterAddr string) *bool {
|
||||
return s.getClusterCapability(ctx, clusterAddr, "supports_custom_ports")
|
||||
}
|
||||
|
||||
// GetClusterRequireSubdomain returns whether any active proxy in the cluster
|
||||
// requires a subdomain. Returns nil when no proxy reported the capability.
|
||||
func (s *SqlStore) GetClusterRequireSubdomain(ctx context.Context, clusterAddr string) *bool {
|
||||
return s.getClusterCapability(ctx, clusterAddr, "require_subdomain")
|
||||
}
|
||||
|
||||
// GetClusterSupportsPrivate reports whether any active proxy in the cluster
|
||||
// has the private capability (nil = unreported).
|
||||
func (s *SqlStore) GetClusterSupportsPrivate(ctx context.Context, clusterAddr string) *bool {
|
||||
return s.getClusterCapability(ctx, clusterAddr, "private")
|
||||
}
|
||||
|
||||
// GetClusterSupportsCrowdSec returns whether all active proxies in the cluster
|
||||
// have CrowdSec configured. Returns nil when no proxy reported the capability.
|
||||
// Unlike other capabilities that use ANY-true (for rolling upgrades), CrowdSec
|
||||
// requires unanimous support: a single unconfigured proxy would let requests
|
||||
// bypass reputation checks.
|
||||
func (s *SqlStore) GetClusterSupportsCrowdSec(ctx context.Context, clusterAddr string) *bool {
|
||||
return s.getClusterUnanimousCapability(ctx, clusterAddr, "supports_crowdsec")
|
||||
}
|
||||
|
||||
// getClusterUnanimousCapability returns an aggregated boolean capability
|
||||
// requiring all active proxies in the cluster to report true.
|
||||
func (s *SqlStore) getClusterUnanimousCapability(ctx context.Context, clusterAddr, column string) *bool {
|
||||
if _, ok := validCapabilityColumns[column]; !ok {
|
||||
log.WithContext(ctx).Errorf("invalid capability column: %s", column)
|
||||
return nil
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Total int64
|
||||
Reported int64
|
||||
AllTrue bool
|
||||
}
|
||||
|
||||
// All active proxies must have reported the capability (no NULLs) and all
|
||||
// must report true. A single unreported or false proxy means the cluster
|
||||
// does not unanimously support the capability.
|
||||
err := s.db.WithContext(ctx).
|
||||
Model(&proxy.Proxy{}).
|
||||
Select("COUNT(*) AS total, "+
|
||||
"COUNT(CASE WHEN "+column+" IS NOT NULL THEN 1 END) AS reported, "+
|
||||
"COUNT(*) > 0 AND COUNT(*) = COUNT(CASE WHEN "+column+" = true THEN 1 END) AS all_true").
|
||||
Where("cluster_address = ? AND status = ? AND last_seen > ?",
|
||||
clusterAddr, "connected", time.Now().Add(-proxyActiveThreshold)).
|
||||
Scan(&result).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("query cluster capability %s for %s: %v", column, clusterAddr, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
if result.Total == 0 || result.Reported == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// If any proxy has not reported (NULL), we can't confirm unanimous support.
|
||||
if result.Reported < result.Total {
|
||||
v := false
|
||||
return &v
|
||||
}
|
||||
|
||||
return &result.AllTrue
|
||||
}
|
||||
|
||||
// getClusterCapability returns an aggregated boolean capability for the given
|
||||
// cluster. It checks active (connected, recently seen) proxies and returns:
|
||||
// - *true if any proxy in the cluster has the capability set to true,
|
||||
// - *false if at least one proxy reported but none set it to true,
|
||||
// - nil if no proxy reported the capability at all.
|
||||
func (s *SqlStore) getClusterCapability(ctx context.Context, clusterAddr, column string) *bool {
|
||||
if _, ok := validCapabilityColumns[column]; !ok {
|
||||
log.WithContext(ctx).Errorf("invalid capability column: %s", column)
|
||||
return nil
|
||||
}
|
||||
|
||||
var result struct {
|
||||
HasCapability bool
|
||||
AnyTrue bool
|
||||
}
|
||||
|
||||
err := s.db.
|
||||
WithContext(ctx).
|
||||
Model(&proxy.Proxy{}).
|
||||
Select("COUNT(CASE WHEN "+column+" IS NOT NULL THEN 1 END) > 0 AS has_capability, "+
|
||||
"COALESCE(MAX(CASE WHEN "+column+" = true THEN 1 ELSE 0 END), 0) = 1 AS any_true").
|
||||
Where("cluster_address = ? AND status = ? AND last_seen > ?",
|
||||
clusterAddr, "connected", time.Now().Add(-proxyActiveThreshold)).
|
||||
Scan(&result).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("query cluster capability %s for %s: %v", column, clusterAddr, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
if !result.HasCapability {
|
||||
return nil
|
||||
}
|
||||
|
||||
return &result.AnyTrue
|
||||
}
|
||||
|
||||
// CleanupStaleProxies deletes proxies that haven't sent heartbeat in the specified duration
|
||||
func (s *SqlStore) CleanupStaleProxies(ctx context.Context, inactivityDuration time.Duration) error {
|
||||
cutoffTime := time.Now().Add(-inactivityDuration)
|
||||
|
||||
result := s.db.
|
||||
Where("last_seen < ?", cutoffTime).
|
||||
Delete(&proxy.Proxy{})
|
||||
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to cleanup stale proxies: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to cleanup stale proxies")
|
||||
}
|
||||
|
||||
if result.RowsAffected > 0 {
|
||||
log.WithContext(ctx).Infof("Cleaned up %d stale proxies", result.RowsAffected)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// GetProxyAccessTokenByHashedToken retrieves a proxy access token by its hashed value.
|
||||
func (s *SqlStore) GetProxyAccessTokenByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken types.HashedProxyToken) (*types.ProxyAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var token types.ProxyAccessToken
|
||||
result := tx.Take(&token, "hashed_token = ?", hashedToken)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "proxy access token not found")
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "get proxy access token: %v", result.Error)
|
||||
}
|
||||
|
||||
return &token, nil
|
||||
}
|
||||
|
||||
// GetAllProxyAccessTokens retrieves all proxy access tokens.
|
||||
func (s *SqlStore) GetAllProxyAccessTokens(ctx context.Context, lockStrength LockingStrength) ([]*types.ProxyAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var tokens []*types.ProxyAccessToken
|
||||
result := tx.Find(&tokens)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "get proxy access tokens: %v", result.Error)
|
||||
}
|
||||
|
||||
return tokens, nil
|
||||
}
|
||||
|
||||
// SaveProxyAccessToken saves a proxy access token to the database.
|
||||
func (s *SqlStore) SaveProxyAccessToken(ctx context.Context, token *types.ProxyAccessToken) error {
|
||||
if result := s.db.Create(token); result.Error != nil {
|
||||
return status.Errorf(status.Internal, "save proxy access token: %v", result.Error)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RevokeProxyAccessToken revokes a proxy access token by its ID.
|
||||
func (s *SqlStore) RevokeProxyAccessToken(ctx context.Context, tokenID string) error {
|
||||
result := s.db.Model(&types.ProxyAccessToken{}).Where(idQueryCondition, tokenID).Update("revoked", true)
|
||||
if result.Error != nil {
|
||||
return status.Errorf(status.Internal, "revoke proxy access token: %v", result.Error)
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "proxy access token not found")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetProxyAccessTokensByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.ProxyAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var tokens []*types.ProxyAccessToken
|
||||
result := tx.Where("account_id = ?", accountID).Find(&tokens)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "get proxy access tokens by account: %v", result.Error)
|
||||
}
|
||||
|
||||
return tokens, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) IsProxyAccessTokenValid(ctx context.Context, tokenID string) (bool, error) {
|
||||
token, err := s.GetProxyAccessTokenByID(ctx, LockingStrengthNone, tokenID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return token.IsValid(), nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetProxyAccessTokenByID(ctx context.Context, lockStrength LockingStrength, tokenID string) (*types.ProxyAccessToken, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var token types.ProxyAccessToken
|
||||
result := tx.Take(&token, idQueryCondition, tokenID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "proxy access token not found")
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "get proxy access token by ID: %v", result.Error)
|
||||
}
|
||||
|
||||
return &token, nil
|
||||
}
|
||||
|
||||
// MarkProxyAccessTokenUsed updates the last used timestamp for a proxy access token.
|
||||
func (s *SqlStore) MarkProxyAccessTokenUsed(ctx context.Context, tokenID string) error {
|
||||
result := s.db.Model(&types.ProxyAccessToken{}).
|
||||
Where(idQueryCondition, tokenID).
|
||||
Update("last_used", time.Now().UTC())
|
||||
if result.Error != nil {
|
||||
return status.Errorf(status.Internal, "mark proxy access token as used: %v", result.Error)
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "proxy access token not found")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) getRoutes(ctx context.Context, accountID string) ([]route.Route, error) {
|
||||
const query = `SELECT id, account_id, public_id, network, domains, keep_route, net_id, description, peer, peer_groups, network_type, masquerade, metric, enabled, groups, access_control_groups, skip_auto_apply FROM routes WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
routes, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (route.Route, error) {
|
||||
var r route.Route
|
||||
var network, domains, peerGroups, groups, accessGroups []byte
|
||||
var keepRoute, masquerade, enabled, skipAutoApply sql.NullBool
|
||||
var metric sql.NullInt64
|
||||
err := row.Scan(&r.ID, &r.AccountID, &r.PublicID, &network, &domains, &keepRoute, &r.NetID, &r.Description, &r.Peer, &peerGroups, &r.NetworkType, &masquerade, &metric, &enabled, &groups, &accessGroups, &skipAutoApply)
|
||||
if err == nil {
|
||||
if keepRoute.Valid {
|
||||
r.KeepRoute = keepRoute.Bool
|
||||
}
|
||||
if masquerade.Valid {
|
||||
r.Masquerade = masquerade.Bool
|
||||
}
|
||||
if enabled.Valid {
|
||||
r.Enabled = enabled.Bool
|
||||
}
|
||||
if skipAutoApply.Valid {
|
||||
r.SkipAutoApply = skipAutoApply.Bool
|
||||
}
|
||||
if metric.Valid {
|
||||
r.Metric = int(metric.Int64)
|
||||
}
|
||||
if network != nil {
|
||||
_ = json.Unmarshal(network, &r.Network)
|
||||
}
|
||||
if domains != nil {
|
||||
_ = json.Unmarshal(domains, &r.Domains)
|
||||
}
|
||||
if peerGroups != nil {
|
||||
_ = json.Unmarshal(peerGroups, &r.PeerGroups)
|
||||
}
|
||||
if groups != nil {
|
||||
_ = json.Unmarshal(groups, &r.Groups)
|
||||
}
|
||||
if accessGroups != nil {
|
||||
_ = json.Unmarshal(accessGroups, &r.AccessControlGroups)
|
||||
}
|
||||
}
|
||||
return r, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return routes, nil
|
||||
}
|
||||
|
||||
// GetAccountRoutes retrieves network routes for an account.
|
||||
func (s *SqlStore) GetAccountRoutes(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*route.Route, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var routes []*route.Route
|
||||
result := tx.Find(&routes, accountIDCondition, accountID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get routes from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get routes from store")
|
||||
}
|
||||
|
||||
return routes, nil
|
||||
}
|
||||
|
||||
// GetRouteByID retrieves a route by its ID and account ID.
|
||||
func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrength, accountID string, routeID string) (*route.Route, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var route *route.Route
|
||||
result := tx.Take(&route, accountAndIDQueryCondition, accountID, routeID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewRouteNotFoundError(routeID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get route from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get route from store")
|
||||
}
|
||||
|
||||
return route, nil
|
||||
}
|
||||
|
||||
// GetRouteByIDOrPublicID retrieves a route by either its ID or its PublicID. See
|
||||
// GetPolicyByIDOrPublicID for why peer-reported references need both.
|
||||
func (s *SqlStore) GetRouteByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID string, routeID string) (*route.Route, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var route *route.Route
|
||||
result := tx.Take(&route, accountAndAnyIDQueryCondition, accountID, routeID, routeID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewRouteNotFoundError(routeID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get route from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get route from store")
|
||||
}
|
||||
|
||||
return route, nil
|
||||
}
|
||||
|
||||
// SaveRoute saves a route to the database.
|
||||
func (s *SqlStore) SaveRoute(ctx context.Context, route *route.Route) error {
|
||||
result := s.db.Save(route)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save route to the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to save route to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteRoute deletes a route from the database.
|
||||
func (s *SqlStore) DeleteRoute(ctx context.Context, accountID, routeID string) error {
|
||||
result := s.db.Delete(&route.Route{}, accountAndIDQueryCondition, accountID, routeID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete route from the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to delete route from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewRouteNotFoundError(routeID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetAccountRoutes(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "retrieve routes by existing account ID",
|
||||
accountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "non-existing account ID",
|
||||
accountID: "nonexistent",
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "empty account ID",
|
||||
accountID: "",
|
||||
expectedCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
routes, err := store.GetAccountRoutes(context.Background(), LockingStrengthNone, tt.accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, routes, tt.expectedCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetRouteByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
tests := []struct {
|
||||
name string
|
||||
routeID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing route",
|
||||
routeID: "ct03t427qv97vmtmglog",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing route",
|
||||
routeID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty route ID",
|
||||
routeID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
route, err := store.GetRouteByID(context.Background(), LockingStrengthNone, accountID, tt.routeID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, route)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, route)
|
||||
require.Equal(t, tt.routeID, string(route.ID))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetRouteByIDOrPublicID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
routeID := "ct03t427qv97vmtmglog"
|
||||
|
||||
route, err := store.GetRouteByID(context.Background(), LockingStrengthNone, accountID, routeID)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, route.PublicID)
|
||||
|
||||
for _, id := range []string{routeID, route.PublicID} {
|
||||
route, err := store.GetRouteByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, routeID, string(route.ID))
|
||||
}
|
||||
|
||||
route, err = store.GetRouteByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, route)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveRoute(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
route := &nbroute.Route{
|
||||
ID: "route-id",
|
||||
AccountID: accountID,
|
||||
Network: netip.MustParsePrefix("10.10.0.0/16"),
|
||||
NetID: "netID",
|
||||
PeerGroups: []string{"routeA"},
|
||||
NetworkType: nbroute.IPv4Network,
|
||||
Masquerade: true,
|
||||
Metric: 9999,
|
||||
Enabled: true,
|
||||
Groups: []string{"groupA"},
|
||||
AccessControlGroups: []string{},
|
||||
}
|
||||
err = store.SaveRoute(context.Background(), route)
|
||||
require.NoError(t, err)
|
||||
|
||||
saveRoute, err := store.GetRouteByID(context.Background(), LockingStrengthNone, accountID, string(route.ID))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, route, saveRoute)
|
||||
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteRoute(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
routeID := "ct03t427qv97vmtmglog"
|
||||
|
||||
err = store.DeleteRoute(context.Background(), accountID, routeID)
|
||||
require.NoError(t, err)
|
||||
|
||||
route, err := store.GetRouteByID(context.Background(), LockingStrengthNone, accountID, routeID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, route)
|
||||
}
|
||||
@@ -0,0 +1,452 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// serviceSelectColumns and targetSelectColumns are the column lists the Postgres
|
||||
// pgx read path scans. They must stay in sync with the rpservice.Service and
|
||||
// rpservice.Target gorm models; TestPgxServiceColumnsMatchGorm enforces this.
|
||||
const serviceSelectColumns = `id, account_id, name, domain, enabled, auth, restrictions,
|
||||
meta_created_at, meta_certificate_issued_at, meta_last_renewed_at, meta_status, proxy_cluster,
|
||||
pass_host_header, rewrite_redirects, session_private_key, session_public_key,
|
||||
mode, listen_port, port_auto_assigned, source, source_peer, terminated,
|
||||
private, access_groups`
|
||||
|
||||
func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpservice.Service, error) {
|
||||
const serviceQuery = `SELECT ` + serviceSelectColumns + ` FROM services WHERE account_id = $1`
|
||||
|
||||
serviceRows, err := s.pool.Query(ctx, serviceQuery, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
services, err := pgx.CollectRows(serviceRows, scanService)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(services) == 0 {
|
||||
return services, nil
|
||||
}
|
||||
|
||||
serviceIDs := make([]string, len(services))
|
||||
serviceMap := make(map[string]*rpservice.Service)
|
||||
for i, svc := range services {
|
||||
serviceIDs[i] = svc.ID
|
||||
serviceMap[svc.ID] = svc
|
||||
}
|
||||
|
||||
targets, err := s.getServiceTargets(ctx, serviceIDs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for _, target := range targets {
|
||||
if service, ok := serviceMap[target.ServiceID]; ok {
|
||||
service.Targets = append(service.Targets, target)
|
||||
}
|
||||
}
|
||||
|
||||
return services, nil
|
||||
}
|
||||
|
||||
func scanService(row pgx.CollectableRow) (*rpservice.Service, error) {
|
||||
var s rpservice.Service
|
||||
var auth []byte
|
||||
var restrictions []byte
|
||||
var accessGroups []byte
|
||||
var createdAt, certIssuedAt, lastRenewedAt sql.NullTime
|
||||
var status, proxyCluster, sessionPrivateKey, sessionPublicKey sql.NullString
|
||||
var mode, source, sourcePeer sql.NullString
|
||||
var terminated, portAutoAssigned, private sql.NullBool
|
||||
var listenPort sql.NullInt64
|
||||
err := row.Scan(
|
||||
&s.ID,
|
||||
&s.AccountID,
|
||||
&s.Name,
|
||||
&s.Domain,
|
||||
&s.Enabled,
|
||||
&auth,
|
||||
&restrictions,
|
||||
&createdAt,
|
||||
&certIssuedAt,
|
||||
&lastRenewedAt,
|
||||
&status,
|
||||
&proxyCluster,
|
||||
&s.PassHostHeader,
|
||||
&s.RewriteRedirects,
|
||||
&sessionPrivateKey,
|
||||
&sessionPublicKey,
|
||||
&mode,
|
||||
&listenPort,
|
||||
&portAutoAssigned,
|
||||
&source,
|
||||
&sourcePeer,
|
||||
&terminated,
|
||||
&private,
|
||||
&accessGroups,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if auth != nil {
|
||||
if err := json.Unmarshal(auth, &s.Auth); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
if len(restrictions) > 0 {
|
||||
if err := json.Unmarshal(restrictions, &s.Restrictions); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal restrictions: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if len(accessGroups) > 0 {
|
||||
if err := json.Unmarshal(accessGroups, &s.AccessGroups); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal access_groups: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if private.Valid {
|
||||
s.Private = private.Bool
|
||||
}
|
||||
|
||||
s.Meta = serviceMetaFromRow(createdAt, certIssuedAt, lastRenewedAt, status)
|
||||
if proxyCluster.Valid {
|
||||
s.ProxyCluster = proxyCluster.String
|
||||
}
|
||||
if sessionPrivateKey.Valid {
|
||||
s.SessionPrivateKey = sessionPrivateKey.String
|
||||
}
|
||||
if sessionPublicKey.Valid {
|
||||
s.SessionPublicKey = sessionPublicKey.String
|
||||
}
|
||||
if mode.Valid {
|
||||
s.Mode = mode.String
|
||||
}
|
||||
if source.Valid {
|
||||
s.Source = source.String
|
||||
}
|
||||
if sourcePeer.Valid {
|
||||
s.SourcePeer = sourcePeer.String
|
||||
}
|
||||
if terminated.Valid {
|
||||
s.Terminated = terminated.Bool
|
||||
}
|
||||
if portAutoAssigned.Valid {
|
||||
s.PortAutoAssigned = portAutoAssigned.Bool
|
||||
}
|
||||
if listenPort.Valid {
|
||||
if listenPort.Int64 < 0 || listenPort.Int64 > math.MaxUint16 {
|
||||
return nil, fmt.Errorf("listen_port %d out of range", listenPort.Int64)
|
||||
}
|
||||
s.ListenPort = uint16(listenPort.Int64)
|
||||
}
|
||||
s.Targets = []*rpservice.Target{}
|
||||
return &s, nil
|
||||
}
|
||||
|
||||
func serviceMetaFromRow(createdAt, certIssuedAt, lastRenewedAt sql.NullTime, status sql.NullString) rpservice.Meta {
|
||||
meta := rpservice.Meta{}
|
||||
if createdAt.Valid {
|
||||
meta.CreatedAt = createdAt.Time
|
||||
}
|
||||
if certIssuedAt.Valid {
|
||||
t := certIssuedAt.Time
|
||||
meta.CertificateIssuedAt = &t
|
||||
}
|
||||
if lastRenewedAt.Valid {
|
||||
t := lastRenewedAt.Time
|
||||
meta.LastRenewedAt = &t
|
||||
}
|
||||
if status.Valid {
|
||||
meta.Status = status.String
|
||||
}
|
||||
return meta
|
||||
}
|
||||
|
||||
func (s *SqlStore) CreateService(ctx context.Context, service *rpservice.Service) error {
|
||||
serviceCopy := service.Copy()
|
||||
if err := serviceCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return fmt.Errorf("encrypt service data: %w", err)
|
||||
}
|
||||
result := s.db.Create(serviceCopy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to create service to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to create service to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) UpdateService(ctx context.Context, service *rpservice.Service) error {
|
||||
serviceCopy := service.Copy()
|
||||
if err := serviceCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return fmt.Errorf("encrypt service data: %w", err)
|
||||
}
|
||||
|
||||
// Create target type instance outside transaction to avoid variable shadowing
|
||||
targetType := &rpservice.Target{}
|
||||
|
||||
// Use a transaction to ensure atomic updates of the service and its targets
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Delete existing targets
|
||||
if err := tx.Where("service_id = ?", serviceCopy.ID).Delete(targetType).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Update the service and create new targets
|
||||
if err := tx.Session(&gorm.Session{FullSaveAssociations: true}).Save(serviceCopy).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to update service to store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to update service to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteService(ctx context.Context, accountID, serviceID string) error {
|
||||
result := s.db.Delete(&rpservice.Service{}, accountAndIDQueryCondition, accountID, serviceID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete service from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete service from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "service %s not found", serviceID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetServiceByID(ctx context.Context, lockStrength LockingStrength, accountID, serviceID string) (*rpservice.Service, error) {
|
||||
tx := s.db.Preload("Targets")
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var service *rpservice.Service
|
||||
result := tx.Take(&service, accountAndIDQueryCondition, accountID, serviceID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "service %s not found", serviceID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get service from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get service from store")
|
||||
}
|
||||
|
||||
if err := service.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt service data: %w", err)
|
||||
}
|
||||
|
||||
return service, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetServiceByDomain(ctx context.Context, domain string) (*rpservice.Service, error) {
|
||||
var service *rpservice.Service
|
||||
result := s.db.Preload("Targets").Where("domain = ?", domain).First(&service)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "service with domain %s not found", domain)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get service by domain from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get service by domain from store")
|
||||
}
|
||||
|
||||
if err := service.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt service data: %w", err)
|
||||
}
|
||||
|
||||
return service, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetServices(ctx context.Context, lockStrength LockingStrength) ([]*rpservice.Service, error) {
|
||||
tx := s.db.Preload("Targets")
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var serviceList []*rpservice.Service
|
||||
result := tx.Find(&serviceList)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get services from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get services from store")
|
||||
}
|
||||
|
||||
for _, service := range serviceList {
|
||||
if err := service.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt service data: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return serviceList, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountServices(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*rpservice.Service, error) {
|
||||
tx := s.db.Preload("Targets")
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var serviceList []*rpservice.Service
|
||||
result := tx.Find(&serviceList, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get services from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get services from store")
|
||||
}
|
||||
|
||||
for _, service := range serviceList {
|
||||
if err := service.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt service data: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return serviceList, nil
|
||||
}
|
||||
|
||||
// RenewEphemeralService updates the last_renewed_at timestamp for an ephemeral service.
|
||||
func (s *SqlStore) RenewEphemeralService(ctx context.Context, accountID, peerID, serviceID string) error {
|
||||
result := s.db.Model(&rpservice.Service{}).
|
||||
Where("id = ? AND account_id = ? AND source_peer = ? AND source = ?", serviceID, accountID, peerID, rpservice.SourceEphemeral).
|
||||
Update("meta_last_renewed_at", time.Now())
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to renew ephemeral service: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "renew ephemeral service")
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "no active expose session for service %s", serviceID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetExpiredEphemeralServices returns ephemeral services whose last renewal exceeds the given TTL.
|
||||
// Only the fields needed for reaping are selected. The limit parameter caps the batch size to
|
||||
// avoid loading too many rows in a single tick. Rows with empty source_peer are excluded to
|
||||
// skip malformed legacy data.
|
||||
func (s *SqlStore) GetExpiredEphemeralServices(ctx context.Context, ttl time.Duration, limit int) ([]*rpservice.Service, error) {
|
||||
cutoff := time.Now().Add(-ttl)
|
||||
var services []*rpservice.Service
|
||||
result := s.db.
|
||||
Select("id", "account_id", "source_peer", "domain").
|
||||
Where("source = ? AND source_peer <> '' AND meta_last_renewed_at < ?", rpservice.SourceEphemeral, cutoff).
|
||||
Limit(limit).
|
||||
Find(&services)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get expired ephemeral services: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "get expired ephemeral services")
|
||||
}
|
||||
return services, nil
|
||||
}
|
||||
|
||||
// CountEphemeralServicesByPeer returns the count of ephemeral services for a specific peer.
|
||||
// Use LockingStrengthUpdate inside a transaction to serialize concurrent create operations.
|
||||
// The locking is applied via a row-level SELECT ... FOR UPDATE (not on the aggregate) to
|
||||
// stay compatible with Postgres, which disallows FOR UPDATE on COUNT(*).
|
||||
func (s *SqlStore) CountEphemeralServicesByPeer(ctx context.Context, lockStrength LockingStrength, accountID, peerID string) (int64, error) {
|
||||
if lockStrength == LockingStrengthNone {
|
||||
var count int64
|
||||
result := s.db.Model(&rpservice.Service{}).
|
||||
Where("account_id = ? AND source_peer = ? AND source = ?", accountID, peerID, rpservice.SourceEphemeral).
|
||||
Count(&count)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count ephemeral services: %v", result.Error)
|
||||
return 0, status.Errorf(status.Internal, "count ephemeral services")
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
var ids []string
|
||||
result := s.db.Model(&rpservice.Service{}).
|
||||
Clauses(clause.Locking{Strength: string(lockStrength)}).
|
||||
Select("id").
|
||||
Where("account_id = ? AND source_peer = ? AND source = ?", accountID, peerID, rpservice.SourceEphemeral).
|
||||
Pluck("id", &ids)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count ephemeral services: %v", result.Error)
|
||||
return 0, status.Errorf(status.Internal, "count ephemeral services")
|
||||
}
|
||||
return int64(len(ids)), nil
|
||||
}
|
||||
|
||||
// EphemeralServiceExists checks if an ephemeral service exists for the given peer and domain.
|
||||
// Use LockingStrengthUpdate inside a transaction to serialize concurrent create operations.
|
||||
func (s *SqlStore) EphemeralServiceExists(ctx context.Context, lockStrength LockingStrength, accountID, peerID, domain string) (bool, error) {
|
||||
if lockStrength == LockingStrengthNone {
|
||||
var count int64
|
||||
result := s.db.Model(&rpservice.Service{}).
|
||||
Where("account_id = ? AND source_peer = ? AND domain = ? AND source = ?", accountID, peerID, domain, rpservice.SourceEphemeral).
|
||||
Count(&count)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to check ephemeral service existence: %v", result.Error)
|
||||
return false, status.Errorf(status.Internal, "check ephemeral service existence")
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
var id string
|
||||
result := s.db.Model(&rpservice.Service{}).
|
||||
Clauses(clause.Locking{Strength: string(lockStrength)}).
|
||||
Select("id").
|
||||
Where("account_id = ? AND source_peer = ? AND domain = ? AND source = ?", accountID, peerID, domain, rpservice.SourceEphemeral).
|
||||
Limit(1).
|
||||
Pluck("id", &id)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to check ephemeral service existence: %v", result.Error)
|
||||
return false, status.Errorf(status.Internal, "check ephemeral service existence")
|
||||
}
|
||||
return id != "", nil
|
||||
}
|
||||
|
||||
// GetServicesByClusterAndPort returns services matching the given proxy cluster, mode, and listen port.
|
||||
func (s *SqlStore) GetServicesByClusterAndPort(ctx context.Context, lockStrength LockingStrength, proxyCluster string, mode string, listenPort uint16) ([]*rpservice.Service, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var services []*rpservice.Service
|
||||
result := tx.Where("proxy_cluster = ? AND mode = ? AND listen_port = ?", proxyCluster, mode, listenPort).Find(&services)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "query services by cluster and port")
|
||||
}
|
||||
|
||||
return services, nil
|
||||
}
|
||||
|
||||
// GetServicesByCluster returns all services for the given proxy cluster.
|
||||
func (s *SqlStore) GetServicesByCluster(ctx context.Context, lockStrength LockingStrength, proxyCluster string) ([]*rpservice.Service, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var services []*rpservice.Service
|
||||
result := tx.Where("proxy_cluster = ?", proxyCluster).Find(&services)
|
||||
if result.Error != nil {
|
||||
return nil, status.Errorf(status.Internal, "query services by cluster")
|
||||
}
|
||||
return services, nil
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
const targetSelectColumns = `id, account_id, service_id, path, host, port, protocol,
|
||||
target_id, target_type, enabled, proxy_protocol,
|
||||
skip_tls_verify, request_timeout, session_idle_timeout, path_rewrite, custom_headers,
|
||||
direct_upstream, middlewares, capture_max_request_bytes, capture_max_response_bytes,
|
||||
capture_content_types, agent_network, disable_access_log`
|
||||
|
||||
func (s *SqlStore) getServiceTargets(ctx context.Context, serviceIDs []string) ([]*rpservice.Target, error) {
|
||||
const targetsQuery = `SELECT ` + targetSelectColumns + ` FROM targets WHERE service_id = ANY($1)`
|
||||
|
||||
rows, err := s.pool.Query(ctx, targetsQuery, serviceIDs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return pgx.CollectRows(rows, scanTarget)
|
||||
}
|
||||
|
||||
func scanTarget(row pgx.CollectableRow) (*rpservice.Target, error) {
|
||||
var t rpservice.Target
|
||||
var path sql.NullString
|
||||
var pathRewrite sql.NullString
|
||||
var proxyProtocol, skipTLSVerify, directUpstream, agentNetwork, disableAccessLog sql.NullBool
|
||||
var requestTimeout, sessionIdleTimeout, captureMaxRequestBytes, captureMaxResponseBytes sql.NullInt64
|
||||
var customHeaders, middlewares, captureContentTypes []byte
|
||||
err := row.Scan(
|
||||
&t.ID,
|
||||
&t.AccountID,
|
||||
&t.ServiceID,
|
||||
&path,
|
||||
&t.Host,
|
||||
&t.Port,
|
||||
&t.Protocol,
|
||||
&t.TargetId,
|
||||
&t.TargetType,
|
||||
&t.Enabled,
|
||||
&proxyProtocol,
|
||||
&skipTLSVerify,
|
||||
&requestTimeout,
|
||||
&sessionIdleTimeout,
|
||||
&pathRewrite,
|
||||
&customHeaders,
|
||||
&directUpstream,
|
||||
&middlewares,
|
||||
&captureMaxRequestBytes,
|
||||
&captureMaxResponseBytes,
|
||||
&captureContentTypes,
|
||||
&agentNetwork,
|
||||
&disableAccessLog,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if path.Valid {
|
||||
t.Path = &path.String
|
||||
}
|
||||
|
||||
t.ProxyProtocol = proxyProtocol.Bool
|
||||
t.Options.SkipTLSVerify = skipTLSVerify.Bool
|
||||
t.Options.RequestTimeout = time.Duration(requestTimeout.Int64)
|
||||
t.Options.SessionIdleTimeout = time.Duration(sessionIdleTimeout.Int64)
|
||||
t.Options.PathRewrite = rpservice.PathRewriteMode(pathRewrite.String)
|
||||
t.Options.DirectUpstream = directUpstream.Bool
|
||||
t.Options.CaptureMaxRequestBytes = captureMaxRequestBytes.Int64
|
||||
t.Options.CaptureMaxResponseBytes = captureMaxResponseBytes.Int64
|
||||
t.Options.AgentNetwork = agentNetwork.Bool
|
||||
t.Options.DisableAccessLog = disableAccessLog.Bool
|
||||
|
||||
if len(customHeaders) > 0 {
|
||||
if err := json.Unmarshal(customHeaders, &t.Options.CustomHeaders); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal custom_headers: %w", err)
|
||||
}
|
||||
}
|
||||
if len(middlewares) > 0 {
|
||||
if err := json.Unmarshal(middlewares, &t.Options.Middlewares); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal middlewares: %w", err)
|
||||
}
|
||||
}
|
||||
if len(captureContentTypes) > 0 {
|
||||
if err := json.Unmarshal(captureContentTypes, &t.Options.CaptureContentTypes); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal capture_content_types: %w", err)
|
||||
}
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteTarget(ctx context.Context, accountID string, serviceID string, targetID uint) error {
|
||||
result := s.db.Delete(&rpservice.Target{}, "account_id = ? AND service_id = ? AND id = ?", accountID, serviceID, targetID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete target from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete target from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.Errorf(status.NotFound, "target not found for service %s", serviceID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteServiceTargets(ctx context.Context, accountID string, serviceID string) error {
|
||||
result := s.db.Delete(&rpservice.Target{}, "account_id = ? AND service_id = ?", accountID, serviceID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete targets from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete targets from store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetTargetsByServiceID retrieves all targets for a given service
|
||||
func (s *SqlStore) GetTargetsByServiceID(ctx context.Context, lockStrength LockingStrength, accountID string, serviceID string) ([]*rpservice.Target, error) {
|
||||
var targets []*rpservice.Target
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
result := tx.Where("account_id = ? AND service_id = ?", accountID, serviceID).Find(&targets)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get targets from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get targets from store")
|
||||
}
|
||||
|
||||
return targets, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetServiceTargetByTargetID(ctx context.Context, lockStrength LockingStrength, accountID string, targetID string) (*rpservice.Target, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var target *rpservice.Target
|
||||
result := tx.Take(&target, "account_id = ? AND target_id = ?", accountID, targetID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "service target with ID %s not found", targetID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get service target from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get service target from store")
|
||||
}
|
||||
|
||||
return target, nil
|
||||
}
|
||||
@@ -0,0 +1,219 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) GetAccountBySetupKey(ctx context.Context, setupKey string) (*types.Account, error) {
|
||||
var key types.SetupKey
|
||||
result := s.db.Select("account_id").Take(&key, GetKeyQueryCondition(s), setupKey)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewSetupKeyNotFoundError(setupKey)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get account by setup key from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get account by setup key from store")
|
||||
}
|
||||
|
||||
if key.AccountID == "" {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
|
||||
return s.GetAccount(ctx, key.AccountID)
|
||||
}
|
||||
|
||||
func (s *SqlStore) getSetupKeys(ctx context.Context, accountID string) ([]types.SetupKey, error) {
|
||||
const query = `SELECT id, account_id, key, key_secret, name, type, created_at, expires_at, updated_at,
|
||||
revoked, used_times, last_used, auto_groups, usage_limit, ephemeral, allow_extra_dns_labels FROM setup_keys WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
keys, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (types.SetupKey, error) {
|
||||
var sk types.SetupKey
|
||||
var autoGroups []byte
|
||||
var skCreatedAt, expiresAt, updatedAt, lastUsed sql.NullTime
|
||||
var revoked, ephemeral, allowExtraDNSLabels sql.NullBool
|
||||
var usedTimes, usageLimit sql.NullInt64
|
||||
|
||||
err := row.Scan(&sk.Id, &sk.AccountID, &sk.Key, &sk.KeySecret, &sk.Name, &sk.Type, &skCreatedAt,
|
||||
&expiresAt, &updatedAt, &revoked, &usedTimes, &lastUsed, &autoGroups, &usageLimit, &ephemeral, &allowExtraDNSLabels)
|
||||
|
||||
if err == nil {
|
||||
if expiresAt.Valid {
|
||||
sk.ExpiresAt = &expiresAt.Time
|
||||
}
|
||||
if skCreatedAt.Valid {
|
||||
sk.CreatedAt = skCreatedAt.Time
|
||||
}
|
||||
if updatedAt.Valid {
|
||||
sk.UpdatedAt = updatedAt.Time
|
||||
if sk.UpdatedAt.IsZero() {
|
||||
sk.UpdatedAt = sk.CreatedAt
|
||||
}
|
||||
}
|
||||
if lastUsed.Valid {
|
||||
sk.LastUsed = &lastUsed.Time
|
||||
}
|
||||
if revoked.Valid {
|
||||
sk.Revoked = revoked.Bool
|
||||
}
|
||||
if usedTimes.Valid {
|
||||
sk.UsedTimes = int(usedTimes.Int64)
|
||||
}
|
||||
if usageLimit.Valid {
|
||||
sk.UsageLimit = int(usageLimit.Int64)
|
||||
}
|
||||
if ephemeral.Valid {
|
||||
sk.Ephemeral = ephemeral.Bool
|
||||
}
|
||||
if allowExtraDNSLabels.Valid {
|
||||
sk.AllowExtraDNSLabels = allowExtraDNSLabels.Bool
|
||||
}
|
||||
if autoGroups != nil {
|
||||
_ = json.Unmarshal(autoGroups, &sk.AutoGroups)
|
||||
} else {
|
||||
sk.AutoGroups = []string{}
|
||||
}
|
||||
}
|
||||
return sk, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return keys, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string) (string, error) {
|
||||
var accountID string
|
||||
result := s.db.Model(&types.SetupKey{}).Select("account_id").Where(GetKeyQueryCondition(s), setupKey).Take(&accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return "", status.NewSetupKeyNotFoundError(setupKey)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get account ID by setup key from store: %v", result.Error)
|
||||
return "", status.Errorf(status.Internal, "failed to get account ID by setup key from store")
|
||||
}
|
||||
|
||||
if accountID == "" {
|
||||
return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
|
||||
return accountID, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*types.SetupKey, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var setupKey types.SetupKey
|
||||
result := tx.
|
||||
Take(&setupKey, GetKeyQueryCondition(s), key)
|
||||
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.PreconditionFailed, "setup key not found")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get setup key by secret from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get setup key by secret from store")
|
||||
}
|
||||
return &setupKey, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error {
|
||||
result := s.db.Model(&types.SetupKey{}).
|
||||
Where(idQueryCondition, setupKeyID).
|
||||
Updates(map[string]interface{}{
|
||||
"used_times": gorm.Expr("used_times + 1"),
|
||||
"last_used": time.Now(),
|
||||
})
|
||||
|
||||
if result.Error != nil {
|
||||
return status.Errorf(status.Internal, "issue incrementing setup key usage count: %s", result.Error)
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewSetupKeyNotFoundError(setupKeyID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAccountSetupKeys retrieves setup keys for an account.
|
||||
func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.SetupKey, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var setupKeys []*types.SetupKey
|
||||
result := tx.
|
||||
Find(&setupKeys, accountIDCondition, accountID)
|
||||
if err := result.Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get setup keys from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get setup keys from store")
|
||||
}
|
||||
|
||||
return setupKeys, nil
|
||||
}
|
||||
|
||||
// GetSetupKeyByID retrieves a setup key by its ID and account ID.
|
||||
func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*types.SetupKey, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var setupKey *types.SetupKey
|
||||
result := tx.Take(&setupKey, accountAndIDQueryCondition, accountID, setupKeyID)
|
||||
if err := result.Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewSetupKeyNotFoundError(setupKeyID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get setup key from the store: %s", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get setup key from store")
|
||||
}
|
||||
|
||||
return setupKey, nil
|
||||
}
|
||||
|
||||
// SaveSetupKey saves a setup key to the database.
|
||||
func (s *SqlStore) SaveSetupKey(ctx context.Context, setupKey *types.SetupKey) error {
|
||||
result := s.db.Save(setupKey)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save setup key to store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save setup key to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteSetupKey deletes a setup key from the database.
|
||||
func (s *SqlStore) DeleteSetupKey(ctx context.Context, accountID, keyID string) error {
|
||||
result := s.db.Delete(&types.SetupKey{}, accountAndIDQueryCondition, accountID, keyID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete setup key from store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete setup key from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewSetupKeyNotFoundError(keyID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
b64 "encoding/base64"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
func TestSqlite_GetSetupKeyBySecret(t *testing.T) {
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
plainKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
|
||||
hashedKey := sha256.Sum256([]byte(plainKey))
|
||||
encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
|
||||
|
||||
_, err = store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
setupKey, err := store.GetSetupKeyBySecret(context.Background(), LockingStrengthNone, encodedHashedKey)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, encodedHashedKey, setupKey.Key)
|
||||
assert.Equal(t, types.HiddenKey(plainKey, 4), setupKey.KeySecret)
|
||||
assert.Equal(t, "bf1c8084-ba50-4ce7-9439-34653001fc3b", setupKey.AccountID)
|
||||
assert.Equal(t, "Default key", setupKey.Name)
|
||||
}
|
||||
|
||||
func TestSqlite_incrementSetupKeyUsage(t *testing.T) {
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
plainKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
|
||||
hashedKey := sha256.Sum256([]byte(plainKey))
|
||||
encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
|
||||
|
||||
_, err = store.GetAccount(context.Background(), existingAccountID)
|
||||
require.NoError(t, err)
|
||||
|
||||
setupKey, err := store.GetSetupKeyBySecret(context.Background(), LockingStrengthNone, encodedHashedKey)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, setupKey.UsedTimes)
|
||||
|
||||
err = store.IncrementSetupKeyUsage(context.Background(), setupKey.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
setupKey, err = store.GetSetupKeyBySecret(context.Background(), LockingStrengthNone, encodedHashedKey)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, setupKey.UsedTimes)
|
||||
|
||||
err = store.IncrementSetupKeyUsage(context.Background(), setupKey.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
setupKey, err = store.GetSetupKeyBySecret(context.Background(), LockingStrengthNone, encodedHashedKey)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, setupKey.UsedTimes)
|
||||
}
|
||||
|
||||
func Test_DeleteSetupKeySuccessfully(t *testing.T) {
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
setupKeyID := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
|
||||
|
||||
err = store.DeleteSetupKey(context.Background(), accountID, setupKeyID)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = store.GetSetupKeyByID(context.Background(), LockingStrengthNone, setupKeyID, accountID)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func Test_DeleteSetupKeyFailsForNonExistingKey(t *testing.T) {
|
||||
t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
nonExistingKeyID := "non-existing-key-id"
|
||||
|
||||
err = store.DeleteSetupKey(context.Background(), accountID, nonExistingKeyID)
|
||||
require.Error(t, err)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,272 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// SaveUsers saves the given list of users to the database.
|
||||
func (s *SqlStore) SaveUsers(ctx context.Context, users []*types.User) error {
|
||||
if len(users) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
usersCopy := make([]*types.User, len(users))
|
||||
for i, user := range users {
|
||||
userCopy := user.Copy()
|
||||
userCopy.Email = user.Email
|
||||
userCopy.Name = user.Name
|
||||
if err := userCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return fmt.Errorf("encrypt user: %w", err)
|
||||
}
|
||||
usersCopy[i] = userCopy
|
||||
}
|
||||
|
||||
result := s.db.Clauses(clause.OnConflict{UpdateAll: true}).Create(&usersCopy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save users to store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save users to store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveUser saves the given user to the database.
|
||||
func (s *SqlStore) SaveUser(ctx context.Context, user *types.User) error {
|
||||
userCopy := user.Copy()
|
||||
userCopy.Email = user.Email
|
||||
userCopy.Name = user.Name
|
||||
|
||||
if err := userCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return fmt.Errorf("encrypt user: %w", err)
|
||||
}
|
||||
|
||||
result := s.db.Save(userCopy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save user to store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save user to store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetUserByPATID(ctx context.Context, lockStrength LockingStrength, patID string) (*types.User, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var user types.User
|
||||
result := tx.
|
||||
Joins("JOIN personal_access_tokens ON personal_access_tokens.user_id = users.id").
|
||||
Where("personal_access_tokens.id = ?", patID).Take(&user)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewPATNotFoundError(patID)
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get token user from the store: %s", result.Error)
|
||||
return nil, status.NewGetUserFromStoreError()
|
||||
}
|
||||
|
||||
if err := user.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt user: %w", err)
|
||||
}
|
||||
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (*types.User, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var user types.User
|
||||
result := tx.Take(&user, idQueryCondition, userID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewUserNotFoundError(userID)
|
||||
}
|
||||
return nil, status.NewGetUserFromStoreError()
|
||||
}
|
||||
|
||||
if err := user.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt user: %w", err)
|
||||
}
|
||||
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteUser(ctx context.Context, accountID, userID string) error {
|
||||
err := s.transaction(func(tx *gorm.DB) error {
|
||||
result := tx.Delete(&types.PersonalAccessToken{}, "user_id = ?", userID)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
|
||||
return tx.Delete(&types.User{}, accountAndIDQueryCondition, accountID, userID).Error
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete user from the store: %s", err)
|
||||
return status.Errorf(status.Internal, "failed to delete user from store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountUsers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.User, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var users []*types.User
|
||||
result := tx.Find(&users, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("error when getting users from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "issue getting users from store")
|
||||
}
|
||||
|
||||
for _, user := range users {
|
||||
if err := user.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt user: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return users, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountOwner(ctx context.Context, lockStrength LockingStrength, accountID string) (*types.User, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var user types.User
|
||||
result := tx.Take(&user, "account_id = ? AND role = ?", accountID, types.UserRoleOwner)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "account owner not found: index lookup failed")
|
||||
}
|
||||
return nil, status.Errorf(status.Internal, "failed to get account owner from the store")
|
||||
}
|
||||
|
||||
if err := user.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt user: %w", err)
|
||||
}
|
||||
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) getUsers(ctx context.Context, accountID string) ([]types.User, error) {
|
||||
const query = `SELECT id, account_id, role, is_service_user, non_deletable, service_user_name, auto_groups, blocked, pending_approval, last_login, created_at, issued, integration_ref_id, integration_ref_integration_type, email, name FROM users WHERE account_id = $1`
|
||||
rows, err := s.pool.Query(ctx, query, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
users, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (types.User, error) {
|
||||
var u types.User
|
||||
var autoGroups []byte
|
||||
var lastLogin, createdAt sql.NullTime
|
||||
var isServiceUser, nonDeletable, blocked, pendingApproval sql.NullBool
|
||||
err := row.Scan(&u.Id, &u.AccountID, &u.Role, &isServiceUser, &nonDeletable, &u.ServiceUserName, &autoGroups, &blocked, &pendingApproval, &lastLogin, &createdAt, &u.Issued, &u.IntegrationReference.ID, &u.IntegrationReference.IntegrationType, &u.Email, &u.Name)
|
||||
if err == nil {
|
||||
if lastLogin.Valid {
|
||||
u.LastLogin = &lastLogin.Time
|
||||
}
|
||||
if createdAt.Valid {
|
||||
u.CreatedAt = createdAt.Time
|
||||
}
|
||||
if isServiceUser.Valid {
|
||||
u.IsServiceUser = isServiceUser.Bool
|
||||
}
|
||||
if nonDeletable.Valid {
|
||||
u.NonDeletable = nonDeletable.Bool
|
||||
}
|
||||
if blocked.Valid {
|
||||
u.Blocked = blocked.Bool
|
||||
}
|
||||
if pendingApproval.Valid {
|
||||
u.PendingApproval = pendingApproval.Bool
|
||||
}
|
||||
if autoGroups != nil {
|
||||
_ = json.Unmarshal(autoGroups, &u.AutoGroups)
|
||||
} else {
|
||||
u.AutoGroups = []string{}
|
||||
}
|
||||
}
|
||||
return u, err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return users, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountByUser(ctx context.Context, userID string) (*types.Account, error) {
|
||||
var user types.User
|
||||
result := s.db.Select("account_id").Take(&user, idQueryCondition, userID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
return nil, status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
if user.AccountID == "" {
|
||||
return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
|
||||
return s.GetAccount(ctx, user.AccountID)
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountIDByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (string, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var accountID string
|
||||
result := tx.Model(&types.User{}).
|
||||
Select("account_id").Where(idQueryCondition, userID).Take(&accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
|
||||
}
|
||||
return "", status.NewGetAccountFromStoreError(result.Error)
|
||||
}
|
||||
|
||||
return accountID, nil
|
||||
}
|
||||
|
||||
// SaveUserLastLogin stores the last login time for a user in DB.
|
||||
func (s *SqlStore) SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error {
|
||||
var user types.User
|
||||
result := s.db.Take(&user, accountAndIDQueryCondition, accountID, userID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return status.NewUserNotFoundError(userID)
|
||||
}
|
||||
return status.NewGetUserFromStoreError()
|
||||
}
|
||||
|
||||
if !lastLogin.IsZero() {
|
||||
user.LastLogin = &lastLogin
|
||||
return s.db.Save(&user).Error
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// SaveUserInvite saves a user invite to the database
|
||||
func (s *SqlStore) SaveUserInvite(ctx context.Context, invite *types.UserInviteRecord) error {
|
||||
inviteCopy := invite.Copy()
|
||||
if err := inviteCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return fmt.Errorf("encrypt invite: %w", err)
|
||||
}
|
||||
|
||||
result := s.db.Save(inviteCopy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save user invite to store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save user invite to store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetUserInviteByID retrieves a user invite by its ID and account ID
|
||||
func (s *SqlStore) GetUserInviteByID(ctx context.Context, lockStrength LockingStrength, accountID, inviteID string) (*types.UserInviteRecord, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var invite types.UserInviteRecord
|
||||
result := tx.Where("account_id = ?", accountID).Take(&invite, idQueryCondition, inviteID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "user invite not found")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get user invite from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get user invite from store")
|
||||
}
|
||||
|
||||
if err := invite.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt invite: %w", err)
|
||||
}
|
||||
|
||||
return &invite, nil
|
||||
}
|
||||
|
||||
// GetUserInviteByHashedToken retrieves a user invite by its hashed token
|
||||
func (s *SqlStore) GetUserInviteByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken string) (*types.UserInviteRecord, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var invite types.UserInviteRecord
|
||||
result := tx.Take(&invite, "hashed_token = ?", hashedToken)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "user invite not found")
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get user invite from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get user invite from store")
|
||||
}
|
||||
|
||||
if err := invite.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt invite: %w", err)
|
||||
}
|
||||
|
||||
return &invite, nil
|
||||
}
|
||||
|
||||
// GetUserInviteByEmail retrieves a user invite by account ID and email.
|
||||
// Since email is encrypted with random IVs, we fetch all invites for the account
|
||||
// and compare emails in memory after decryption.
|
||||
func (s *SqlStore) GetUserInviteByEmail(ctx context.Context, lockStrength LockingStrength, accountID, email string) (*types.UserInviteRecord, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var invites []*types.UserInviteRecord
|
||||
result := tx.Find(&invites, "account_id = ?", accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get user invites from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get user invites from store")
|
||||
}
|
||||
|
||||
for _, invite := range invites {
|
||||
if err := invite.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt invite: %w", err)
|
||||
}
|
||||
if strings.EqualFold(invite.Email, email) {
|
||||
return invite, nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, status.Errorf(status.NotFound, "user invite not found for email")
|
||||
}
|
||||
|
||||
// GetAccountUserInvites retrieves all user invites for an account
|
||||
func (s *SqlStore) GetAccountUserInvites(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.UserInviteRecord, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var invites []*types.UserInviteRecord
|
||||
result := tx.Find(&invites, "account_id = ?", accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get user invites from store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get user invites from store")
|
||||
}
|
||||
|
||||
for _, invite := range invites {
|
||||
if err := invite.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
return nil, fmt.Errorf("decrypt invite: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return invites, nil
|
||||
}
|
||||
|
||||
// DeleteUserInvite deletes a user invite by its ID
|
||||
func (s *SqlStore) DeleteUserInvite(ctx context.Context, inviteID string) error {
|
||||
result := s.db.Delete(&types.UserInviteRecord{}, idQueryCondition, inviteID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete user invite from store: %s", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete user invite from store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,343 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/management/server/util"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
"github.com/netbirdio/netbird/util/crypt"
|
||||
)
|
||||
|
||||
func TestSqlStore_GetAccountUsers(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
account, err := store.GetAccount(context.Background(), accountID)
|
||||
require.NoError(t, err)
|
||||
users, err := store.GetAccountUsers(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, users, len(account.Users))
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserByUserID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
userID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing user",
|
||||
userID: "edafee4e-63fb-11ec-90d6-0242ac120003",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing user",
|
||||
userID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty user ID",
|
||||
userID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
user, err := store.GetUserByUserID(context.Background(), LockingStrengthNone, tt.userID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, user)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, user)
|
||||
require.Equal(t, tt.userID, user.Id)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetUserByPATID(t *testing.T) {
|
||||
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
|
||||
t.Cleanup(cleanUp)
|
||||
assert.NoError(t, err)
|
||||
|
||||
id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
user, err := store.GetUserByPATID(context.Background(), LockingStrengthNone, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "f4f6d672-63fb-11ec-90d6-0242ac120003", user.Id)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveUser(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
user := &types.User{
|
||||
Id: "user-id",
|
||||
AccountID: accountID,
|
||||
Role: types.UserRoleAdmin,
|
||||
IsServiceUser: false,
|
||||
AutoGroups: []string{"groupA", "groupB"},
|
||||
Blocked: false,
|
||||
LastLogin: util.ToPtr(time.Now().UTC()),
|
||||
CreatedAt: time.Now().UTC().Add(-time.Hour),
|
||||
Issued: types.UserIssuedIntegration,
|
||||
}
|
||||
err = store.SaveUser(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
saveUser, err := store.GetUserByUserID(context.Background(), LockingStrengthNone, user.Id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, user.Id, saveUser.Id)
|
||||
require.Equal(t, user.AccountID, saveUser.AccountID)
|
||||
require.Equal(t, user.Role, saveUser.Role)
|
||||
require.Equal(t, user.AutoGroups, saveUser.AutoGroups)
|
||||
require.WithinDurationf(t, user.GetLastLogin(), saveUser.LastLogin.UTC(), time.Millisecond, "LastLogin should be equal")
|
||||
require.WithinDurationf(t, user.CreatedAt, saveUser.CreatedAt.UTC(), time.Millisecond, "CreatedAt should be equal")
|
||||
require.Equal(t, user.Issued, saveUser.Issued)
|
||||
require.Equal(t, user.Blocked, saveUser.Blocked)
|
||||
require.Equal(t, user.IsServiceUser, saveUser.IsServiceUser)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveUsers(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
accountUsers, err := store.GetAccountUsers(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accountUsers, 2)
|
||||
|
||||
users := []*types.User{
|
||||
{
|
||||
Id: "user-1",
|
||||
AccountID: accountID,
|
||||
Issued: "api",
|
||||
AutoGroups: []string{"groupA", "groupB"},
|
||||
},
|
||||
{
|
||||
Id: "user-2",
|
||||
AccountID: accountID,
|
||||
Issued: "integration",
|
||||
AutoGroups: []string{"groupA"},
|
||||
},
|
||||
}
|
||||
err = store.SaveUsers(context.Background(), users)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountUsers, err = store.GetAccountUsers(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accountUsers, 4)
|
||||
|
||||
users[1].AutoGroups = []string{"groupA", "groupC"}
|
||||
err = store.SaveUsers(context.Background(), users)
|
||||
require.NoError(t, err)
|
||||
|
||||
user, err := store.GetUserByUserID(context.Background(), LockingStrengthNone, users[1].Id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, users[1].AutoGroups, user.AutoGroups)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveUserWithEncryption(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Enable encryption
|
||||
key, err := crypt.GenerateKey()
|
||||
require.NoError(t, err)
|
||||
fieldEncrypt, err := crypt.NewFieldEncrypt(key)
|
||||
require.NoError(t, err)
|
||||
store.SetFieldEncrypt(fieldEncrypt)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
// rawUser is used to read raw (potentially encrypted) data from the database
|
||||
// without any gorm hooks or automatic decryption
|
||||
type rawUser struct {
|
||||
Id string
|
||||
Email string
|
||||
Name string
|
||||
}
|
||||
|
||||
t.Run("save user with empty email and name", func(t *testing.T) {
|
||||
user := &types.User{
|
||||
Id: "user-empty-fields",
|
||||
AccountID: accountID,
|
||||
Role: types.UserRoleUser,
|
||||
Email: "",
|
||||
Name: "",
|
||||
AutoGroups: []string{"groupA"},
|
||||
}
|
||||
err = store.SaveUser(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify using direct database query that empty strings remain empty (not encrypted)
|
||||
var raw rawUser
|
||||
err = store.(*SqlStore).db.Table("users").Select("id, email, name").Where("id = ?", user.Id).First(&raw).Error
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", raw.Email, "empty email should remain empty in database")
|
||||
require.Equal(t, "", raw.Name, "empty name should remain empty in database")
|
||||
|
||||
// Verify manual decryption returns empty strings
|
||||
decryptedEmail, err := fieldEncrypt.Decrypt(raw.Email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", decryptedEmail)
|
||||
|
||||
decryptedName, err := fieldEncrypt.Decrypt(raw.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", decryptedName)
|
||||
})
|
||||
|
||||
t.Run("save user with email and name", func(t *testing.T) {
|
||||
user := &types.User{
|
||||
Id: "user-with-fields",
|
||||
AccountID: accountID,
|
||||
Role: types.UserRoleAdmin,
|
||||
Email: "test@example.com",
|
||||
Name: "Test User",
|
||||
AutoGroups: []string{"groupB"},
|
||||
}
|
||||
err = store.SaveUser(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify using direct database query that the data is encrypted (not plaintext)
|
||||
var raw rawUser
|
||||
err = store.(*SqlStore).db.Table("users").Select("id, email, name").Where("id = ?", user.Id).First(&raw).Error
|
||||
require.NoError(t, err)
|
||||
require.NotEqual(t, "test@example.com", raw.Email, "email should be encrypted in database")
|
||||
require.NotEqual(t, "Test User", raw.Name, "name should be encrypted in database")
|
||||
|
||||
// Verify manual decryption returns correct values
|
||||
decryptedEmail, err := fieldEncrypt.Decrypt(raw.Email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "test@example.com", decryptedEmail)
|
||||
|
||||
decryptedName, err := fieldEncrypt.Decrypt(raw.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "Test User", decryptedName)
|
||||
})
|
||||
|
||||
t.Run("save multiple users with mixed fields", func(t *testing.T) {
|
||||
users := []*types.User{
|
||||
{
|
||||
Id: "batch-user-1",
|
||||
AccountID: accountID,
|
||||
Email: "",
|
||||
Name: "",
|
||||
},
|
||||
{
|
||||
Id: "batch-user-2",
|
||||
AccountID: accountID,
|
||||
Email: "batch@example.com",
|
||||
Name: "Batch User",
|
||||
},
|
||||
}
|
||||
err = store.SaveUsers(context.Background(), users)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify first user (empty fields) using direct database query
|
||||
var raw1 rawUser
|
||||
err = store.(*SqlStore).db.Table("users").Select("id, email, name").Where("id = ?", "batch-user-1").First(&raw1).Error
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", raw1.Email, "empty email should remain empty in database")
|
||||
require.Equal(t, "", raw1.Name, "empty name should remain empty in database")
|
||||
|
||||
// Verify second user (with fields) using direct database query
|
||||
var raw2 rawUser
|
||||
err = store.(*SqlStore).db.Table("users").Select("id, email, name").Where("id = ?", "batch-user-2").First(&raw2).Error
|
||||
require.NoError(t, err)
|
||||
require.NotEqual(t, "batch@example.com", raw2.Email, "email should be encrypted in database")
|
||||
require.NotEqual(t, "Batch User", raw2.Name, "name should be encrypted in database")
|
||||
|
||||
// Verify manual decryption returns empty strings for first user
|
||||
decryptedEmail1, err := fieldEncrypt.Decrypt(raw1.Email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", decryptedEmail1)
|
||||
|
||||
decryptedName1, err := fieldEncrypt.Decrypt(raw1.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", decryptedName1)
|
||||
|
||||
// Verify manual decryption returns correct values for second user
|
||||
decryptedEmail2, err := fieldEncrypt.Decrypt(raw2.Email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "batch@example.com", decryptedEmail2)
|
||||
|
||||
decryptedName2, err := fieldEncrypt.Decrypt(raw2.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "Batch User", decryptedName2)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteUser(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
userID := "f4f6d672-63fb-11ec-90d6-0242ac120003"
|
||||
|
||||
err = store.DeleteUser(context.Background(), accountID, userID)
|
||||
require.NoError(t, err)
|
||||
|
||||
user, err := store.GetUserByUserID(context.Background(), LockingStrengthNone, userID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, user)
|
||||
|
||||
userPATs, err := store.GetUserPATs(context.Background(), LockingStrengthNone, userID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, userPATs, 0)
|
||||
}
|
||||
|
||||
func TestSqlStore_SaveUsers_LargeBatch(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
accountUsers, err := store.GetAccountUsers(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accountUsers, 2)
|
||||
|
||||
usersToSave := make([]*types.User, 0)
|
||||
|
||||
for i := 1; i <= 8000; i++ {
|
||||
usersToSave = append(usersToSave, &types.User{
|
||||
Id: fmt.Sprintf("user-%d", i),
|
||||
AccountID: accountID,
|
||||
Role: types.UserRoleUser,
|
||||
})
|
||||
}
|
||||
|
||||
err = store.SaveUsers(context.Background(), usersToSave)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountUsers, err = store.GetAccountUsers(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 8002, len(accountUsers))
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func (s *SqlStore) CreateZone(ctx context.Context, zone *zones.Zone) error {
|
||||
result := s.db.Create(zone)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to create zone to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to create zone to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) UpdateZone(ctx context.Context, zone *zones.Zone) error {
|
||||
result := s.db.Select("*").Save(zone)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to update zone to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to update zone to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteZone(ctx context.Context, accountID, zoneID string) error {
|
||||
result := s.db.Delete(&zones.Zone{}, accountAndIDQueryCondition, accountID, zoneID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete zone from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete zone from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewZoneNotFoundError(zoneID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetZoneByID(ctx context.Context, lockStrength LockingStrength, accountID, zoneID string) (*zones.Zone, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var zone *zones.Zone
|
||||
result := tx.Preload("Records").Take(&zone, accountAndIDQueryCondition, accountID, zoneID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewZoneNotFoundError(zoneID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get zone from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get zone from store")
|
||||
}
|
||||
|
||||
return zone, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetZoneByDomain(ctx context.Context, accountID, domain string) (*zones.Zone, error) {
|
||||
var zone *zones.Zone
|
||||
result := s.db.Where("account_id = ? AND domain = ?", accountID, domain).First(&zone)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewZoneNotFoundError(domain)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get zone by domain from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get zone by domain from store")
|
||||
}
|
||||
|
||||
return zone, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountZones(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*zones.Zone, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var zones []*zones.Zone
|
||||
result := tx.Preload("Records").Find(&zones, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get zones from the store: %s", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get zones from store")
|
||||
}
|
||||
|
||||
return zones, nil
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
func TestSqlStore_CreateZone(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
savedZone, err := store.GetZoneByID(context.Background(), LockingStrengthNone, accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, savedZone)
|
||||
assert.Equal(t, zone.ID, savedZone.ID)
|
||||
assert.Equal(t, zone.Name, savedZone.Name)
|
||||
assert.Equal(t, zone.Domain, savedZone.Domain)
|
||||
assert.Equal(t, zone.Enabled, savedZone.Enabled)
|
||||
assert.Equal(t, zone.EnableSearchDomain, savedZone.EnableSearchDomain)
|
||||
assert.Equal(t, zone.DistributionGroups, savedZone.DistributionGroups)
|
||||
}
|
||||
|
||||
func TestSqlStore_GetZoneByID(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
zoneID string
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing zone",
|
||||
accountID: accountID,
|
||||
zoneID: zone.ID,
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing zone",
|
||||
accountID: accountID,
|
||||
zoneID: "non-existing",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty zone ID",
|
||||
accountID: accountID,
|
||||
zoneID: "",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
savedZone, err := store.GetZoneByID(context.Background(), LockingStrengthNone, tt.accountID, tt.zoneID)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
require.Nil(t, savedZone)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, savedZone)
|
||||
assert.Equal(t, tt.zoneID, savedZone.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_GetAccountZones(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone1 := zones.NewZone(accountID, "Zone 1", "example1.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone1)
|
||||
require.NoError(t, err)
|
||||
|
||||
zone2 := zones.NewZone(accountID, "Zone 2", "example2.com", true, true, []string{"group1", "group2"})
|
||||
err = store.CreateZone(context.Background(), zone2)
|
||||
require.NoError(t, err)
|
||||
|
||||
allZones, err := store.GetAccountZones(context.Background(), LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, allZones)
|
||||
assert.GreaterOrEqual(t, len(allZones), 2)
|
||||
|
||||
zoneIDs := make(map[string]bool)
|
||||
for _, z := range allZones {
|
||||
zoneIDs[z.ID] = true
|
||||
}
|
||||
assert.True(t, zoneIDs[zone1.ID])
|
||||
assert.True(t, zoneIDs[zone2.ID])
|
||||
}
|
||||
|
||||
func TestSqlStore_GetZoneByDomain(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
otherAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3c"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID string
|
||||
domain string
|
||||
expectError bool
|
||||
errorType status.Type
|
||||
}{
|
||||
{
|
||||
name: "retrieve existing zone by domain",
|
||||
accountID: accountID,
|
||||
domain: "example.com",
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "retrieve non-existing zone domain",
|
||||
accountID: accountID,
|
||||
domain: "non-existing.com",
|
||||
expectError: true,
|
||||
errorType: status.NotFound,
|
||||
},
|
||||
{
|
||||
name: "retrieve with empty domain",
|
||||
accountID: accountID,
|
||||
domain: "",
|
||||
expectError: true,
|
||||
errorType: status.NotFound,
|
||||
},
|
||||
{
|
||||
name: "retrieve with different account ID",
|
||||
accountID: otherAccountID,
|
||||
domain: "example.com",
|
||||
expectError: true,
|
||||
errorType: status.NotFound,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
savedZone, err := store.GetZoneByDomain(context.Background(), tt.accountID, tt.domain)
|
||||
if tt.expectError {
|
||||
require.Error(t, err)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, tt.errorType, sErr.Type())
|
||||
require.Nil(t, savedZone)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, savedZone)
|
||||
assert.Equal(t, tt.domain, savedZone.Domain)
|
||||
assert.Equal(t, zone.ID, savedZone.ID)
|
||||
assert.Equal(t, zone.Name, savedZone.Name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSqlStore_UpdateZone(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
zone.Name = "Updated Zone"
|
||||
zone.Domain = "updated.com"
|
||||
zone.Enabled = false
|
||||
zone.EnableSearchDomain = true
|
||||
zone.DistributionGroups = []string{"group2", "group3"}
|
||||
|
||||
err = store.UpdateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
updatedZone, err := store.GetZoneByID(context.Background(), LockingStrengthNone, accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, updatedZone)
|
||||
assert.Equal(t, "Updated Zone", updatedZone.Name)
|
||||
assert.Equal(t, "updated.com", updatedZone.Domain)
|
||||
assert.False(t, updatedZone.Enabled)
|
||||
assert.True(t, updatedZone.EnableSearchDomain)
|
||||
assert.Equal(t, []string{"group2", "group3"}, updatedZone.DistributionGroups)
|
||||
}
|
||||
|
||||
func TestSqlStore_DeleteZone(t *testing.T) {
|
||||
store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
|
||||
t.Cleanup(cleanup)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
|
||||
|
||||
zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
|
||||
err = store.CreateZone(context.Background(), zone)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = store.DeleteZone(context.Background(), accountID, zone.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
deletedZone, err := store.GetZoneByID(context.Background(), LockingStrengthNone, accountID, zone.ID)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, deletedZone)
|
||||
sErr, ok := status.FromError(err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, sErr.Type(), status.NotFound)
|
||||
}
|
||||
Reference in New Issue
Block a user