From 80916ffb537e0f9b032440ddb6273306cf178601 Mon Sep 17 00:00:00 2001
From: Egon Elbre <egonelbre@gmail.com>
Date: Mon, 18 Mar 2019 16:29:54 +0200
Subject: [PATCH] storagenode/pieces: ensure we can call Commit or Cancel only
 once (#1511)

---
 storage/filestore/blob.go        | 15 ++++++++++++---
 storage/filestore/store_test.go  |  6 ++++++
 storagenode/pieces/readwrite.go  | 12 +++++++++++-
 storagenode/pieces/store_test.go |  7 ++++++-
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/storage/filestore/blob.go b/storage/filestore/blob.go
index f84ce9e7b..fffb87697 100644
--- a/storage/filestore/blob.go
+++ b/storage/filestore/blob.go
@@ -32,18 +32,23 @@ func (blob *blobReader) Size() (int64, error) {
 
 // blobWriter implements writing blobs
 type blobWriter struct {
-	ref   storage.BlobRef
-	store *Store
+	ref    storage.BlobRef
+	store  *Store
+	closed bool
 
 	*os.File
 }
 
 func newBlobWriter(ref storage.BlobRef, store *Store, file *os.File) *blobWriter {
-	return &blobWriter{ref, store, file}
+	return &blobWriter{ref, store, false, file}
 }
 
 // Cancel discards the blob.
 func (blob *blobWriter) Cancel() error {
+	if blob.closed {
+		return nil
+	}
+	blob.closed = true
 	err := blob.File.Close()
 	removeErr := os.Remove(blob.File.Name())
 	return Error.Wrap(errs.Combine(err, removeErr))
@@ -51,6 +56,10 @@ func (blob *blobWriter) Cancel() error {
 
 // Commit moves the file to the target location.
 func (blob *blobWriter) Commit() error {
+	if blob.closed {
+		return Error.New("already closed")
+	}
+	blob.closed = true
 	err := blob.store.dir.Commit(blob.File, blob.ref)
 	return Error.Wrap(err)
 }
diff --git a/storage/filestore/store_test.go b/storage/filestore/store_test.go
index d768b2446..2d6572ac6 100644
--- a/storage/filestore/store_test.go
+++ b/storage/filestore/store_test.go
@@ -58,6 +58,10 @@ func TestStoreLoad(t *testing.T) {
 		require.Equal(t, n, len(data))
 
 		require.NoError(t, writer.Commit())
+		// after committing we should be able to call cancel without an error
+		require.NoError(t, writer.Cancel())
+		// two commits should fail
+		require.Error(t, writer.Commit())
 	}
 
 	namespace = randomValue()
@@ -114,6 +118,8 @@ func TestStoreLoad(t *testing.T) {
 		require.Equal(t, n, len(data))
 
 		require.NoError(t, writer.Cancel())
+		// commit after cancel should return an error
+		require.Error(t, writer.Commit())
 
 		_, err = store.Open(ctx, ref)
 		require.Error(t, err)
diff --git a/storagenode/pieces/readwrite.go b/storagenode/pieces/readwrite.go
index c591083a2..d4e12718f 100644
--- a/storagenode/pieces/readwrite.go
+++ b/storagenode/pieces/readwrite.go
@@ -20,6 +20,8 @@ type Writer struct {
 	hash hash.Hash
 	blob storage.BlobWriter
 	size int64
+
+	closed bool
 }
 
 // NewWriter creates a new writer for storage.BlobWriter.
@@ -47,14 +49,22 @@ func (w *Writer) Hash() []byte { return w.hash.Sum(nil) }
 
 // Commit commits piece to permanent storage.
 func (w *Writer) Commit() error {
+	if w.closed {
+		return Error.New("already closed")
+	}
+	w.closed = true
 	if err := w.buf.Flush(); err != nil {
-		return Error.Wrap(errs.Combine(err, w.Cancel()))
+		return Error.Wrap(errs.Combine(err, w.blob.Cancel()))
 	}
 	return Error.Wrap(w.blob.Commit())
 }
 
 // Cancel deletes any temporarily written data.
 func (w *Writer) Cancel() error {
+	if w.closed {
+		return nil
+	}
+	w.closed = true
 	w.buf.Reset(nil)
 	return Error.Wrap(w.blob.Cancel())
 }
diff --git a/storagenode/pieces/store_test.go b/storagenode/pieces/store_test.go
index 5261abb29..509aa3601 100644
--- a/storagenode/pieces/store_test.go
+++ b/storagenode/pieces/store_test.go
@@ -56,6 +56,8 @@ func TestPieces(t *testing.T) {
 
 		// commit
 		require.NoError(t, writer.Commit())
+		// after commit we should be able to call cancel without an error
+		require.NoError(t, writer.Cancel())
 	}
 
 	{ // valid reads
@@ -72,6 +74,8 @@ func TestPieces(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, int(length), n)
 
+			require.NoError(t, reader.Close())
+
 			return data
 		}
 
@@ -99,10 +103,11 @@ func TestPieces(t *testing.T) {
 
 		// cancel writing
 		require.NoError(t, writer.Cancel())
+		// commit should not fail
+		require.Error(t, writer.Commit())
 
 		// read should fail
 		_, err = store.Reader(ctx, satelliteID, cancelledPieceID)
 		assert.Error(t, err)
 	}
-
 }