From de573f87484df3278d9b7412da1964313aaeb90a Mon Sep 17 00:00:00 2001
From: Geert-Johan Riemer <gjr19912@gmail.com>
Date: Mon, 19 Jan 2015 14:39:33 +0100
Subject: [PATCH] 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>
---
 src/archive/zip/writer.go      | 22 +++++++++++++++++++++-
 src/archive/zip/writer_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/archive/zip/writer.go b/src/archive/zip/writer.go
index 6a71887c63..c2e04163cf 100644
--- 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 184a7d96a7..8be86408d8 100644
--- 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})
-- 
2.48.1