From 1e1760eebf5ae90e63d713fc8b2ee69d7eb885d0 Mon Sep 17 00:00:00 2001
From: Brandon Westcott <brandon@quorumcontrol.com>
Date: Thu, 2 Apr 2020 11:26:44 -0400
Subject: [PATCH] Ensure zlib buffer is closed before writing objects to
 storage

---
 storage/chaintree/object.go | 17 +----------------
 storage/object.go           | 24 ++++++++++++++++++++++++
 storage/object_test.go      | 25 +++++++++++++++++++++++++
 storage/siaskynet/net.go    | 17 ++---------------
 4 files changed, 52 insertions(+), 31 deletions(-)
 create mode 100644 storage/object_test.go

diff --git a/storage/chaintree/object.go b/storage/chaintree/object.go
index d5e2e05..f31d46f 100644
--- a/storage/chaintree/object.go
+++ b/storage/chaintree/object.go
@@ -116,25 +116,10 @@ func (s *ObjectStorage) SetEncodedObjectTxn(o plumbing.EncodedObject) (*transact
 		return nil, plumbing.ErrInvalidType
 	}
 
-	buf := bytes.NewBuffer(nil)
-
-	writer := objfile.NewWriter(buf)
-	defer writer.Close()
-
-	reader, err := o.Reader()
+	buf, err := storage.ZlibBufferForObject(o)
 	if err != nil {
 		return nil, err
 	}
-	defer reader.Close()
-
-	if err := writer.WriteHeader(o.Type(), o.Size()); err != nil {
-		return nil, err
-	}
-
-	if _, err = io.Copy(writer, reader); err != nil {
-		return nil, err
-	}
-	writer.Close()
 
 	objectBytes, err := ioutil.ReadAll(buf)
 	if err != nil {
diff --git a/storage/object.go b/storage/object.go
index d6e0aae..56cef49 100644
--- a/storage/object.go
+++ b/storage/object.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/go-git/go-git/v5/plumbing"
+	"github.com/go-git/go-git/v5/plumbing/format/objfile"
 	"github.com/go-git/go-git/v5/plumbing/format/packfile"
 	"github.com/go-git/go-git/v5/plumbing/storer"
 	logging "github.com/ipfs/go-log"
@@ -37,6 +38,29 @@ func ObjectWritePath(h plumbing.Hash) string {
 	return strings.Join(ObjectReadPath(h)[2:], "/")
 }
 
+func ZlibBufferForObject(o plumbing.EncodedObject) (*bytes.Buffer, error) {
+	buf := bytes.NewBuffer(nil)
+
+	writer := objfile.NewWriter(buf)
+	defer writer.Close()
+
+	reader, err := o.Reader()
+	if err != nil {
+		return nil, err
+	}
+	defer reader.Close()
+
+	if err := writer.WriteHeader(o.Type(), o.Size()); err != nil {
+		return nil, err
+	}
+
+	if _, err = io.Copy(writer, reader); err != nil {
+		return nil, err
+	}
+
+	return buf, err
+}
+
 func (s *ChaintreeObjectStorage) Chaintree() *chaintree.ChainTree {
 	return s.ChainTree.ChainTree
 }
diff --git a/storage/object_test.go b/storage/object_test.go
new file mode 100644
index 0000000..9667e8c
--- /dev/null
+++ b/storage/object_test.go
@@ -0,0 +1,25 @@
+package storage
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/go-git/go-git/v5/plumbing"
+	"github.com/stretchr/testify/require"
+)
+
+func TestZlibBufferForObject(t *testing.T) {
+	o := &plumbing.MemoryObject{}
+	o.SetType(plumbing.BlobObject)
+	o.SetSize(14)
+	_, err := o.Write([]byte("Hello, World!\n"))
+	require.Nil(t, err)
+	require.Equal(t, o.Hash().String(), "8ab686eafeb1f44702738c8b0f24f2567c36da6d")
+
+	buf, err := ZlibBufferForObject(o)
+	require.Nil(t, err)
+
+	fmt.Printf("%v", buf.Bytes())
+
+	require.Equal(t, buf.Bytes(), []byte{120, 156, 74, 202, 201, 79, 82, 48, 52, 97, 240, 72, 205, 201, 201, 215, 81, 8, 207, 47, 202, 73, 81, 228, 2, 4, 0, 0, 255, 255, 78, 21, 6, 152})
+}
diff --git a/storage/siaskynet/net.go b/storage/siaskynet/net.go
index 65cc222..e77138b 100644
--- a/storage/siaskynet/net.go
+++ b/storage/siaskynet/net.go
@@ -1,13 +1,13 @@
 package siaskynet
 
 import (
-	"bytes"
 	"io"
 	"sync"
 
 	"github.com/NebulousLabs/go-skynet"
 	"github.com/go-git/go-git/v5/plumbing"
 	"github.com/go-git/go-git/v5/plumbing/format/objfile"
+	"github.com/quorumcontrol/dgit/storage"
 	"go.uber.org/zap"
 )
 
@@ -47,24 +47,11 @@ func InitSkynet(uploaderCount, downloaderCount int) *Skynet {
 }
 
 func (s *Skynet) uploadObject(o plumbing.EncodedObject) (string, error) {
-	buf := bytes.NewBuffer(nil)
-
-	writer := objfile.NewWriter(buf)
-	defer writer.Close()
-
-	reader, err := o.Reader()
+	buf, err := storage.ZlibBufferForObject(o)
 	if err != nil {
 		return "", err
 	}
 
-	if err = writer.WriteHeader(o.Type(), o.Size()); err != nil {
-		return "", err
-	}
-
-	if _, err = io.Copy(writer, reader); err != nil {
-		return "", err
-	}
-
 	uploadData := make(skynet.UploadData)
 	uploadData[o.Hash().String()] = buf