5 files changed, 67 insertions, 2 deletions
diff --git a/formats/objfile/common.go b/formats/objfile/common.go
index 7389086..2f0585f 100644
--- a/formats/objfile/common.go
+++ b/formats/objfile/common.go
@@ -13,6 +13,8 @@ var (
 	ErrClosed = errors.New("objfile: already closed")
 	// ErrHeader is returned when the objfile has an invalid header.
 	ErrHeader = errors.New("objfile: invalid header")
+	// ErrNegativeSize is returned when a negative object size is declared.
+	ErrNegativeSize = errors.New("objfile: negative object size")
 )
 
 type header struct {
@@ -38,7 +40,11 @@ func (h *header) Read(r io.Reader) error {
 
 	h.size, err = strconv.ParseInt(string(size), 10, 64)
 	if err != nil {
-		return err
+		return ErrHeader
+	}
+
+	if h.size < 0 {
+		return ErrNegativeSize
 	}
 
 	return nil
diff --git a/formats/objfile/common_test.go b/formats/objfile/common_test.go
index 4727685..0c5a4cd 100644
--- a/formats/objfile/common_test.go
+++ b/formats/objfile/common_test.go
@@ -1,6 +1,7 @@
 package objfile
 
 import (
+	"bytes"
 	"encoding/base64"
 	"testing"
 
@@ -67,3 +68,33 @@ var objfileFixtures = []objfileFixture{
 }
 
 func Test(t *testing.T) { TestingT(t) }
+
+type SuiteCommon struct{}
+
+var _ = Suite(&SuiteCommon{})
+
+func (s *SuiteCommon) TestHeaderReadEmpty(c *C) {
+	var h header
+	c.Assert(h.Read(new(bytes.Buffer)), Equals, ErrHeader)
+}
+
+func (s *SuiteCommon) TestHeaderReadGarbage(c *C) {
+	var h header
+	c.Assert(h.Read(bytes.NewBuffer([]byte{1, 2, 3, 4, 5})), Equals, ErrHeader)
+	c.Assert(h.Read(bytes.NewBuffer([]byte{1, 2, 3, 4, 5, '0'})), Equals, ErrHeader)
+}
+
+func (s *SuiteCommon) TestHeaderReadInvalidType(c *C) {
+	var h header
+	c.Assert(h.Read(bytes.NewBuffer([]byte{1, 2, ' ', 4, 5, 0})), Equals, core.ErrInvalidType)
+}
+
+func (s *SuiteCommon) TestHeaderReadInvalidSize(c *C) {
+	var h header
+	c.Assert(h.Read(bytes.NewBuffer([]byte{'b', 'l', 'o', 'b', ' ', 'a', 0})), Equals, ErrHeader)
+}
+
+func (s *SuiteCommon) TestHeaderReadNegativeSize(c *C) {
+	var h header
+	c.Assert(h.Read(bytes.NewBuffer([]byte{'b', 'l', 'o', 'b', ' ', '-', '1', 0})), Equals, ErrNegativeSize)
+}
diff --git a/formats/objfile/reader_test.go b/formats/objfile/reader_test.go
index 871eefe..caebb60 100644
--- a/formats/objfile/reader_test.go
+++ b/formats/objfile/reader_test.go
@@ -35,9 +35,12 @@ func testReader(c *C, source io.Reader, hash core.Hash, typ core.ObjectType, con
 	rc, err := ioutil.ReadAll(r)
 	c.Assert(err, IsNil)
 	c.Assert(rc, DeepEquals, content, Commentf("%scontent=%s, expected=%s", base64.StdEncoding.EncodeToString(rc), base64.StdEncoding.EncodeToString(content)))
+	c.Assert(r.Size(), Equals, int64(len(content)))
 	c.Assert(r.Hash(), Equals, hash) // Test Hash() before close
 	c.Assert(r.Close(), IsNil)
 	c.Assert(r.Hash(), Equals, hash) // Test Hash() after close
+	_, err = r.Read(make([]byte, 0, 1))
+	c.Assert(err, Equals, ErrClosed)
 }
 
 func (s *SuiteReader) TestReadEmptyObjfile(c *C) {
diff --git a/formats/objfile/writer.go b/formats/objfile/writer.go
index d80256c..d9d40f0 100644
--- a/formats/objfile/writer.go
+++ b/formats/objfile/writer.go
@@ -39,9 +39,18 @@ type Writer struct {
 // size and type information. Any errors encountered in that process will be
 // returned in err.
 //
+// If an invalid t is provided, core.ErrInvalidType is returned. If a negative
+// size is provided, ErrNegativeSize is returned.
+//
 // The returned Writer implements io.WriteCloser. Close should be called when
 // finished with the Writer. Close will not close the underlying io.Writer.
 func NewWriter(w io.Writer, t core.ObjectType, size int64) (*Writer, error) {
+	if !t.Valid() {
+		return nil, core.ErrInvalidType
+	}
+	if size < 0 {
+		return nil, ErrNegativeSize
+	}
 	writer := &Writer{
 		header: header{t: t, size: size},
 	}
@@ -58,7 +67,7 @@ func (w *Writer) init(output io.Writer) (err error) {
 
 	err = w.header.Write(w.compressor)
 	if err != nil {
-		defer w.compressor.Close()
+		w.compressor.Close()
 		return
 	}
 
diff --git a/formats/objfile/writer_test.go b/formats/objfile/writer_test.go
index 03b8370..0061f3f 100644
--- a/formats/objfile/writer_test.go
+++ b/formats/objfile/writer_test.go
@@ -38,9 +38,12 @@ func testWriter(c *C, dest io.Writer, hash core.Hash, typ core.ObjectType, conte
 	written, err := io.Copy(w, bytes.NewReader(content))
 	c.Assert(err, IsNil)
 	c.Assert(written, Equals, length)
+	c.Assert(w.Size(), Equals, int64(len(content)))
 	c.Assert(w.Hash(), Equals, hash) // Test Hash() before close
 	c.Assert(w.Close(), IsNil)
 	c.Assert(w.Hash(), Equals, hash) // Test Hash() after close
+	_, err = w.Write([]byte{1})
+	c.Assert(err, Equals, ErrClosed)
 }
 
 func (s *SuiteWriter) TestWriteOverflow(c *C) {
@@ -51,3 +54,16 @@ func (s *SuiteWriter) TestWriteOverflow(c *C) {
 	_, err = w.Write([]byte("56789"))
 	c.Assert(err, Equals, ErrOverflow)
 }
+
+func (s *SuiteWriter) TestNewWriterInvalidType(c *C) {
+	var t core.ObjectType
+	_, err := NewWriter(new(bytes.Buffer), t, 8)
+	c.Assert(err, Equals, core.ErrInvalidType)
+}
+
+func (s *SuiteWriter) TestNewWriterInvalidSize(c *C) {
+	_, err := NewWriter(new(bytes.Buffer), core.BlobObject, -1)
+	c.Assert(err, Equals, ErrNegativeSize)
+	_, err = NewWriter(new(bytes.Buffer), core.BlobObject, -1651860)
+	c.Assert(err, Equals, ErrNegativeSize)
+}