Merge pull request #9 from tschaub/body-hash

Hash the body too

Author: tschaub

internal/cache/cache.go

@@ -37,17 +37,37 @@ func NewItem(cacheDir string, req *http.Request) (*Item, error) {
 
 	dir, file := filepath.Split(req.URL.EscapedPath())
 
-	var extra string
+	extra := false
+	hasher := sha256.New()
 	if req.URL.RawQuery != "" {
-		hasher := sha256.New()
 		hasher.Write([]byte(req.URL.RawQuery))
-		extra = base64.URLEncoding.EncodeToString(hasher.Sum(nil))
+		extra = true
 	}
 
-	// TODO: hash the body too
+	if req.Body != nil {
+		body := req.Body
+		buffer := &bytes.Buffer{}
+		if _, err := io.Copy(buffer, body); err != nil {
+			req.Body = struct {
+				io.Closer
+				io.Reader
+			}{
+				Closer: body,
+				Reader: io.MultiReader(buffer, body),
+			}
+			return nil, err
+		}
+
+		defer body.Close()
+		req.Body = io.NopCloser(buffer)
+		if buffer.Len() > 0 {
+			hasher.Write(buffer.Bytes())
+			extra = true
+		}
+	}
 
-	if extra != "" {
-		file += "?" + extra
+	if extra {
+		file += "?" + base64.URLEncoding.EncodeToString(hasher.Sum(nil))
 		if len(file) > maxFileNameLength {
 			file = file[:maxFileNameLength]
 		}
@@ -167,6 +187,14 @@ func (i *Item) Move(dir string) (*Item, error) {
 	return newItem, nil
 }
 
+func (i *Item) Rebase(dir string) *Item {
+	return &Item{
+		baseDir:     dir,
+		relBodyPath: i.relBodyPath,
+		relMetaPath: i.relMetaPath,
+	}
+}
+
 type Meta struct {
 	Header     http.Header
 	StatusCode int

internal/cache/cache_test.go

@@ -1,8 +1,10 @@
 package cache_test
 
 import (
+	"io"
 	"net/http"
 	"net/url"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -13,6 +15,7 @@ import (
 func TestNewItem(t *testing.T) {
 	cases := []struct {
 		method string
+		body   string
 		dir    string
 		url    string
 		key    string
@@ -24,6 +27,13 @@ func TestNewItem(t *testing.T) {
 			url:    "https://example.com/foo/bar",
 			key:    "base/GET/https/example.com/foo/#body#bar",
 		},
+		{
+			method: http.MethodPost,
+			body:   "example post body",
+			dir:    "base",
+			url:    "https://example.com/foo/bar",
+			key:    "base/POST/https/example.com/foo/#body#bar?alAlQ_ChTDxBRLzeTmqgXqkBuHIvGo1fj1rl_IHFCn8=",
+		},
 		{
 			method: http.MethodHead,
 			dir:    "base",
@@ -86,6 +96,9 @@ func TestNewItem(t *testing.T) {
 			}
 
 			request := &http.Request{URL: u, Method: c.method}
+			if c.body != "" {
+				request.Body = io.NopCloser(strings.NewReader(c.body))
+			}
 			item, err := cache.NewItem(c.dir, request)
 			if err != nil {
 				require.EqualError(t, err, c.err)

internal/proxy/proxy.go

@@ -81,6 +81,7 @@ func (p *Proxy) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.R
 		p.logger.Error("failed to get new cache item", "error", err, "url", req.URL, "method", req.Method)
 		return req, nil
 	}
+	ctx.UserData = cacheItem
 
 	exists, err := cacheItem.Exists()
 	if err != nil {
@@ -116,9 +117,11 @@ func (p *Proxy) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *http
 	}
 	req := resp.Request
 
-	cacheItem, err := cache.NewItem(p.cacheDir(), req)
-	if err != nil {
-		p.logger.Error("failed to get cache item", "error", err, "url", req.URL, "method", req.Method)
+	var cacheItem *cache.Item
+	if i, ok := ctx.UserData.(*cache.Item); ok && i != nil {
+		cacheItem = i
+	} else {
+		p.logger.Error("missing cache item in response handler user data", "url", req.URL)
 		return resp
 	}
 
@@ -142,7 +145,7 @@ func (p *Proxy) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *http
 	}
 
 	if resp.StatusCode == http.StatusPartialContent {
-		go p.download(req)
+		go p.download(req, cacheItem)
 		return resp
 	}
 
@@ -165,11 +168,8 @@ func (p *Proxy) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *http
 	return clonedResponse
 }
 
-func (p *Proxy) download(req *http.Request) {
-	tempCacheItem, err := cache.NewItem(p.tempDir(), req)
-	if err != nil {
-		p.logger.Error("download failed", "error", err, "url", req.URL, "method", req.Method)
-	}
+func (p *Proxy) download(req *http.Request, cacheItem *cache.Item) {
+	tempCacheItem := cacheItem.Rebase(p.tempDir())
 
 	cacheKey := tempCacheItem.Key()
 	if _, alreadyDownloading := p.downloading.LoadOrStore(cacheKey, true); alreadyDownloading {
@@ -181,7 +181,7 @@ func (p *Proxy) download(req *http.Request) {
 	clonedRequest := req.Clone(context.Background())
 	clonedRequest.Header.Del("Range")
 
-	p.logger.Debug("starting download", "url", clonedRequest.URL, "method", req.Method)
+	p.logger.Info("starting download", "url", clonedRequest.URL, "method", req.Method)
 	resp, err := http.DefaultClient.Do(clonedRequest)
 	if err != nil {
 		p.logger.Error("download request failed", "error", err, "url", clonedRequest.URL, "method", req.Method)
@@ -202,7 +202,7 @@ func (p *Proxy) download(req *http.Request) {
 	if _, err := tempCacheItem.Move(p.cacheDir()); err != nil {
 		p.logger.Error("failed to rename download", "error", err, "url", clonedRequest.URL, "method", req.Method)
 	}
-	p.logger.Debug("download complete", "url", clonedRequest.URL, "method", req.Method)
+	p.logger.Info("download complete", "url", clonedRequest.URL, "method", req.Method)
 }
 
 func hostMatches(u *url.URL, hosts []string) bool {