archive/zip: support "end of central directory record comment"

This change added support "end of central directory record comemnt" to the Writer.

There is a new exported field Writer.Comment in this change.
If invalid size of comment was set, Close returns error without closing resources.

Fixes #21634

Change-Id: Ifb62bc6c7f81b9257ac83eb570ad9915de727f8c
Reviewed-on: https://go-review.googlesource.com/59310
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
Run-TryBot: Joe Tsai <thebrokentoaster@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/archive/zip/writer.go b/src/archive/zip/writer.go
index 079917cadc36d12275f4b6d6455afc9b2f8d0085..1aca8518ca3603a576f486782b6433f9822f3f87 100644 (file)
--- a/src/archive/zip/writer.go
+++ b/src/archive/zip/writer.go
@@ -30,6 +30,9 @@ type Writer struct {
        // testHookCloseSizeOffset if non-nil is called with the size
        // of offset of the central directory at Close.
        testHookCloseSizeOffset func(size, offset uint64)
+
+       // Comment is the central directory comment and must be set before Close is called.
+       Comment string
 }
 
 type header struct {
@@ -62,6 +65,10 @@ func (w *Writer) Flush() error {
 // Close finishes writing the zip file by writing the central directory.
 // It does not (and cannot) close the underlying writer.
 func (w *Writer) Close() error {
+       if len(w.Comment) > uint16max {
+               return errors.New("zip: Writer.Comment too long")
+       }
+
        if w.last != nil && !w.last.closed {
                if err := w.last.close(); err != nil {
                        return err
@@ -177,15 +184,18 @@ func (w *Writer) Close() error {
        var buf [directoryEndLen]byte
        b := writeBuf(buf[:])
        b.uint32(uint32(directoryEndSignature))
-       b = b[4:]                 // skip over disk number and first disk number (2x uint16)
-       b.uint16(uint16(records)) // number of entries this disk
-       b.uint16(uint16(records)) // number of entries total
-       b.uint32(uint32(size))    // size of directory
-       b.uint32(uint32(offset))  // start of directory
-       // skipped size of comment (always zero)
+       b = b[4:]                        // skip over disk number and first disk number (2x uint16)
+       b.uint16(uint16(records))        // number of entries this disk
+       b.uint16(uint16(records))        // number of entries total
+       b.uint32(uint32(size))           // size of directory
+       b.uint32(uint32(offset))         // start of directory
+       b.uint16(uint16(len(w.Comment))) // byte size of EOCD comment
        if _, err := w.cw.Write(buf[:]); err != nil {
                return err
        }
+       if _, err := io.WriteString(w.cw, w.Comment); err != nil {
+               return err
+       }
 
        return w.cw.w.(*bufio.Writer).Flush()
 }

diff --git a/src/archive/zip/writer_test.go b/src/archive/zip/writer_test.go
index 92fb6ecf0ed11e7decce45b72b77058994259e73..8db159f232d8aae2a54843a95b95991c7a1fb5b1 100644 (file)
--- a/src/archive/zip/writer_test.go
+++ b/src/archive/zip/writer_test.go
@@ -10,6 +10,7 @@ import (
        "io/ioutil"
        "math/rand"
        "os"
+       "strings"
        "testing"
 )
 
@@ -87,6 +88,48 @@ func TestWriter(t *testing.T) {
        }
 }
 
+// TestWriterComment is test for EOCD comment read/write.
+func TestWriterComment(t *testing.T) {
+       var tests = []struct {
+               comment string
+               ok      bool
+       }{
+               {"hi, hello", true},
+               {"hi, こんにちわ", true},
+               {strings.Repeat("a", uint16max), true},
+               {strings.Repeat("a", uint16max+1), false},
+       }
+
+       for _, test := range tests {
+               // write a zip file
+               buf := new(bytes.Buffer)
+               w := NewWriter(buf)
+               w.Comment = test.comment
+
+               if err := w.Close(); test.ok == (err != nil) {
+                       t.Fatal(err)
+               }
+
+               if w.closed != test.ok {
+                       t.Fatalf("Writer.closed: got %v, want %v", w.closed, test.ok)
+               }
+
+               // skip read test in failure cases
+               if !test.ok {
+                       continue
+               }
+
+               // read it back
+               r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
+               if err != nil {
+                       t.Fatal(err)
+               }
+               if r.Comment != test.comment {
+                       t.Fatalf("Reader.Comment: got %v, want %v", r.Comment, test.comment)
+               }
+       }
+}
+
 func TestWriterUTF8(t *testing.T) {
        var utf8Tests = []struct {
                name    string