Repository: buchgr/bazel-remote
Pull request #107 (2019-10-18 18:30:29)

Check PUT request content size
Ensure that the PUT requests size is the same as the actual size of the data received.

This was discovered in some cases bazel is in the middle of uploading the file, the user terminates bazel, the cache server blindly accepts the file and bazel now redistributes incomplete data.

buchgr: Good catch @allada !

Patch:
diff --git a/cache/disk/disk.go b/cache/disk/disk.go
index 7504862..99be912 100644
--- a/cache/disk/disk.go
+++ b/cache/disk/disk.go
@@ -145,7 +145,7 @@ func (c *diskCache) loadExistingFiles() error {
 	return nil
 }
 
-func (c *diskCache) Put(kind cache.EntryKind, hash string, size int64, r io.Reader) (err error) {
+func (c *diskCache) Put(kind cache.EntryKind, hash string, expectedSize int64, r io.Reader) (err error) {
 	c.mux.Lock()
 
 	key := cacheKey(kind, hash)
@@ -164,7 +164,7 @@ func (c *diskCache) Put(kind cache.EntryKind, hash string, size int64, r io.Read
 
 	// Try to add the item to the LRU.
 	newItem := &lruItem{
-		size:      size,
+		size:      expectedSize,
 		committed: false,
 	}
 	ok := c.lru.Add(key, newItem)
@@ -201,11 +201,14 @@ func (c *diskCache) Put(kind cache.EntryKind, hash string, size int64, r io.Read
 			// Only delete the temp file if moving it didn't succeed.
 			os.Remove(f.Name())
 		}
+		// Just in case we didn't already close it.  No need to check errors.
+		f.Close()
 	}()
 
+	var bytesCopied int64 = 0
 	if kind == cache.CAS {
 		hasher := sha256.New()
-		if _, err = io.Copy(io.MultiWriter(f, hasher), r); err != nil {
+		if bytesCopied, err = io.Copy(io.MultiWriter(f, hasher), r); err != nil {
 			return
 		}
 		actualHash := hex.EncodeToString(hasher.Sum(nil))
@@ -215,7 +218,7 @@ func (c *diskCache) Put(kind cache.EntryKind, hash string, size int64, r io.Read
 			return
 		}
 	} else {
-		if _, err = io.Copy(f, r); err != nil {
+		if bytesCopied, err = io.Copy(f, r); err != nil {
 			return
 		}
 	}
@@ -228,6 +231,12 @@ func (c *diskCache) Put(kind cache.EntryKind, hash string, size int64, r io.Read
 		return err
 	}
 
+	if bytesCopied != expectedSize {
+		err = fmt.Errorf(
+			"sizes don't match. Expected %d, found %d", expectedSize, bytesCopied)
+		return err
+	}
+
 	// Rename to the final path
 	filePath := cacheFilePath(kind, c.dir, hash)
 	err = os.Rename(f.Name(), filePath)
diff --git a/cache/http/http.go b/cache/http/http.go
index 3f1c466..9f1408f 100644
--- a/cache/http/http.go
+++ b/cache/http/http.go
@@ -85,12 +85,15 @@ func logResponse(log cache.Logger, method string, code int, url string) {
 	log.Printf("%4s %d %15s %s", method, code, "", url)
 }
 
-func (r *remoteHTTPProxyCache) Put(kind cache.EntryKind, hash string, size int64, data io.Reader) (err error) {
+func (r *remoteHTTPProxyCache) Put(kind cache.EntryKind, hash string, size int64, data io.Reader) (error) {
 	if r.local.Contains(kind, hash) {
 		io.Copy(ioutil.Discard, data)
 		return nil
 	}
-	r.local.Put(kind, hash, size, data)
+	err := r.local.Put(kind, hash, size, data)
+	if err != nil {
+		return err
+	}
 
 	select {
 	case r.uploadQueue <- &uploadReq{
@@ -100,7 +103,7 @@ func (r *remoteHTTPProxyCache) Put(kind cache.EntryKind, hash string, size int64
 	default:
 		r.errorLogger.Printf("too many uploads queued")
 	}
-	return
+	return err
 }
 
 func (r *remoteHTTPProxyCache) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {


Test patch:
diff --git a/cache/disk/disk_test.go b/cache/disk/disk_test.go
index 4f07a38..b282d05 100644
--- a/cache/disk/disk_test.go
+++ b/cache/disk/disk_test.go
@@ -113,7 +113,8 @@ func TestCacheEviction(t *testing.T) {
 	}
 
 	for i, thisExp := range expectedSizesNumItems {
-		err := testCache.Put(cache.AC, fmt.Sprintf("aa-%d", i), int64(i), strings.NewReader("hello"))
+		strReader := strings.NewReader(strings.Repeat("a", i))
+		err := testCache.Put(cache.AC, fmt.Sprintf("aa-%d", i), int64(i), strReader)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -122,6 +123,17 @@ func TestCacheEviction(t *testing.T) {
 	}
 }
 
+func TestCachePutWrongSize(t *testing.T) {
+        cacheDir := tempDir(t)
+        defer os.RemoveAll(cacheDir)
+        testCache := New(cacheDir, 100)
+
+        err := testCache.Put(cache.AC, "aa-aa", int64(10), strings.NewReader("hello"))
+        if err == nil {
+            t.Fatal("Expected error due to size being different")
+        }
+}
+
 func expectContentEquals(t *testing.T, data io.ReadCloser, sizeBytes int64, expectedContent []byte) {
 	if data == nil {
 		t.Fatal("expected the item to exist")
diff --git a/cache/http/http_test.go b/cache/http/http_test.go
index 7c68c79..689d42c 100644
--- a/cache/http/http_test.go
+++ b/cache/http/http_test.go
@@ -192,3 +192,39 @@ func TestProxyWriteErrorsAreNotPropagated(t *testing.T) {
 		t.Error("Expected the blob to be stored locally")
 	}
 }
+
+func TestProxyLocalPutFailuresNotRelayed(t *testing.T) {
+	// Test that when there is an error writing to the remote proxy
+	// then the error is not propagated to the client. This is because
+	// the writes to the proxy happen asynchronously and on a best effort
+	// basis.
+
+	data := []byte("hello world")
+	hashBytes := sha256.Sum256(data)
+	hash := hex.EncodeToString(hashBytes[:])
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Error("Size miss-matched Put request should not have been forwarded")
+	}))
+
+	cacheDir := testutils.TempDir(t)
+	defer os.RemoveAll(cacheDir)
+	diskCache := disk.New(cacheDir, 100)
+
+	baseURL, err := url.Parse(ts.URL)
+	if err != nil {
+		t.Error(err)
+	}
+
+	proxy := New(baseURL, diskCache, &http.Client{}, testutils.NewSilentLogger(),
+		testutils.NewSilentLogger())
+
+	err = proxy.Put(cache.AC, hash, int64(len(data) + 1), bytes.NewReader(data))
+	if err == nil {
+		t.Error("Expected Put to error with size miss-match")
+	}
+
+	if diskCache.Contains(cache.AC, hash) {
+		t.Error("Expected not to be stored locally")
+	}
+}