From 0b664f94fb7a9def9de28823c28da8df808b0b13 Mon Sep 17 00:00:00 2001
From: Nimi Wariboko Jr <nimiwaribokoj@gmail.com>
Date: Thu, 11 Oct 2018 10:56:22 -0700
Subject: [PATCH] Retry calls to vault on token renew and token lookup. Fixes
 #65

---
 Gopkg.lock    |  14 +++++-
 gatekeeper.go | 119 ++++++++++++++++++++++++++++----------------------
 2 files changed, 81 insertions(+), 52 deletions(-)

diff --git a/Gopkg.lock b/Gopkg.lock
index af7138b..1fadc58 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -38,6 +38,12 @@
   revision = "bc3f534c19ffdf835e524e11f0f825b3eaf541c3"
   version = "v1.14.31"
 
+[[projects]]
+  name = "github.com/cenkalti/backoff"
+  packages = ["."]
+  revision = "2ea60e5f094469f9e65adb9cd103795b73ae743e"
+  version = "v2.0.0"
+
 [[projects]]
   branch = "master"
   name = "github.com/dsnet/compress"
@@ -217,6 +223,12 @@
   packages = ["ssh/terminal"]
   revision = "a2144134853fc9a27a7b1e3eb4f19f1a76df13c9"
 
+[[projects]]
+  branch = "master"
+  name = "golang.org/x/net"
+  packages = ["context"]
+  revision = "49bb7cea24b1df9410e1712aa6433dae904ff66a"
+
 [[projects]]
   branch = "master"
   name = "golang.org/x/sys"
@@ -254,6 +266,6 @@
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "9048eb7a83b6106108531a77fa0c5808767dbbbe67efbe56ee3846d71f38815c"
+  inputs-digest = "bed642cfe4467d987157c0e31e00dd96397c58ee4d4ab12a0836578341817471"
   solver-name = "gps-cdcl"
   solver-version = 1
diff --git a/gatekeeper.go b/gatekeeper.go
index c530f0a..739ef7a 100644
--- a/gatekeeper.go
+++ b/gatekeeper.go
@@ -14,6 +14,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/cenkalti/backoff"
 	"github.com/franela/goreq"
 	gkClient "github.com/nemosupremo/vault-gatekeeper/gatekeeper"
 	"github.com/nemosupremo/vault-gatekeeper/policy"
@@ -75,6 +76,8 @@ type Config struct {
 	Version string
 
 	SkipPolicyLoading bool
+
+	Backoff *backoff.ExponentialBackOff
 }
 
 type peer struct {
@@ -282,6 +285,10 @@ func NewGatekeeper(conf Config) (*Gatekeeper, error) {
 		}
 	}
 
+	g.config.Backoff = backoff.NewExponentialBackOff()
+	g.config.Backoff.MaxElapsedTime = 1 * time.Minute
+	g.config.Backoff.MaxInterval = 15 * time.Second
+
 	return g, nil
 }
 
@@ -602,39 +609,45 @@ func (g *Gatekeeper) TokenTtl() (time.Duration, error) {
 	g.RLock()
 	token := g.Token
 	g.RUnlock()
-	r, err := vault.Request{goreq.Request{
-		Uri:             vault.Path("v1/auth/token/lookup-self"),
-		MaxRedirects:    10,
-		RedirectHeaders: true,
-		Method:          "GET",
-	}.WithHeader("X-Vault-Token", token)}.Do()
-	if err == nil {
-		defer r.Body.Close()
-		switch r.StatusCode {
-		case 200:
-			var resp struct {
-				Data struct {
-					Ttl int `json:"ttl"`
-				} `json:"data"`
-			}
-			if err := r.Body.FromJsonTo(&resp); err == nil {
-				return time.Duration(resp.Data.Ttl) * time.Second, nil
-			} else {
-				return 0, err
-			}
-		default:
-			var e vault.Error
-			e.Code = r.StatusCode
-			if err := r.Body.FromJsonTo(&e); err == nil {
-				return 0, e
-			} else {
-				e.Errors = []string{"communication error."}
-				return 0, e
+	ttl := time.Duration(0)
+	f := func() error {
+		r, err := vault.Request{goreq.Request{
+			Uri:             vault.Path("v1/auth/token/lookup-self"),
+			MaxRedirects:    10,
+			RedirectHeaders: true,
+			Method:          "GET",
+		}.WithHeader("X-Vault-Token", token)}.Do()
+		if err == nil {
+			defer r.Body.Close()
+			switch r.StatusCode {
+			case 200:
+				var resp struct {
+					Data struct {
+						Ttl int `json:"ttl"`
+					} `json:"data"`
+				}
+				if err := r.Body.FromJsonTo(&resp); err == nil {
+					ttl = time.Duration(resp.Data.Ttl) * time.Second
+					return nil
+				} else {
+					return err
+				}
+			default:
+				var e vault.Error
+				e.Code = r.StatusCode
+				if err := r.Body.FromJsonTo(&e); err == nil {
+					return backoff.Permanent(e)
+				} else {
+					e.Errors = []string{"communication error."}
+					return e
+				}
 			}
+		} else {
+			return err
 		}
-	} else {
-		return 0, err
 	}
+	err := backoff.Retry(f, g.config.Backoff)
+	return ttl, err
 }
 
 func (g *Gatekeeper) RenewToken() error {
@@ -644,30 +657,34 @@ func (g *Gatekeeper) RenewToken() error {
 	g.RLock()
 	token := g.Token
 	g.RUnlock()
-	r, err := vault.Request{goreq.Request{
-		Uri:             vault.Path("v1/auth/token/renew-self"),
-		MaxRedirects:    10,
-		RedirectHeaders: true,
-		Method:          "POST",
-	}.WithHeader("X-Vault-Token", token)}.Do()
-	if err == nil {
-		defer r.Body.Close()
-		switch r.StatusCode {
-		case 200, 204:
-			return nil
-		default:
-			var e vault.Error
-			e.Code = r.StatusCode
-			if err := r.Body.FromJsonTo(&e); err == nil {
-				return e
-			} else {
-				e.Errors = []string{"communication error."}
-				return e
+	f := func() error {
+		r, err := vault.Request{goreq.Request{
+			Uri:             vault.Path("v1/auth/token/renew-self"),
+			MaxRedirects:    10,
+			RedirectHeaders: true,
+			Method:          "POST",
+		}.WithHeader("X-Vault-Token", token)}.Do()
+		if err == nil {
+			defer r.Body.Close()
+			switch r.StatusCode {
+			case 200, 204:
+				return nil
+			default:
+				var e vault.Error
+				e.Code = r.StatusCode
+				if err := r.Body.FromJsonTo(&e); err == nil {
+					return backoff.Permanent(e)
+				} else {
+					e.Errors = []string{"communication error."}
+					return e
+				}
 			}
+		} else {
+			return err
 		}
-	} else {
-		return err
 	}
+	err := backoff.Retry(f, g.config.Backoff)
+	return err
 }
 
 func (g *Gatekeeper) Peers() []peer {