[iamclient] Lock public & protected grpc calls if connection lost

Signed-off-by: Mykola Kobets <mykola_kobets@epam.com>

Author: mykola-kobets-epam

iamclient/iamcertprovider.go

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright (C) 2025 EPAM Systems, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iamclient
+
+import (
+	"context"
+
+	"github.com/aosedge/aos_common/aoserrors"
+	pb "github.com/aosedge/aos_common/api/iamanager"
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+)
+
+/***********************************************************************************************************************
+ * Consts
+ **********************************************************************************************************************/
+
+/***********************************************************************************************************************
+ * Types
+ **********************************************************************************************************************/
+
+// IAM certificate provider.
+type IAMCertProvider struct {
+	connection *grpc.ClientConn
+}
+
+/***********************************************************************************************************************
+ * Public
+ **********************************************************************************************************************/
+
+// New creates new certificate provider instance.
+func NewCertProvider(connection *grpc.ClientConn) *IAMCertProvider {
+	return &IAMCertProvider{connection: connection}
+}
+
+// GetCertificate gets certificate by issuer.
+func (provider *IAMCertProvider) GetCertificate(
+	certType string, issuer []byte, serial string,
+) (certURL, keyURL string, err error) {
+	log.Debug("Get IAM certificate")
+
+	ctx, cancel := context.WithTimeout(context.Background(), iamRequestTimeout)
+	defer cancel()
+
+	publicService := pb.NewIAMPublicServiceClient(provider.connection)
+	response, err := publicService.GetCert(
+		ctx, &pb.GetCertRequest{Type: certType, Issuer: issuer, Serial: serial})
+	if err != nil {
+		return "", "", aoserrors.Wrap(err)
+	}
+
+	return response.GetCertUrl(), response.GetKeyUrl(), nil
+}

iamclient/iamclient.go

