feat(lib/store): verify download blob digest when moving to cache

lib/store/ca_download_store.go

@@ -21,17 +21,17 @@ import (
 	"github.com/uber-go/tally"
 	"github.com/uber/kraken/lib/store/base"
 	"github.com/uber/kraken/lib/store/metadata"
+	"github.com/uber/kraken/utils/closers"
 )
 
 // CADownloadStore allows simultaneously downloading and uploading
 // content-adddressable files.
 type CADownloadStore struct {
+	config        CADownloadStoreConfig
 	backend       base.FileStore
 	downloadState base.FileState
 	cacheState    base.FileState
 	cleanup       *cleanupManager
-	readPartSize  int
-	writePartSize int
 }
 
 // NewCADownloadStore creates a new CADownloadStore.
@@ -61,12 +61,11 @@ func NewCADownloadStore(config CADownloadStoreConfig, stats tally.Scope) (*CADow
 		backend.NewFileOp().AcceptState(cacheState))
 
 	return &CADownloadStore{
+		config:        config,
 		backend:       backend,
 		downloadState: downloadState,
 		cacheState:    cacheState,
 		cleanup:       cleanup,
-		readPartSize:  config.ReadPartSize,
-		writePartSize: config.WritePartSize,
 	}, nil
 }
 
@@ -80,13 +79,33 @@ func (s *CADownloadStore) CreateDownloadFile(name string, length int64) error {
 	return s.backend.NewFileOp().CreateFile(name, s.downloadState, length)
 }
 
+// GetDownloadFileReader returns a FileReader for name.
+func (s *CADownloadStore) getDownloadFileReader(name string) (FileReader, error) {
+	return s.backend.NewFileOp().AcceptState(s.downloadState).GetFileReader(name, s.config.ReadPartSize)
+}
+
 // GetDownloadFileReadWriter returns a FileReadWriter for name.
 func (s *CADownloadStore) GetDownloadFileReadWriter(name string) (FileReadWriter, error) {
-	return s.backend.NewFileOp().AcceptState(s.downloadState).GetFileReadWriter(name, s.readPartSize, s.writePartSize)
+	return s.backend.NewFileOp().AcceptState(s.downloadState).GetFileReadWriter(name, s.config.ReadPartSize, s.config.WritePartSize)
 }
 
-// MoveDownloadFileToCache moves a download file to the cache.
+// MoveDownloadFileToCache moves a download file to the cache
+// after verifying its digest.
 func (s *CADownloadStore) MoveDownloadFileToCache(name string) error {
+	f, err := s.getDownloadFileReader(name)
+	if err != nil {
+		if s.InCacheError(err) {
+			// Inform the caller the file has already been moved,
+			// matching the behaviour of MoveFile(...).
+			return os.ErrExist
+		}
+		return fmt.Errorf("get file reader %s: %w", name, err)
+	}
+	defer closers.Close(f)
+
+	if err := verifyDigest(f, name, s.config.SkipHashVerification); err != nil {
+		return fmt.Errorf("verify digest: %w", err)
+	}
 	return s.backend.NewFileOp().AcceptState(s.downloadState).MoveFile(name, s.cacheState)
 }
 
@@ -158,7 +177,7 @@ func (s *CADownloadStore) Any() *CADownloadStoreScope {
 
 // GetFileReader returns a reader for name.
 func (a *CADownloadStoreScope) GetFileReader(name string) (FileReader, error) {
-	return a.op.GetFileReader(name, a.store.readPartSize)
+	return a.op.GetFileReader(name, a.store.config.ReadPartSize)
 }
 
 // GetFileStat returns file info for name.

lib/store/ca_download_store_test.go

@@ -14,8 +14,10 @@
 package store
 
 import (
+	"errors"
 	"os"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/uber/kraken/core"
@@ -31,13 +33,21 @@ func TestCADownloadStoreDownloadAndDeleteFiles(t *testing.T) {
 
 	var names []string
 	var wg sync.WaitGroup
-	for i := 0; i < 100; i++ {
-		name := core.DigestFixture().Hex()
+	for range 100 {
+		blob := core.NewBlobFixture()
+		name := blob.Digest.Hex()
 		names = append(names, name)
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			require.NoError(s.CreateDownloadFile(name, 1))
+
+			require.NoError(s.CreateDownloadFile(name, blob.Length()))
+			w, err := s.GetDownloadFileReadWriter(name)
+			require.NoError(err)
+			_, err = w.Write(blob.Content)
+			require.NoError(err)
+			require.NoError(w.Close())
+
 			require.NoError(s.MoveDownloadFileToCache(name))
 			require.NoError(s.Cache().DeleteFile(name))
 		}()
@@ -49,3 +59,125 @@ func TestCADownloadStoreDownloadAndDeleteFiles(t *testing.T) {
 		require.True(os.IsNotExist(err))
 	}
 }
+
+func TestCADownloadStoreMoveDownloadFileToCacheConcurrent(t *testing.T) {
+	require := require.New(t)
+
+	s, cleanup := CADownloadStoreFixture()
+	defer cleanup()
+
+	blob := core.NewBlobFixture()
+	name := blob.Digest.Hex()
+
+	require.NoError(s.CreateDownloadFile(name, blob.Length()))
+	w, err := s.GetDownloadFileReadWriter(name)
+	require.NoError(err)
+	_, err = w.Write(blob.Content)
+	require.NoError(err)
+	require.NoError(w.Close())
+
+	var success atomic.Int32
+	var exists atomic.Int32
+	var errMu sync.Mutex
+	var errs error
+	var wg sync.WaitGroup
+	for range 10 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			err := s.MoveDownloadFileToCache(name)
+			switch err {
+			case nil:
+				success.Add(1)
+			case os.ErrExist:
+				exists.Add(1)
+			default:
+				errMu.Lock()
+				errs = errors.Join(errs, err)
+				errMu.Unlock()
+			}
+		}()
+	}
+	wg.Wait()
+
+	require.NoError(errs)
+	require.Equal(int32(1), success.Load())
+	require.Equal(int32(9), exists.Load())
+}
+
+func TestCADownloadStoreMoveDownloadFileToCacheNonMatchingDigest(t *testing.T) {
+	require := require.New(t)
+
+	s, cleanup := CADownloadStoreFixture()
+	defer cleanup()
+
+	blob := core.NewBlobFixture()
+	name := blob.Digest.Hex()
+
+	require.NoError(s.CreateDownloadFile(name, blob.Length()))
+	w, err := s.GetDownloadFileReadWriter(name)
+	require.NoError(err)
+	corrupted := make([]byte, len(blob.Content))
+	copy(corrupted, blob.Content)
+	corrupted[0]++
+	_, err = w.Write(corrupted)
+	require.NoError(err)
+	require.NoError(w.Close())
+
+	err = s.MoveDownloadFileToCache(name)
+	require.ErrorContains(err, "verify digest: computed digest")
+
+	_, err = s.Download().GetFileStat(name)
+	require.NoError(err)
+	_, err = s.Cache().GetFileStat(name)
+	require.True(s.InDownloadError(err))
+}
+
+func TestCADownloadStoreMoveDownloadFileToCacheSkipVerification(t *testing.T) {
+	require := require.New(t)
+
+	s, cleanup := CADownloadStoreFixture()
+	s.config.SkipHashVerification = true
+	defer cleanup()
+
+	blob := core.NewBlobFixture()
+	name := blob.Digest.Hex()
+
+	require.NoError(s.CreateDownloadFile(name, blob.Length()))
+	w, err := s.GetDownloadFileReadWriter(name)
+	require.NoError(err)
+	emptyBlob := make([]byte, len(blob.Content))
+	_, err = w.Write(emptyBlob)
+	require.NoError(err)
+	require.NoError(w.Close())
+
+	err = s.MoveDownloadFileToCache(name)
+	require.NoError(err)
+
+	_, err = s.Cache().GetFileStat(name)
+	require.NoError(err)
+}
+
+func TestCADownloadStoreMoveDownloadFileToCacheAlreadyInCache(t *testing.T) {
+	require := require.New(t)
+
+	s, cleanup := CADownloadStoreFixture()
+	defer cleanup()
+
+	blob := core.NewBlobFixture()
+	name := blob.Digest.Hex()
+
+	require.NoError(s.CreateDownloadFile(name, blob.Length()))
+	w, err := s.GetDownloadFileReadWriter(name)
+	require.NoError(err)
+	_, err = w.Write(blob.Content)
+	require.NoError(err)
+	require.NoError(w.Close())
+
+	err = s.MoveDownloadFileToCache(name)
+	require.NoError(err)
+
+	err = s.MoveDownloadFileToCache(name)
+	require.True(os.IsExist(err))
+}

lib/store/ca_store.go

@@ -180,8 +180,8 @@ func (s *CAStore) MoveUploadFileToCache(uploadName, cacheName string) error {
 		return fmt.Errorf("get file reader %s: %s", uploadName, err)
 	}
 	defer closers.Close(f)
-	if err := s.verify(f, cacheName); err != nil {
-		return fmt.Errorf("verify digest: %s", err)
+	if err := verifyDigest(f, cacheName, s.config.SkipHashVerification); err != nil {
+		return fmt.Errorf("verify digest: %w", err)
 	}
 
 	return s.cacheStore.newFileOp().MoveFileFrom(cacheName, s.cacheStore.state, uploadPath)
@@ -330,28 +330,6 @@ func (s *CAStore) addItemForDiskSync(item *drainItem) {
 	s.drain.queue.PushBack(item)
 }
 
-// verify verifies that name is a valid SHA256 digest, and checks if the given
-// blob content matches the digset unless explicitly skipped.
-func (s *CAStore) verify(r io.Reader, name string) error {
-	// Verify that expected name is a valid SHA256 digest.
-	expected, err := core.NewSHA256DigestFromHex(name)
-	if err != nil {
-		return fmt.Errorf("new digest from file name: %s", err)
-	}
-
-	if !s.config.SkipHashVerification {
-		digester := core.NewDigester()
-		computed, err := digester.FromReader(r)
-		if err != nil {
-			return fmt.Errorf("calculate digest: %s", err)
-		}
-		if computed != expected {
-			return fmt.Errorf("computed digest %s doesn't match expected value %s", computed, expected)
-		}
-	}
-	return nil
-}
-
 func (s *CAStore) memoryCacheCleanupWorker() {
 	defer s.ttlWg.Done()
 

lib/store/config.go

@@ -95,4 +95,6 @@ type CADownloadStoreConfig struct {
 	ReadPartSize int `yaml:"read_part_size"`
 	// Part size limit for each file write. 0 means no limit.
 	WritePartSize int `yaml:"write_part_size"`
+
+	SkipHashVerification bool `yaml:"skip_hash_verification"`
 }

lib/store/utils.go

@@ -15,7 +15,12 @@ package store
 
 import (
 	"bytes"
+	"fmt"
+	"io"
 	"os"
+
+	"github.com/uber/kraken/core"
+	"github.com/uber/kraken/utils/log"
 )
 
 func createOrUpdateSymlink(sourcePath, targetPath string) error {
@@ -42,6 +47,25 @@ func createOrUpdateSymlink(sourcePath, targetPath string) error {
 	return nil
 }
 
+func verifyDigest(r io.Reader, name string, skipHashVerification bool) error {
+	expected, err := core.NewSHA256DigestFromHex(name)
+	if err != nil {
+		return fmt.Errorf("new digest from file name: %w", err)
+	}
+
+	if !skipHashVerification {
+		computed, err := core.NewDigester().FromReader(r)
+		if err != nil {
+			return fmt.Errorf("calculate digest: %w", err)
+		}
+		if computed != expected {
+			log.With("name", name, "expected", expected, "computed", computed).Error("Digest verification did not match")
+			return fmt.Errorf("computed digest %s doesn't match expected value %s", computed, expected)
+		}
+	}
+	return nil
+}
+
 type bufferFileReader struct {
 	*bytes.Reader
 }

lib/torrent/storage/agentstorage/torrent_test.go

@@ -258,7 +258,10 @@ func (w *coordinatedWriter) Write(b []byte) (int, error) {
 func TestTorrentWritePieceConflictsDoNotBlock(t *testing.T) {
 	require := require.New(t)
 
-	blob := core.SizedBlobFixture(1, 1)
+	// Use two pieces to prevent triggering a blob cache move after piece
+	// completion, which would fail as a mocked store writer is used.
+	blob := core.SizedBlobFixture(2, 1)
+	piece := blob.Content[:1]
 
 	f, cleanup := store.NewMockFileReadWriter([]byte{})
 	defer cleanup()
@@ -278,18 +281,18 @@ func TestTorrentWritePieceConflictsDoNotBlock(t *testing.T) {
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
-		require.NoError(tor.WritePiece(piecereader.NewBuffer(blob.Content), 0))
+		require.NoError(tor.WritePiece(piecereader.NewBuffer(piece), 0))
 	}()
 
 	// Writing while another goroutine is mid-write should not block.
 	<-w.startWriting
-	require.Equal(errWritePieceConflict, tor.WritePiece(piecereader.NewBuffer(blob.Content), 0))
+	require.Equal(errWritePieceConflict, tor.WritePiece(piecereader.NewBuffer(piece), 0))
 	w.stopWriting <- true
 
 	<-done
 
 	// Duplicate write should detect piece is complete.
-	require.Equal(storage.ErrPieceComplete, tor.WritePiece(piecereader.NewBuffer(blob.Content), 0))
+	require.Equal(storage.ErrPieceComplete, tor.WritePiece(piecereader.NewBuffer(piece), 0))
 }
 
 func TestTorrentWritePieceFailuresRemoveDirtyStatus(t *testing.T) {
@@ -303,7 +306,10 @@ func TestTorrentWritePieceFailuresRemoveDirtyStatus(t *testing.T) {
 	cads, cleanup := store.CADownloadStoreFixture()
 	defer cleanup()
 
-	blob := core.SizedBlobFixture(1, 1)
+	// Use two pieces to prevent triggering a blob cache move after piece
+	// completion, which would fail as a mocked store writer is used.
+	blob := core.SizedBlobFixture(2, 1)
+	piece := blob.Content[:1]
 
 	prepareStore(cads, blob.MetaInfo)
 
@@ -312,12 +318,12 @@ func TestTorrentWritePieceFailuresRemoveDirtyStatus(t *testing.T) {
 	gomock.InOrder(
 		// First write fails.
 		w.EXPECT().Seek(int64(0), 0).Return(int64(0), nil),
-		w.EXPECT().Write(blob.Content).Return(0, errors.New("first write error")),
+		w.EXPECT().Write(piece).Return(0, errors.New("first write error")),
 		w.EXPECT().Close().Return(nil),
 
 		// Second write succeeds.
 		w.EXPECT().Seek(int64(0), 0).Return(int64(0), nil),
-		w.EXPECT().Write(blob.Content).Return(len(blob.Content), nil),
+		w.EXPECT().Write(piece).Return(len(piece), nil),
 		w.EXPECT().Close().Return(nil),
 	)
 
@@ -326,8 +332,8 @@ func TestTorrentWritePieceFailuresRemoveDirtyStatus(t *testing.T) {
 
 	// After the first write fails, the dirty bit should be flipped to empty,
 	// allowing future writes to succeed.
-	require.Error(tor.WritePiece(piecereader.NewBuffer(blob.Content), 0))
-	require.NoError(tor.WritePiece(piecereader.NewBuffer(blob.Content), 0))
+	require.Error(tor.WritePiece(piecereader.NewBuffer(piece), 0))
+	require.NoError(tor.WritePiece(piecereader.NewBuffer(piece), 0))
 }
 
 func TestTorrentRestoreCompletedTorrent(t *testing.T) {