archive/zip: add Writer.Flush

This is needed for callers to be able to keep track of the
writing position within a zip file. Otherwise it's not
possible to compute the size of headers, and the TOC isn't
written until the very end.

LGTM=adg
R=adg
CC=golang-codereviews
https://golang.org/cl/134210043

2 files changed, 25 insertions, 0 deletions

diff --git a/src/pkg/archive/zip/writer.go b/src/pkg/archive/zip/writer.go
index 6c9800a78f..170beec0ee 100644
--- a/src/pkg/archive/zip/writer.go
+++ b/src/pkg/archive/zip/writer.go
@@ -34,6 +34,12 @@ func NewWriter(w io.Writer) *Writer {
 	return &Writer{cw: &countWriter{w: bufio.NewWriter(w)}}
 }
 
+// Flush flushes any buffered data to the underlying writer.
+// Calling Flush is not normally necessary; calling Close is sufficient.
+func (w *Writer) Flush() error {
+	return w.cw.w.(*bufio.Writer).Flush()
+}
+
 // Close finishes writing the zip file by writing the central directory.
 // It does not (and can not) close the underlying writer.
 func (w *Writer) Close() error {
diff --git a/src/pkg/archive/zip/writer_test.go b/src/pkg/archive/zip/writer_test.go
index 4bfa870809..184a7d96a7 100644
--- a/src/pkg/archive/zip/writer_test.go
+++ b/src/pkg/archive/zip/writer_test.go
@@ -6,6 +6,7 @@ package zip
 
 import (
 	"bytes"
+	"io"
 	"io/ioutil"
 	"math/rand"
 	"os"
@@ -86,6 +87,24 @@ func TestWriter(t *testing.T) {
 	}
 }
 
+func TestWriterFlush(t *testing.T) {
+	var buf bytes.Buffer
+	w := NewWriter(struct{ io.Writer }{&buf})
+	_, err := w.Create("foo")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if buf.Len() > 0 {
+		t.Fatalf("Unexpected %d bytes already in buffer", buf.Len())
+	}
+	if err := w.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if buf.Len() == 0 {
+		t.Fatal("No bytes written after Flush")
+	}
+}
+
 func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
 	header := &FileHeader{
 		Name:   wt.Name,