Handle vault redirects

vault/app-id_strategy.go

@@ -28,20 +28,31 @@ func NewAppIDAuthStrategy() *AppIDAuthStrategy {
 	return nil
 }
 
-// GetToken - log in to the app-id auth backend and return the client token
-func (a *AppIDAuthStrategy) GetToken(addr *url.URL) (string, error) {
+func (a *AppIDAuthStrategy) GetHttpClient() *http.Client {
 	if a.hc == nil {
 		a.hc = &http.Client{Timeout: time.Second * 5}
 	}
-	client := a.hc
+	return a.hc
+}
+
+func (a *AppIDAuthStrategy) SetToken(req *http.Request) {
+	// no-op
+}
 
+func (a *AppIDAuthStrategy) Do(req *http.Request) (*http.Response, error) {
+	hc := a.GetHttpClient()
+	return hc.Do(req)
+}
+
+// GetToken - log in to the app-id auth backend and return the client token
+func (a *AppIDAuthStrategy) GetToken(addr *url.URL) (string, error) {
 	buf := new(bytes.Buffer)
 	json.NewEncoder(buf).Encode(&a)
 
 	u := &url.URL{}
 	*u = *addr
 	u.Path = "/v1/auth/app-id/login"
-	res, err := client.Post(u.String(), "application/json; charset=utf-8", buf)
+	res, err := requestAndFollow(a, "POST", u, buf.Bytes())
 	if err != nil {
 		return "", err
 	}

vault/client.go

@@ -54,6 +54,28 @@ func getAuthStrategy() AuthStrategy {
 	return nil
 }
 
+func (c *Client) GetHttpClient() *http.Client {
+	if c.hc == nil {
+		c.hc = &http.Client{
+			Timeout: time.Second * 5,
+			CheckRedirect: func(req *http.Request, via []*http.Request) error {
+				req.Header.Set("X-Vault-Token", c.token)
+				return nil
+			},
+		}
+	}
+	return c.hc
+}
+
+func (c *Client) SetToken(req *http.Request) {
+	req.Header.Set("X-Vault-Token", c.token)
+}
+
+func (c *Client) Do(req *http.Request) (*http.Response, error) {
+	hc := c.GetHttpClient()
+	return hc.Do(req)
+}
+
 // Login - log in to Vault with the discovered auth backend and save the token
 func (c *Client) Login() error {
 	token, err := c.Auth.GetToken(c.Addr)
@@ -72,17 +94,12 @@ func (c *Client) RevokeToken() {
 		return
 	}
 
-	if c.hc == nil {
-		c.hc = &http.Client{Timeout: time.Second * 5}
-	}
-
 	u := &url.URL{}
 	*u = *c.Addr
 	u.Path = "/v1/auth/token/revoke-self"
-	req, _ := http.NewRequest("POST", u.String(), nil)
-	req.Header.Set("X-Vault-Token", c.token)
 
-	res, err := c.hc.Do(req)
+	res, err := requestAndFollow(c, "POST", u, nil)
+
 	if err != nil {
 		log.Println("Error while revoking Vault Token", err)
 	}
@@ -94,32 +111,23 @@ func (c *Client) RevokeToken() {
 
 func (c *Client) Read(path string) ([]byte, error) {
 	path = normalizeURLPath(path)
-	if c.hc == nil {
-		c.hc = &http.Client{Timeout: time.Second * 5}
-	}
 
 	u := &url.URL{}
 	*u = *c.Addr
 	u.Path = "/v1" + path
-	req, err := http.NewRequest("GET", u.String(), nil)
-	if err != nil {
-		return nil, err
-	}
-	req.Header.Set("X-Vault-Token", c.token)
 
-	res, err := c.hc.Do(req)
+	res, err := requestAndFollow(c, "GET", u, nil)
 	if err != nil {
 		return nil, err
 	}
-
 	body, err := ioutil.ReadAll(res.Body)
 	res.Body.Close()
 	if err != nil {
 		return nil, err
 	}
 
 	if res.StatusCode != 200 {
-		err = fmt.Errorf("Unexpected HTTP status %d on Read from %s: %s", res.StatusCode, u, body)
+		err = fmt.Errorf("Unexpected HTTP status %d on Read from %s: %s", res.StatusCode, path, body)
 		return nil, err
 	}
 
@@ -131,7 +139,7 @@ func (c *Client) Read(path string) ([]byte, error) {
 	}
 
 	if _, ok := response["data"]; !ok {
-		return nil, fmt.Errorf("Unexpected HTTP body on Read for %s: %s", u, body)
+		return nil, fmt.Errorf("Unexpected HTTP body on Read for %s: %s", path, body)
 	}
 
 	return json.Marshal(response["data"])

vault/http.go

@@ -0,0 +1,47 @@
+package vault
+
+import (
+	"bytes"
+	"net/http"
+	"net/url"
+)
+
+// httpClient
+type httpClient interface {
+	GetHttpClient() *http.Client
+	SetToken(req *http.Request)
+	Do(req *http.Request) (*http.Response, error)
+}
+
+func requestAndFollow(hc httpClient, method string, u *url.URL, body []byte) (*http.Response, error) {
+	var res *http.Response
+	var err error
+	for attempts := 0; attempts < 2; attempts++ {
+		reader := bytes.NewReader(body)
+		req, err := http.NewRequest(method, u.String(), reader)
+
+		if err != nil {
+			return nil, err
+		}
+		hc.SetToken(req)
+		if method == "POST" {
+			req.Header.Set("Content-Type", "application/json; charset=utf-8")
+		}
+
+		res, err = hc.Do(req)
+		if err != nil {
+			return nil, err
+		}
+		if res.StatusCode == 307 {
+			res.Body.Close()
+			location, errLocation := res.Location()
+			if errLocation != nil {
+				return nil, errLocation
+			}
+			u.Host = location.Host
+		} else {
+			break
+		}
+	}
+	return res, err
+}