directory.go

Documentation: github.com/sassoftware/relic/lib/zipslicer

     1  //
     2  // Copyright (c) SAS Institute Inc.
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  //
    16  
    17  package zipslicer
    18  
    19  import (
    20  	"bufio"
    21  	"bytes"
    22  	"encoding/binary"
    23  	"errors"
    24  	"fmt"
    25  	"io"
    26  	"io/ioutil"
    27  )
    28  
    29  type Directory struct {
    30  	File   []*File
    31  	Size   int64
    32  	DirLoc int64
    33  	r      io.ReaderAt
    34  	end64  zip64End
    35  	loc64  zip64Loc
    36  	end    zipEndRecord
    37  }
    38  
    39  // Return the offset of the zip central directory
    40  func FindDirectory(r io.ReaderAt, size int64) (int64, error) {
    41  	pos := size - directoryEndLen - directory64LocLen
    42  	var endb [directoryEndLen + directory64LocLen]byte
    43  	if _, err := r.ReadAt(endb[:], pos); err != nil {
    44  		return 0, err
    45  	}
    46  	re := bytes.NewReader(endb[:])
    47  	var loc64 zip64Loc
    48  	var end zipEndRecord
    49  	binary.Read(re, binary.LittleEndian, &loc64)
    50  	binary.Read(re, binary.LittleEndian, &end)
    51  	if end.Signature != directoryEndSignature {
    52  		return 0, errors.New("zip central directory not found")
    53  	}
    54  	if end.TotalCDCount == uint16Max || end.CDSize == uint32Max || end.CDOffset == uint32Max {
    55  		if loc64.Signature != directory64LocSignature {
    56  			return 0, errors.New("expected ZIP64 locator")
    57  		}
    58  		// ZIP64
    59  		var end64b [directory64EndLen]byte
    60  		if _, err := r.ReadAt(end64b[:], int64(loc64.Offset)); err != nil {
    61  			return 0, err
    62  		}
    63  		var end64 zip64End
    64  		binary.Read(bytes.NewReader(end64b[:]), binary.LittleEndian, &end64)
    65  		if end64.Signature != directory64EndSignature {
    66  			return 0, errors.New("zip central directory not found")
    67  		}
    68  		return int64(end64.CDOffset), nil
    69  	}
    70  	return int64(end.CDOffset), nil
    71  }
    72  
    73  // Read a zip from a ReaderAt, with a separate copy of the central directory
    74  func ReadWithDirectory(r io.ReaderAt, size int64, cd []byte) (*Directory, error) {
    75  	dirLoc := size - int64(len(cd))
    76  	files := make([]*File, 0)
    77  	for {
    78  		if binary.LittleEndian.Uint32(cd) != directoryHeaderSignature {
    79  			break
    80  		}
    81  		var hdr zipCentralDir
    82  		binary.Read(bytes.NewReader(cd), binary.LittleEndian, &hdr)
    83  		f := &File{
    84  			CreatorVersion:   hdr.CreatorVersion,
    85  			ReaderVersion:    hdr.ReaderVersion,
    86  			Flags:            hdr.Flags,
    87  			Method:           hdr.Method,
    88  			ModifiedTime:     hdr.ModifiedTime,
    89  			ModifiedDate:     hdr.ModifiedDate,
    90  			CRC32:            hdr.CRC32,
    91  			CompressedSize:   uint64(hdr.CompressedSize),
    92  			UncompressedSize: uint64(hdr.UncompressedSize),
    93  			InternalAttrs:    hdr.InternalAttrs,
    94  			ExternalAttrs:    hdr.ExternalAttrs,
    95  			Offset:           uint64(hdr.Offset),
    96  
    97  			r:  r,
    98  			rs: size,
    99  		}
   100  		f.raw = make([]byte, directoryHeaderLen+int(hdr.FilenameLen)+int(hdr.ExtraLen)+int(hdr.CommentLen))
   101  		copy(f.raw, cd)
   102  		cd = cd[directoryHeaderLen:]
   103  		f.Name, cd = string(cd[:int(hdr.FilenameLen)]), cd[int(hdr.FilenameLen):]
   104  		f.Extra, cd = cd[:int(hdr.ExtraLen)], cd[int(hdr.ExtraLen):]
   105  		f.Comment, cd = cd[:int(hdr.CommentLen)], cd[int(hdr.CommentLen):]
   106  		needUSize := f.UncompressedSize == uint32Max
   107  		needCSize := f.CompressedSize == uint32Max
   108  		needOffset := f.Offset == uint32Max
   109  		extra := f.Extra
   110  		for len(extra) >= 4 {
   111  			tag := binary.LittleEndian.Uint16(extra[:2])
   112  			size := binary.LittleEndian.Uint16(extra[2:4])
   113  			if int(size) > len(extra)-4 {
   114  				break
   115  			}
   116  			if tag == zip64ExtraID {
   117  				e := extra[4 : 4+size]
   118  				if needUSize && size >= 8 {
   119  					f.UncompressedSize = binary.LittleEndian.Uint64(e)
   120  					needUSize = false
   121  				}
   122  				if needCSize && size >= 16 {
   123  					f.CompressedSize = binary.LittleEndian.Uint64(e[8:])
   124  					needCSize = false
   125  				}
   126  				if needOffset && size >= 24 {
   127  					f.Offset = binary.LittleEndian.Uint64(e[16:])
   128  					needOffset = false
   129  				}
   130  				break
   131  			}
   132  			extra = extra[4+size:]
   133  		}
   134  		if needCSize || needOffset {
   135  			return nil, errors.New("missing ZIP64 header")
   136  		}
   137  		files = append(files, f)
   138  	}
   139  	d := &Directory{
   140  		File:   files,
   141  		Size:   size,
   142  		DirLoc: dirLoc,
   143  		r:      r,
   144  	}
   145  	rd := bytes.NewReader(cd)
   146  	switch binary.LittleEndian.Uint32(cd) {
   147  	case directory64EndSignature:
   148  		binary.Read(rd, binary.LittleEndian, &d.end64)
   149  		binary.Read(rd, binary.LittleEndian, &d.loc64)
   150  	case directoryEndSignature:
   151  	default:
   152  		return nil, errors.New("expected end record")
   153  	}
   154  	binary.Read(rd, binary.LittleEndian, &d.end)
   155  	return d, nil
   156  }
   157  
   158  // Read a zip from a ReaderAt
   159  func Read(r io.ReaderAt, size int64) (*Directory, error) {
   160  	loc, err := FindDirectory(r, size)
   161  	if err != nil {
   162  		return nil, err
   163  	}
   164  	cd := make([]byte, size-loc)
   165  	if _, err := r.ReadAt(cd, loc); err != nil {
   166  		return nil, err
   167  	}
   168  	return ReadWithDirectory(r, size, cd)
   169  }
   170  
   171  // Read a zip from a stream, using a separate copy of the central directory.
   172  // Contents must be read in zip order or an error will be raised.
   173  func ReadStream(r io.Reader, size int64, cd []byte) (*Directory, error) {
   174  	ra := &streamReaderAt{r: r}
   175  	return ReadWithDirectory(ra, size, cd)
   176  }
   177  
   178  // Serialize a zip file with all of the files up to, but not including, the
   179  // given index. The contents and central directory are written to separate
   180  // writers, which may be the same writer.
   181  func (d *Directory) Truncate(n int, body, dir io.Writer) error {
   182  	if body != nil {
   183  		for i := 0; i < n; i++ {
   184  			f := d.File[i]
   185  			fs, err := f.GetTotalSize()
   186  			if err != nil {
   187  				return err
   188  			}
   189  			if _, err := io.Copy(body, io.NewSectionReader(d.r, int64(f.Offset), fs)); err != nil {
   190  				return err
   191  			}
   192  		}
   193  	}
   194  	cdOffset := d.File[n].Offset
   195  	var size uint64
   196  	for i := 0; i < n; i++ {
   197  		blob, err := d.File[i].GetDirectoryHeader()
   198  		if err != nil {
   199  			return err
   200  		}
   201  		dir.Write(blob)
   202  		size += uint64(len(blob))
   203  	}
   204  	end := d.end
   205  	if d.end64.Signature != 0 {
   206  		end64 := d.end64
   207  		end64.DiskCDCount = uint64(n)
   208  		end64.TotalCDCount = uint64(n)
   209  		end64.CDSize = size
   210  		end64.CDOffset = cdOffset
   211  		binary.Write(dir, binary.LittleEndian, end64)
   212  		loc := d.loc64
   213  		loc.Offset = cdOffset + size
   214  		binary.Write(dir, binary.LittleEndian, loc)
   215  	} else {
   216  		if cdOffset >= uint32Max || n >= uint16Max {
   217  			return errors.New("file too big for 32-bit ZIP")
   218  		}
   219  		end.DiskCDCount = uint16(n)
   220  		end.TotalCDCount = uint16(n)
   221  		end.CDSize = uint32(size)
   222  		end.CDOffset = uint32(cdOffset)
   223  	}
   224  	binary.Write(dir, binary.LittleEndian, end)
   225  	return nil
   226  }
   227  
   228  // Get the original central directory and end-of-directory from a previously-read file.
   229  //
   230  // If trim is true, then the end-of-directory will be updated to skip over any
   231  // non-ZIP data between the last file's contents and the first central
   232  // directory entry.
   233  func (d *Directory) GetOriginalDirectory(trim bool) (cdEntries, endOfDir []byte, err error) {
   234  	if d.end.Signature == 0 {
   235  		return nil, nil, errors.New("new zipfile, can't produce original directory")
   236  	}
   237  	var wcd, weod bytes.Buffer
   238  	if err := d.WriteDirectory(&wcd, nil, false); err != nil {
   239  		return nil, nil, err
   240  	}
   241  	end64 := d.end64
   242  	loc64 := d.loc64
   243  	end := d.end
   244  	if trim {
   245  		contentEnd, err := d.NextFileOffset()
   246  		if err != nil {
   247  			return nil, nil, err
   248  		}
   249  		delta := d.DirLoc - contentEnd
   250  		if delta < 0 || delta > uint32Max {
   251  			return nil, nil, errors.New("non-ZIP data out of bounds")
   252  		}
   253  		if end64.Signature != 0 {
   254  			end64.CDOffset -= uint64(delta)
   255  		}
   256  		if loc64.Signature != 0 {
   257  			loc64.Offset -= uint64(delta)
   258  		}
   259  		if end.CDOffset != uint32Max || loc64.Signature == 0 {
   260  			end.CDOffset -= uint32(delta)
   261  		}
   262  	}
   263  	binary.Write(&weod, binary.LittleEndian, end64)
   264  	binary.Write(&weod, binary.LittleEndian, loc64)
   265  	binary.Write(&weod, binary.LittleEndian, end)
   266  	return wcd.Bytes(), weod.Bytes(), nil
   267  }
   268  
   269  // Serialize a zip central directory to file. The file entries will be written
   270  // to wcd, and the end-of-directory markers will be written to weod.
   271  //
   272  // If forceZip64 is true then a ZIP64 end-of-directory marker will always be
   273  // written; otherwise it is only done if ZIP64 features are required.
   274  func (d *Directory) WriteDirectory(wcd, weod io.Writer, forceZip64 bool) error {
   275  	buf := bufio.NewWriter(wcd)
   276  	cdoff := d.DirLoc
   277  	var count, size uint64
   278  	minVersion := uint16(zip20)
   279  	for _, f := range d.File {
   280  		if f.ReaderVersion > minVersion {
   281  			minVersion = f.ReaderVersion
   282  		}
   283  		blob, err := f.GetDirectoryHeader()
   284  		if err != nil {
   285  			return err
   286  		}
   287  		if _, err := buf.Write(blob); err != nil {
   288  			return err
   289  		}
   290  		count++
   291  		size += uint64(len(blob))
   292  	}
   293  	if wcd != weod {
   294  		if err := buf.Flush(); err != nil {
   295  			return err
   296  		}
   297  		buf.Reset(weod)
   298  	} else if weod == nil {
   299  		return nil
   300  	}
   301  	var end zipEndRecord
   302  	if count >= uint16Max || size >= uint32Max || cdoff >= uint32Max || forceZip64 {
   303  		minVersion = zip45
   304  	}
   305  	if minVersion == zip45 {
   306  		end64off := cdoff + int64(size)
   307  		end64 := zip64End{
   308  			Signature:      directory64EndSignature,
   309  			RecordSize:     directory64EndLen - 12,
   310  			CreatorVersion: zip45,
   311  			ReaderVersion:  minVersion,
   312  			DiskCDCount:    count,
   313  			TotalCDCount:   count,
   314  			CDSize:         size,
   315  			CDOffset:       uint64(cdoff),
   316  		}
   317  		if err := binary.Write(buf, binary.LittleEndian, end64); err != nil {
   318  			return err
   319  		}
   320  		loc64 := zip64Loc{
   321  			Signature: directory64LocSignature,
   322  			Offset:    uint64(end64off),
   323  			DiskCount: 1,
   324  		}
   325  		if err := binary.Write(buf, binary.LittleEndian, loc64); err != nil {
   326  			return err
   327  		}
   328  		end = zipEndRecord{
   329  			Signature:    directoryEndSignature,
   330  			DiskCDCount:  uint16Max,
   331  			TotalCDCount: uint16Max,
   332  			CDSize:       uint32Max,
   333  			CDOffset:     uint32Max,
   334  		}
   335  	} else {
   336  		end = zipEndRecord{
   337  			Signature:    directoryEndSignature,
   338  			DiskCDCount:  uint16(count),
   339  			TotalCDCount: uint16(count),
   340  			CDSize:       uint32(size),
   341  			CDOffset:     uint32(cdoff),
   342  		}
   343  	}
   344  	if err := binary.Write(buf, binary.LittleEndian, end); err != nil {
   345  		return err
   346  	}
   347  	return buf.Flush()
   348  }
   349  
   350  type streamReaderAt struct {
   351  	r   io.Reader
   352  	pos int64
   353  }
   354  
   355  func (r *streamReaderAt) ReadAt(d []byte, p int64) (int, error) {
   356  	if p > r.pos {
   357  		if _, err := io.CopyN(ioutil.Discard, r.r, p-r.pos); err != nil {
   358  			return 0, err
   359  		}
   360  		r.pos = p
   361  	} else if p < r.pos {
   362  		return 0, fmt.Errorf("attempted to seek backwards: at %d, to %d", r.pos, p)
   363  	}
   364  	n, err := r.r.Read(d)
   365  	r.pos += int64(n)
   366  	return n, err
   367  }
   368  
   369  // Add a file to the central directory. Its contents are assumed to be already
   370  // located after the last added file.
   371  func (d *Directory) AddFile(f *File) (*File, error) {
   372  	size, err := f.GetTotalSize()
   373  	if err != nil {
   374  		return nil, err
   375  	}
   376  	offset := uint64(d.DirLoc)
   377  	if f.Offset != offset {
   378  		f.raw = nil
   379  	}
   380  	f.Offset = offset
   381  	d.DirLoc += size
   382  	d.File = append(d.File, f)
   383  	return f, nil
   384  }
   385  
   386  // Get the offset immediately following the last file's contents. This is the
   387  // same as DirLoc unless there is non-zip data in between.
   388  func (d *Directory) NextFileOffset() (int64, error) {
   389  	if len(d.File) == 0 {
   390  		return 0, nil
   391  	}
   392  	lastFile := d.File[len(d.File)-1]
   393  	size, err := lastFile.GetTotalSize()
   394  	if err != nil {
   395  		return 0, err
   396  	}
   397  	return int64(lastFile.Offset) + size, nil
   398  }
   399
View as plain text