Merge branch 'main' into client_server_test

netbirdio · Feb 8, 2022 · 95ab5d9 · 95ab5d9
2 parents a789083 + 3c959bb
commit 95ab5d9
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 49 deletions.
diff --git a/client/cmd/login.go b/client/cmd/login.go
@@ -26,55 +26,78 @@ var (
 		RunE: func(cmd *cobra.Command, args []string) error {
 			SetFlagsFromEnvVars()
 
-			err := util.InitLog(logLevel, logFile)
-			if err != nil {
-				log.Errorf("failed initializing log %v", err)
-				return err
-			}
-
-			config, err := internal.GetConfig(managementURL, configPath, preSharedKey)
-			if err != nil {
-				log.Errorf("failed getting config %s %v", configPath, err)
-				return err
+			var backOff = &backoff.ExponentialBackOff{
+				InitialInterval:     time.Second,
+				RandomizationFactor: backoff.DefaultRandomizationFactor,
+				Multiplier:          backoff.DefaultMultiplier,
+				MaxInterval:         2 * time.Second,
+				MaxElapsedTime:      time.Second * 10,
+				Stop:                backoff.Stop,
+				Clock:               backoff.SystemClock,
 			}
 
-			//validate our peer's Wireguard PRIVATE key
-			myPrivateKey, err := wgtypes.ParseKey(config.PrivateKey)
-			if err != nil {
-				log.Errorf("failed parsing Wireguard key %s: [%s]", config.PrivateKey, err.Error())
-				return err
-			}
-
-			ctx := context.Background()
-
-			mgmTlsEnabled := false
-			if config.ManagementURL.Scheme == "https" {
-				mgmTlsEnabled = true
-			}
-
-			log.Debugf("connecting to Management Service %s", config.ManagementURL.String())
-			mgmClient, err := mgm.NewClient(ctx, config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled)
-			if err != nil {
-				log.Errorf("failed connecting to Management Service %s %v", config.ManagementURL.String(), err)
-				return err
-			}
-			log.Debugf("connected to anagement Service %s", config.ManagementURL.String())
-
-			serverKey, err := mgmClient.GetServerPublicKey()
-			if err != nil {
-				log.Errorf("failed while getting Management Service public key: %v", err)
-				return err
-			}
-
-			_, err = loginPeer(*serverKey, mgmClient, setupKey)
-			if err != nil {
-				log.Errorf("failed logging-in peer on Management Service : %v", err)
-				return err
+			loginOp := func() error {
+
+				err := util.InitLog(logLevel, logFile)
+				if err != nil {
+					log.Errorf("failed initializing log %v", err)
+					return err
+				}
+
+				config, err := internal.GetConfig(managementURL, configPath, preSharedKey)
+				if err != nil {
+					log.Errorf("failed getting config %s %v", configPath, err)
+					return err
+				}
+
+				//validate our peer's Wireguard PRIVATE key
+				myPrivateKey, err := wgtypes.ParseKey(config.PrivateKey)
+				if err != nil {
+					log.Errorf("failed parsing Wireguard key %s: [%s]", config.PrivateKey, err.Error())
+					return err
+				}
+
+				ctx := context.Background()
+
+				mgmTlsEnabled := false
+				if config.ManagementURL.Scheme == "https" {
+					mgmTlsEnabled = true
+				}
+
+				log.Debugf("connecting to Management Service %s", config.ManagementURL.String())
+				mgmClient, err := mgm.NewClient(ctx, config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled)
+				if err != nil {
+					log.Errorf("failed connecting to Management Service %s %v", config.ManagementURL.String(), err)
+					return err
+				}
+				log.Debugf("connected to management Service %s", config.ManagementURL.String())
+
+				serverKey, err := mgmClient.GetServerPublicKey()
+				if err != nil {
+					log.Errorf("failed while getting Management Service public key: %v", err)
+					return err
+				}
+
+				_, err = loginPeer(*serverKey, mgmClient, setupKey)
+				if err != nil {
+					log.Errorf("failed logging-in peer on Management Service : %v", err)
+					return err
+				}
+
+				err = mgmClient.Close()
+				if err != nil {
+					log.Errorf("failed closing Management Service client: %v", err)
+					return err
+				}
+
+				return nil
 			}
 
-			err = mgmClient.Close()
+			err := backoff.RetryNotify(loginOp, backOff, func(err error, duration time.Duration) {
+				log.Warnf("retrying Login to the Management service in %v due to error %v", duration, err)
+			})
 			if err != nil {
-				log.Errorf("failed closing Management Service client: %v", err)
+				log.Errorf("exiting login retry loop due to unrecoverable error: %v", err)
 				return err
 			}
 

diff --git a/management/client/grpc.go b/management/client/grpc.go
@@ -36,7 +36,7 @@ func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsE
 		transportOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{}))
 	}
 
-	mgmCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	mgmCtx, cancel := context.WithTimeout(ctx, time.Second*3)
 	defer cancel()
 	conn, err := grpc.DialContext(
 		mgmCtx,
@@ -185,7 +185,7 @@ func (c *GrpcClient) GetServerPublicKey() (*wgtypes.Key, error) {
 		return nil, fmt.Errorf("no connection to management")
 	}
 
-	mgmCtx, cancel := context.WithTimeout(c.ctx, 5*time.Second) //todo make a general setting
+	mgmCtx, cancel := context.WithTimeout(c.ctx, time.Second*2)
 	defer cancel()
 	resp, err := c.realClient.GetServerKey(mgmCtx, &proto.Empty{})
 	if err != nil {
@@ -209,7 +209,7 @@ func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*pro
 		log.Errorf("failed to encrypt message: %s", err)
 		return nil, err
 	}
-	mgmCtx, cancel := context.WithTimeout(c.ctx, 5*time.Second) //todo make a general setting
+	mgmCtx, cancel := context.WithTimeout(c.ctx, time.Second*2)
 	defer cancel()
 	resp, err := c.realClient.Login(mgmCtx, &proto.EncryptedMessage{
 		WgPubKey: c.key.PublicKey().String(),

diff --git a/signal/client/grpc.go b/signal/client/grpc.go
@@ -58,7 +58,7 @@ func NewClient(ctx context.Context, addr string, key wgtypes.Key, tlsEnabled boo
 		transportOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{}))
 	}
 
-	sigCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	sigCtx, cancel := context.WithTimeout(ctx, time.Second*3)
 	defer cancel()
 	conn, err := grpc.DialContext(
 		sigCtx,
@@ -291,7 +291,7 @@ func (c *GrpcClient) Send(msg *proto.Message) error {
 		return err
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 	defer cancel()
 	_, err = c.realClient.Send(ctx, encryptedMessage)
 	if err != nil {