package client import ( "context" "errors" "fmt" "io" "sync" "time" "google.golang.org/grpc/codes" gstatus "google.golang.org/grpc/status" "github.com/cenkalti/backoff/v4" log "github.com/sirupsen/logrus" "golang.zx2c4.com/wireguard/wgctrl/wgtypes" "google.golang.org/grpc" "google.golang.org/grpc/connectivity" nbgrpc "github.com/netbirdio/netbird/client/grpc" "github.com/netbirdio/netbird/client/system" "github.com/netbirdio/netbird/encryption" "github.com/netbirdio/netbird/shared/management/domain" "github.com/netbirdio/netbird/shared/management/proto" "github.com/netbirdio/netbird/util/wsproxy" ) // Custom management client errors that abstract away gRPC error codes var ( // ErrPermissionDenied is returned when the server denies access to a resource ErrPermissionDenied = errors.New("permission denied") // ErrInvalidArgument is returned when the request contains invalid arguments ErrInvalidArgument = errors.New("invalid argument") // ErrUnauthenticated is returned when authentication is required ErrUnauthenticated = errors.New("unauthenticated") // ErrNotFound is returned when the requested resource is not found ErrNotFound = errors.New("not found") // ErrUnimplemented is returned when the operation is not implemented ErrUnimplemented = errors.New("not implemented") ) const ConnectTimeout = 10 * time.Second const ( errMsgMgmtPublicKey = "failed getting Management Service public key: %s" errMsgNoMgmtConnection = "no connection to management" ) // ConnStateNotifier is a wrapper interface of the status recorders type ConnStateNotifier interface { MarkManagementDisconnected(error) MarkManagementConnected() } type GrpcClient struct { key wgtypes.Key realClient proto.ManagementServiceClient conn *grpc.ClientConn connStateCallback ConnStateNotifier connStateCallbackLock sync.RWMutex addr string tlsEnabled bool reconnectMutex sync.Mutex } // NewClient creates a new client to Management service // The client is not connected after creation - call Connect to establish the connection func NewClient(addr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) *GrpcClient { return &GrpcClient{ key: ourPrivateKey, addr: addr, tlsEnabled: tlsEnabled, connStateCallbackLock: sync.RWMutex{}, reconnectMutex: sync.Mutex{}, } } // Connect establishes a connection to the Management Service with retry logic // Retries connection attempts with exponential backoff on failure func (c *GrpcClient) Connect(ctx context.Context) error { var conn *grpc.ClientConn operation := func() error { var err error conn, err = nbgrpc.CreateConnection(ctx, c.addr, c.tlsEnabled, wsproxy.ManagementComponent) if err != nil { log.Warnf("failed to connect to Management Service: %v", err) return err } return nil } if err := backoff.Retry(operation, defaultBackoff(ctx)); err != nil { log.Errorf("failed creating connection to Management Service after retries: %v", err) return fmt.Errorf("create connection: %w", err) } c.conn = conn c.realClient = proto.NewManagementServiceClient(conn) log.Infof("connected to the Management Service at %s", c.addr) return nil } // ConnectWithoutRetry establishes a connection to the Management Service without retry logic // Performs a single connection attempt - callers should implement their own retry logic if needed func (c *GrpcClient) ConnectWithoutRetry(ctx context.Context) error { conn, err := nbgrpc.CreateConnection(ctx, c.addr, c.tlsEnabled, wsproxy.ManagementComponent) if err != nil { log.Warnf("failed to connect to Management Service: %v", err) return fmt.Errorf("create connection: %w", err) } c.conn = conn c.realClient = proto.NewManagementServiceClient(conn) log.Debugf("connected to the Management Service at %s", c.addr) return nil } // Close closes connection to the Management Service func (c *GrpcClient) Close() error { return c.conn.Close() } // SetConnStateListener set the ConnStateNotifier func (c *GrpcClient) SetConnStateListener(notifier ConnStateNotifier) { c.connStateCallbackLock.Lock() defer c.connStateCallbackLock.Unlock() c.connStateCallback = notifier } // ready indicates whether the client is okay and ready to be used // for now it just checks whether gRPC connection to the service is ready func (c *GrpcClient) ready() bool { return c.conn.GetState() == connectivity.Ready || c.conn.GetState() == connectivity.Idle } // Sync wraps the real client's Sync endpoint call and takes care of retries and encryption/decryption of messages // Blocking request. The result will be sent via msgHandler callback function func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error { operation := func() error { log.Debugf("management connection state %v", c.conn.GetState()) connState := c.conn.GetState() if connState == connectivity.Shutdown { return backoff.Permanent(fmt.Errorf("connection to management has been shut down")) } else if !(connState == connectivity.Ready || connState == connectivity.Idle) { c.conn.WaitForStateChange(ctx, connState) return fmt.Errorf("connection to management is not ready and in %s state", connState) } serverPubKey, err := c.getServerPublicKey(ctx) if err != nil { log.Debugf(errMsgMgmtPublicKey, err) return err } return c.handleStream(ctx, *serverPubKey, sysInfo, msgHandler) } err := backoff.Retry(operation, defaultBackoff(ctx)) if err != nil { log.Warnf("exiting the Management service connection retry loop due to the unrecoverable error: %s", err) } return err } func (c *GrpcClient) handleStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error { ctx, cancelStream := context.WithCancel(ctx) defer cancelStream() stream, err := c.connectToStream(ctx, serverPubKey, sysInfo) if err != nil { log.Debugf("failed to open Management Service stream: %s", err) if s, ok := gstatus.FromError(err); ok && s.Code() == codes.PermissionDenied { return backoff.Permanent(err) // unrecoverable error, propagate to the upper layer } return err } log.Infof("connected to the Management Service stream") c.notifyConnected() // blocking until error err = c.receiveEvents(stream, serverPubKey, msgHandler) if err != nil { c.notifyDisconnected(err) s, _ := gstatus.FromError(err) switch s.Code() { case codes.PermissionDenied: return backoff.Permanent(err) // unrecoverable error, propagate to the upper layer case codes.Canceled: log.Debugf("management connection context has been canceled, this usually indicates shutdown") return nil default: log.Warnf("disconnected from the Management service but will retry silently. Reason: %v", err) return err } } return nil } // GetNetworkMap return with the network map func (c *GrpcClient) GetNetworkMap(ctx context.Context, sysInfo *system.Info) (*proto.NetworkMap, error) { serverPubKey, err := c.getServerPublicKey(ctx) if err != nil { log.Debugf("failed getting Management Service public key: %s", err) return nil, err } stream, err := c.connectToStream(ctx, *serverPubKey, sysInfo) if err != nil { log.Debugf("failed to open Management Service stream: %s", err) return nil, err } defer func() { _ = stream.CloseSend() }() update, err := stream.Recv() if err == io.EOF { log.Debugf("Management stream has been closed by server: %s", err) return nil, err } if err != nil { log.Debugf("disconnected from Management Service sync stream: %v", err) return nil, err } decryptedResp := &proto.SyncResponse{} err = encryption.DecryptMessage(*serverPubKey, c.key, update.Body, decryptedResp) if err != nil { log.Errorf("failed decrypting update message from Management Service: %s", err) return nil, err } if decryptedResp.GetNetworkMap() == nil { return nil, fmt.Errorf("invalid msg, required network map") } return decryptedResp.GetNetworkMap(), nil } func (c *GrpcClient) connectToStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info) (proto.ManagementService_SyncClient, error) { req := &proto.SyncRequest{Meta: infoToMetaData(sysInfo)} myPrivateKey := c.key myPublicKey := myPrivateKey.PublicKey() encryptedReq, err := encryption.EncryptMessage(serverPubKey, myPrivateKey, req) if err != nil { log.Errorf("failed encrypting message: %s", err) return nil, err } syncReq := &proto.EncryptedMessage{WgPubKey: myPublicKey.String(), Body: encryptedReq} sync, err := c.realClient.Sync(ctx, syncReq) if err != nil { return nil, err } return sync, nil } func (c *GrpcClient) receiveEvents(stream proto.ManagementService_SyncClient, serverPubKey wgtypes.Key, msgHandler func(msg *proto.SyncResponse) error) error { for { update, err := stream.Recv() if err == io.EOF { log.Debugf("Management stream has been closed by server: %s", err) return err } if err != nil { log.Debugf("disconnected from Management Service sync stream: %v", err) return err } log.Debugf("got an update message from Management Service") decryptedResp := &proto.SyncResponse{} err = encryption.DecryptMessage(serverPubKey, c.key, update.Body, decryptedResp) if err != nil { log.Errorf("failed decrypting update message from Management Service: %s", err) return err } if err := msgHandler(decryptedResp); err != nil { log.Errorf("failed handling an update message received from Management Service: %v", err.Error()) } } } // getServerPublicKey returns server's WireGuard public key (used later for encrypting messages sent to the server) // This is a simple operation without retry logic - callers should handle retries at the operation level func (c *GrpcClient) getServerPublicKey(ctx context.Context) (*wgtypes.Key, error) { if !c.ready() { return nil, errors.New(errMsgNoMgmtConnection) } mgmCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() resp, err := c.realClient.GetServerKey(mgmCtx, &proto.Empty{}) if err != nil { log.Errorf("failed while getting Management Service public key: %v", err) return nil, fmt.Errorf("failed while getting Management Service public key") } key, err := wgtypes.ParseKey(resp.Key) if err != nil { return nil, err } return &key, nil } // IsHealthy probes the gRPC connection and returns false on errors func (c *GrpcClient) IsHealthy(ctx context.Context) bool { switch c.conn.GetState() { case connectivity.TransientFailure: return false case connectivity.Connecting: return true case connectivity.Shutdown: return true case connectivity.Idle: case connectivity.Ready: } healthCtx, cancel := context.WithTimeout(ctx, 1*time.Second) defer cancel() _, err := c.realClient.GetServerKey(healthCtx, &proto.Empty{}) if err != nil { c.notifyDisconnected(err) log.Warnf("health check returned: %s", err) return false } c.notifyConnected() return true } // HealthCheck verifies connectivity to the management server // Returns an error if the server is not reachable // Internally uses getServerPublicKey to verify the connection func (c *GrpcClient) HealthCheck(ctx context.Context) error { _, err := c.getServerPublicKey(ctx) return err } func (c *GrpcClient) login(ctx context.Context, req *proto.LoginRequest) (*proto.LoginResponse, error) { if !c.ready() { return nil, errors.New(errMsgNoMgmtConnection) } serverKey, err := c.getServerPublicKey(ctx) if err != nil { log.Debugf(errMsgMgmtPublicKey, err) return nil, err } loginReq, err := encryption.EncryptMessage(*serverKey, c.key, req) if err != nil { log.Errorf("failed to encrypt message: %s", err) return nil, err } var resp *proto.EncryptedMessage operation := func() error { mgmCtx, cancel := context.WithTimeout(ctx, ConnectTimeout) defer cancel() var err error resp, err = c.realClient.Login(mgmCtx, &proto.EncryptedMessage{ WgPubKey: c.key.PublicKey().String(), Body: loginReq, }) if err != nil { // retry only on context canceled if s, ok := gstatus.FromError(err); ok && s.Code() == codes.Canceled { return err } return backoff.Permanent(err) } return nil } err = backoff.Retry(operation, nbgrpc.Backoff(ctx)) if err != nil { log.Errorf("failed to login to Management Service: %v", err) return nil, err } loginResp := &proto.LoginResponse{} err = encryption.DecryptMessage(*serverKey, c.key, resp.Body, loginResp) if err != nil { log.Errorf("failed to decrypt login response: %s", err) return nil, err } return loginResp, nil } // Register registers peer on Management Server. It actually calls a Login endpoint with a provided setup key // Takes care of encrypting and decrypting messages. // This method will also collect system info and send it with the request (e.g. hostname, os, etc) // Returns custom errors: ErrPermissionDenied, ErrInvalidArgument, ErrUnauthenticated func (c *GrpcClient) Register(ctx context.Context, setupKey string, jwtToken string, sysInfo *system.Info, pubSSHKey []byte, dnsLabels domain.List) error { keys := &proto.PeerKeys{ SshPubKey: pubSSHKey, WgPubKey: []byte(c.key.PublicKey().String()), } _, err := c.login(ctx, &proto.LoginRequest{SetupKey: setupKey, Meta: infoToMetaData(sysInfo), JwtToken: jwtToken, PeerKeys: keys, DnsLabels: dnsLabels.ToPunycodeList()}) return wrapGRPCError(err) } // Login attempts login to Management Server. Takes care of encrypting and decrypting messages. // Returns custom errors: ErrPermissionDenied, ErrInvalidArgument, ErrUnauthenticated func (c *GrpcClient) Login(ctx context.Context, sysInfo *system.Info, pubSSHKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) { keys := &proto.PeerKeys{ SshPubKey: pubSSHKey, WgPubKey: []byte(c.key.PublicKey().String()), } resp, err := c.login(ctx, &proto.LoginRequest{Meta: infoToMetaData(sysInfo), PeerKeys: keys, DnsLabels: dnsLabels.ToPunycodeList()}) return resp, wrapGRPCError(err) } // GetDeviceAuthorizationFlow returns a device authorization flow information. // It also takes care of encrypting and decrypting messages. // It automatically retries with backoff and reconnection on connection errors. // Returns custom errors: ErrNotFound, ErrUnimplemented func (c *GrpcClient) GetDeviceAuthorizationFlow(ctx context.Context) (*proto.DeviceAuthorizationFlow, error) { var flowInfoResp *proto.DeviceAuthorizationFlow err := c.withRetry(ctx, func() error { if !c.ready() { return fmt.Errorf("no connection to management in order to get device authorization flow") } serverKey, err := c.getServerPublicKey(ctx) if err != nil { log.Debugf(errMsgMgmtPublicKey, err) return err } mgmCtx, cancel := context.WithTimeout(ctx, time.Second*2) defer cancel() message := &proto.DeviceAuthorizationFlowRequest{} encryptedMSG, err := encryption.EncryptMessage(*serverKey, c.key, message) if err != nil { return err } resp, err := c.realClient.GetDeviceAuthorizationFlow(mgmCtx, &proto.EncryptedMessage{ WgPubKey: c.key.PublicKey().String(), Body: encryptedMSG}, ) if err != nil { return err } flowInfo := &proto.DeviceAuthorizationFlow{} err = encryption.DecryptMessage(*serverKey, c.key, resp.Body, flowInfo) if err != nil { errWithMSG := fmt.Errorf("failed to decrypt device authorization flow message: %s", err) log.Error(errWithMSG) return errWithMSG } flowInfoResp = flowInfo return nil }) return flowInfoResp, wrapGRPCError(err) } // GetPKCEAuthorizationFlow returns a pkce authorization flow information. // It also takes care of encrypting and decrypting messages. // It automatically retries with backoff and reconnection on connection errors. // Returns custom errors: ErrNotFound, ErrUnimplemented func (c *GrpcClient) GetPKCEAuthorizationFlow(ctx context.Context) (*proto.PKCEAuthorizationFlow, error) { var flowInfoResp *proto.PKCEAuthorizationFlow err := c.withRetry(ctx, func() error { if !c.ready() { return fmt.Errorf("no connection to management in order to get pkce authorization flow") } serverKey, err := c.getServerPublicKey(ctx) if err != nil { log.Debugf(errMsgMgmtPublicKey, err) return err } mgmCtx, cancel := context.WithTimeout(ctx, time.Second*2) defer cancel() message := &proto.PKCEAuthorizationFlowRequest{} encryptedMSG, err := encryption.EncryptMessage(*serverKey, c.key, message) if err != nil { return err } resp, err := c.realClient.GetPKCEAuthorizationFlow(mgmCtx, &proto.EncryptedMessage{ WgPubKey: c.key.PublicKey().String(), Body: encryptedMSG, }) if err != nil { return err } flowInfo := &proto.PKCEAuthorizationFlow{} err = encryption.DecryptMessage(*serverKey, c.key, resp.Body, flowInfo) if err != nil { errWithMSG := fmt.Errorf("failed to decrypt pkce authorization flow message: %s", err) log.Error(errWithMSG) return errWithMSG } flowInfoResp = flowInfo return nil }) return flowInfoResp, wrapGRPCError(err) } // SyncMeta sends updated system metadata to the Management Service. // It should be used if there is changes on peer posture check after initial sync. func (c *GrpcClient) SyncMeta(ctx context.Context, sysInfo *system.Info) error { if !c.ready() { return errors.New(errMsgNoMgmtConnection) } serverPubKey, err := c.getServerPublicKey(ctx) if err != nil { log.Debugf(errMsgMgmtPublicKey, err) return err } syncMetaReq, err := encryption.EncryptMessage(*serverPubKey, c.key, &proto.SyncMetaRequest{Meta: infoToMetaData(sysInfo)}) if err != nil { log.Errorf("failed to encrypt message: %s", err) return err } mgmCtx, cancel := context.WithTimeout(ctx, ConnectTimeout) defer cancel() _, err = c.realClient.SyncMeta(mgmCtx, &proto.EncryptedMessage{ WgPubKey: c.key.PublicKey().String(), Body: syncMetaReq, }) return err } func (c *GrpcClient) notifyDisconnected(err error) { c.connStateCallbackLock.RLock() defer c.connStateCallbackLock.RUnlock() if c.connStateCallback == nil { return } c.connStateCallback.MarkManagementDisconnected(err) } func (c *GrpcClient) notifyConnected() { c.connStateCallbackLock.RLock() defer c.connStateCallbackLock.RUnlock() if c.connStateCallback == nil { return } c.connStateCallback.MarkManagementConnected() } func (c *GrpcClient) Logout(ctx context.Context) error { serverKey, err := c.getServerPublicKey(ctx) if err != nil { return fmt.Errorf("get server public key: %w", err) } mgmCtx, cancel := context.WithTimeout(ctx, time.Second*15) defer cancel() message := &proto.Empty{} encryptedMSG, err := encryption.EncryptMessage(*serverKey, c.key, message) if err != nil { return fmt.Errorf("encrypt logout message: %w", err) } _, err = c.realClient.Logout(mgmCtx, &proto.EncryptedMessage{ WgPubKey: c.key.PublicKey().String(), Body: encryptedMSG, }) if err != nil { return fmt.Errorf("logout: %w", err) } return nil } // reconnect closes the current connection and creates a new one func (c *GrpcClient) reconnect(ctx context.Context) error { c.reconnectMutex.Lock() defer c.reconnectMutex.Unlock() // Close existing connection if c.conn != nil { if err := c.conn.Close(); err != nil { log.Debugf("error closing old connection: %v", err) } } // Create new connection log.Debugf("reconnecting to Management Service %s", c.addr) conn, err := nbgrpc.CreateConnection(ctx, c.addr, c.tlsEnabled, wsproxy.ManagementComponent) if err != nil { log.Errorf("failed reconnecting to Management Service %s: %v", c.addr, err) return fmt.Errorf("reconnect: create connection: %w", err) } c.conn = conn c.realClient = proto.NewManagementServiceClient(conn) log.Debugf("successfully reconnected to Management service %s", c.addr) return nil } // withRetry wraps an operation with exponential backoff retry logic // It automatically reconnects on connection errors func (c *GrpcClient) withRetry(ctx context.Context, operation func() error) error { backoffSettings := &backoff.ExponentialBackOff{ InitialInterval: 500 * time.Millisecond, RandomizationFactor: 0.5, Multiplier: 1.5, MaxInterval: 10 * time.Second, MaxElapsedTime: 2 * time.Minute, Stop: backoff.Stop, Clock: backoff.SystemClock, } backoffSettings.Reset() return backoff.RetryNotify( func() error { err := operation() if err == nil { return nil } // If it's a connection error, attempt reconnection if isConnectionError(err) { log.Warnf("connection error detected, attempting reconnection: %v", err) if reconnectErr := c.reconnect(ctx); reconnectErr != nil { log.Errorf("reconnection failed: %v", reconnectErr) return reconnectErr } // Return the original error to trigger retry with the new connection return err } // For authentication errors (InvalidArgument, PermissionDenied), don't retry if isAuthenticationError(err) { return backoff.Permanent(err) } return err }, backoff.WithContext(backoffSettings, ctx), func(err error, duration time.Duration) { log.Warnf("operation failed, retrying in %v: %v", duration, err) }, ) } // defaultBackoff is a basic backoff mechanism for general issues func defaultBackoff(ctx context.Context) backoff.BackOff { return backoff.WithContext(&backoff.ExponentialBackOff{ InitialInterval: 800 * time.Millisecond, RandomizationFactor: 1, Multiplier: 1.7, MaxInterval: 10 * time.Second, MaxElapsedTime: 3 * 30 * 24 * time.Hour, // 3 months Stop: backoff.Stop, Clock: backoff.SystemClock, }, ctx) } // isConnectionError checks if the error is a connection-related error that should trigger reconnection func isConnectionError(err error) bool { if err == nil { return false } s, ok := gstatus.FromError(err) if !ok { return false } // These error codes indicate connection issues return s.Code() == codes.Unavailable || s.Code() == codes.DeadlineExceeded || s.Code() == codes.Canceled || s.Code() == codes.Internal } // isAuthenticationError checks if the error is an authentication-related error that should not be retried func isAuthenticationError(err error) bool { if err == nil { return false } s, ok := gstatus.FromError(err) if !ok { return false } return s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied } // wrapGRPCError converts gRPC errors to custom management client errors func wrapGRPCError(err error) error { if err == nil { return nil } // Check if it's already a custom error if errors.Is(err, ErrPermissionDenied) || errors.Is(err, ErrInvalidArgument) || errors.Is(err, ErrUnauthenticated) || errors.Is(err, ErrNotFound) || errors.Is(err, ErrUnimplemented) { return err } // Convert gRPC status errors to custom errors s, ok := gstatus.FromError(err) if !ok { return err } switch s.Code() { case codes.PermissionDenied: return fmt.Errorf("%w: %s", ErrPermissionDenied, s.Message()) case codes.InvalidArgument: return fmt.Errorf("%w: %s", ErrInvalidArgument, s.Message()) case codes.Unauthenticated: return fmt.Errorf("%w: %s", ErrUnauthenticated, s.Message()) case codes.NotFound: return fmt.Errorf("%w: %s", ErrNotFound, s.Message()) case codes.Unimplemented: return fmt.Errorf("%w: %s", ErrUnimplemented, s.Message()) default: return err } } func infoToMetaData(info *system.Info) *proto.PeerSystemMeta { if info == nil { return nil } addresses := make([]*proto.NetworkAddress, 0, len(info.NetworkAddresses)) for _, addr := range info.NetworkAddresses { addresses = append(addresses, &proto.NetworkAddress{ NetIP: addr.NetIP.String(), Mac: addr.Mac, }) } files := make([]*proto.File, 0, len(info.Files)) for _, file := range info.Files { files = append(files, &proto.File{ Path: file.Path, Exist: file.Exist, ProcessIsRunning: file.ProcessIsRunning, }) } return &proto.PeerSystemMeta{ Hostname: info.Hostname, GoOS: info.GoOS, OS: info.OS, Core: info.OSVersion, OSVersion: info.OSVersion, Platform: info.Platform, Kernel: info.Kernel, NetbirdVersion: info.NetbirdVersion, UiVersion: info.UIVersion, KernelVersion: info.KernelVersion, NetworkAddresses: addresses, SysSerialNumber: info.SystemSerialNumber, SysManufacturer: info.SystemManufacturer, SysProductName: info.SystemProductName, Environment: &proto.Environment{ Cloud: info.Environment.Cloud, Platform: info.Environment.Platform, }, Files: files, Flags: &proto.Flags{ RosenpassEnabled: info.RosenpassEnabled, RosenpassPermissive: info.RosenpassPermissive, ServerSSHAllowed: info.ServerSSHAllowed, DisableClientRoutes: info.DisableClientRoutes, DisableServerRoutes: info.DisableServerRoutes, DisableDNS: info.DisableDNS, DisableFirewall: info.DisableFirewall, BlockLANAccess: info.BlockLANAccess, BlockInbound: info.BlockInbound, LazyConnectionEnabled: info.LazyConnectionEnabled, }, } }