agent/caching: add X-Cache and Age headers (#6394)

* agent/caching: add X-Cache and Age headers, update Date header on cached resp

* Update command/agent/cache/lease_cache.go

Co-Authored-By: calvn <cleung2010@gmail.com>

* Update command/agent/cache/proxy.go

Co-Authored-By: calvn <cleung2010@gmail.com>

command/agent/cache/api_proxy.go

@@ -1,10 +1,8 @@
 package cache
 
 import (
-	"bytes"
 	"context"
 	"fmt"
-	"io/ioutil"
 
 	hclog "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/vault/api"
@@ -46,32 +44,15 @@ func (ap *APIProxy) Send(ctx context.Context, req *SendRequest) (*SendResponse,
 	// Make the request to Vault and get the response
 	ap.logger.Info("forwarding request", "path", req.Request.URL.Path, "method", req.Request.Method)
 
-	var sendResponse *SendResponse
 	resp, err := client.RawRequestWithContext(ctx, fwReq)
-	if resp != nil {
-		sendResponse = &SendResponse{Response: resp}
-	}
-	if err != nil {
-		// Bubble back the api.Response as well for error checking/handling at the handler layer.
-		return sendResponse, err
+
+	// Before error checking from the request call, we'd want to initialize a SendResponse to
+	// potentially return
+	sendResponse, newErr := NewSendResponse(resp, nil)
+	if newErr != nil {
+		return nil, newErr
 	}
 
-	// Set SendResponse.ResponseBody if the response body is non-nil
-	if resp.Body != nil {
-		respBody, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			ap.logger.Error("failed to read request body", "error", err)
-			return nil, err
-		}
-		// Close the old body
-		resp.Body.Close()
-
-		// Re-set the response body for potential consumption on the way back up the
-		// Proxier middleware chain.
-		resp.Body = ioutil.NopCloser(bytes.NewReader(respBody))
-
-		sendResponse.ResponseBody = respBody
-	}
-
-	return sendResponse, nil
+	// Bubble back the api.Response as well for error checking/handling at the handler layer.
+	return sendResponse, err
 }

command/agent/cache/handler.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"time"
 
 	"github.com/hashicorp/errwrap"
 	hclog "github.com/hashicorp/go-hclog"
@@ -66,13 +67,42 @@ func Handler(ctx context.Context, logger hclog.Logger, proxier Proxier, inmemSin
 
 		defer resp.Response.Body.Close()
 
-		copyHeader(w.Header(), resp.Response.Header)
-		w.WriteHeader(resp.Response.StatusCode)
+		// Set headers
+		setHeaders(w, resp)
+
+		// Set response body
 		io.Copy(w, resp.Response.Body)
 		return
 	})
 }
 
+// setHeaders is a helper that sets the header values based on SendResponse. It
+// copies over the headers from the original response and also includes any
+// cache-related headers.
+func setHeaders(w http.ResponseWriter, resp *SendResponse) {
+	// Set header values
+	copyHeader(w.Header(), resp.Response.Header)
+	if resp.CacheMeta != nil {
+		xCacheVal := "MISS"
+
+		if resp.CacheMeta.Hit {
+			xCacheVal = "HIT"
+
+			// If this is a cache hit, we also set the Age header
+			age := fmt.Sprintf("%.0f", resp.CacheMeta.Age.Seconds())
+			w.Header().Set("Age", age)
+
+			// Update the date value
+			w.Header().Set("Date", time.Now().Format(http.TimeFormat))
+		}
+
+		w.Header().Set("X-Cache", xCacheVal)
+	}
+
+	// Set status code
+	w.WriteHeader(resp.Response.StatusCode)
+}
+
 // processTokenLookupResponse checks if the request was one of token
 // lookup-self. If the auto-auth token was used to perform lookup-self, the
 // identifier of the token and its accessor same will be stripped off of the

command/agent/cache/lease_cache.go

@@ -13,6 +13,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/hashicorp/errwrap"
 	hclog "github.com/hashicorp/go-hclog"
@@ -141,12 +142,21 @@ func (c *LeaseCache) checkCacheForRequest(id string) (*SendResponse, error) {
 		return nil, err
 	}
 
