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