Add replaceKey variable

config/configs.go

@@ -25,6 +25,8 @@ var CaCertificatePEM []byte
 var VereignCaCertificatePEM []byte
 var VereignCaKeyPEM []byte
 
+var ReplaceKey bool
+
 var MaxMessageSize int
 
 var GrpcListenAddress string
@@ -42,6 +44,8 @@ func SetConfigValues(configFile, etcdURL string) {
 	viper.SetDefault("dataStorageUrl", "data-storage-agent:7777")
 	viper.SetDefault("entitiesManagerUrl", "entities-management-agent:7779")
 
+	viper.SetDefault("replaceKey", false)
+
 	viper.SetDefault("viamUUID", "viam-system")
 	viper.SetDefault("viamSession", "viam-session")
 
@@ -127,6 +131,8 @@ func SetConfigValues(configFile, etcdURL string) {
 	DataStorageUrl = viper.GetString("dataStorageUrl")
 	EntitiesManagerUrl = viper.GetString("entitiesManagerUrl")
 
+	ReplaceKey = viper.GetBool("replaceKey")
+
 	SystemAuth.Uuid = viper.GetString("viamUUID")
 	SystemAuth.Session = viper.GetString("viamSession")
 

handler/handler.go

@@ -22,6 +22,7 @@ import (
 	"errors"
 	"log"
 	"strings"
+
 	"code.vereign.com/code/viam-apis/clientutils"
 
 	keyutils "code.vereign.com/code/key-storage-agent/utils"
@@ -113,7 +114,7 @@ func (s *KeyStorageServerImpl) GetKey(ctx context.Context, in *api.GetKeyRequest
 		getKeyResponse.StatusList = utils.AddStatus(getKeyResponse.StatusList,
 			"500", api.StatusType_ERROR, errorsString)
 	}
-	
+
 	if !hasData {
 		log.Println("No such key " + in.Uuid)
 		getKeyResponse.Key = nil
@@ -149,18 +150,20 @@ func (s *KeyStorageServerImpl) SetKey(ctx context.Context, in *api.SetKeyRequest
 
 	key := &api.Key{}
 
-	_, _, err := client.GetData("keys", in.Uuid+"/"+api.KeyType.String(in.KeyType), key)
-	if err != nil {
-		log.Printf("grpc call GetData to DataStorage failed: %s", err)
-		setKeyResponse.StatusList = utils.AddStatus(setKeyResponse.StatusList,
-			"500", api.StatusType_ERROR, err.Error())
-		return setKeyResponse, nil
-	}
+	if ReplaceKey == false {
+		_, _, err := client.GetData("keys", in.Uuid+"/"+api.KeyType.String(in.KeyType), key)
+		if err != nil {
+			log.Printf("grpc call GetData to DataStorage failed: %s", err)
+			setKeyResponse.StatusList = utils.AddStatus(setKeyResponse.StatusList,
+				"500", api.StatusType_ERROR, err.Error())
+			return setKeyResponse, nil
+		}
 
-	if len(key.Content) > 0 {
-		setKeyResponse.StatusList = utils.AddStatus(setKeyResponse.StatusList,
-			"400", api.StatusType_ERROR, "Key is already set")
-		return setKeyResponse, nil
+		if len(key.Content) > 0 {
+			setKeyResponse.StatusList = utils.AddStatus(setKeyResponse.StatusList,
+				"400", api.StatusType_ERROR, "Key is already set")
+			return setKeyResponse, nil
+		}
 	}
 
 	result, errors, err := client.PutData("keys", in.Uuid+"/"+api.KeyType.String(in.KeyType), in.Key)
@@ -268,7 +271,7 @@ func (s *KeyStorageServerImpl) GetKeyId(ctx context.Context, in *api.GetKeyIdByK
 			log.Printf("Error: %s", errorsString)
 			return nil, errors.New(errorsString)
 		}
-	
+
 		if !hasData {
 			log.Println("No such checkID " + checkID)
 			return nil, errors.New("No such checkID " + checkID)