From 63cf3860c6c922b6bd5be40253e69972d8ab87fb Mon Sep 17 00:00:00 2001
From: Calvin Leung Huang <cleung2010@gmail.com>
Date: Mon, 27 Apr 2020 17:20:25 -0700
Subject: [PATCH] database/mongodb: revert to old retry behavior (#8863)

* database/mongodb: revert to old retry behavior

* add a default case for non-EOF errors
---
 plugins/database/mongodb/mongodb.go | 67 +++++++++++++++--------------
 1 file changed, 34 insertions(+), 33 deletions(-)

diff --git a/plugins/database/mongodb/mongodb.go b/plugins/database/mongodb/mongodb.go
index 1d47082aa67ae..aafa42ba57bc6 100644
--- a/plugins/database/mongodb/mongodb.go
+++ b/plugins/database/mongodb/mongodb.go
@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"strings"
 	"time"
 
 	"github.com/hashicorp/vault/api"
@@ -95,11 +97,6 @@ func (m *MongoDB) CreateUser(ctx context.Context, statements dbplugin.Statements
 		return "", "", dbutil.ErrEmptyCreationStatement
 	}
 
-	client, err := m.getConnection(ctx)
-	if err != nil {
-		return "", "", err
-	}
-
 	username, err = m.GenerateUsername(usernameConfig)
 	if err != nil {
 		return "", "", err
@@ -132,7 +129,7 @@ func (m *MongoDB) CreateUser(ctx context.Context, statements dbplugin.Statements
 		Roles:    mongoCS.Roles.toStandardRolesArray(),
 	}
 
-	if err := runCommandWithRetry(ctx, client, mongoCS.DB, createUserCmd); err != nil {
+	if err := m.runCommandWithRetry(ctx, mongoCS.DB, createUserCmd); err != nil {
 		return "", "", err
 	}
 
@@ -150,11 +147,6 @@ func (m *MongoDB) SetCredentials(ctx context.Context, statements dbplugin.Statem
 	m.Lock()
 	defer m.Unlock()
 
-	client, err := m.getConnection(ctx)
-	if err != nil {
-		return "", "", err
-	}
-
 	username = staticUser.Username
 	password = staticUser.Password
 
@@ -167,7 +159,7 @@ func (m *MongoDB) SetCredentials(ctx context.Context, statements dbplugin.Statem
 	if err != nil {
 		return "", "", err
 	}
-	if err := runCommandWithRetry(ctx, client, cs.Database, changeUserCmd); err != nil {
+	if err := m.runCommandWithRetry(ctx, cs.Database, changeUserCmd); err != nil {
 		return "", "", err
 	}
 
@@ -188,11 +180,6 @@ func (m *MongoDB) RevokeUser(ctx context.Context, statements dbplugin.Statements
 
 	statements = dbutil.StatementCompatibilityHelper(statements)
 
-	client, err := m.getConnection(ctx)
-	if err != nil {
-		return err
-	}
-
 	// If no revocation statements provided, pass in empty JSON
 	var revocationStatement string
 	switch len(statements.Revocation) {
@@ -206,7 +193,7 @@ func (m *MongoDB) RevokeUser(ctx context.Context, statements dbplugin.Statements
 
 	// Unmarshal revocation statements into mongodbRoles
 	var mongoCS mongoDBStatement
-	err = json.Unmarshal([]byte(revocationStatement), &mongoCS)
+	err := json.Unmarshal([]byte(revocationStatement), &mongoCS)
 	if err != nil {
 		return err
 	}
@@ -222,7 +209,7 @@ func (m *MongoDB) RevokeUser(ctx context.Context, statements dbplugin.Statements
 		WriteConcern: writeconcern.New(writeconcern.WMajority()),
 	}
 
-	return runCommandWithRetry(ctx, client, db, dropUserCmd)
+	return m.runCommandWithRetry(ctx, db, dropUserCmd)
 }
 
 // RotateRootCredentials is not currently supported on MongoDB
@@ -230,22 +217,36 @@ func (m *MongoDB) RotateRootCredentials(ctx context.Context, statements []string
 	return nil, errors.New("root credential rotation is not currently implemented in this database secrets engine")
 }
 
-// runCommandWithRetry runs a command with retry.
-func runCommandWithRetry(ctx context.Context, client *mongo.Client, db string, cmd interface{}) error {
-	timeout := time.Now().Add(1 * time.Minute)
-	backoffTime := 3
-	for {
-		// Run command
-		result := client.Database(db).RunCommand(ctx, cmd, nil)
-		if result.Err() == nil {
-			break
-		}
+// runCommandWithRetry runs a command and retries once more if there's a failure
+// on the first attempt. This should be called with the lock held
+func (m *MongoDB) runCommandWithRetry(ctx context.Context, db string, cmd interface{}) error {
+	// Get the client
+	client, err := m.getConnection(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Run command
+	result := client.Database(db).RunCommand(ctx, cmd, nil)
 
-		if time.Now().After(timeout) {
-			return result.Err()
+	// Error check on the first attempt
+	err = result.Err()
+	switch {
+	case err == nil:
+		return nil
+	case err == io.EOF, strings.Contains(err.Error(), "EOF"):
+		// Call getConnection to reset and retry query if we get an EOF error on first attempt.
+		client, err = m.getConnection(ctx)
+		if err != nil {
+			return err
+		}
+		result = client.Database(db).RunCommand(ctx, cmd, nil)
+		if err := result.Err(); err != nil {
+			return err
 		}
-		time.Sleep(time.Duration(backoffTime) * time.Second)
-		backoffTime += backoffTime
+	default:
+		return err
 	}
+
 	return nil
 }