archive/zip: enable overriding (de)compressors per file

Implement setting the compression level for a zip archive by registering
a per-Writer compressor through Writer.RegisterCompressor.  If no
compressors are registered, fall back to the ones registered at the
package level.  Also implements per-Reader decompressors.

Fixes #8359

Change-Id: I93b27c81947b0f817b42e0067aa610ff267fdb21
Reviewed-on: https://go-review.googlesource.com/16669
Reviewed-by: Joe Tsai <joetsai@digital-static.net>
Run-TryBot: Joe Tsai <joetsai@digital-static.net>
Reviewed-by: Klaus Post <klauspost@gmail.com>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

1 files changed, 27 insertions, 5 deletions

diff --git a/src/archive/zip/reader.go b/src/archive/zip/reader.go
index 0f7086081a..9aa77d9c43 100644
--- a/src/archive/zip/reader.go
+++ b/src/archive/zip/reader.go
@@ -22,9 +22,10 @@ var (
 )
 
 type Reader struct {
-	r       io.ReaderAt
-	File    []*File
-	Comment string
+	r             io.ReaderAt
+	File          []*File
+	Comment       string
+	decompressors map[uint16]Decompressor
 }
 
 type ReadCloser struct {
@@ -34,6 +35,7 @@ type ReadCloser struct {
 
 type File struct {
 	FileHeader
+	zip          *Reader
 	zipr         io.ReaderAt
 	zipsize      int64
 	headerOffset int64
@@ -95,7 +97,7 @@ func (z *Reader) init(r io.ReaderAt, size int64) error {
 	// a bad one, and then only report a ErrFormat or UnexpectedEOF if
 	// the file count modulo 65536 is incorrect.
 	for {
-		f := &File{zipr: r, zipsize: size}
+		f := &File{zip: z, zipr: r, zipsize: size}
 		err = readDirectoryHeader(f, buf)
 		if err == ErrFormat || err == io.ErrUnexpectedEOF {
 			break
@@ -113,6 +115,26 @@ func (z *Reader) init(r io.ReaderAt, size int64) error {
 	return nil
 }
 
+// RegisterDecompressor registers or overrides a custom decompressor for a
+// specific method ID. If a decompressor for a given method is not found,
+// Reader will default to looking up the decompressor at the package level.
+//
+// Must not be called concurrently with Open on any Files in the Reader.
+func (z *Reader) RegisterDecompressor(method uint16, dcomp Decompressor) {
+	if z.decompressors == nil {
+		z.decompressors = make(map[uint16]Decompressor)
+	}
+	z.decompressors[method] = dcomp
+}
+
+func (z *Reader) decompressor(method uint16) Decompressor {
+	dcomp := z.decompressors[method]
+	if dcomp == nil {
+		dcomp = decompressor(method)
+	}
+	return dcomp
+}
+
 // Close closes the Zip file, rendering it unusable for I/O.
 func (rc *ReadCloser) Close() error {
 	return rc.f.Close()
@@ -140,7 +162,7 @@ func (f *File) Open() (rc io.ReadCloser, err error) {
 	}
 	size := int64(f.CompressedSize64)
 	r := io.NewSectionReader(f.zipr, f.headerOffset+bodyOffset, size)
-	dcomp := decompressor(f.Method)
+	dcomp := f.zip.decompressor(f.Method)
 	if dcomp == nil {
 		err = ErrAlgorithm
 		return