@@ -74,7 +74,7 @@ type Client struct {
 	isMainNode bool
 
 	publicConnection    *grpchelpers.GRPCConn
-	protectedConnection *grpc.ClientConn
+	protectedConnection *grpchelpers.GRPCConn
 
 	publicService            pb.IAMPublicServiceClient
 	identService             pb.IAMPublicIdentityServiceClient
@@ -133,13 +133,14 @@ func New(
 	cryptocontext *cryptutils.CryptoContext, insecure bool,
 ) (client *Client, err error) {
 	localClient := &Client{
-		publicURL:        publicURL,
-		protectedURL:     protectedURL,
-		certStorage:      certStorage,
-		sender:           sender,
-		cryptocontext:    cryptocontext,
-		insecure:         insecure,
-		publicConnection: grpchelpers.NewGRPCConn(),
+		publicURL:           publicURL,
+		protectedURL:        protectedURL,
+		certStorage:         certStorage,
+		sender:              sender,
+		cryptocontext:       cryptocontext,
+		insecure:            insecure,
+		publicConnection:    grpchelpers.NewGRPCConn(),
+		protectedConnection: grpchelpers.NewGRPCConn(),
 		nodeInfoSubs: &nodeInfoChangeSub{
 			listeners: make([]chan cloudprotocol.NodeInfo, 0),
 		},
@@ -613,7 +614,6 @@ func (client *Client) Close() error {
 	client.isReconnecting.Store(true)
 
 	client.closeGRPCConnection()
-	client.publicConnection.Close()
 
 	log.Debug("Disconnected from IAM")
 
@@ -657,7 +657,8 @@ func (client *Client) SubscribeCertChanged(certType string) (<-chan *pb.CertInfo
 	ch := make(chan *pb.CertInfo, 1)
 
 	if _, ok := client.certChangeSub[certType]; !ok {
-		grpcStream, err := client.subscribeCertChange(certType)
+		grpcStream, err := client.publicService.SubscribeCertChanged(
+			context.Background(), &pb.SubscribeCertChangedRequest{Type: certType})
 		if err != nil {
 			return nil, aoserrors.Wrap(err)
 		}
@@ -839,38 +840,55 @@ func (client *Client) openGRPCConnection() (err error) {
 	log.Debug("Connecting to IAM...")
 
 	var publicConn *grpc.ClientConn
+	var protectedConn *grpc.ClientConn
 
 	publicConn, err = grpchelpers.CreatePublicConnection(
 		client.publicURL, client.cryptocontext, client.insecure)
 	if err != nil {
 		return aoserrors.Wrap(err)
 	}
 
-	if err := client.publicConnection.Start(publicConn); err != nil {
-		return aoserrors.Wrap(err)
-	}
+	client.publicConnection.Set(publicConn)
+
+	defer func() {
+		if err == nil {
+			client.publicConnection.Start()
+		} else {
+			publicConn.Close()
+		}
+	}()
 
 	client.publicService = pb.NewIAMPublicServiceClient(client.publicConnection)
 	client.identService = pb.NewIAMPublicIdentityServiceClient(client.publicConnection)
 	client.publicNodesService = pb.NewIAMPublicNodesServiceClient(client.publicConnection)
 	client.publicPermissionsService = pb.NewIAMPublicPermissionsServiceClient(client.publicConnection)
 
-	if err = client.restoreCertInfoSubs(); err != nil {
-		log.Error("Failed subscribe on CertInfo change")
-
-		return aoserrors.Wrap(err)
+	if err = client.restorePublicSubs(); err != nil {
+		return err
 	}
 
 	if !client.isProtectedConnEnabled() {
 		return nil
 	}
 
-	client.protectedConnection, err = grpchelpers.CreateProtectedConnection(client.certStorage,
-		client.protectedURL, client.cryptocontext, client, client.insecure)
+	certProvider := NewCertProvider(publicConn)
+
+	protectedConn, err = grpchelpers.CreateProtectedConnection(client.certStorage,
+		client.protectedURL, client.cryptocontext, certProvider, client.insecure)
 	if err != nil {
 		return aoserrors.Wrap(err)
 	}
 
+	client.protectedConnection.Set(protectedConn)
+
+	defer func() {
+		if err == nil {
+			client.protectedConnection.Start()
+		} else {
+			protectedConn.Close()
+		}
+	}()
+
 	client.certificateService = pb.NewIAMCertificateServiceClient(client.protectedConnection)
 	client.provisioningService = pb.NewIAMProvisioningServiceClient(client.protectedConnection)
 	client.nodesService = pb.NewIAMNodesServiceClient(client.protectedConnection)
@@ -889,7 +907,7 @@ func (client *Client) closeGRPCConnection() {
 	}
 
 	if client.protectedConnection != nil {
-		client.protectedConnection.Close()
+		client.protectedConnection.Stop()
 	}
 
 	for _, sub := range client.certChangeSub {
@@ -944,20 +962,6 @@ func (client *Client) subscribeUnitSubjectsChange() error {
 	return nil
 }
 
-func (client *Client) subscribeCertChange(certType string) (
-	listener pb.IAMPublicService_SubscribeCertChangedClient, err error,
-) {
-	listener, err = client.publicService.SubscribeCertChanged(context.Background(),
-		&pb.SubscribeCertChangedRequest{Type: certType})
-	if err != nil {
-		log.WithField("error", err).Error("Can't subscribe on CertChange event")
-
-		return nil, aoserrors.Wrap(err)
-	}
-
-	return listener, aoserrors.Wrap(err)
-}
-
 func (client *Client) processNodeInfoChange(sub *nodeInfoChangeSub) {
 	defer sub.stopWG.Done()
 
@@ -1182,9 +1186,10 @@ func (client *Client) finishProvisioning(nodeID, password string) (errorInfo *cl
 
 func (client *Client) restoreCertInfoSubs() error {
 	for certType, sub := range client.certChangeSub {
-		grpcStream, err := client.subscribeCertChange(certType)
+		grpcStream, err := client.publicService.SubscribeCertChanged(
+			context.Background(), &pb.SubscribeCertChangedRequest{Type: certType})
 		if err != nil {
-			return err
+			return aoserrors.Wrap(err)
 		}
 
 		sub.grpcStream = &grpcStream
@@ -1197,6 +1202,32 @@ func (client *Client) restoreCertInfoSubs() error {
 	return nil
 }
 
+func (client *Client) restorePublicSubs() error {
+	var err error
+
+	if err = client.restoreCertInfoSubs(); err != nil {
+		log.Error("Failed subscribe on CertInfo change")
+
+		return aoserrors.Wrap(err)
+	}
+
+	if client.isMainNode {
+		if err = client.subscribeNodeInfoChange(); err != nil {
+			log.Error("Failed subscribe on NodeInfo change")
+
+			return aoserrors.Wrap(err)
+		}
+
+		if err = client.subscribeUnitSubjectsChange(); err != nil {
+			log.Error("Failed subscribe on UnitSubject change")
+
+			return aoserrors.Wrap(err)
+		}
+	}
+
+	return nil
+}
+
 func (client *Client) onConnectionLost() {
 	select {
 	case <-client.closeChannel: