diff --git a/lib/store/ca_download_store.go b/lib/store/ca_download_store.go index 28dbe63ce..19a2a75d5 100644 --- a/lib/store/ca_download_store.go +++ b/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. diff --git a/lib/store/ca_download_store_test.go b/lib/store/ca_download_store_test.go index 3f7103c7b..c44c5d422 100644 --- a/lib/store/ca_download_store_test.go +++ b/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)) +} diff --git a/lib/store/ca_store.go b/lib/store/ca_store.go index f2c2f7653..63dd2d4a9 100644 --- a/lib/store/ca_store.go +++ b/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() diff --git a/lib/store/config.go b/lib/store/config.go index 085e648f6..045ebd6a4 100644 --- a/lib/store/config.go +++ b/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"` } diff --git a/lib/store/utils.go b/lib/store/utils.go index 68ee86f4e..142795068 100644 --- a/lib/store/utils.go +++ b/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 } diff --git a/lib/torrent/storage/agentstorage/torrent_test.go b/lib/torrent/storage/agentstorage/torrent_test.go index 7894a484d..32fe3b018 100644 --- a/lib/torrent/storage/agentstorage/torrent_test.go +++ b/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) {