From f493c37052e6859790177082ef5ab5c276d55551 Mon Sep 17 00:00:00 2001
From: Jay Mundrawala <jay@mondoo.com>
Date: Mon, 29 Jan 2024 15:03:07 -0600
Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B9=20Retry=20login=20command=20when?=
 =?UTF-8?q?=20concurrent=20IAM=20updates=20are=20detected=20(#3151)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 🧹 Retry login command when concurrent IAM updates are detected

If the server returns aborted for RegisterAgent, that means it aborted
the change because of concurrent writes to the IAM policy, likely
meaning concurrent registrations. Clients are safe to backoff and
retry in this case.

* 🧹 improve logout error message on cli

---------

Co-authored-by: Christoph Hartmann <chris@lollyrock.com>
---
 apps/cnquery/cmd/login.go  | 33 +++++++++++++++++++++++++++++++--
 apps/cnquery/cmd/logout.go |  2 +-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/apps/cnquery/cmd/login.go b/apps/cnquery/cmd/login.go
index 01f4f16097..2699654ef8 100644
--- a/apps/cnquery/cmd/login.go
+++ b/apps/cnquery/cmd/login.go
@@ -5,6 +5,7 @@ package cmd
 
 import (
 	"context"
+	"math/rand"
 	"strings"
 	"time"
 
@@ -19,7 +20,9 @@ import (
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/sysinfo"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/upstream"
 	"go.mondoo.com/ranger-rpc"
+	"go.mondoo.com/ranger-rpc/codes"
 	"go.mondoo.com/ranger-rpc/plugins/authentication/statictoken"
+	"go.mondoo.com/ranger-rpc/status"
 )
 
 func init() {
@@ -116,7 +119,7 @@ func register(token string, annotations map[string]string) error {
 			name = sysInfo.Hostname
 		}
 
-		confirmation, err := client.RegisterAgent(context.Background(), &upstream.AgentRegistrationRequest{
+		confirmation, err := registerAgent(context.Background(), client, &upstream.AgentRegistrationRequest{
 			Token: token,
 			Name:  name,
 			AgentInfo: &upstream.AgentInfo{
@@ -193,7 +196,7 @@ func register(token string, annotations map[string]string) error {
 				name = sysInfo.Hostname
 			}
 
-			confirmation, err := client.RegisterAgent(context.Background(), &upstream.AgentRegistrationRequest{
+			confirmation, err := registerAgent(context.Background(), client, &upstream.AgentRegistrationRequest{
 				Name: name,
 				AgentInfo: &upstream.AgentInfo{
 					Mrn:              opts.AgentMrn,
@@ -247,3 +250,29 @@ func register(token string, annotations map[string]string) error {
 	log.Info().Msgf("client %s has logged in successfully", viper.Get("agent_mrn"))
 	return nil
 }
+
+func registerAgent(ctx context.Context, client *upstream.AgentManagerClient, req *upstream.AgentRegistrationRequest) (*upstream.AgentRegistrationConfirmation, error) {
+	const maxRetries = 3
+	try := 0
+	for {
+		confirmation, err := client.RegisterAgent(ctx, req)
+		if err != nil {
+			if status.Code(err) == codes.Aborted {
+				jitter := time.Duration(rand.Intn(5000)) * time.Millisecond
+				sleepTime := 5*(1<<try)*time.Second + jitter
+
+				try++
+				if try > maxRetries {
+					return nil, errors.Wrap(err, "failed to log in client due to concurrent IAM changes")
+				}
+
+				log.Warn().Err(err).Msgf("failed to log in client due to concurrent IAM changes, retrying (%d/%d) in %dms", try, maxRetries, sleepTime.Milliseconds())
+				time.Sleep(sleepTime)
+			} else {
+				return nil, errors.Wrap(err, "failed to log in client")
+			}
+		} else {
+			return confirmation, nil
+		}
+	}
+}
diff --git a/apps/cnquery/cmd/logout.go b/apps/cnquery/cmd/logout.go
index e572fdbc6b..08dd8c0045 100644
--- a/apps/cnquery/cmd/logout.go
+++ b/apps/cnquery/cmd/logout.go
@@ -80,7 +80,7 @@ ensure the credentials cannot be used in the future.
 
 		if !viper.GetBool("force") {
 			log.Info().Msg("are you sure you want to revoke client access to Mondoo Platform? Use --force if you are sure")
-			return cli_errors.ExitCode1WithoutError
+			return cli_errors.NewCommandError(errors.New("--force is required to logout"), ConfigurationErrorCode)
 		}
 
 		// try to load config into credentials struct