archive/zip: prevent writing data for a directory

When creating a directory, Writer.Create now returns a dummy
io.Writer that always returns an error on Write.

Fixes #24043

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

2 files changed, 46 insertions, 20 deletions

diff --git a/src/archive/zip/writer.go b/src/archive/zip/writer.go
index 12675d60a0..f3abe8770c 100644
--- a/src/archive/zip/writer.go
+++ b/src/archive/zip/writer.go
@@ -11,6 +11,7 @@ import (
 	"hash"
 	"hash/crc32"
 	"io"
+	"strings"
 	"unicode/utf8"
 )
 
@@ -320,35 +321,43 @@ func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
 		fh.Extra = append(fh.Extra, mbuf[:]...)
 	}
 
-	fw := &fileWriter{
-		zipw:      w.cw,
-		compCount: &countWriter{w: w.cw},
-		crc32:     crc32.NewIEEE(),
-	}
-	comp := w.compressor(fh.Method)
-	if comp == nil {
-		return nil, ErrAlgorithm
-	}
-	var err error
-	fw.comp, err = comp(fw.compCount)
-	if err != nil {
-		return nil, err
-	}
-	fw.rawCount = &countWriter{w: fw.comp}
-
+	var (
+		ow io.Writer
+		fw *fileWriter
+	)
 	h := &header{
 		FileHeader: fh,
 		offset:     uint64(w.cw.count),
 	}
-	w.dir = append(w.dir, h)
-	fw.header = h
 
+	if strings.HasSuffix(fh.Name, "/") {
+		ow = dirWriter{}
+	} else {
+		fw = &fileWriter{
+			zipw:      w.cw,
+			compCount: &countWriter{w: w.cw},
+			crc32:     crc32.NewIEEE(),
+		}
+		comp := w.compressor(fh.Method)
+		if comp == nil {
+			return nil, ErrAlgorithm
+		}
+		var err error
+		fw.comp, err = comp(fw.compCount)
+		if err != nil {
+			return nil, err
+		}
+		fw.rawCount = &countWriter{w: fw.comp}
+		fw.header = h
+		ow = fw
+	}
+	w.dir = append(w.dir, h)
 	if err := writeHeader(w.cw, fh); err != nil {
 		return nil, err
 	}
-
+	// If we're creating a directory, fw is nil.
 	w.last = fw
-	return fw, nil
+	return ow, nil
 }
 
 func writeHeader(w io.Writer, h *FileHeader) error {
@@ -401,6 +410,12 @@ func (w *Writer) compressor(method uint16) Compressor {
 	return comp
 }
 
+type dirWriter struct{}
+
+func (dirWriter) Write([]byte) (int, error) {
+	return 0, errors.New("zip: write to directory")
+}
+
 type fileWriter struct {
 	*header
 	zipw      io.Writer
diff --git a/src/archive/zip/writer_test.go b/src/archive/zip/writer_test.go
index 38f32296fa..271a36729c 100644
--- a/src/archive/zip/writer_test.go
+++ b/src/archive/zip/writer_test.go
@@ -299,6 +299,17 @@ func TestWriterFlush(t *testing.T) {
 	}
 }
 
+func TestWriterDir(t *testing.T) {
+	w := NewWriter(ioutil.Discard)
+	dw, err := w.Create("dir/")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := dw.Write([]byte("hello")); err == nil {
+		t.Error("Write to directory: got nil error, want non-nil")
+	}
+}
+
 func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
 	header := &FileHeader{
 		Name:   wt.Name,