From 27c655ef672332d79bdbafa4759015769820a2e9 Mon Sep 17 00:00:00 2001 From: Calvin Leung Huang Date: Tue, 12 Mar 2019 13:21:02 -0700 Subject: [PATCH] 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 * Update command/agent/cache/proxy.go Co-Authored-By: calvn --- command/agent/cache/api_proxy.go | 35 ++++-------------- command/agent/cache/handler.go | 34 ++++++++++++++++- command/agent/cache/lease_cache.go | 22 ++++++++--- command/agent/cache/proxy.go | 59 ++++++++++++++++++++++++++++-- command/agent/cache/testing.go | 4 +- 5 files changed, 115 insertions(+), 39 deletions(-) diff --git a/command/agent/cache/api_proxy.go b/command/agent/cache/api_proxy.go index 709851ad04..5a321f8a3e 100644 --- a/command/agent/cache/api_proxy.go +++ b/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 } diff --git a/command/agent/cache/handler.go b/command/agent/cache/handler.go index abacb844b8..24ceeeda5c 100644 --- a/command/agent/cache/handler.go +++ b/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 diff --git a/command/agent/cache/lease_cache.go b/command/agent/cache/lease_cache.go index c88c105ff1..d9cd55f233 100644 --- a/command/agent/cache/lease_cache.go +++ b/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, diff --git a/command/agent/cache/proxy.go b/command/agent/cache/proxy.go index 4637590917..609cfa2972 100644 --- a/command/agent/cache/proxy.go +++ b/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 +} diff --git a/command/agent/cache/testing.go b/command/agent/cache/testing.go index 9db9934ddd..444f8ae8fe 100644 --- a/command/agent/cache/testing.go +++ b/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") }