[client, management] Support DNS Labels for Peer Addressing (#3252)

* [client] Support Extra DNS Labels for Peer Addressing

* [management] Support Extra DNS Labels for Peer Addressing

---------

Co-authored-by: Viktor Liu <17948409+lixmal@users.noreply.github.com>

management/server/account.go

@@ -63,7 +63,7 @@ type AccountManager interface {
 	GetOrCreateAccountByUser(ctx context.Context, userId, domain string) (*types.Account, error)
 	GetAccount(ctx context.Context, accountID string) (*types.Account, error)
 	CreateSetupKey(ctx context.Context, accountID string, keyName string, keyType types.SetupKeyType, expiresIn time.Duration,
-		autoGroups []string, usageLimit int, userID string, ephemeral bool) (*types.SetupKey, error)
+		autoGroups []string, usageLimit int, userID string, ephemeral bool, allowExtraDNSLabels bool) (*types.SetupKey, error)
 	SaveSetupKey(ctx context.Context, accountID string, key *types.SetupKey, userID string) (*types.SetupKey, error)
 	CreateUser(ctx context.Context, accountID, initiatorUserID string, key *types.UserInfo) (*types.UserInfo, error)
 	DeleteUser(ctx context.Context, accountID, initiatorUserID string, targetUserID string) error

management/server/account_test.go

@@ -1080,7 +1080,7 @@ func TestAccountManager_AddPeer(t *testing.T) {
 
 	serial := account.Network.CurrentSerial() // should be 0
 
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userID, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userID, false, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 		return
@@ -1456,7 +1456,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userID, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userID, false, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 		return
@@ -2948,7 +2948,7 @@ func setupNetworkMapTest(t *testing.T) (*DefaultAccountManager, *types.Account,
 		t.Fatal(err)
 	}
 
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userID, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userID, false, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 	}

management/server/grpcserver.go

@@ -481,6 +481,7 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 		UserID:          userID,
 		SetupKey:        loginReq.GetSetupKey(),
 		ConnectionIP:    realIP,
+		ExtraDNSLabels:  loginReq.GetDnsLabels(),
 	})
 	if err != nil {
 		log.WithContext(ctx).Warnf("failed logging in peer %s: %s", peerKey, err)

management/server/http/api/openapi.yml

@@ -361,6 +361,12 @@ components:
               description: System serial number
               type: string
               example: "C02XJ0J0JGH7"
+            extra_dns_labels:
+              description: Extra DNS labels added to the peer
+              type: array
+              items:
+                type: string
+                example: "stage-host-1"
           required:
             - city_name
             - connected
@@ -384,6 +390,7 @@ components:
             - ui_version
             - approval_required
             - serial_number
+            - extra_dns_labels
     AccessiblePeer:
       allOf:
         - $ref: '#/components/schemas/PeerMinimum'
@@ -503,6 +510,10 @@ components:
           description: Indicate that the peer will be ephemeral or not
           type: boolean
           example: true
+        allow_extra_dns_labels:
+          description: Allow extra DNS labels to be added to the peer
+          type: boolean
+          example: true
       required:
         - id
         - key
@@ -518,6 +529,7 @@ components:
         - updated_at
         - usage_limit
         - ephemeral
+        - allow_extra_dns_labels
     SetupKeyClear:
       allOf:
         - $ref: '#/components/schemas/SetupKeyBase'
@@ -587,6 +599,10 @@ components:
           description: Indicate that the peer will be ephemeral or not
           type: boolean
           example: true
+        allow_extra_dns_labels:
+          description: Allow extra DNS labels to be added to the peer
+          type: boolean
+          example: true
       required:
         - name
         - type

management/server/http/api/types.gen.go

@@ -297,6 +297,9 @@ type CountryCode = string
 
 // CreateSetupKeyRequest defines model for CreateSetupKeyRequest.
 type CreateSetupKeyRequest struct {
+	// AllowExtraDnsLabels Allow extra DNS labels to be added to the peer
+	AllowExtraDnsLabels *bool `json:"allow_extra_dns_labels,omitempty"`
+
 	// AutoGroups List of group IDs to auto-assign to peers registered with this key
 	AutoGroups []string `json:"auto_groups"`
 
@@ -689,6 +692,9 @@ type Peer struct {
 	// DnsLabel Peer's DNS label is the parsed peer name for domain resolution. It is used to form an FQDN by appending the account's domain to the peer label. e.g. peer-dns-label.netbird.cloud
 	DnsLabel string `json:"dns_label"`
 
+	// ExtraDnsLabels Extra DNS labels added to the peer
+	ExtraDnsLabels []string `json:"extra_dns_labels"`
+
 	// GeonameId Unique identifier from the GeoNames database for a specific geographical location.
 	GeonameId int `json:"geoname_id"`
 
@@ -767,6 +773,9 @@ type PeerBatch struct {
 	// DnsLabel Peer's DNS label is the parsed peer name for domain resolution. It is used to form an FQDN by appending the account's domain to the peer label. e.g. peer-dns-label.netbird.cloud
 	DnsLabel string `json:"dns_label"`
 
+	// ExtraDnsLabels Extra DNS labels added to the peer
+	ExtraDnsLabels []string `json:"extra_dns_labels"`
+
 	// GeonameId Unique identifier from the GeoNames database for a specific geographical location.
 	GeonameId int `json:"geoname_id"`
 
@@ -1230,6 +1239,9 @@ type RulePortRange struct {
 
 // SetupKey defines model for SetupKey.
 type SetupKey struct {
+	// AllowExtraDnsLabels Allow extra DNS labels to be added to the peer
+	AllowExtraDnsLabels bool `json:"allow_extra_dns_labels"`
+
 	// AutoGroups List of group IDs to auto-assign to peers registered with this key
 	AutoGroups []string `json:"auto_groups"`
 
@@ -1275,6 +1287,9 @@ type SetupKey struct {
 
 // SetupKeyBase defines model for SetupKeyBase.
 type SetupKeyBase struct {
+	// AllowExtraDnsLabels Allow extra DNS labels to be added to the peer
+	AllowExtraDnsLabels bool `json:"allow_extra_dns_labels"`
+
 	// AutoGroups List of group IDs to auto-assign to peers registered with this key
 	AutoGroups []string `json:"auto_groups"`
 
@@ -1317,6 +1332,9 @@ type SetupKeyBase struct {
 
 // SetupKeyClear defines model for SetupKeyClear.
 type SetupKeyClear struct {
+	// AllowExtraDnsLabels Allow extra DNS labels to be added to the peer
+	AllowExtraDnsLabels bool `json:"allow_extra_dns_labels"`
+
 	// AutoGroups List of group IDs to auto-assign to peers registered with this key
 	AutoGroups []string `json:"auto_groups"`
 

management/server/http/handlers/peers/peers_handler.go

@@ -338,6 +338,7 @@ func toSinglePeerResponse(peer *nbpeer.Peer, groupsInfo []api.GroupMinimum, dnsD
 		UserId:                      peer.UserID,
 		UiVersion:                   peer.Meta.UIVersion,
 		DnsLabel:                    fqdn(peer, dnsDomain),
+		ExtraDnsLabels:              fqdnList(peer.ExtraDNSLabels, dnsDomain),
 		LoginExpirationEnabled:      peer.LoginExpirationEnabled,
 		LastLogin:                   peer.GetLastLogin(),
 		LoginExpired:                peer.Status.LoginExpired,
@@ -372,6 +373,7 @@ func toPeerListItemResponse(peer *nbpeer.Peer, groupsInfo []api.GroupMinimum, dn
 		UserId:                 peer.UserID,
 		UiVersion:              peer.Meta.UIVersion,
 		DnsLabel:               fqdn(peer, dnsDomain),
+		ExtraDnsLabels:         fqdnList(peer.ExtraDNSLabels, dnsDomain),
 		LoginExpirationEnabled: peer.LoginExpirationEnabled,
 		LastLogin:              peer.GetLastLogin(),
 		LoginExpired:           peer.Status.LoginExpired,
@@ -392,3 +394,11 @@ func fqdn(peer *nbpeer.Peer, dnsDomain string) string {
 		return fqdn
 	}
 }
+func fqdnList(extraLabels []string, dnsDomain string) []string {
+	fqdnList := make([]string, 0, len(extraLabels))
+	for _, label := range extraLabels {
+		fqdn := fmt.Sprintf("%s.%s", label, dnsDomain)
+		fqdnList = append(fqdnList, fqdn)
+	}
+	return fqdnList
+}

management/server/http/handlers/routes/routes_handler.go

@@ -2,11 +2,8 @@ package routes
 
 import (
 	"encoding/json"
-	"fmt"
 	"net/http"
 	"net/netip"
-	"regexp"
-	"strings"
 	"unicode/utf8"
 
 	"github.com/gorilla/mux"
@@ -21,7 +18,6 @@ import (
 	"github.com/netbirdio/netbird/route"
 )
 
-const maxDomains = 32
 const failedToConvertRoute = "failed to convert route to response: %v"
 
 // handler is the routes handler of the account
@@ -102,7 +98,7 @@ func (h *handler) createRoute(w http.ResponseWriter, r *http.Request) {
 	var networkType route.NetworkType
 	var newPrefix netip.Prefix
 	if req.Domains != nil {
-		d, err := validateDomains(*req.Domains)
+		d, err := domain.ValidateDomains(*req.Domains)
 		if err != nil {
 			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid domains: %v", err), w)
 			return
@@ -225,7 +221,7 @@ func (h *handler) updateRoute(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if req.Domains != nil {
-		d, err := validateDomains(*req.Domains)
+		d, err := domain.ValidateDomains(*req.Domains)
 		if err != nil {
 			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid domains: %v", err), w)
 			return
@@ -350,34 +346,3 @@ func toRouteResponse(serverRoute *route.Route) (*api.Route, error) {
 	}
 	return route, nil
 }
-
-// validateDomains checks if each domain in the list is valid and returns a punycode-encoded DomainList.
-func validateDomains(domains []string) (domain.List, error) {
-	if len(domains) == 0 {
-		return nil, fmt.Errorf("domains list is empty")
-	}
-	if len(domains) > maxDomains {
-		return nil, fmt.Errorf("domains list exceeds maximum allowed domains: %d", maxDomains)
-	}
-
-	domainRegex := regexp.MustCompile(`^(?:\*\.)?(?:(?:xn--)?[a-zA-Z0-9_](?:[a-zA-Z0-9-_]{0,61}[a-zA-Z0-9])?\.)*(?:xn--)?[a-zA-Z0-9](?:[a-zA-Z0-9-_]{0,61}[a-zA-Z0-9])?$`)
-
-	var domainList domain.List
-
-	for _, d := range domains {
-		d := strings.ToLower(d)
-
-		// handles length and idna conversion
-		punycode, err := domain.FromString(d)
-		if err != nil {
-			return domainList, fmt.Errorf("failed to convert domain to punycode: %s: %v", d, err)
-		}
-
-		if !domainRegex.MatchString(string(punycode)) {
-			return domainList, fmt.Errorf("invalid domain format: %s", d)
-		}
-
-		domainList = append(domainList, punycode)
-	}
-	return domainList, nil
-}

management/server/http/handlers/routes/routes_handler_test.go

@@ -561,96 +561,6 @@ func TestRoutesHandlers(t *testing.T) {
 	}
 }
 
-func TestValidateDomains(t *testing.T) {
-	tests := []struct {
-		name     string
-		domains  []string
-		expected domain.List
-		wantErr  bool
-	}{
-		{
-			name:     "Empty list",
-			domains:  nil,
-			expected: nil,
-			wantErr:  true,
-		},
-		{
-			name:     "Valid ASCII domain",
-			domains:  []string{"sub.ex-ample.com"},
-			expected: domain.List{"sub.ex-ample.com"},
-			wantErr:  false,
-		},
-		{
-			name:     "Valid Unicode domain",
-			domains:  []string{"münchen.de"},
-			expected: domain.List{"xn--mnchen-3ya.de"},
-			wantErr:  false,
-		},
-		{
-			name:     "Valid Unicode, all labels",
-			domains:  []string{"中国.中国.中国"},
-			expected: domain.List{"xn--fiqs8s.xn--fiqs8s.xn--fiqs8s"},
-			wantErr:  false,
-		},
-		{
-			name:     "With underscores",
-			domains:  []string{"_jabber._tcp.gmail.com"},
-			expected: domain.List{"_jabber._tcp.gmail.com"},
-			wantErr:  false,
-		},
-		{
-			name:     "Invalid domain format",
-			domains:  []string{"-example.com"},
-			expected: nil,
-			wantErr:  true,
-		},
-		{
-			name:     "Invalid domain format 2",
-			domains:  []string{"example.com-"},
-			expected: nil,
-			wantErr:  true,
-		},
-		{
-			name:     "Multiple domains valid and invalid",
-			domains:  []string{"google.com", "invalid,nbdomain.com", "münchen.de"},
-			expected: domain.List{"google.com"},
-			wantErr:  true,
-		},
-		{
-			name:     "Valid wildcard domain",
-			domains:  []string{"*.example.com"},
-			expected: domain.List{"*.example.com"},
-			wantErr:  false,
-		},
-		{
-			name:     "Wildcard with dot domain",
-			domains:  []string{".*.example.com"},
-			expected: nil,
-			wantErr:  true,
-		},
-		{
-			name:     "Wildcard with dot domain",
-			domains:  []string{".*.example.com"},
-			expected: nil,
-			wantErr:  true,
-		},
-		{
-			name:     "Invalid wildcard domain",
-			domains:  []string{"a.*.example.com"},
-			expected: nil,
-			wantErr:  true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := validateDomains(tt.domains)
-			assert.Equal(t, tt.wantErr, err != nil)
-			assert.Equal(t, got, tt.expected)
-		})
-	}
-}
-
 func toApiRoute(t *testing.T, r *route.Route) *api.Route {
 	t.Helper()
 

management/server/http/handlers/setup_keys/setupkeys_handler.go

@@ -3,6 +3,7 @@ package setup_keys
 import (
 	"context"
 	"encoding/json"
+
 	"net/http"
 	"time"
 
@@ -86,8 +87,13 @@ func (h *handler) createSetupKey(w http.ResponseWriter, r *http.Request) {
 		ephemeral = *req.Ephemeral
 	}
 
+	var allowExtraDNSLabels bool
+	if req.AllowExtraDnsLabels != nil {
+		allowExtraDNSLabels = *req.AllowExtraDnsLabels
+	}
+
 	setupKey, err := h.accountManager.CreateSetupKey(r.Context(), accountID, req.Name, types.SetupKeyType(req.Type), expiresIn,
-		req.AutoGroups, req.UsageLimit, userID, ephemeral)
+		req.AutoGroups, req.UsageLimit, userID, ephemeral, allowExtraDNSLabels)
 	if err != nil {
 		util.WriteError(r.Context(), err, w)
 		return
@@ -237,19 +243,20 @@ func ToResponseBody(key *types.SetupKey) *api.SetupKey {
 	}
 
 	return &api.SetupKey{
-		Id:         key.Id,
-		Key:        key.KeySecret,
-		Name:       key.Name,
-		Expires:    key.GetExpiresAt(),
-		Type:       string(key.Type),
-		Valid:      key.IsValid(),
-		Revoked:    key.Revoked,
-		UsedTimes:  key.UsedTimes,
-		LastUsed:   key.GetLastUsed(),
-		State:      state,
-		AutoGroups: key.AutoGroups,
-		UpdatedAt:  key.UpdatedAt,
-		UsageLimit: key.UsageLimit,
-		Ephemeral:  key.Ephemeral,
+		Id:                  key.Id,
+		Key:                 key.KeySecret,
+		Name:                key.Name,
+		Expires:             key.GetExpiresAt(),
+		Type:                string(key.Type),
+		Valid:               key.IsValid(),
+		Revoked:             key.Revoked,
+		UsedTimes:           key.UsedTimes,
+		LastUsed:            key.GetLastUsed(),
+		State:               state,
+		AutoGroups:          key.AutoGroups,
+		UpdatedAt:           key.UpdatedAt,
+		UsageLimit:          key.UsageLimit,
+		Ephemeral:           key.Ephemeral,
+		AllowExtraDnsLabels: key.AllowExtraDNSLabels,
 	}
 }

management/server/http/handlers/setup_keys/setupkeys_handler_test.go

@@ -37,11 +37,12 @@ func initSetupKeysTestMetaData(defaultKey *types.SetupKey, newKey *types.SetupKe
 				return claims.AccountId, claims.UserId, nil
 			},
 			CreateSetupKeyFunc: func(_ context.Context, _ string, keyName string, typ types.SetupKeyType, _ time.Duration, _ []string,
-				_ int, _ string, ephemeral bool,
+				_ int, _ string, ephemeral bool, allowExtraDNSLabels bool,
 			) (*types.SetupKey, error) {
 				if keyName == newKey.Name || typ != newKey.Type {
 					nk := newKey.Copy()
 					nk.Ephemeral = ephemeral
+					nk.AllowExtraDNSLabels = allowExtraDNSLabels
 					return nk, nil
 				}
 				return nil, fmt.Errorf("failed creating setup key")
@@ -94,7 +95,7 @@ func TestSetupKeysHandlers(t *testing.T) {
 	adminUser := types.NewAdminUser("test_user")
 
 	newSetupKey, plainKey := types.GenerateSetupKey(newSetupKeyName, types.SetupKeyReusable, 0, []string{"group-1"},
-		types.SetupKeyUnlimitedUsage, true)
+		types.SetupKeyUnlimitedUsage, true, false)
 	newSetupKey.Key = plainKey
 	updatedDefaultSetupKey := defaultSetupKey.Copy()
 	updatedDefaultSetupKey.AutoGroups = []string{"group-1"}

management/server/management_proto_test.go

@@ -714,7 +714,7 @@ func Test_LoginPerformance(t *testing.T) {
 						return
 					}
 
-					setupKey, err := am.CreateSetupKey(context.Background(), account.Id, fmt.Sprintf("key-%d", j), types.SetupKeyReusable, time.Hour, nil, 0, fmt.Sprintf("user-%d", j), false)
+					setupKey, err := am.CreateSetupKey(context.Background(), account.Id, fmt.Sprintf("key-%d", j), types.SetupKeyReusable, time.Hour, nil, 0, fmt.Sprintf("user-%d", j), false, false)
 					if err != nil {
 						t.Logf("error creating setup key: %v", err)
 						return

management/server/mock_server/account_mock.go

@@ -25,7 +25,7 @@ type MockAccountManager struct {
 	GetOrCreateAccountByUserFunc func(ctx context.Context, userId, domain string) (*types.Account, error)
 	GetAccountFunc               func(ctx context.Context, accountID string) (*types.Account, error)
 	CreateSetupKeyFunc           func(ctx context.Context, accountId string, keyName string, keyType types.SetupKeyType,
-		expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool) (*types.SetupKey, error)
+		expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool, allowExtraDNSLabels bool) (*types.SetupKey, error)
 	GetSetupKeyFunc                     func(ctx context.Context, accountID, userID, keyID string) (*types.SetupKey, error)
 	AccountExistsFunc                   func(ctx context.Context, accountID string) (bool, error)
 	GetAccountIDByUserIdFunc            func(ctx context.Context, userId, domain string) (string, error)
@@ -205,9 +205,10 @@ func (am *MockAccountManager) CreateSetupKey(
 	usageLimit int,
 	userID string,
 	ephemeral bool,
+	allowExtraDNSLabels bool,
 ) (*types.SetupKey, error) {
 	if am.CreateSetupKeyFunc != nil {
-		return am.CreateSetupKeyFunc(ctx, accountID, keyName, keyType, expiresIn, autoGroups, usageLimit, userID, ephemeral)
+		return am.CreateSetupKeyFunc(ctx, accountID, keyName, keyType, expiresIn, autoGroups, usageLimit, userID, ephemeral, allowExtraDNSLabels)
 	}
 	return nil, status.Errorf(codes.Unimplemented, "method CreateSetupKey is not implemented")
 }

management/server/peer.go

@@ -15,6 +15,7 @@ import (
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/exp/maps"
 
+	"github.com/netbirdio/netbird/management/domain"
 	"github.com/netbirdio/netbird/management/server/geolocation"
 
 	"github.com/netbirdio/netbird/management/server/idp"
@@ -53,6 +54,9 @@ type PeerLogin struct {
 	SetupKey string
 	// ConnectionIP is the real IP of the peer
 	ConnectionIP net.IP
+
+	// ExtraDNSLabels is a list of extra DNS labels that the peer wants to use
+	ExtraDNSLabels []string
 }
 
 // GetPeers returns a list of peers under the given account filtering out peers that do not belong to a user if
@@ -502,6 +506,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		var setupKeyName string
 		var ephemeral bool
 		var groupsToAdd []string
+		var allowExtraDNSLabels bool
 		if addedByUser {
 			user, err := transaction.GetUserByUserID(ctx, store.LockingStrengthUpdate, userID)
 			if err != nil {
@@ -527,6 +532,11 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 			ephemeral = sk.Ephemeral
 			setupKeyID = sk.Id
 			setupKeyName = sk.Name
+			allowExtraDNSLabels = sk.AllowExtraDNSLabels
+
+			if !sk.AllowExtraDNSLabels && len(peer.ExtraDNSLabels) > 0 {
+				return status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key doesn't allow extra DNS labels")
+			}
 		}
 
 		if (strings.ToLower(peer.Meta.Hostname) == "iphone" || strings.ToLower(peer.Meta.Hostname) == "ipad") && userID != "" {
@@ -567,6 +577,8 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 			Ephemeral:                   ephemeral,
 			Location:                    peer.Location,
 			InactivityExpirationEnabled: addedByUser,
+			ExtraDNSLabels:              peer.ExtraDNSLabels,
+			AllowExtraDNSLabels:         allowExtraDNSLabels,
 		}
 		opEvent.TargetID = newPeer.ID
 		opEvent.Meta = newPeer.EventMeta(am.GetDNSDomain())
@@ -860,6 +872,20 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 			shouldStorePeer = true
 		}
 
+		if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
+			return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
+		}
+
+		extraLabels, err := domain.ValidateDomainsStrSlice(login.ExtraDNSLabels)
+		if err != nil {
+			return status.Errorf(status.InvalidArgument, "invalid extra DNS labels: %v", err)
+		}
+
+		if !slices.Equal(peer.ExtraDNSLabels, extraLabels) {
+			peer.ExtraDNSLabels = extraLabels
+			shouldStorePeer = true
+		}
+
 		if shouldStorePeer {
 			if err = transaction.SavePeer(ctx, store.LockingStrengthUpdate, accountID, peer); err != nil {
 				return err

management/server/peer/peer.go

@@ -49,6 +49,11 @@ type Peer struct {
 	Ephemeral bool `gorm:"index"`
 	// Geo location based on connection IP
 	Location Location `gorm:"embedded;embeddedPrefix:location_"`
+
+	// ExtraDNSLabels is a list of additional DNS labels that can be used to resolve the peer
+	ExtraDNSLabels []string `gorm:"serializer:json"`
+	// AllowExtraDNSLabels indicates whether the peer allows extra DNS labels to be used for resolving the peer
+	AllowExtraDNSLabels bool
 }
 
 type PeerStatus struct { //nolint:revive
@@ -202,6 +207,8 @@ func (p *Peer) Copy() *Peer {
 		Ephemeral:                   p.Ephemeral,
 		Location:                    p.Location,
 		InactivityExpirationEnabled: p.InactivityExpirationEnabled,
+		ExtraDNSLabels:              slices.Clone(p.ExtraDNSLabels),
+		AllowExtraDNSLabels:         p.AllowExtraDNSLabels,
 	}
 }
 

management/server/peer_test.go

@@ -168,7 +168,7 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userId, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userId, false, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 		return
@@ -417,7 +417,7 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userId, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, userId, false, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 		return
@@ -489,7 +489,7 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 	}
 
 	// two peers one added by a regular user and one with a setup key
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, adminUser, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, nil, 999, adminUser, false, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 		return

management/server/setupkey.go

@@ -52,7 +52,7 @@ type SetupKeyUpdateOperation struct {
 // CreateSetupKey generates a new setup key with a given name, type, list of groups IDs to auto-assign to peers registered with this key,
 // and adds it to the specified account. A list of autoGroups IDs can be empty.
 func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID string, keyName string, keyType types.SetupKeyType,
-	expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool) (*types.SetupKey, error) {
+	expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool, allowExtraDNSLabels bool) (*types.SetupKey, error) {
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
@@ -78,7 +78,7 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 			return status.Errorf(status.InvalidArgument, "invalid auto groups: %v", err)
 		}
 
-		setupKey, plainKey = types.GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
+		setupKey, plainKey = types.GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral, allowExtraDNSLabels)
 		setupKey.AccountID = accountID
 
 		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, autoGroups, nil, setupKey)

management/server/setupkey_test.go

@@ -50,7 +50,7 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 	keyName := "my-test-key"
 
 	key, err := manager.CreateSetupKey(context.Background(), account.Id, keyName, types.SetupKeyReusable, expiresIn, []string{},
-		types.SetupKeyUnlimitedUsage, userID, false)
+		types.SetupKeyUnlimitedUsage, userID, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -168,7 +168,7 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 	for _, tCase := range []testCase{testCase1, testCase2, testCase3} {
 		t.Run(tCase.name, func(t *testing.T) {
 			key, err := manager.CreateSetupKey(context.Background(), account.Id, tCase.expectedKeyName, types.SetupKeyReusable, expiresIn,
-				tCase.expectedGroups, types.SetupKeyUnlimitedUsage, userID, false)
+				tCase.expectedGroups, types.SetupKeyUnlimitedUsage, userID, false, false)
 
 			if tCase.expectedFailure {
 				if err == nil {
@@ -210,7 +210,7 @@ func TestGetSetupKeys(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	plainKey, err := manager.CreateSetupKey(context.Background(), account.Id, "key1", types.SetupKeyReusable, time.Hour, nil, types.SetupKeyUnlimitedUsage, userID, false)
+	plainKey, err := manager.CreateSetupKey(context.Background(), account.Id, "key1", types.SetupKeyReusable, time.Hour, nil, types.SetupKeyUnlimitedUsage, userID, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -275,7 +275,7 @@ func TestGenerateSetupKey(t *testing.T) {
 	expectedUpdatedAt := time.Now().UTC()
 	var expectedAutoGroups []string
 
-	key, plain := types.GenerateSetupKey(expectedName, types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	key, plain := types.GenerateSetupKey(expectedName, types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 
 	assertKey(t, key, expectedName, expectedRevoke, expectedType, expectedUsedTimes, expectedCreatedAt,
 		expectedExpiresAt, strconv.Itoa(int(types.Hash(plain))), expectedUpdatedAt, expectedAutoGroups, true)
@@ -283,33 +283,33 @@ func TestGenerateSetupKey(t *testing.T) {
 }
 
 func TestSetupKey_IsValid(t *testing.T) {
-	validKey, _ := types.GenerateSetupKey("valid key", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	validKey, _ := types.GenerateSetupKey("valid key", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 	if !validKey.IsValid() {
 		t.Errorf("expected key to be valid, got invalid %v", validKey)
 	}
 
 	// expired
-	expiredKey, _ := types.GenerateSetupKey("invalid key", types.SetupKeyOneOff, -time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	expiredKey, _ := types.GenerateSetupKey("invalid key", types.SetupKeyOneOff, -time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 	if expiredKey.IsValid() {
 		t.Errorf("expected key to be invalid due to expiration, got valid %v", expiredKey)
 	}
 
 	// revoked
-	revokedKey, _ := types.GenerateSetupKey("invalid key", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	revokedKey, _ := types.GenerateSetupKey("invalid key", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 	revokedKey.Revoked = true
 	if revokedKey.IsValid() {
 		t.Errorf("expected revoked key to be invalid, got valid %v", revokedKey)
 	}
 
 	// overused
-	overUsedKey, _ := types.GenerateSetupKey("invalid key", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	overUsedKey, _ := types.GenerateSetupKey("invalid key", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 	overUsedKey.UsedTimes = 1
 	if overUsedKey.IsValid() {
 		t.Errorf("expected overused key to be invalid, got valid %v", overUsedKey)
 	}
 
 	// overused
-	reusableKey, _ := types.GenerateSetupKey("valid key", types.SetupKeyReusable, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	reusableKey, _ := types.GenerateSetupKey("valid key", types.SetupKeyReusable, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 	reusableKey.UsedTimes = 99
 	if !reusableKey.IsValid() {
 		t.Errorf("expected reusable key to be valid when used many times, got valid %v", reusableKey)
@@ -388,7 +388,7 @@ func isValidBase64SHA256(encodedKey string) bool {
 
 func TestSetupKey_Copy(t *testing.T) {
 
-	key, _ := types.GenerateSetupKey("key name", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false)
+	key, _ := types.GenerateSetupKey("key name", types.SetupKeyOneOff, time.Hour, []string{}, types.SetupKeyUnlimitedUsage, false, false)
 	keyCopy := key.Copy()
 
 	assertKey(t, keyCopy, key.Name, key.Revoked, string(key.Type), key.UsedTimes, key.CreatedAt, key.GetExpiresAt(), key.Id,
@@ -436,7 +436,7 @@ func TestSetupKeyAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		setupKey, err = manager.CreateSetupKey(context.Background(), account.Id, "key1", types.SetupKeyReusable, time.Hour, nil, 999, userID, false)
+		setupKey, err = manager.CreateSetupKey(context.Background(), account.Id, "key1", types.SetupKeyReusable, time.Hour, nil, 999, userID, false, false)
 		assert.NoError(t, err)
 
 		select {
@@ -477,7 +477,7 @@ func TestDefaultAccountManager_CreateSetupKey_ShouldNotAllowToUpdateRevokedKey(t
 		t.Fatal(err)
 	}
 
-	key, err := manager.CreateSetupKey(context.Background(), account.Id, "testName", types.SetupKeyReusable, time.Hour, nil, types.SetupKeyUnlimitedUsage, userID, false)
+	key, err := manager.CreateSetupKey(context.Background(), account.Id, "testName", types.SetupKeyReusable, time.Hour, nil, types.SetupKeyUnlimitedUsage, userID, false, false)
 	assert.NoError(t, err)
 
 	// revoke the key

management/server/types/account.go

@@ -459,8 +459,23 @@ func (a *Account) GetPeersCustomZone(ctx context.Context, dnsDomain string) nbdn
 			TTL:   defaultTTL,
 			RData: peer.IP.String(),
 		})
-
 		sb.Reset()
+
+		for _, extraLabel := range peer.ExtraDNSLabels {
+			sb.Grow(len(extraLabel) + len(domainSuffix))
+			sb.WriteString(extraLabel)
+			sb.WriteString(domainSuffix)
+
+			customZone.Records = append(customZone.Records, nbdns.SimpleRecord{
+				Name:  sb.String(),
+				Type:  int(dns.TypeA),
+				Class: nbdns.DefaultClass,
+				TTL:   defaultTTL,
+				RData: peer.IP.String(),
+			})
+			sb.Reset()
+		}
+
 	}
 
 	go func() {

management/server/types/setupkey.go

@@ -10,6 +10,7 @@ import (
 	"unicode/utf8"
 
 	"github.com/google/uuid"
+
 	"github.com/netbirdio/netbird/management/server/util"
 )
 
@@ -54,6 +55,8 @@ type SetupKey struct {
 	UsageLimit int
 	// Ephemeral indicate if the peers will be ephemeral or not
 	Ephemeral bool
+	// AllowExtraDNSLabels indicates if the key allows extra DNS labels
+	AllowExtraDNSLabels bool
 }
 
 // Copy copies SetupKey to a new object
@@ -64,21 +67,22 @@ func (key *SetupKey) Copy() *SetupKey {
 		key.UpdatedAt = key.CreatedAt
 	}
 	return &SetupKey{
-		Id:         key.Id,
-		AccountID:  key.AccountID,
-		Key:        key.Key,
-		KeySecret:  key.KeySecret,
-		Name:       key.Name,
-		Type:       key.Type,
-		CreatedAt:  key.CreatedAt,
-		ExpiresAt:  key.ExpiresAt,
-		UpdatedAt:  key.UpdatedAt,
-		Revoked:    key.Revoked,
-		UsedTimes:  key.UsedTimes,
-		LastUsed:   key.LastUsed,
-		AutoGroups: autoGroups,
-		UsageLimit: key.UsageLimit,
-		Ephemeral:  key.Ephemeral,
+		Id:                  key.Id,
+		AccountID:           key.AccountID,
+		Key:                 key.Key,
+		KeySecret:           key.KeySecret,
+		Name:                key.Name,
+		Type:                key.Type,
+		CreatedAt:           key.CreatedAt,
+		ExpiresAt:           key.ExpiresAt,
+		UpdatedAt:           key.UpdatedAt,
+		Revoked:             key.Revoked,
+		UsedTimes:           key.UsedTimes,
+		LastUsed:            key.LastUsed,
+		AutoGroups:          autoGroups,
+		UsageLimit:          key.UsageLimit,
+		Ephemeral:           key.Ephemeral,
+		AllowExtraDNSLabels: key.AllowExtraDNSLabels,
 	}
 }
 
@@ -150,7 +154,7 @@ func (key *SetupKey) IsOverUsed() bool {
 
 // GenerateSetupKey generates a new setup key
 func GenerateSetupKey(name string, t SetupKeyType, validFor time.Duration, autoGroups []string,
-	usageLimit int, ephemeral bool) (*SetupKey, string) {
+	usageLimit int, ephemeral bool, allowExtraDNSLabels bool) (*SetupKey, string) {
 	key := strings.ToUpper(uuid.New().String())
 	limit := usageLimit
 	if t == SetupKeyOneOff {
@@ -166,26 +170,27 @@ func GenerateSetupKey(name string, t SetupKeyType, validFor time.Duration, autoG
 	encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
 
 	return &SetupKey{
-		Id:         strconv.Itoa(int(Hash(key))),
-		Key:        encodedHashedKey,
-		KeySecret:  HiddenKey(key, 4),
-		Name:       name,
-		Type:       t,
-		CreatedAt:  time.Now().UTC(),
-		ExpiresAt:  expiresAt,
-		UpdatedAt:  time.Now().UTC(),
-		Revoked:    false,
-		UsedTimes:  0,
-		AutoGroups: autoGroups,
-		UsageLimit: limit,
-		Ephemeral:  ephemeral,
+		Id:                  strconv.Itoa(int(Hash(key))),
+		Key:                 encodedHashedKey,
+		KeySecret:           HiddenKey(key, 4),
+		Name:                name,
+		Type:                t,
+		CreatedAt:           time.Now().UTC(),
+		ExpiresAt:           expiresAt,
+		UpdatedAt:           time.Now().UTC(),
+		Revoked:             false,
+		UsedTimes:           0,
+		AutoGroups:          autoGroups,
+		UsageLimit:          limit,
+		Ephemeral:           ephemeral,
+		AllowExtraDNSLabels: allowExtraDNSLabels,
 	}, key
 }
 
 // GenerateDefaultSetupKey generates a default reusable setup key with an unlimited usage and 30 days expiration
 func GenerateDefaultSetupKey() (*SetupKey, string) {
 	return GenerateSetupKey(DefaultSetupKeyName, SetupKeyReusable, DefaultSetupKeyDuration, []string{},
-		SetupKeyUnlimitedUsage, false)
+		SetupKeyUnlimitedUsage, false, false)
 }
 
 func Hash(s string) uint32 {