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>

2 files changed, 59 insertions, 6 deletions

diff --git a/src/archive/zip/writer.go b/src/archive/zip/writer.go
index 079917cadc..1aca8518ca 100644
--- 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 92fb6ecf0e..8db159f232 100644
--- 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