From 553a13588bea8ac9a70c5001379a6ef6ed68fbc2 Mon Sep 17 00:00:00 2001 From: Misha Bragin Date: Thu, 1 Sep 2022 18:28:45 +0200 Subject: [PATCH] Free up gRPC client resources on errors (#448) --- client/internal/config.go | 14 ++++----- client/internal/connect.go | 61 ++++++++++++++++++++------------------ client/internal/login.go | 14 ++++++--- 3 files changed, 49 insertions(+), 40 deletions(-) diff --git a/client/internal/config.go b/client/internal/config.go index c573160b8..d02fb2167 100644 --- a/client/internal/config.go +++ b/client/internal/config.go @@ -226,7 +226,13 @@ func GetDeviceAuthorizationFlowInfo(ctx context.Context, config *Config) (Device log.Errorf("failed connecting to Management Service %s %v", config.ManagementURL.String(), err) return DeviceAuthorizationFlow{}, err } - log.Debugf("connected to management Service %s", config.ManagementURL.String()) + log.Debugf("connected to the Management service %s", config.ManagementURL.String()) + defer func() { + err = mgmClient.Close() + if err != nil { + log.Warnf("failed to close the Management service client %v", err) + } + }() serverKey, err := mgmClient.GetServerPublicKey() if err != nil { @@ -245,12 +251,6 @@ func GetDeviceAuthorizationFlowInfo(ctx context.Context, config *Config) (Device } } - err = mgmClient.Close() - if err != nil { - log.Errorf("failed closing Management Service client: %v", err) - return DeviceAuthorizationFlow{}, err - } - return DeviceAuthorizationFlow{ Provider: protoDeviceAuthorizationFlow.Provider.String(), diff --git a/client/internal/connect.go b/client/internal/connect.go index 6d35113ce..719a43bd8 100644 --- a/client/internal/connect.go +++ b/client/internal/connect.go @@ -79,9 +79,21 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *nbStatus.Sta cancel() }() + log.Debugf("conecting to the Management service %s", config.ManagementURL.Host) + mgmClient, err := mgm.NewClient(engineCtx, config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled) + if err != nil { + return wrapErr(gstatus.Errorf(codes.FailedPrecondition, "failed connecting to Management Service : %s", err)) + } + log.Debugf("connected to the Management service %s", config.ManagementURL.Host) + defer func() { + err = mgmClient.Close() + if err != nil { + log.Warnf("failed to close the Management service client %v", err) + } + }() + // connect (just a connection, no stream yet) and login to Management Service to get an initial global Wiretrustee config - mgmClient, loginResp, err := connectToManagement(engineCtx, config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled, - publicSSHKey) + loginResp, err := loginToManagement(engineCtx, mgmClient, publicSSHKey) if err != nil { log.Debug(err) if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.PermissionDenied) { @@ -114,6 +126,12 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *nbStatus.Sta log.Error(err) return wrapErr(err) } + defer func() { + err = signalClient.Close() + if err != nil { + log.Warnf("failed closing Signal service client %v", err) + } + }() statusRecorder.MarkSignalConnected(signalURL) @@ -139,18 +157,6 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *nbStatus.Sta backOff.Reset() - err = mgmClient.Close() - if err != nil { - log.Errorf("failed closing Management Service client %v", err) - return wrapErr(err) - } - - err = signalClient.Close() - if err != nil { - log.Errorf("failed closing Signal Service client %v", err) - return wrapErr(err) - } - err = engine.Stop() if err != nil { log.Errorf("failed stopping engine %v", err) @@ -215,34 +221,26 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.WiretrusteeConfig, return signalClient, nil } -// connectToManagement creates Management Services client, establishes a connection, logs-in and gets a global Wiretrustee config (signal, turn, stun hosts, etc) -func connectToManagement(ctx context.Context, managementAddr string, ourPrivateKey wgtypes.Key, tlsEnabled bool, pubSSHKey []byte) (*mgm.GrpcClient, *mgmProto.LoginResponse, error) { - log.Debugf("connecting to Management Service %s", managementAddr) - client, err := mgm.NewClient(ctx, managementAddr, ourPrivateKey, tlsEnabled) - if err != nil { - return nil, nil, gstatus.Errorf(codes.FailedPrecondition, "failed connecting to Management Service : %s", err) - } - log.Debugf("connected to management server %s", managementAddr) +// loginToManagement creates Management Services client, establishes a connection, logs-in and gets a global Wiretrustee config (signal, turn, stun hosts, etc) +func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte) (*mgmProto.LoginResponse, error) { serverPublicKey, err := client.GetServerPublicKey() if err != nil { - return nil, nil, gstatus.Errorf(codes.FailedPrecondition, "failed while getting Management Service public key: %s", err) + return nil, gstatus.Errorf(codes.FailedPrecondition, "failed while getting Management Service public key: %s", err) } sysInfo := system.GetInfo(ctx) loginResp, err := client.Login(*serverPublicKey, sysInfo, pubSSHKey) if err != nil { - return nil, nil, err + return nil, err } - log.Debugf("peer logged in to Management Service %s", managementAddr) - - return client, loginResp, nil + return loginResp, nil } -// NB: hardcoded from github.com/netbirdio/netbird/management/cmd to avoid import // ManagementLegacyPort is the port that was used before by the Management gRPC server. // It is used for backward compatibility now. +// NB: hardcoded from github.com/netbirdio/netbird/management/cmd to avoid import const ManagementLegacyPort = 33073 // UpdateOldManagementPort checks whether client can switch to the new Management port 443. @@ -286,7 +284,12 @@ func UpdateOldManagementPort(ctx context.Context, config *Config, configPath str log.Infof("couldn't switch to the new Management %s", newURL.String()) return config, err } - defer client.Close() //nolint + defer func() { + err = client.Close() + if err != nil { + log.Warnf("failed to close the Management service client %v", err) + } + }() // gRPC check _, err = client.GetServerPublicKey() diff --git a/client/internal/login.go b/client/internal/login.go index 2064a8b7e..3946cfda5 100644 --- a/client/internal/login.go +++ b/client/internal/login.go @@ -26,13 +26,19 @@ func Login(ctx context.Context, config *Config, setupKey string, jwtToken string mgmTlsEnabled = true } - log.Debugf("connecting to Management Service %s", config.ManagementURL.String()) + log.Debugf("connecting to the 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) + log.Errorf("failed connecting to the Management service %s %v", config.ManagementURL.String(), err) return err } - log.Debugf("connected to management Service %s", config.ManagementURL.String()) + log.Debugf("connected to the Management service %s", config.ManagementURL.String()) + defer func() { + err = mgmClient.Close() + if err != nil { + log.Warnf("failed to close the Management service client %v", err) + } + }() serverKey, err := mgmClient.GetServerPublicKey() if err != nil { @@ -53,7 +59,7 @@ func Login(ctx context.Context, config *Config, setupKey string, jwtToken string err = mgmClient.Close() if err != nil { - log.Errorf("failed closing Management Service client: %v", err) + log.Errorf("failed to close the Management service client: %v", err) return err }