feat: stream write to zip directly (#863)

1 files changed, 48 insertions, 26 deletions

diff --git a/file.go b/file.go
index 36b1a42..495718a 100644
--- a/file.go
+++ b/file.go
@@ -87,17 +87,55 @@ func (f *File) Write(w io.Writer) error {
 
 // WriteTo implements io.WriterTo to write the file.
 func (f *File) WriteTo(w io.Writer) (int64, error) {
-	buf, err := f.WriteToBuffer()
-	if err != nil {
+	if f.options != nil && f.options.Password != "" {
+		buf, err := f.WriteToBuffer()
+		if err != nil {
+			return 0, err
+		}
+		return buf.WriteTo(w)
+	}
+	if err := f.writeDirectToWriter(w); err != nil {
 		return 0, err
 	}
-	return buf.WriteTo(w)
+	return 0, nil
 }
 
-// WriteToBuffer provides a function to get bytes.Buffer from the saved file.
+// WriteToBuffer provides a function to get bytes.Buffer from the saved file. And it allocate space in memory. Be careful when the file size is large.
 func (f *File) WriteToBuffer() (*bytes.Buffer, error) {
 	buf := new(bytes.Buffer)
 	zw := zip.NewWriter(buf)
+
+	if err := f.writeToZip(zw); err != nil {
+		return buf, zw.Close()
+	}
+
+	if f.options != nil && f.options.Password != "" {
+		if err := zw.Close(); err != nil {
+			return buf, err
+		}
+		b, err := Encrypt(buf.Bytes(), f.options)
+		if err != nil {
+			return buf, err
+		}
+		buf.Reset()
+		buf.Write(b)
+		return buf, nil
+	}
+	return buf, zw.Close()
+}
+
+// writeDirectToWriter provides a function to write to io.Writer.
+func (f *File) writeDirectToWriter(w io.Writer) error {
+	zw := zip.NewWriter(w)
+	if err := f.writeToZip(zw); err != nil {
+		zw.Close()
+		return err
+	}
+	return zw.Close()
+}
+
+// writeToZip provides a function to write to zip.Writer
+func (f *File) writeToZip(zw *zip.Writer) error {
 	f.calcChainWriter()
 	f.commentsWriter()
 	f.contentTypesWriter()
@@ -112,19 +150,17 @@ func (f *File) WriteToBuffer() (*bytes.Buffer, error) {
 	for path, stream := range f.streams {
 		fi, err := zw.Create(path)
 		if err != nil {
-			zw.Close()
-			return buf, err
+			return err
 		}
 		var from io.Reader
 		from, err = stream.rawData.Reader()
 		if err != nil {
 			stream.rawData.Close()
-			return buf, err
+			return err
 		}
 		_, err = io.Copy(fi, from)
 		if err != nil {
-			zw.Close()
-			return buf, err
+			return err
 		}
 		stream.rawData.Close()
 	}
@@ -135,27 +171,13 @@ func (f *File) WriteToBuffer() (*bytes.Buffer, error) {
 		}
 		fi, err := zw.Create(path)
 		if err != nil {
-			zw.Close()
-			return buf, err
+			return err
 		}
 		_, err = fi.Write(content)
 		if err != nil {
-			zw.Close()
-			return buf, err
+			return err
 		}
 	}
 
-	if f.options != nil && f.options.Password != "" {
-		if err := zw.Close(); err != nil {
-			return buf, err
-		}
-		b, err := Encrypt(buf.Bytes(), f.options)
-		if err != nil {
-			return buf, err
-		}
-		buf.Reset()
-		buf.Write(b)
-		return buf, nil
-	}
-	return buf, zw.Close()
+	return nil
 }