archive/tar: make Writer error handling consistent

The Writer logic was not consistent about when an IO error would
persist across multiple calls on Writer's methods.

Thus, to make the error handling more consistent we always check
the persistent state of the error prior to every exported method
call, and return an error if set. Otherwise, it is the responsibility
of every exported method to persist any fatal errors that may occur.

As a simplification, we can remove the close field since that
information can be represented by simply storing ErrWriteAfterClose
in the err field.

Change-Id: I8746ca36b3739803e0373253450db69b3bd12f38
Reviewed-on: https://go-review.googlesource.com/55590
Run-TryBot: Joe Tsai <joetsai@digital-static.net>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

src/archive/tar/writer.go

@@ -4,9 +4,6 @@
 
 package tar
 
-// TODO(dsymonds):
-// - catch more errors (no first header, etc.)
-
 import (
 	"bytes"
 	"fmt"
@@ -23,13 +20,15 @@ import (
 // writing at most hdr.Size bytes in total.
 type Writer struct {
 	w   io.Writer
-	err    error
 	nb  int64  // number of unwritten bytes for current file entry
 	pad int64  // amount of padding to write after current file entry
-	closed bool
-
 	hdr Header // Shallow copy of Header that is safe for mutations
 	blk block  // Buffer to use as temporary local storage
+
+	// err is a persistent error.
+	// It is only the responsibility of every exported method of Writer to
+	// ensure that this error is sticky.
+	err error
 }
 
 // NewWriter creates a new Writer writing to w.
@@ -41,10 +40,12 @@ func NewWriter(w io.Writer) *Writer { return &Writer{w: w} }
 // Deprecated: This is unecessary as the next call to WriteHeader or Close
 // will implicitly flush out the file's padding.
 func (tw *Writer) Flush() error {
-	if tw.nb > 0 {
-		tw.err = fmt.Errorf("tar: missed writing %d bytes", tw.nb)
+	if tw.err != nil {
 		return tw.err
 	}
+	if tw.nb > 0 {
+		return fmt.Errorf("archive/tar: missed writing %d bytes", tw.nb)
+	}
 	if _, tw.err = tw.w.Write(zeroBlock[:tw.pad]); tw.err != nil {
 		return tw.err
 	}
@@ -63,13 +64,16 @@ func (tw *Writer) WriteHeader(hdr *Header) error {
 	tw.hdr = *hdr // Shallow copy of Header
 	switch allowedFormats, paxHdrs := tw.hdr.allowedFormats(); {
 	case allowedFormats&formatUSTAR != 0:
-		return tw.writeUSTARHeader(&tw.hdr)
+		tw.err = tw.writeUSTARHeader(&tw.hdr)
+		return tw.err
 	case allowedFormats&formatPAX != 0:
-		return tw.writePAXHeader(&tw.hdr, paxHdrs)
+		tw.err = tw.writePAXHeader(&tw.hdr, paxHdrs)
+		return tw.err
 	case allowedFormats&formatGNU != 0:
-		return tw.writeGNUHeader(&tw.hdr)
+		tw.err = tw.writeGNUHeader(&tw.hdr)
+		return tw.err
 	default:
-		return ErrHeader
+		return ErrHeader // Non-fatal error
 	}
 }
 
@@ -273,45 +277,46 @@ func splitUSTARPath(name string) (prefix, suffix string, ok bool) {
 
 // Write writes to the current entry in the tar archive.
 // Write returns the error ErrWriteTooLong if more than
-// hdr.Size bytes are written after WriteHeader.
-func (tw *Writer) Write(b []byte) (n int, err error) {
-	if tw.closed {
-		err = ErrWriteAfterClose
-		return
-	}
-	overwrite := false
-	if int64(len(b)) > tw.nb {
-		b = b[0:tw.nb]
-		overwrite = true
-	}
-	n, err = tw.w.Write(b)
+// Header.Size bytes are written after WriteHeader.
+//
+// Calling Write on special types like TypeLink, TypeSymLink, TypeChar,
+// TypeBlock, TypeDir, and TypeFifo returns (0, ErrWriteTooLong) regardless
+// of what the Header.Size claims.
+func (tw *Writer) Write(b []byte) (int, error) {
+	if tw.err != nil {
+		return 0, tw.err
+	}
+
+	overwrite := int64(len(b)) > tw.nb
+	if overwrite {
+		b = b[:tw.nb]
+	}
+	n, err := tw.w.Write(b)
 	tw.nb -= int64(n)
 	if err == nil && overwrite {
-		err = ErrWriteTooLong
-		return
+		return n, ErrWriteTooLong // Non-fatal error
 	}
 	tw.err = err
-	return
+	return n, err
 }
 
 // Close closes the tar archive, flushing any unwritten
 // data to the underlying writer.
 func (tw *Writer) Close() error {
-	if tw.err != nil || tw.closed {
-		return tw.err
+	if tw.err == ErrWriteAfterClose {
+		return nil
 	}
-	tw.Flush()
-	tw.closed = true
 	if tw.err != nil {
 		return tw.err
 	}
 
-	// trailer: two zero blocks
-	for i := 0; i < 2; i++ {
-		_, tw.err = tw.w.Write(zeroBlock[:])
-		if tw.err != nil {
-			break
-		}
+	// Trailer: two zero blocks.
+	err := tw.Flush()
+	for i := 0; i < 2 && err == nil; i++ {
+		_, err = tw.w.Write(zeroBlock[:])
 	}
-	return tw.err
+
+	// Ensure all future actions are invalid.
+	tw.err = ErrWriteAfterClose
+	return err // Report IO errors
 }

src/archive/tar/writer_test.go

@@ -576,7 +576,19 @@ func TestValidTypeflagWithPAXHeader(t *testing.T) {
 	}
 }
 
-func TestWriteHeaderOnly(t *testing.T) {
+// failOnceWriter fails exactly once and then always reports success.
+type failOnceWriter bool
+
+func (w *failOnceWriter) Write(b []byte) (int, error) {
+	if !*w {
+		return 0, io.ErrShortWrite
+	}
+	*w = true
+	return len(b), nil
+}
+
+func TestWriterErrors(t *testing.T) {
+	t.Run("HeaderOnly", func(t *testing.T) {
 		tw := NewWriter(new(bytes.Buffer))
 		hdr := &Header{Name: "dir/", Typeflag: TypeDir}
 		if err := tw.WriteHeader(hdr); err != nil {
@@ -585,31 +597,83 @@ func TestWriteHeaderOnly(t *testing.T) {
 		if _, err := tw.Write([]byte{0x00}); err != ErrWriteTooLong {
 			t.Fatalf("Write() = %v, want %v", err, ErrWriteTooLong)
 		}
-}
+	})
 
-func TestWriteNegativeSize(t *testing.T) {
+	t.Run("NegativeSize", func(t *testing.T) {
 		tw := NewWriter(new(bytes.Buffer))
 		hdr := &Header{Name: "small.txt", Size: -1}
 		if err := tw.WriteHeader(hdr); err != ErrHeader {
 			t.Fatalf("WriteHeader() = nil, want %v", ErrHeader)
 		}
-}
-
-func TestWriteAfterClose(t *testing.T) {
-	var buffer bytes.Buffer
-	tw := NewWriter(&buffer)
+	})
 
-	hdr := &Header{
-		Name: "small.txt",
-		Size: 5,
+	t.Run("BeforeHeader", func(t *testing.T) {
+		tw := NewWriter(new(bytes.Buffer))
+		if _, err := tw.Write([]byte("Kilts")); err != ErrWriteTooLong {
+			t.Fatalf("Write() = %v, want %v", err, ErrWriteTooLong)
 		}
+	})
+
+	t.Run("AfterClose", func(t *testing.T) {
+		tw := NewWriter(new(bytes.Buffer))
+		hdr := &Header{Name: "small.txt"}
 		if err := tw.WriteHeader(hdr); err != nil {
-		t.Fatalf("Failed to write header: %s", err)
+			t.Fatalf("WriteHeader() = %v, want nil", err)
+		}
+		if err := tw.Close(); err != nil {
+			t.Fatalf("Close() = %v, want nil", err)
 		}
-	tw.Close()
 		if _, err := tw.Write([]byte("Kilts")); err != ErrWriteAfterClose {
-		t.Fatalf("Write: got %v; want ErrWriteAfterClose", err)
+			t.Fatalf("Write() = %v, want %v", err, ErrWriteAfterClose)
+		}
+		if err := tw.Flush(); err != ErrWriteAfterClose {
+			t.Fatalf("Flush() = %v, want %v", err, ErrWriteAfterClose)
 		}
+		if err := tw.Close(); err != nil {
+			t.Fatalf("Close() = %v, want nil", err)
+		}
+	})
+
+	t.Run("PrematureFlush", func(t *testing.T) {
+		tw := NewWriter(new(bytes.Buffer))
+		hdr := &Header{Name: "small.txt", Size: 5}
+		if err := tw.WriteHeader(hdr); err != nil {
+			t.Fatalf("WriteHeader() = %v, want nil", err)
+		}
+		if err := tw.Flush(); err == nil {
+			t.Fatalf("Flush() = %v, want non-nil error", err)
+		}
+	})
+
+	t.Run("PrematureClose", func(t *testing.T) {
+		tw := NewWriter(new(bytes.Buffer))
+		hdr := &Header{Name: "small.txt", Size: 5}
+		if err := tw.WriteHeader(hdr); err != nil {
+			t.Fatalf("WriteHeader() = %v, want nil", err)
+		}
+		if err := tw.Close(); err == nil {
+			t.Fatalf("Close() = %v, want non-nil error", err)
+		}
+	})
+
+	t.Run("Persistence", func(t *testing.T) {
+		tw := NewWriter(new(failOnceWriter))
+		if err := tw.WriteHeader(&Header{}); err != io.ErrShortWrite {
+			t.Fatalf("WriteHeader() = %v, want %v", err, io.ErrShortWrite)
+		}
+		if err := tw.WriteHeader(&Header{Name: "small.txt"}); err == nil {
+			t.Errorf("WriteHeader() = got %v, want non-nil error", err)
+		}
+		if _, err := tw.Write(nil); err == nil {
+			t.Errorf("Write() = %v, want non-nil error", err)
+		}
+		if err := tw.Flush(); err == nil {
+			t.Errorf("Flush() = %v, want non-nil error", err)
+		}
+		if err := tw.Close(); err == nil {
+			t.Errorf("Close() = %v, want non-nil error", err)
+		}
+	})
 }
 
 func TestSplitUSTARPath(t *testing.T) {