From 93d3973f34e85718ce41b345cab1f00a284e3e05 Mon Sep 17 00:00:00 2001
From: Artem Glazychev <artem.glazychev@xored.com>
Date: Wed, 15 Feb 2023 20:05:32 +0700
Subject: [PATCH] Revert back label values in tokenClient chain element

Signed-off-by: Artem Glazychev <artem.glazychev@xored.com>
---
 .../common/token/multitoken/client.go            | 16 +++++++++++++---
 .../common/token/multitoken/client_test.go       | 10 ++++++----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/pkg/networkservice/common/token/multitoken/client.go b/pkg/networkservice/common/token/multitoken/client.go
index b7ea3ebd..27d73370 100644
--- a/pkg/networkservice/common/token/multitoken/client.go
+++ b/pkg/networkservice/common/token/multitoken/client.go
@@ -1,6 +1,8 @@
-// Copyright (c) 2020-2022 Doc.ai and/or its affiliates.
+// Copyright (c) 2020-2023 Doc.ai and/or its affiliates.
 //
-// Copyright (c) 2021-2022 Nordix Foundation.
+// Copyright (c) 2021-2023 Nordix Foundation.
+//
+// Copyright (c) 2023 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -58,8 +60,10 @@ func (c *tokenClient) Request(ctx context.Context, request *networkservice.Netwo
 	isEstablished := c.config.get(request.GetConnection()) != ""
 
 	var tokenID string
+	var tokenName string
 	if labels := request.GetConnection().GetLabels(); labels != nil {
-		if tokenName, ok := labels[sriovTokenLabel]; ok {
+		var ok bool
+		if tokenName, ok = labels[sriovTokenLabel]; ok {
 			tokenID = c.config.assign(tokenName, request.GetConnection())
 			if tokenID == "" {
 				return nil, errors.Errorf("no free token for the name: %v", tokenName)
@@ -83,6 +87,12 @@ func (c *tokenClient) Request(ctx context.Context, request *networkservice.Netwo
 		c.config.release(request.GetConnection())
 	}
 
+	if tokenName != "" {
+		// Set the previous values in the labels. We need them for healing
+		delete(conn.GetLabels(), serviceDomainLabel)
+		conn.GetLabels()[sriovTokenLabel] = tokenName
+	}
+
 	return conn, err
 }
 
diff --git a/pkg/networkservice/common/token/multitoken/client_test.go b/pkg/networkservice/common/token/multitoken/client_test.go
index 0a31d1f9..663c00fc 100644
--- a/pkg/networkservice/common/token/multitoken/client_test.go
+++ b/pkg/networkservice/common/token/multitoken/client_test.go
@@ -1,6 +1,8 @@
-// Copyright (c) 2020-2022 Doc.ai and/or its affiliates.
+// Copyright (c) 2020-2023 Doc.ai and/or its affiliates.
 //
-// Copyright (c) 2021-2022 Nordix Foundation.
+// Copyright (c) 2021-2023 Nordix Foundation.
+//
+// Copyright (c) 2023 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -73,12 +75,12 @@ func TestTokenClient_Request(t *testing.T) {
 		token.NewClient(),
 		&validateClient{t},
 	)
-	_, err = client.Request(context.TODO(), request)
+	conn, err := client.Request(context.Background(), request)
 	require.NoError(t, err)
 
 	require.Equal(t, map[string]string{
 		sriovTokenLabel: tokenName,
-	}, request.GetConnection().GetLabels())
+	}, conn.GetLabels())
 }
 
 type validateClient struct {