archive/zip: add NewWriterWithOptions

When appending zip data to existing data such as a binary file the
zip headers must use the correct offset. NewWriterWithOptions
allows creating a Writer that uses the provided offset in the zip
headers.

Fixes #8669

Change-Id: I6ec64f1e816cc57b6fc8bb9e8a0918e586fc56b0
Reviewed-on: https://go-review.googlesource.com/2978
Reviewed-by: Andrew Gerrand <adg@golang.org>

diff --git a/src/archive/zip/writer.go b/src/archive/zip/writer.go
index 6a71887c6360c8c30a6d7f3d86a0f47b0193dc13..c2e04163cf5a8cbe15c5e1b6a8ff185bbe0edb58 100644 (file)
--- a/src/archive/zip/writer.go
+++ b/src/archive/zip/writer.go
@@ -29,9 +29,29 @@ type header struct {
        offset uint64
 }
 
+// WriterOptions contains configuration options for a zip.Writer.
+type WriterOptions struct {
+       // Offset modifies the initial zip offset.
+       // This is useful when the zip is appended to other data such as a binary executable.
+       Offset int64
+}
+
 // NewWriter returns a new Writer writing a zip file to w.
 func NewWriter(w io.Writer) *Writer {
-       return &Writer{cw: &countWriter{w: bufio.NewWriter(w)}}
+       return NewWriterWithOptions(w, nil)
+}
+
+// NewWriterWithOptions returns a new Writer writing a zip file to w and uses the given options.
+func NewWriterWithOptions(w io.Writer, options *WriterOptions) *Writer {
+       writer := &Writer{
+               cw: &countWriter{
+                       w: bufio.NewWriter(w),
+               },
+       }
+       if options != nil {
+               writer.cw.count = options.Offset
+       }
+       return writer
 }
 
 // Flush flushes any buffered data to the underlying writer.

diff --git a/src/archive/zip/writer_test.go b/src/archive/zip/writer_test.go
index 184a7d96a7f2bb331d35a77932de8e4189dd38a5..8be86408d80c3c721cf7bc7f35762c6f5e4dd4c5 100644 (file)
--- a/src/archive/zip/writer_test.go
+++ b/src/archive/zip/writer_test.go
@@ -87,6 +87,40 @@ func TestWriter(t *testing.T) {
        }
 }
 
+func TestWriterOffsetOption(t *testing.T) {
+       largeData := make([]byte, 1<<17)
+       for i := range largeData {
+               largeData[i] = byte(rand.Int())
+       }
+       writeTests[1].Data = largeData
+       defer func() {
+               writeTests[1].Data = nil
+       }()
+
+       // write a zip file
+       buf := new(bytes.Buffer)
+       existingData := []byte{1, 2, 3, 1, 2, 3, 1, 2, 3}
+       n, _ := buf.Write(existingData)
+       w := NewWriterWithOptions(buf, &WriterOptions{Offset: int64(n)})
+
+       for _, wt := range writeTests {
+               testCreate(t, w, &wt)
+       }
+
+       if err := w.Close(); err != nil {
+               t.Fatal(err)
+       }
+
+       // read it back
+       r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
+       if err != nil {
+               t.Fatal(err)
+       }
+       for i, wt := range writeTests {
+               testReadFile(t, r.File[i], &wt)
+       }
+}
+
 func TestWriterFlush(t *testing.T) {
        var buf bytes.Buffer
        w := NewWriter(struct{ io.Writer }{&buf})