1
2
3
4
5 package zip
6
7 import (
8 "bufio"
9 "encoding/binary"
10 "errors"
11 "hash"
12 "hash/crc32"
13 "io"
14 "io/fs"
15 "strings"
16 "unicode/utf8"
17 )
18
19 var (
20 errLongName = errors.New("zip: FileHeader.Name too long")
21 errLongExtra = errors.New("zip: FileHeader.Extra too long")
22 )
23
24
25 type Writer struct {
26 cw *countWriter
27 dir []*header
28 last *fileWriter
29 closed bool
30 compressors map[uint16]Compressor
31 comment string
32
33
34
35 testHookCloseSizeOffset func(size, offset uint64)
36 }
37
38 type header struct {
39 *FileHeader
40 offset uint64
41 raw bool
42 }
43
44
45 func NewWriter(w io.Writer) *Writer {
46 return &Writer{cw: &countWriter{w: bufio.NewWriter(w)}}
47 }
48
49
50
51
52
53 func (w *Writer) SetOffset(n int64) {
54 if w.cw.count != 0 {
55 panic("zip: SetOffset called after data was written")
56 }
57 w.cw.count = n
58 }
59
60
61
62 func (w *Writer) Flush() error {
63 return w.cw.w.(*bufio.Writer).Flush()
64 }
65
66
67
68 func (w *Writer) SetComment(comment string) error {
69 if len(comment) > uint16max {
70 return errors.New("zip: Writer.Comment too long")
71 }
72 w.comment = comment
73 return nil
74 }
75
76
77
78 func (w *Writer) Close() error {
79 if w.last != nil && !w.last.closed {
80 if err := w.last.close(); err != nil {
81 return err
82 }
83 w.last = nil
84 }
85 if w.closed {
86 return errors.New("zip: writer closed twice")
87 }
88 w.closed = true
89
90
91 start := w.cw.count
92 for _, h := range w.dir {
93 var buf [directoryHeaderLen]byte
94 b := writeBuf(buf[:])
95 b.uint32(uint32(directoryHeaderSignature))
96 b.uint16(h.CreatorVersion)
97 b.uint16(h.ReaderVersion)
98 b.uint16(h.Flags)
99 b.uint16(h.Method)
100 b.uint16(h.ModifiedTime)
101 b.uint16(h.ModifiedDate)
102 b.uint32(h.CRC32)
103 if h.isZip64() || h.offset >= uint32max {
104
105
106
107 b.uint32(uint32max)
108 b.uint32(uint32max)
109
110
111 var buf [28]byte
112 eb := writeBuf(buf[:])
113 eb.uint16(zip64ExtraID)
114 eb.uint16(24)
115 eb.uint64(h.UncompressedSize64)
116 eb.uint64(h.CompressedSize64)
117 eb.uint64(h.offset)
118 h.Extra = append(h.Extra, buf[:]...)
119 } else {
120 b.uint32(h.CompressedSize)
121 b.uint32(h.UncompressedSize)
122 }
123
124 b.uint16(uint16(len(h.Name)))
125 b.uint16(uint16(len(h.Extra)))
126 b.uint16(uint16(len(h.Comment)))
127 b = b[4:]
128 b.uint32(h.ExternalAttrs)
129 if h.offset > uint32max {
130 b.uint32(uint32max)
131 } else {
132 b.uint32(uint32(h.offset))
133 }
134 if _, err := w.cw.Write(buf[:]); err != nil {
135 return err
136 }
137 if _, err := io.WriteString(w.cw, h.Name); err != nil {
138 return err
139 }
140 if _, err := w.cw.Write(h.Extra); err != nil {
141 return err
142 }
143 if _, err := io.WriteString(w.cw, h.Comment); err != nil {
144 return err
145 }
146 }
147 end := w.cw.count
148
149 records := uint64(len(w.dir))
150 size := uint64(end - start)
151 offset := uint64(start)
152
153 if f := w.testHookCloseSizeOffset; f != nil {
154 f(size, offset)
155 }
156
157 if records >= uint16max || size >= uint32max || offset >= uint32max {
158 var buf [directory64EndLen + directory64LocLen]byte
159 b := writeBuf(buf[:])
160
161
162 b.uint32(directory64EndSignature)
163 b.uint64(directory64EndLen - 12)
164 b.uint16(zipVersion45)
165 b.uint16(zipVersion45)
166 b.uint32(0)
167 b.uint32(0)
168 b.uint64(records)
169 b.uint64(records)
170 b.uint64(size)
171 b.uint64(offset)
172
173
174 b.uint32(directory64LocSignature)
175 b.uint32(0)
176 b.uint64(uint64(end))
177 b.uint32(1)
178
179 if _, err := w.cw.Write(buf[:]); err != nil {
180 return err
181 }
182
183
184
185 records = uint16max
186 size = uint32max
187 offset = uint32max
188 }
189
190
191 var buf [directoryEndLen]byte
192 b := writeBuf(buf[:])
193 b.uint32(uint32(directoryEndSignature))
194 b = b[4:]
195 b.uint16(uint16(records))
196 b.uint16(uint16(records))
197 b.uint32(uint32(size))
198 b.uint32(uint32(offset))
199 b.uint16(uint16(len(w.comment)))
200 if _, err := w.cw.Write(buf[:]); err != nil {
201 return err
202 }
203 if _, err := io.WriteString(w.cw, w.comment); err != nil {
204 return err
205 }
206
207 return w.cw.w.(*bufio.Writer).Flush()
208 }
209
210
211
212
213
214
215
216
217
218
219 func (w *Writer) Create(name string) (io.Writer, error) {
220 header := &FileHeader{
221 Name: name,
222 Method: Deflate,
223 }
224 return w.CreateHeader(header)
225 }
226
227
228
229
230 func detectUTF8(s string) (valid, require bool) {
231 for i := 0; i < len(s); {
232 r, size := utf8.DecodeRuneInString(s[i:])
233 i += size
234
235
236
237
238
239
240 if r < 0x20 || r > 0x7d || r == 0x5c {
241 if !utf8.ValidRune(r) || (r == utf8.RuneError && size == 1) {
242 return false, false
243 }
244 require = true
245 }
246 }
247 return true, require
248 }
249
250
251
252 func (w *Writer) prepare(fh *FileHeader) error {
253 if w.last != nil && !w.last.closed {
254 if err := w.last.close(); err != nil {
255 return err
256 }
257 }
258 if len(w.dir) > 0 && w.dir[len(w.dir)-1].FileHeader == fh {
259
260 return errors.New("archive/zip: invalid duplicate FileHeader")
261 }
262 return nil
263 }
264
265
266
267
268
269
270
271
272 func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
273 if err := w.prepare(fh); err != nil {
274 return nil, err
275 }
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293 utf8Valid1, utf8Require1 := detectUTF8(fh.Name)
294 utf8Valid2, utf8Require2 := detectUTF8(fh.Comment)
295 switch {
296 case fh.NonUTF8:
297 fh.Flags &^= 0x800
298 case (utf8Require1 || utf8Require2) && (utf8Valid1 && utf8Valid2):
299 fh.Flags |= 0x800
300 }
301
302 fh.CreatorVersion = fh.CreatorVersion&0xff00 | zipVersion20
303 fh.ReaderVersion = zipVersion20
304
305
306 if !fh.Modified.IsZero() {
307
308
309
310
311
312
313
314
315 fh.ModifiedDate, fh.ModifiedTime = timeToMsDosTime(fh.Modified)
316
317
318
319
320
321
322
323 var mbuf [9]byte
324 mt := uint32(fh.Modified.Unix())
325 eb := writeBuf(mbuf[:])
326 eb.uint16(extTimeExtraID)
327 eb.uint16(5)
328 eb.uint8(1)
329 eb.uint32(mt)
330 fh.Extra = append(fh.Extra, mbuf[:]...)
331 }
332
333 var (
334 ow io.Writer
335 fw *fileWriter
336 )
337 h := &header{
338 FileHeader: fh,
339 offset: uint64(w.cw.count),
340 }
341
342 if strings.HasSuffix(fh.Name, "/") {
343
344
345
346
347 fh.Method = Store
348 fh.Flags &^= 0x8
349
350
351 fh.CompressedSize = 0
352 fh.CompressedSize64 = 0
353 fh.UncompressedSize = 0
354 fh.UncompressedSize64 = 0
355
356 ow = dirWriter{}
357 } else {
358 fh.Flags |= 0x8
359
360 fw = &fileWriter{
361 zipw: w.cw,
362 compCount: &countWriter{w: w.cw},
363 crc32: crc32.NewIEEE(),
364 }
365 comp := w.compressor(fh.Method)
366 if comp == nil {
367 return nil, ErrAlgorithm
368 }
369 var err error
370 fw.comp, err = comp(fw.compCount)
371 if err != nil {
372 return nil, err
373 }
374 fw.rawCount = &countWriter{w: fw.comp}
375 fw.header = h
376 ow = fw
377 }
378 w.dir = append(w.dir, h)
379 if err := writeHeader(w.cw, h); err != nil {
380 return nil, err
381 }
382
383 w.last = fw
384 return ow, nil
385 }
386
387 func writeHeader(w io.Writer, h *header) error {
388 const maxUint16 = 1<<16 - 1
389 if len(h.Name) > maxUint16 {
390 return errLongName
391 }
392 if len(h.Extra) > maxUint16 {
393 return errLongExtra
394 }
395
396 var buf [fileHeaderLen]byte
397 b := writeBuf(buf[:])
398 b.uint32(uint32(fileHeaderSignature))
399 b.uint16(h.ReaderVersion)
400 b.uint16(h.Flags)
401 b.uint16(h.Method)
402 b.uint16(h.ModifiedTime)
403 b.uint16(h.ModifiedDate)
404
405
406
407 if h.raw && !h.hasDataDescriptor() {
408 b.uint32(h.CRC32)
409 b.uint32(uint32(min(h.CompressedSize64, uint32max)))
410 b.uint32(uint32(min(h.UncompressedSize64, uint32max)))
411 } else {
412
413
414 b.uint32(0)
415 b.uint32(0)
416 b.uint32(0)
417 }
418 b.uint16(uint16(len(h.Name)))
419 b.uint16(uint16(len(h.Extra)))
420 if _, err := w.Write(buf[:]); err != nil {
421 return err
422 }
423 if _, err := io.WriteString(w, h.Name); err != nil {
424 return err
425 }
426 _, err := w.Write(h.Extra)
427 return err
428 }
429
430
431
432
433
434
435
436
437
438
439
440 func (w *Writer) CreateRaw(fh *FileHeader) (io.Writer, error) {
441 if err := w.prepare(fh); err != nil {
442 return nil, err
443 }
444
445 fh.CompressedSize = uint32(min(fh.CompressedSize64, uint32max))
446 fh.UncompressedSize = uint32(min(fh.UncompressedSize64, uint32max))
447
448 h := &header{
449 FileHeader: fh,
450 offset: uint64(w.cw.count),
451 raw: true,
452 }
453 w.dir = append(w.dir, h)
454 if err := writeHeader(w.cw, h); err != nil {
455 return nil, err
456 }
457
458 if strings.HasSuffix(fh.Name, "/") {
459 w.last = nil
460 return dirWriter{}, nil
461 }
462
463 fw := &fileWriter{
464 header: h,
465 zipw: w.cw,
466 }
467 w.last = fw
468 return fw, nil
469 }
470
471
472
473 func (w *Writer) Copy(f *File) error {
474 r, err := f.OpenRaw()
475 if err != nil {
476 return err
477 }
478
479
480 fh := f.FileHeader
481 fw, err := w.CreateRaw(&fh)
482 if err != nil {
483 return err
484 }
485 _, err = io.Copy(fw, r)
486 return err
487 }
488
489
490
491
492 func (w *Writer) RegisterCompressor(method uint16, comp Compressor) {
493 if w.compressors == nil {
494 w.compressors = make(map[uint16]Compressor)
495 }
496 w.compressors[method] = comp
497 }
498
499
500
501
502 func (w *Writer) AddFS(fsys fs.FS) error {
503 return fs.WalkDir(fsys, ".", func(name string, d fs.DirEntry, err error) error {
504 if err != nil {
505 return err
506 }
507 if d.IsDir() {
508 return nil
509 }
510 info, err := d.Info()
511 if err != nil {
512 return err
513 }
514 if !info.Mode().IsRegular() {
515 return errors.New("zip: cannot add non-regular file")
516 }
517 h, err := FileInfoHeader(info)
518 if err != nil {
519 return err
520 }
521 h.Name = name
522 h.Method = Deflate
523 fw, err := w.CreateHeader(h)
524 if err != nil {
525 return err
526 }
527 f, err := fsys.Open(name)
528 if err != nil {
529 return err
530 }
531 defer f.Close()
532 _, err = io.Copy(fw, f)
533 return err
534 })
535 }
536
537 func (w *Writer) compressor(method uint16) Compressor {
538 comp := w.compressors[method]
539 if comp == nil {
540 comp = compressor(method)
541 }
542 return comp
543 }
544
545 type dirWriter struct{}
546
547 func (dirWriter) Write(b []byte) (int, error) {
548 if len(b) == 0 {
549 return 0, nil
550 }
551 return 0, errors.New("zip: write to directory")
552 }
553
554 type fileWriter struct {
555 *header
556 zipw io.Writer
557 rawCount *countWriter
558 comp io.WriteCloser
559 compCount *countWriter
560 crc32 hash.Hash32
561 closed bool
562 }
563
564 func (w *fileWriter) Write(p []byte) (int, error) {
565 if w.closed {
566 return 0, errors.New("zip: write to closed file")
567 }
568 if w.raw {
569 return w.zipw.Write(p)
570 }
571 w.crc32.Write(p)
572 return w.rawCount.Write(p)
573 }
574
575 func (w *fileWriter) close() error {
576 if w.closed {
577 return errors.New("zip: file closed twice")
578 }
579 w.closed = true
580 if w.raw {
581 return w.writeDataDescriptor()
582 }
583 if err := w.comp.Close(); err != nil {
584 return err
585 }
586
587
588 fh := w.header.FileHeader
589 fh.CRC32 = w.crc32.Sum32()
590 fh.CompressedSize64 = uint64(w.compCount.count)
591 fh.UncompressedSize64 = uint64(w.rawCount.count)
592
593 if fh.isZip64() {
594 fh.CompressedSize = uint32max
595 fh.UncompressedSize = uint32max
596 fh.ReaderVersion = zipVersion45
597 } else {
598 fh.CompressedSize = uint32(fh.CompressedSize64)
599 fh.UncompressedSize = uint32(fh.UncompressedSize64)
600 }
601
602 return w.writeDataDescriptor()
603 }
604
605 func (w *fileWriter) writeDataDescriptor() error {
606 if !w.hasDataDescriptor() {
607 return nil
608 }
609
610
611
612
613
614 var buf []byte
615 if w.isZip64() {
616 buf = make([]byte, dataDescriptor64Len)
617 } else {
618 buf = make([]byte, dataDescriptorLen)
619 }
620 b := writeBuf(buf)
621 b.uint32(dataDescriptorSignature)
622 b.uint32(w.CRC32)
623 if w.isZip64() {
624 b.uint64(w.CompressedSize64)
625 b.uint64(w.UncompressedSize64)
626 } else {
627 b.uint32(w.CompressedSize)
628 b.uint32(w.UncompressedSize)
629 }
630 _, err := w.zipw.Write(buf)
631 return err
632 }
633
634 type countWriter struct {
635 w io.Writer
636 count int64
637 }
638
639 func (w *countWriter) Write(p []byte) (int, error) {
640 n, err := w.w.Write(p)
641 w.count += int64(n)
642 return n, err
643 }
644
645 type nopCloser struct {
646 io.Writer
647 }
648
649 func (w nopCloser) Close() error {
650 return nil
651 }
652
653 type writeBuf []byte
654
655 func (b *writeBuf) uint8(v uint8) {
656 (*b)[0] = v
657 *b = (*b)[1:]
658 }
659
660 func (b *writeBuf) uint16(v uint16) {
661 binary.LittleEndian.PutUint16(*b, v)
662 *b = (*b)[2:]
663 }
664
665 func (b *writeBuf) uint32(v uint32) {
666 binary.LittleEndian.PutUint32(*b, v)
667 *b = (*b)[4:]
668 }
669
670 func (b *writeBuf) uint64(v uint64) {
671 binary.LittleEndian.PutUint64(*b, v)
672 *b = (*b)[8:]
673 }
674
View as plain text