-	return &SendResponse{
-		Response: &api.Response{
-			Response: resp,
-		},
-		ResponseBody: index.Response,
-	}, nil
+	sendResp, err := NewSendResponse(&api.Response{Response: resp}, index.Response)
+	if err != nil {
+		c.logger.Error("failed to create new send response", "error", err)
+		return nil, err
+	}
+	sendResp.CacheMeta.Hit = true
+
+	respTime, err := http.ParseTime(resp.Header.Get("Date"))
+	if err != nil {
+		c.logger.Error("failed to parse cached response date", "error", err)
+		return nil, err
+	}
+	sendResp.CacheMeta.Age = time.Now().Sub(respTime)
+
+	return sendResp, nil
 }
 
 // Send performs a cache lookup on the incoming request. If it's a cache hit,

command/agent/cache/proxy.go

@@ -1,23 +1,42 @@
 package cache
 
 import (
+	"bytes"
 	"context"
+	"fmt"
+	"io/ioutil"
 	"net/http"
+	"time"
 
 	"github.com/hashicorp/vault/api"
 )
 
 // SendRequest is the input for Proxier.Send.
 type SendRequest struct {
-	Token       string
-	Request     *http.Request
+	Token   string
+	Request *http.Request
+
+	// RequestBody is the stored body bytes from Request.Body. It is set here to
+	// avoid reading and re-setting the stream multiple times.
 	RequestBody []byte
 }
 
 // SendResponse is the output from Proxier.Send.
 type SendResponse struct {
-	Response     *api.Response
+	Response *api.Response
+
+	// ResponseBody is the stored body bytes from Response.Body. It is set here to
+	// avoid reading and re-setting the stream multiple times.
 	ResponseBody []byte
+	CacheMeta    *CacheMeta
+}
+
+// CacheMeta contains metadata information about the response,
+// such as whether it was a cache hit or miss, and the age of the
+// cached entry.
+type CacheMeta struct {
+	Hit bool
+	Age time.Duration
 }
 
 // Proxier is the interface implemented by different components that are
@@ -26,3 +45,37 @@ type SendResponse struct {
 type Proxier interface {
 	Send(ctx context.Context, req *SendRequest) (*SendResponse, error)
 }
+
+// NewSendResponse creates a new SendResponse and takes care of initializing its
+// fields properly.
+func NewSendResponse(apiResponse *api.Response, responseBody []byte) (*SendResponse, error) {
+	if apiResponse == nil {
+		return nil, fmt.Errorf("nil api response provided")
+	}
+
+	resp := &SendResponse{
+		Response:  apiResponse,
+		CacheMeta: &CacheMeta{},
+	}
+
+	// If a response body is separately provided we set that as the SendResponse.ResponseBody,
+	// otherwise we will do an ioutil.ReadAll to extract the response body from apiResponse.
+	switch {
+	case len(responseBody) > 0:
+		resp.ResponseBody = responseBody
+	case apiResponse.Body != nil:
+		respBody, err := ioutil.ReadAll(apiResponse.Body)
+		if err != nil {
+			return nil, err
+		}
+		// Close the old body
+		apiResponse.Body.Close()
+
+		// Re-set the response body after reading from the Reader
+		apiResponse.Body = ioutil.NopCloser(bytes.NewReader(respBody))
+
+		resp.ResponseBody = respBody
+	}
+
+	return resp, nil
+}

command/agent/cache/testing.go

@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/hashicorp/vault/api"
 )
@@ -46,9 +47,11 @@ func newTestSendResponse(status int, body string) *SendResponse {
 		Response: &api.Response{
 			Response: &http.Response{
 				StatusCode: status,
+				Header:     http.Header{},
 			},
 		},
 	}
+	resp.Response.Header.Set("Date", time.Now().Format(http.TimeFormat))
 
 	if body != "" {
 		resp.Response.Body = ioutil.NopCloser(strings.NewReader(body))
@@ -56,7 +59,6 @@ func newTestSendResponse(status int, body string) *SendResponse {
 	}
 
 	if json.Valid([]byte(body)) {
-		resp.Response.Header = http.Header{}
 		resp.Response.Header.Set("content-type", "application/json")
 	}