code review 162062: Robustness fixes for tar/archive.

src/pkg/archive/tar/writer.go

 // Copyright 2009 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 package tar
 
 // TODO(dsymonds):
 // - catch more errors (no first header, write after close, etc.)
 
 import (
         "io";
         "os";
         "strconv";
         "strings";
 )
 
 var (
-        ErrWriteTooLong = os.NewError("write too long");
-        ErrFieldTooLong = os.NewError("header field too long");
+        ErrWriteTooLong         = os.NewError("write too long");
+        ErrFieldTooLong         = os.NewError("header field too long");
+        ErrWriteAfterClose      = os.NewError("write after close");
 )
 
 // A Writer provides sequential writing of a tar archive in POSIX.1 format.
 // A tar archive consists of a sequence of files.
 // Call WriteHeader to begin a new file, and then call Write to supply that file's data,
 // writing at most hdr.Size bytes in total.
 //
 // Example:
 //      tw := tar.NewWriter(w);
 //      hdr := new(Header);
@@ skipping 71 matching lines @@
         tw.usedBinary = true;
         for i := len(b) - 1; x > 0 && i >= 0; i-- {
                 b[i] = byte(x);
                 x >>= 8;
         }
         b[0] |= 0x80;   // highest bit indicates binary format
 }
 
 // WriteHeader writes hdr and prepares to accept the file's contents.
 // WriteHeader calls Flush if it is not the first header.
+// Calling after a Close will return ErrWriteAfterClose.
 func (tw *Writer) WriteHeader(hdr *Header) os.Error {
+        if tw.closed {
+                return ErrWriteAfterClose
+        }
         if tw.err == nil {
                 tw.Flush()
         }
         if tw.err != nil {
                 return tw.err
         }
 
         tw.nb = int64(hdr.Size);
         tw.pad = -tw.nb & (blockSize - 1);      // blockSize is a power of two
 
@@ skipping 35 matching lines @@
 
         _, tw.err = tw.w.Write(header);
 
         return tw.err;
 }
 
 // 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 os.Error) {
+        if tw.closed {
+                err = ErrWriteTooLong;
+                return;
+        }
         overwrite := false;
         if int64(len(b)) > tw.nb {
                 b = b[0:tw.nb];
                 overwrite = true;
         }
         n, err = tw.w.Write(b);
         tw.nb -= int64(n);
         if err == nil && overwrite {
-                err = ErrWriteTooLong
+                err = ErrWriteTooLong;
+                return;
         }
         tw.err = err;
         return;
 }
 
+// Close closes the tar archive, flushing any unwritten
+// data to the underlying writer.
 func (tw *Writer) Close() os.Error {
         if tw.err != nil || tw.closed {
                 return tw.err
         }
         tw.Flush();
         tw.closed = true;
 
         // trailer: two zero blocks
         for i := 0; i < 2; i++ {
                 _, tw.err = tw.w.Write(zeroBlock);
                 if tw.err != nil {
                         break
                 }
         }
         return tw.err;
 }