diff --git a/management/internals/network_map_db/db_store.go b/management/internals/network_map_db/db_store.go index 5b60e3e3a..a9325ee06 100644 --- a/management/internals/network_map_db/db_store.go +++ b/management/internals/network_map_db/db_store.go @@ -33,6 +33,7 @@ type NetworkMapDBStore interface { GetAccountSettings(ctx context.Context, accountId string) (nmdata.AccountSettingsInfo, error) GetPostureChecks(ctx context.Context, accountId string) ([]nmdata.PostureChecks, error) GetNetworkMapData(ctx context.Context, accountId string) (*networkmap.NetworkMapData, error) + GetAllowedUserIds(ctx context.Context, accountId string) (map[string]struct{}, error) } type NetworkMapDBStoreImpl struct { diff --git a/management/internals/network_map_db/pgsql/allowed_user_ids.go b/management/internals/network_map_db/pgsql/allowed_user_ids.go new file mode 100644 index 000000000..3a4dcf408 --- /dev/null +++ b/management/internals/network_map_db/pgsql/allowed_user_ids.go @@ -0,0 +1,42 @@ +package networkmap_pgsql + +import ( + "context" + + "github.com/jackc/pgx/v5" +) + +const ( + GetAllowedUserIdsQuery = ` + select id + from users + where account_id=$1 and not blocked and not is_service_user + ` +) + +func (pg *PgStore) GetAllowedUserIds(ctx context.Context, accountId string) (map[string]struct{}, error) { + c, err := pg.Pool.Acquire(ctx) + if err != nil { + return nil, err + } + return GetGetAllowedUserIdsViaPgxConnection(ctx, c.Conn(), accountId) +} + +func GetGetAllowedUserIdsViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (map[string]struct{}, error) { + rows, err := con.Query(ctx, GetAllowedUserIdsQuery, accountId) + if err != nil { + return nil, err + } + + toret := make(map[string]struct{}) + var id string + _, err = pgx.ForEachRow(rows, []any{&id}, func() error { + toret[id] = struct{}{} + return nil + }) + if err != nil { + return nil, err + } + + return toret, nil +} diff --git a/management/internals/network_map_db/pgsql/group_test.go b/management/internals/network_map_db/pgsql/group_test.go index e5f3bf0f8..2d0a08f85 100644 --- a/management/internals/network_map_db/pgsql/group_test.go +++ b/management/internals/network_map_db/pgsql/group_test.go @@ -3,13 +3,11 @@ package networkmap_pgsql import ( "context" "fmt" - "strings" "testing" "time" _ "embed" - "github.com/jackc/pgx/v5/pgxpool" "github.com/netbirdio/netbird/shared/management/networkmap/nmdata" "github.com/stretchr/testify/assert" ) @@ -59,7 +57,7 @@ func TestGetPeers(t *testing.T) { // ) } -func TestGetPolocies(t *testing.T) { +func TestGetPolicies(t *testing.T) { ctx := context.TODO() s, err := NewPostgresqlStore(ctx, "postgresql://root:netbird@localhost:5432/netbird") @@ -256,18 +254,15 @@ func TestGetPostureChecks(t *testing.T) { // ) } -func loadSQL(ctx context.Context, pool *pgxpool.Pool, initdb string) error { - queries := strings.Split(string(initdb), ";") +func TestGetGetAllowedUserIds(t *testing.T) { + ctx := context.TODO() - for _, query := range queries { - query = strings.TrimSpace(query) - if query != "" { - _, err := pool.Query(ctx, query) - if err != nil { - return err - } - } - } + s, err := NewPostgresqlStore(ctx, "postgresql://root:netbird@localhost:5432/netbird") + assert.NoError(t, err) + // err = loadSQL(ctx, s.pool, initDb) + //assert.NoError(t, err) - return nil + ids, err := s.GetAllowedUserIds(ctx, "cus73sbl0ubs73cfoo90") //"ckd7ee2fic3c73dtendg") + assert.NoError(t, err) + assert.NotEmpty(t, ids) } diff --git a/management/internals/network_map_db/pgsql/network_map_data.go b/management/internals/network_map_db/pgsql/network_map_data.go index d5a785b47..e94df8108 100644 --- a/management/internals/network_map_db/pgsql/network_map_data.go +++ b/management/internals/network_map_db/pgsql/network_map_data.go @@ -62,6 +62,10 @@ func (pg *PgStore) GetNetworkMapData(ctx context.Context, accountId string) (*ne if err != nil { return rollbackAndReturnError(ctx, tx, err) } + allowedUserIds, err := GetGetAllowedUserIdsViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } resourcePolicies := make(map[string][]*nmdata.Policy) for _, resource := range networkResources { @@ -104,6 +108,7 @@ func (pg *PgStore) GetNetworkMapData(ctx context.Context, accountId string) (*ne NameServerGroups: toSliceOfPtrs(nsGroups), NetworkResources: toSliceOfPtrs(networkResources), PostureChecks: toMap(postureChecks, func(pc nmdata.PostureChecks) string { return pc.ID }), + AllowedUserIDs: allowedUserIds, NetworkXIDToPublicID: networkXIDToPublicID, // TODO (dmitri) maybe we can switch to public ids everywhere? }