Fixes for #4352

client/cmd/service_controller.go

@@ -61,7 +61,11 @@ func (p *program) Start(svc service.Service) error {
 			}
 		}
 
-		serverInstance := server.New(p.ctx, util.FindFirstLogPath(logFiles), configPath, profilesDisabled, updateSettingsDisabled)
+		key, err := getSetupKey()
+	if err != nil {
+		log.Errorf("Error getting setup key: %v", err)
+	}
+		serverInstance := server.New(p.ctx, util.FindFirstLogPath(logFiles), configPath, profilesDisabled, updateSettingsDisabled, key)
 		if err := serverInstance.Start(); err != nil {
 			log.Fatalf("failed to start daemon: %v", err)
 		}

client/cmd/up.go

@@ -196,7 +196,7 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command, activeProf *pr
 	r := peer.NewRecorder(config.ManagementURL.String())
 	r.GetFullStatus()
 
-	connectClient := internal.NewConnectClient(ctx, config, r)
+	connectClient := internal.NewConnectClient(ctx, config, r, providedSetupKey)
 	SetupDebugHandler(ctx, config, r, connectClient, "")
 
 	return connectClient.Run(nil)

client/embed/embed.go

@@ -131,7 +131,7 @@ func (c *Client) Start(startCtx context.Context) error {
 	}
 
 	recorder := peer.NewRecorder(c.config.ManagementURL.String())
-	client := internal.NewConnectClient(ctx, c.config, recorder)
+	client := internal.NewConnectClient(ctx, c.config, recorder, c.setupKey)
 
 	// either startup error (permanent backoff err) or nil err (successful engine up)
 	// TODO: make after-startup backoff err available

client/internal/connect.go

@@ -43,6 +43,7 @@ type ConnectClient struct {
 	statusRecorder *peer.Status
 	engine         *Engine
 	engineMutex    sync.Mutex
+	setupKey       string
 
 	persistSyncResponse bool
 }
@@ -51,13 +52,14 @@ func NewConnectClient(
 	ctx context.Context,
 	config *profilemanager.Config,
 	statusRecorder *peer.Status,
-
+	setupKey string,
 ) *ConnectClient {
 	return &ConnectClient{
 		ctx:            ctx,
 		config:         config,
 		statusRecorder: statusRecorder,
 		engineMutex:    sync.Mutex{},
+		setupKey:       setupKey,
 	}
 }
 
@@ -192,7 +194,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		}()
 
 		// connect (just a connection, no stream yet) and login to Management Service to get an initial global Netbird config
-		loginResp, err := loginToManagement(engineCtx, mgmClient, publicSSHKey, c.config)
+		loginResp, err := loginToManagement(engineCtx, mgmClient, publicSSHKey, c.config, c.setupKey)
 		if err != nil {
 			log.Debug(err)
 			if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.PermissionDenied) {
@@ -485,7 +487,7 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourP
 }
 
 // loginToManagement creates Management ServiceDependencies client, establishes a connection, logs-in and gets a global Netbird config (signal, turn, stun hosts, etc)
-func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte, config *profilemanager.Config) (*mgmProto.LoginResponse, error) {
+func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte, config *profilemanager.Config, setupKey string) (*mgmProto.LoginResponse, error) {
 
 	serverPublicKey, err := client.GetServerPublicKey()
 	if err != nil {
@@ -505,7 +507,12 @@ func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte,
 		config.BlockInbound,
 		config.LazyConnectionEnabled,
 	)
-	loginResp, err := client.Login(*serverPublicKey, sysInfo, pubSSHKey, config.DNSLabels)
+	var loginResp *mgmProto.LoginResponse
+	if setupKey != "" {
+		loginResp, err = client.Register(*serverPublicKey, setupKey, "", sysInfo, pubSSHKey, config.DNSLabels)
+	} else {
+		loginResp, err = client.Login(*serverPublicKey, sysInfo, pubSSHKey, config.DNSLabels)
+	}
 	if err != nil {
 		return nil, err
 	}

client/server/server.go

@@ -58,7 +58,8 @@ type Server struct {
 	rootCtx   context.Context
 	actCancel context.CancelFunc
 
-	logFile string
+	logFile  string
+	setupKey string
 
 	oauthAuthFlow oauthAuthFlow
 
@@ -88,10 +89,11 @@ type oauthAuthFlow struct {
 }
 
 // New server instance constructor.
-func New(ctx context.Context, logFile string, configFile string, profilesDisabled bool, updateSettingsDisabled bool) *Server {
+func New(ctx context.Context, logFile string, configFile string, profilesDisabled bool, updateSettingsDisabled bool, setupKey string) *Server {
 	return &Server{
 		rootCtx:                ctx,
 		logFile:                logFile,
+		setupKey:               setupKey,
 		persistSyncResponse:    true,
 		statusRecorder:         peer.NewRecorder(""),
 		profileManager:         profilemanager.NewServiceManager(configFile),
@@ -172,7 +174,7 @@ func (s *Server) Start() error {
 		return nil
 	}
 
-	go s.connectWithRetryRuns(ctx, config, s.statusRecorder, nil)
+	go s.connectWithRetryRuns(ctx, config, s.statusRecorder, nil, s.setupKey)
 
 	return nil
 }
@@ -205,7 +207,7 @@ func (s *Server) setDefaultConfigIfNotExists(ctx context.Context) error {
 // mechanism to keep the client connected even when the connection is lost.
 // we cancel retry if the client receive a stop or down command, or if disable auto connect is configured.
 func (s *Server) connectWithRetryRuns(ctx context.Context, config *profilemanager.Config, statusRecorder *peer.Status,
-	runningChan chan struct{},
+	runningChan chan struct{}, setupKey string,
 ) {
 	backOff := getConnectWithBackoff(ctx)
 	retryStarted := false
@@ -235,7 +237,7 @@ func (s *Server) connectWithRetryRuns(ctx context.Context, config *profilemanage
 
 	runOperation := func() error {
 		log.Tracef("running client connection")
-		s.connectClient = internal.NewConnectClient(ctx, config, statusRecorder)
+		s.connectClient = internal.NewConnectClient(ctx, config, statusRecorder, setupKey)
 		s.connectClient.SetSyncResponsePersistence(s.persistSyncResponse)
 
 		err := s.connectClient.Run(runningChan)
@@ -711,7 +713,7 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR
 	defer cancel()
 
 	runningChan := make(chan struct{}, 1) // buffered channel to do not lose the signal
-	go s.connectWithRetryRuns(ctx, s.config, s.statusRecorder, runningChan)
+	go s.connectWithRetryRuns(ctx, s.config, s.statusRecorder, runningChan, s.setupKey)
 	for {
 		select {
 		case <-runningChan: