From 73f8811a4b5a9d3ea0fb66561a202ecd5e8768ed Mon Sep 17 00:00:00 2001
From: Jakub <jakub@cuth.sk>
Date: Mon, 8 Aug 2022 13:55:00 +0200
Subject: [PATCH] GODT-1645: Changing timeouts to not send too many login
 attempts.

---
 pkg/pmapi/manager.go      |  2 +-
 pkg/pmapi/manager_auth.go |  4 ++++
 pkg/pmapi/response.go     | 16 +++++++++++++---
 test/context/users.go     | 22 ++++++++++++++--------
 4 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/pkg/pmapi/manager.go b/pkg/pmapi/manager.go
index 09997cf0..e0401f05 100644
--- a/pkg/pmapi/manager.go
+++ b/pkg/pmapi/manager.go
@@ -85,7 +85,7 @@ func newManager(cfg Config) *manager {
 	// The resty is increasing the delay between retries up to 1 minute
 	// (SetRetryMaxWaitTime) so for 10 retries the cumulative delay can be
 	// up to 5min.
-	m.rc.SetRetryCount(5)
+	m.rc.SetRetryCount(3)
 	m.rc.SetRetryMaxWaitTime(time.Minute)
 	m.rc.SetRetryAfter(catchRetryAfter)
 	m.rc.AddRetryCondition(m.shouldRetry)
diff --git a/pkg/pmapi/manager_auth.go b/pkg/pmapi/manager_auth.go
index 77abef50..65d53266 100644
--- a/pkg/pmapi/manager_auth.go
+++ b/pkg/pmapi/manager_auth.go
@@ -62,6 +62,10 @@ func (m *manager) NewClientWithLogin(ctx context.Context, username string, passw
 		return nil, nil, err
 	}
 
+	// Do not retry requests after this point. The ephemeral from auth info
+	// won't be valid any more
+	ctx = ContextWithoutRetry(ctx)
+
 	auth, err := m.auth(ctx, AuthReq{
 		Username:        username,
 		ClientProof:     base64.StdEncoding.EncodeToString(proofs.ClientProof),
diff --git a/pkg/pmapi/response.go b/pkg/pmapi/response.go
index 12146009..8f43cbfb 100644
--- a/pkg/pmapi/response.go
+++ b/pkg/pmapi/response.go
@@ -105,17 +105,27 @@ func logConnReuse(_ *resty.Client, res *resty.Response) error {
 func catchRetryAfter(_ *resty.Client, res *resty.Response) (time.Duration, error) {
 	if res.StatusCode() == http.StatusTooManyRequests {
 		if after := res.Header().Get("Retry-After"); after != "" {
+			l := log.
+				WithField("statusCode", res.StatusCode()).
+				WithField("url", res.Request.URL).
+				WithField("verb", res.Request.Method)
+
 			seconds, err := strconv.Atoi(after)
 			if err != nil {
-				log.WithError(err).Warning("Cannot convert Retry-After to number")
+				l.WithError(err).Warning("Cannot convert Retry-After to number")
 				seconds = 10
 			}
 
 			// To avoid spikes when all clients retry at the same time, we add some random wait.
 			seconds += rand.Intn(10) //nolint:gosec // It is OK to use weak random number generator here.
+			l = l.WithField("seconds", seconds).WithField("start", time.Now().Unix())
 
-			log.Warningf("Retrying %s after %ds induced by http code %d", res.Request.URL, seconds, res.StatusCode())
-			return time.Duration(seconds) * time.Second, nil
+			// Maximum retry time in client is is one minute. But
+			// here wait times can be longer e.g. high API load
+			l.Warn("Retrying after induced by http code. Waiting now...")
+			time.Sleep(time.Duration(seconds) * time.Second)
+			l.Warn("Wait done")
+			return 0, nil
 		}
 	}
 
diff --git a/test/context/users.go b/test/context/users.go
index b6aaafd0..a308c139 100644
--- a/test/context/users.go
+++ b/test/context/users.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
+	"os"
 	"path/filepath"
 	"time"
 
@@ -29,6 +30,7 @@ import (
 	"github.com/ProtonMail/proton-bridge/v2/internal/users"
 	"github.com/ProtonMail/proton-bridge/v2/pkg/pmapi"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -52,14 +54,20 @@ func (ctx *TestContext) LoginUser(username string, password, mailboxPassword []b
 		}
 	}
 
-	userID, err := ctx.users.FinishLogin(client, auth, mailboxPassword)
+	user, err := ctx.users.FinishLogin(client, auth, mailboxPassword)
 	if err != nil {
 		return errors.Wrap(err, "failed to finish login")
 	}
 
-	ctx.addCleanupChecked(func() error {
-		return ctx.bridge.LogoutUser(userID)
-	}, "Logging out user")
+	ctx.addCleanupChecked(
+		func() error {
+			if os.Getenv(EnvName) == EnvLive {
+				logrus.Warn("Pausing user.Logout by two minutes to not hit issues with too many login attempts.")
+				time.Sleep(2 * time.Minute)
+			}
+			return user.Logout()
+		}, "Logging out user",
+	)
 
 	return nil
 }
@@ -75,14 +83,12 @@ func (ctx *TestContext) FinishLogin(client pmapi.Client, mailboxPassword []byte)
 		return errors.New("cannot get current auth tokens from client")
 	}
 
-	userID, err := ctx.users.FinishLogin(client, c.GetCurrentAuth(), mailboxPassword)
+	user, err := ctx.users.FinishLogin(client, c.GetCurrentAuth(), mailboxPassword)
 	if err != nil {
 		return errors.Wrap(err, "failed to finish login")
 	}
 
-	ctx.addCleanupChecked(func() error {
-		return ctx.bridge.LogoutUser(userID)
-	}, "Logging out user")
+	ctx.addCleanupChecked(user.Logout, "Logging out user")
 
 	return nil
 }